ark_ff_macros/
lib.rs

1#![warn(
2    unused,
3    future_incompatible,
4    nonstandard_style,
5    rust_2018_idioms,
6    rust_2021_compatibility
7)]
8#![forbid(unsafe_code)]
9
10use num_bigint::BigUint;
11use proc_macro::TokenStream;
12use syn::{Expr, Item, ItemFn, Lit};
13
14mod montgomery;
15mod unroll;
16
17pub(crate) mod utils;
18
19#[proc_macro]
20pub fn to_sign_and_limbs(input: TokenStream) -> TokenStream {
21    let num = utils::parse_string(input).expect("expected decimal string");
22    let (is_positive, limbs) = utils::str_to_limbs(&num);
23
24    let limbs: String = limbs.join(", ");
25    let limbs_and_sign = format!("({}", is_positive) + ", [" + &limbs + "])";
26    let tuple: Expr = syn::parse_str(&limbs_and_sign).unwrap();
27    quote::quote!(#tuple).into()
28}
29
30/// Derive the `MontConfig` trait.
31///
32/// The attributes available to this macro are
33/// * `modulus`: Specify the prime modulus underlying this prime field.
34/// * `generator`: Specify the generator of the multiplicative subgroup of this
35///   prime field. This value must be a quadratic non-residue in the field.
36/// * `small_subgroup_base` and `small_subgroup_power` (optional): If the field
37///   has insufficient two-adicity, specify an additional subgroup of size
38///   `small_subgroup_base.pow(small_subgroup_power)`.
39// This code was adapted from the `PrimeField` Derive Macro in ff-derive.
40#[proc_macro_derive(
41    MontConfig,
42    attributes(modulus, generator, small_subgroup_base, small_subgroup_power)
43)]
44pub fn mont_config(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
45    // Parse the type definition
46    let ast: syn::DeriveInput = syn::parse(input).unwrap();
47
48    // We're given the modulus p of the prime field
49    let modulus: BigUint = fetch_attr("modulus", &ast.attrs)
50        .expect("Please supply a modulus attribute")
51        .parse()
52        .expect("Modulus should be a number");
53
54    // We may be provided with a generator of p - 1 order. It is required that this
55    // generator be quadratic nonresidue.
56    let generator: BigUint = fetch_attr("generator", &ast.attrs)
57        .expect("Please supply a generator attribute")
58        .parse()
59        .expect("Generator should be a number");
60
61    let small_subgroup_base: Option<u32> = fetch_attr("small_subgroup_base", &ast.attrs)
62        .map(|s| s.parse().expect("small_subgroup_base should be a number"));
63
64    let small_subgroup_power: Option<u32> = fetch_attr("small_subgroup_power", &ast.attrs)
65        .map(|s| s.parse().expect("small_subgroup_power should be a number"));
66
67    montgomery::mont_config_helper(
68        modulus,
69        generator,
70        small_subgroup_base,
71        small_subgroup_power,
72        ast.ident,
73    )
74    .into()
75}
76
77const ARG_MSG: &str = "Failed to parse unroll threshold; must be a positive integer";
78
79/// Attribute used to unroll for loops found inside a function block.
80#[proc_macro_attribute]
81pub fn unroll_for_loops(args: TokenStream, input: TokenStream) -> TokenStream {
82    let unroll_by = match syn::parse2::<syn::Lit>(args.into()).expect(ARG_MSG) {
83        Lit::Int(int) => int.base10_parse().expect(ARG_MSG),
84        _ => panic!("{}", ARG_MSG),
85    };
86
87    let item: Item = syn::parse(input).expect("Failed to parse input.");
88
89    if let Item::Fn(item_fn) = item {
90        let new_block = {
91            let &ItemFn {
92                block: ref box_block,
93                ..
94            } = &item_fn;
95            unroll::unroll_in_block(box_block, unroll_by)
96        };
97        let new_item = Item::Fn(ItemFn {
98            block: Box::new(new_block),
99            ..item_fn
100        });
101        quote::quote! ( #new_item ).into()
102    } else {
103        quote::quote! ( #item ).into()
104    }
105}
106
107/// Fetch an attribute string from the derived struct.
108fn fetch_attr(name: &str, attrs: &[syn::Attribute]) -> Option<String> {
109    for attr in attrs {
110        if let Ok(meta) = attr.parse_meta() {
111            match meta {
112                syn::Meta::NameValue(nv) => {
113                    if nv.path.get_ident().map(|i| i.to_string()) == Some(name.to_string()) {
114                        match nv.lit {
115                            syn::Lit::Str(ref s) => return Some(s.value()),
116                            _ => {
117                                panic!("attribute {} should be a string", name);
118                            },
119                        }
120                    }
121                },
122                _ => {
123                    panic!("attribute {} should be a string", name);
124                },
125            }
126        }
127    }
128
129    None
130}
131
132#[test]
133fn test_str_to_limbs() {
134    let (is_positive, limbs) = utils::str_to_limbs("-5");
135    assert!(!is_positive);
136    assert_eq!(&limbs, &["5u64".to_string()]);
137
138    let (is_positive, limbs) = utils::str_to_limbs("100");
139    assert!(is_positive);
140    assert_eq!(&limbs, &["100u64".to_string()]);
141
142    let large_num = -((1i128 << 64) + 101234001234i128);
143    let (is_positive, limbs) = utils::str_to_limbs(&large_num.to_string());
144    assert!(!is_positive);
145    assert_eq!(&limbs, &["101234001234u64".to_string(), "1u64".to_string()]);
146
147    let num = "80949648264912719408558363140637477264845294720710499478137287262712535938301461879813459410946";
148    let (is_positive, limbs) = utils::str_to_limbs(num);
149    assert!(is_positive);
150    let expected_limbs = [
151        format!("{}u64", 0x8508c00000000002u64),
152        format!("{}u64", 0x452217cc90000000u64),
153        format!("{}u64", 0xc5ed1347970dec00u64),
154        format!("{}u64", 0x619aaf7d34594aabu64),
155        format!("{}u64", 0x9b3af05dd14f6ecu64),
156    ];
157    assert_eq!(&limbs, &expected_limbs);
158}