pint_abi_gen/
lib.rs

1//! Macros for generating items from pint-generated contract ABI JSON.
2//!
3//! The entry points for this crate are:
4//!
5//! - [`from_file!`][crate::from_file!]
6//! - [`from_str!`][crate::from_str!]
7//!
8//! For a given contract, the following items are generated:
9//!
10//! - A `mod` representing `storage`.
11//! - For each `predicate`, a module with the following:
12//!     - An `Args` struct for the predicate's arguments.
13//!
14//! The aim for the generated items is to ease the construction of solutions
15//! including the encoding of keys, values and mutations from higher-level types.
16
17use addr::Addresses;
18use essential_types::{contract::Contract, predicate::Program, PredicateAddress};
19use pint_abi_types::{ContractABI, PredicateABI, StorageVarABI, TupleField, TypeABI, UnionVariant};
20use pint_abi_visit::Nesting;
21use proc_macro::TokenStream;
22use proc_macro2::Span;
23use quote::ToTokens;
24use std::collections::BTreeSet;
25use syn::parse_macro_input;
26
27mod addr;
28mod args;
29mod array;
30mod keys;
31mod macro_args;
32mod map;
33mod mutations;
34mod tuple;
35mod unions;
36mod utils;
37
38/// The name of the root module within the predicate set produced by the compiler.
39const ROOT_MOD_NAME: &str = "";
40
41/// Pint keyed var types that occupy only a single key.
42enum SingleKeyTy {
43    Bool,
44    Int,
45    Real,
46    String,
47    B256,
48    Optional(Box<TypeABI>),
49    Union(String /* Pint path to the union*/),
50}
51
52impl SingleKeyTy {
53    /// The type of the builder method value.
54    fn syn_ty(&self, mod_level: usize) -> syn::Type {
55        match self {
56            SingleKeyTy::Bool => syn::parse_quote!(bool),
57            SingleKeyTy::Int => syn::parse_quote!(i64),
58            SingleKeyTy::Real => syn::parse_quote!(f64),
59            SingleKeyTy::String => syn::parse_quote!(String),
60            SingleKeyTy::B256 => syn::parse_quote!([i64; 4]),
61            SingleKeyTy::Optional(ty) => ty_from_optional(ty, mod_level),
62            SingleKeyTy::Union(name) => unions::ty_from_union(name, mod_level),
63        }
64    }
65}
66
67/// Convert the given pint tuple fields to unnamed Rust fields.
68fn fields_from_tuple_fields(fields: &[TupleField], mod_level: usize) -> Vec<syn::Field> {
69    fields
70        .iter()
71        .map(|TupleField { name, ty }| {
72            // NOTE: Currently we ignore tuple field names.
73            let _name = name;
74            let ty = ty_from_pint_ty(ty, mod_level);
75            syn::parse_quote!(#ty)
76        })
77        .collect()
78}
79
80/// Convert the given pint tuple to an equivalent Rust tuple type.
81fn ty_from_tuple(tuple: &[TupleField], mod_level: usize) -> syn::Type {
82    let fields = fields_from_tuple_fields(tuple, mod_level);
83    syn::parse_quote! {
84        ( #( #fields ),* )
85    }
86}
87
88/// Convert the given pint array to an equivalent Rust array type.
89fn ty_from_array(ty: &TypeABI, size: i64, mod_level: usize) -> syn::Type {
90    let syn_ty = ty_from_pint_ty(ty, mod_level);
91    let len = usize::try_from(size).expect("array size out of range of `usize`");
92    syn::parse_quote! {
93        [#syn_ty; #len]
94    }
95}
96
97/// Convert the given pint tuple to an equivalent Rust tuple type.
98fn ty_from_optional(ty: &TypeABI, mod_level: usize) -> syn::Type {
99    let syn_ty = ty_from_pint_ty(ty, mod_level);
100    syn::parse_quote! {
101        Option<#syn_ty>
102    }
103}
104
105/// Convert the given pint ABI type to an equivalent Rust type.
106fn ty_from_pint_ty(ty: &TypeABI, mod_level: usize) -> syn::Type {
107    match ty {
108        TypeABI::Bool => syn::parse_quote!(bool),
109        TypeABI::Int => syn::parse_quote!(i64),
110        TypeABI::Real => syn::parse_quote!(f64),
111        TypeABI::String => syn::parse_quote!(String),
112        TypeABI::B256 => syn::parse_quote!([i64; 4]),
113        TypeABI::Optional(ty) => ty_from_optional(ty, mod_level),
114        TypeABI::Union { name, .. } => unions::ty_from_union(name, mod_level),
115        TypeABI::Tuple(tuple) => ty_from_tuple(tuple, mod_level),
116        TypeABI::Array { ty, size } => ty_from_array(ty, *size, mod_level),
117        TypeABI::Map { .. } => unreachable!("Maps are not allowed as non-storage types"),
118    }
119}
120
121/// Names for fields are emitted by pint with a `::` prefix.
122/// This function checks for the `::` prefix and strips it if it exists.
123fn strip_colons_prefix(name: &str) -> &str {
124    name.trim_start_matches("::")
125}
126
127/// Var names have the `::` prefix, and sometimes have `.` or `@` or `::` separators for
128/// flattened tuple fields. We strip the `::` prefix, and replace all `.` or `@` or `::`
129/// occurrences with `_`.
130fn field_name_from_var_name(name: &str) -> String {
131    strip_colons_prefix(name)
132        .replace(['.', '@'], "_")
133        .replace("::", "_")
134}
135
136/// Generate all items for the given predicate.
137fn items_from_predicate(
138    predicate: &PredicateABI,
139    addr: Option<&PredicateAddress>,
140) -> Vec<syn::Item> {
141    let mut items = vec![];
142    if let Some(addr) = addr {
143        items.push(addr::predicate_const(&addr.contract, &addr.predicate).into());
144    }
145    if !predicate.params.is_empty() {
146        items.extend(args::items(&predicate.params));
147    }
148    items
149}
150
151/// Generate a module for the given predicate.
152/// Modules are only generated for named predicates.
153fn mod_from_predicate(
154    name: &str,
155    predicate: &PredicateABI,
156    addr: Option<PredicateAddress>,
157) -> syn::ItemMod {
158    let doc_str = format!("Items for the `{name}` predicate.");
159    let ident = syn::Ident::new(name, Span::call_site());
160    let items = items_from_predicate(predicate, addr.as_ref());
161    syn::parse_quote! {
162        #[allow(non_snake_case)]
163        #[doc = #doc_str]
164        pub mod #ident {
165            #(
166                #items
167            )*
168        }
169    }
170}
171
172/// Whether or not the given predicate contains any items.
173fn is_predicate_empty(pred: &PredicateABI) -> bool {
174    pred.params.is_empty()
175}
176
177/// Zip the predicates with their addresses.
178fn predicates_with_addrs<'a>(
179    predicates: &'a [PredicateABI],
180    addrs: Option<&'a Addresses>,
181) -> impl 'a + Iterator<Item = (&'a PredicateABI, Option<PredicateAddress>)> {
182    predicates.iter().enumerate().map(move |(ix, predicate)| {
183        let addr = addrs.map(|addrs| PredicateAddress {
184            contract: addrs.contract.clone(),
185            predicate: addrs.predicates[ix].clone(),
186        });
187        (predicate, addr)
188    })
189}
190
191/// Generate a module for each named predicate.
192fn mods_from_named_predicates(
193    predicates: &[PredicateABI],
194    addrs: Option<&Addresses>,
195) -> Vec<syn::ItemMod> {
196    predicates_with_addrs(predicates, addrs)
197        .filter(|(predicate, addr)| !is_predicate_empty(predicate) || addr.is_some())
198        .filter(|(predicate, _)| predicate.name != ROOT_MOD_NAME)
199        .map(|(predicate, addr)| {
200            let name = strip_colons_prefix(&predicate.name);
201            mod_from_predicate(name, predicate, addr)
202        })
203        .collect()
204}
205
206/// Find the root predicate.
207fn find_root_predicate<'a>(
208    predicates: &'a [PredicateABI],
209    addrs: Option<&'a Addresses>,
210) -> Option<(&'a PredicateABI, Option<PredicateAddress>)> {
211    predicates_with_addrs(predicates, addrs).find(|(predicate, _)| predicate.name == ROOT_MOD_NAME)
212}
213
214/// Given the set of predicates, generate all items.
215///
216/// This includes a module for each named predicate, and types for the root predicate.
217fn items_from_predicates(predicates: &[PredicateABI], addrs: Option<&Addresses>) -> Vec<syn::Item> {
218    let mut items = vec![];
219    // Add the root predicate items.
220    if let Some((root_pred, addr)) = find_root_predicate(predicates, addrs) {
221        // By default, the root predicate has no name. We name its predicate module `root`.
222        let name = "root";
223        items.push(mod_from_predicate(name, root_pred, addr).into());
224    }
225    // Add the named predicate modules.
226    items.extend(
227        mods_from_named_predicates(predicates, addrs)
228            .into_iter()
229            .map(syn::Item::from),
230    );
231    items
232}
233
234/// Given a keyed var `Nesting`, create a Rust expression that results in its value.
235fn nesting_expr(nesting: &[Nesting]) -> syn::ExprArray {
236    let elems = nesting
237        .iter()
238        .map(|n| {
239            let expr: syn::Expr = match n {
240                Nesting::Var { ix } => {
241                    syn::parse_quote!(pint_abi::key::Nesting::Var { ix: #ix })
242                }
243                Nesting::TupleField { ix: _, flat_ix } => {
244                    syn::parse_quote!(pint_abi::key::Nesting::TupleField { flat_ix: #flat_ix })
245                }
246                Nesting::MapEntry { key_size: _ } => {
247                    syn::parse_quote!(pint_abi::key::Nesting::MapEntry)
248                }
249                Nesting::ArrayElem { elem_len, .. } => {
250                    syn::parse_quote!(pint_abi::key::Nesting::ArrayElem { elem_len: #elem_len })
251                }
252            };
253            expr
254        })
255        .collect();
256    syn::ExprArray {
257        attrs: vec![],
258        bracket_token: Default::default(),
259        elems,
260    }
261}
262
263/// A small expression for constructing a key from a keyed var nesting and a
264/// `Mutations` builder's `key_elems` stack.
265fn construct_key_expr() -> syn::Expr {
266    syn::parse_quote! {
267        pint_abi::key::construct(&nesting[..], &self.key_elems[..])
268    }
269}
270
271/// Convert the given type `&[Nesting]` into a string to use for a generated
272/// builder struct name.
273///
274/// E.g. `[Var { ix: 1 }, MapEntry { key_ty: i64 }, TupleField { ix: 3, .. }]`
275/// becomes `Var1_MapEntry_Tuple3` so that it may be appended to a builder type
276/// name, e.g. `Tuple_1_MapEntry_3`.
277fn nesting_ty_str<'a>(nesting: impl IntoIterator<Item = &'a Nesting>) -> String {
278    fn elem_str(nesting: &Nesting) -> String {
279        match nesting {
280            Nesting::Var { ix } => ix.to_string(),
281            Nesting::TupleField { ix, .. } => ix.to_string(),
282            Nesting::MapEntry { .. } => "MapEntry".to_string(),
283            Nesting::ArrayElem { .. } => "ArrayElem".to_string(),
284        }
285    }
286    let mut iter = nesting.into_iter();
287    let mut s = elem_str(iter.next().expect("nesting must contain at least one item"));
288    for n in iter {
289        use std::fmt::Write;
290        write!(&mut s, "_{}", elem_str(n)).expect("failed to fmt nesting ty str");
291    }
292    s
293}
294
295/// Given a type nesting, create a string for presenting the associated key in docs.
296///
297/// E.g. `[0, 1, _, _, _, _, 6, 7]`.
298fn nesting_key_doc_str(nesting: &[Nesting]) -> String {
299    use core::fmt::Write;
300    let partial_key = pint_abi_visit::partial_key_from_nesting(nesting);
301    let mut s = "[".to_string();
302    let mut opts = partial_key.iter();
303    fn write_opt(s: &mut String, opt: &Option<i64>) {
304        match opt {
305            None => write!(s, "_"),
306            Some(u) => write!(s, "{u}"),
307        }
308        .expect("failed to write key element to string")
309    }
310    if let Some(opt) = opts.next() {
311        write_opt(&mut s, opt);
312        for opt in opts {
313            write!(&mut s, ", ").unwrap();
314            write_opt(&mut s, opt);
315        }
316    }
317    write!(&mut s, "]").unwrap();
318    s
319}
320
321/// The `mutations` and `keys` items for the given keyed vars.
322///
323/// This is used for `storage` mod generation.
324fn items_from_keyed_vars(vars: &[StorageVarABI]) -> Vec<syn::Item> {
325    let mut items = vec![];
326
327    // The `mutations` module and re-exports.
328    items.push(mutations::module(vars).into());
329    items.push(syn::parse_quote! {
330        #[doc(inline)]
331        pub use mutations::{mutations, Mutations};
332    });
333
334    // The `keys` module and re-exports.
335    items.push(keys::module(vars).into());
336    items.push(syn::parse_quote! {
337        #[doc(inline)]
338        pub use keys::{keys, Keys};
339    });
340
341    items
342}
343
344/// Create a module with `mutations` and `keys` fns for the given keyed vars.
345///
346/// This is used for `storage` mod generation.
347fn mod_from_keyed_vars(mod_name: &str, vars: &[StorageVarABI]) -> syn::ItemMod {
348    let items = items_from_keyed_vars(vars);
349    let mod_ident = syn::Ident::new(mod_name, Span::call_site());
350    syn::parse_quote! {
351        pub mod #mod_ident {
352            //! Items related to simplifying the process of building sets of
353            //! [`Mutation`][pint_abi::types::essential::solution::Mutation]s and
354            //! [`Key`][pint_abi::types::essential::Key]s for
355            //! [`Solution`][pint_abi::types::essential::solution::Solution]s and queries.
356            //!
357            //!
358            //! See the [`mutations`](./fn.mutations.html) fn to start constructing
359            //! a set of [`Mutations`].
360            //!
361            //! See the [`keys`](./fn.mutations.html) fn to start constructing a set
362            //! of [`Keys`].
363            //!
364            //! The [`Mutations`] and [`Keys`] impls provides a set of builder
365            //! methods that allow for writing `Mutation`s and `Key`s to an
366            //! inner `Vec` from higher-level values.
367            //!
368            //! The final `Vec<Mutation>` or `Vec<Key>` can be produced using the
369            //! `From<Mutations>` or `From<Keys>` conversion impls.
370            #(
371                #items
372            )*
373        }
374    }
375}
376
377/// Given an ABI, generate all items.
378fn items_from_abi_and_addrs(abi: &ContractABI, addrs: Option<&Addresses>) -> Vec<syn::Item> {
379    let mut items = vec![];
380
381    // Collect all the union types encountered in the contract
382    let mut unions: BTreeSet<(Vec<String>, Vec<UnionVariant>)> = BTreeSet::new();
383    abi.storage
384        .iter()
385        .map(|var| var.ty.clone())
386        .chain(
387            abi.predicates
388                .iter()
389                .flat_map(|predicate| predicate.params.iter().map(|param| param.ty.clone())),
390        )
391        .for_each(|ty| unions::collect_unions(&ty, &mut unions));
392
393    items.extend(unions::items_from_unions(&unions));
394
395    items.extend(items_from_predicates(&abi.predicates, addrs));
396    if let Some(addrs) = addrs {
397        items.push(addr::contract_const(&addrs.contract).into());
398    }
399    if !abi.storage.is_empty() {
400        items.push(mod_from_keyed_vars("storage", &abi.storage).into());
401    }
402    items
403}
404
405/// Shorthand for producing tokens given the full deserialized ABI.
406fn tokens_from_abi_and_addrs(abi: &ContractABI, addrs: Option<&Addresses>) -> TokenStream {
407    let items = items_from_abi_and_addrs(abi, addrs);
408    items
409        .into_iter()
410        .map(|item| TokenStream::from(item.into_token_stream()))
411        .collect()
412}
413
414/// Given a path specified as an argument to the `from_file!` macro, resolve
415/// whether it's relative to the `CARGO_MANIFEST_DIR` or absolute.
416fn resolve_path(path: &std::path::Path) -> std::path::PathBuf {
417    if path.is_relative() {
418        let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")
419            .expect("`CARGO_MANIFEST_DIR` not set, but required for relative path expansion");
420        let manifest_dir_path = std::path::Path::new(&manifest_dir);
421        manifest_dir_path.join(path)
422    } else {
423        path.to_path_buf()
424    }
425}
426
427/// Generate all items from a raw ABI JSON string.
428#[proc_macro]
429pub fn from_str(input: TokenStream) -> TokenStream {
430    let input_lit_str = parse_macro_input!(input as syn::LitStr);
431    let string = input_lit_str.value();
432    let abi: ContractABI = serde_json::from_str(&string)
433        .expect("failed to deserialize str from JSON to `ContractABI`");
434    tokens_from_abi_and_addrs(&abi, None)
435}
436
437/// Read and deserialize an instance of type `T` from the JSON file.
438fn read_from_json_file<T>(path: &std::path::Path) -> Result<T, serde_json::Error>
439where
440    T: for<'de> serde::Deserialize<'de>,
441{
442    let file = std::fs::File::open(path).unwrap_or_else(|err| {
443        panic!("failed to open {path:?}: {err}");
444    });
445    let reader = std::io::BufReader::new(file);
446    serde_json::from_reader(reader)
447}
448
449/// Generate all items from an ABI JSON file at the given path.
450///
451/// Also allows for optionally providing a path to the associated contract JSON
452/// for generating `ADDRESS` consts for the contract and its predicates.
453///
454/// NOTE: This macro is designed for use via the `pint-abi` crate, from which
455/// it is re-exported as `pint_abi::gen_from_file!`. Much of the code generation
456/// assumes the `pint-abi` crate is available as a dependency.
457///
458/// ## Usage
459///
460/// ```ignore
461/// pint_abi::gen_from_file! {
462///     abi: "path/to/abi.json",
463///     contract: "path/to/contract.json",
464/// }
465/// ```
466///
467/// ## Supported Paths
468///
469/// The given path must be either:
470///
471/// 1. An absolute path or
472/// 2. A path that is relative to the project's `CARGO_MANIFEST_DIR`.
473///
474/// The limitation around relative paths is due to having no way to know the path
475/// to the source file in which the macro is being invoked in stable Rust.
476#[proc_macro]
477pub fn from_file(input: TokenStream) -> TokenStream {
478    let args = parse_macro_input!(input as macro_args::FromFile);
479
480    // Load the contract ABI.
481    let abi_path = resolve_path(args.abi.value().as_ref());
482    let abi: ContractABI = read_from_json_file(&abi_path).unwrap_or_else(|err| {
483        panic!("failed to deserialize {abi_path:?} from JSON to `ContractABI`: {err}");
484    });
485
486    // Load the contract itself if provided.
487    let contract: Option<Contract> = args
488        .contract
489        .map(|contract| {
490            let contract_path = resolve_path(contract.value().as_ref());
491            read_from_json_file::<(Contract, Vec<Program>)>(&contract_path).unwrap_or_else(|err| {
492                panic!("failed to deserialize {contract_path:?} from JSON to `Contract`: {err}");
493            })
494        })
495        .map(|(contract, _)| contract);
496
497    // Collect the contract and predicate addresses.
498    let addrs = contract.as_ref().map(Addresses::from);
499
500    tokens_from_abi_and_addrs(&abi, addrs.as_ref())
501}