diff --git a/pint-abi-gen-tests/test-pkgs/simple/src/contract.pnt b/pint-abi-gen-tests/test-pkgs/simple/src/contract.pnt index 25f5d4c6..7c3f1049 100644 --- a/pint-abi-gen-tests/test-pkgs/simple/src/contract.pnt +++ b/pint-abi-gen-tests/test-pkgs/simple/src/contract.pnt @@ -1,5 +1,5 @@ -union U = A(int) | B | C(b256); -union WW = D | E(U); +union UU = A(int) | B | C(b256); +union WW = D | E(UU); storage { s0: bool, @@ -12,9 +12,9 @@ storage { my_nested_map0: (int => (int => int)), my_nested_map1: (int => (b256 => { int, { b256, int } })), my_array: int[5], - u1: U, - u2: U, - u3: U, + u1: UU, + u2: UU, + u3: UU, w1: WW, } @@ -85,10 +85,10 @@ predicate Foo { state u3 = mut storage::u3; state w1 = mut storage::w1; - constraint u1' == U::A(69); - constraint u2' == U::B; - constraint u3' == U::C(0x6969696969696969696969696969696969696969696969696969696969696969); - constraint w1' == WW::E(U::A(69)); + constraint u1' == UU::A(69); + constraint u2' == UU::B; + constraint u3' == UU::C(0x6969696969696969696969696969696969696969696969696969696969696969); + constraint w1' == WW::E(UU::A(69)); // Check the state. constraint s0' == true; diff --git a/pint-abi-gen-tests/tests/simple.rs b/pint-abi-gen-tests/tests/simple.rs index 8ce30086..56edf4d5 100644 --- a/pint-abi-gen-tests/tests/simple.rs +++ b/pint-abi-gen-tests/tests/simple.rs @@ -96,14 +96,12 @@ async fn test_solution_foo() { .enumerate() .fold(arr, |arr, (ix, val)| arr.entry(ix, val)) }) - .u1(simple::U::A(69)) - .u2(simple::U::B) - .u3(simple::U::C([0x6969696969696969; 4])) - .w1(simple::WW::E(simple::U::A(69))) + .u1(simple::UU::A(69)) + .u2(simple::UU::B) + .u3(simple::UU::C([0x6969696969696969; 4])) + .w1(simple::WW::E(simple::UU::A(69))) .into(); - dbg!(pint_abi::encode(&simple::WW::E(simple::U::A(69)))); - // Build the same set of keys, so we can ensure they match the mutations. let keys: Vec = simple::storage::keys() .s0() diff --git a/pint-abi-gen/src/array.rs b/pint-abi-gen/src/array.rs index bbad323c..0af77581 100644 --- a/pint-abi-gen/src/array.rs +++ b/pint-abi-gen/src/array.rs @@ -2,7 +2,7 @@ use crate::{ construct_key_expr, keys, map, mutations, nesting_expr, nesting_key_doc_str, nesting_ty_str, - tuple, SingleKeyTy, TypeABI, + tuple, ModLevel, SingleKeyTy, TypeABI, }; use pint_abi_visit::{KeyedVarTree, Nesting, NodeIx}; use proc_macro2::Span; @@ -224,7 +224,7 @@ fn mutation_method_for_tuple(tup_nesting: &[Nesting]) -> syn::ImplItemFn { /// An array mutation builder method for an element with a single-key value. fn mutation_method_for_single_key(ty: &SingleKeyTy, elem_nesting: &[Nesting]) -> syn::ImplItemFn { - let ty = ty.syn_ty(); + let ty = ty.syn_ty(&ModLevel::Mutations); let nesting_expr: syn::ExprArray = nesting_expr(elem_nesting); let construct_key_expr: syn::Expr = construct_key_expr(); let nesting_key_doc_str = nesting_key_doc_str(elem_nesting); diff --git a/pint-abi-gen/src/keys.rs b/pint-abi-gen/src/keys.rs index 5596c79d..c73a3627 100644 --- a/pint-abi-gen/src/keys.rs +++ b/pint-abi-gen/src/keys.rs @@ -3,10 +3,9 @@ //! a solution. use crate::{array, map, tuple}; -use pint_abi_types::{TypeABI, UnionVariant, VarABI}; +use pint_abi_types::{TypeABI, VarABI}; use pint_abi_visit::{KeyedVarTree, Nesting, NodeIx}; use proc_macro2::Span; -use std::collections::BTreeSet; /// Recursively traverse the given keyed var and create a builder struct and /// associated impl for each tuple, map and array. @@ -245,19 +244,9 @@ fn keys_items(vars: &[VarABI]) -> Vec { } /// A `keys` module for all `Keys`-related items. -pub(crate) fn module( - vars: &[VarABI], - unions: &BTreeSet<(String, Vec)>, -) -> syn::ItemMod { +pub(crate) fn module(vars: &[VarABI]) -> syn::ItemMod { let mut items = vec![]; - for (name, _) in unions { - let ident = syn::Ident::new(crate::strip_colons_prefix(name), Span::call_site()); - items.push(syn::parse_quote! { - use super::#ident; - }); - } - items.extend(keys_items(vars)); syn::parse_quote! { diff --git a/pint-abi-gen/src/lib.rs b/pint-abi-gen/src/lib.rs index 6e75afce..32cb1090 100644 --- a/pint-abi-gen/src/lib.rs +++ b/pint-abi-gen/src/lib.rs @@ -51,47 +51,64 @@ enum SingleKeyTy { impl SingleKeyTy { /// The type of the builder method value. - fn syn_ty(&self) -> syn::Type { + fn syn_ty(&self, mod_level: &ModLevel) -> syn::Type { match self { SingleKeyTy::Bool => syn::parse_quote!(bool), SingleKeyTy::Int => syn::parse_quote!(i64), SingleKeyTy::Real => syn::parse_quote!(f64), SingleKeyTy::String => syn::parse_quote!(String), SingleKeyTy::B256 => syn::parse_quote!([i64; 4]), - SingleKeyTy::Union(name) => ty_from_union(name), + SingleKeyTy::Union(name) => ty_from_union(name, mod_level), } } } -fn ty_from_union(name: &str) -> syn::Type { +#[derive(Clone)] +pub(crate) enum ModLevel { + Top, + Predicate, + Mutations, + + #[allow(dead_code)] + Storage, + #[allow(dead_code)] + Keys, +} + +fn ty_from_union(name: &str, mod_level: &ModLevel) -> syn::Type { let ident = syn::Ident::new(strip_colons_prefix(name), Span::call_site()); - syn::parse_quote!(#ident) + + match mod_level { + ModLevel::Top => syn::parse_quote! { #ident }, + ModLevel::Predicate | ModLevel::Storage => syn::parse_quote! { super::#ident }, + ModLevel::Mutations | ModLevel::Keys => syn::parse_quote! { super::super::#ident }, + } } /// Convert the given pint tuple fields to unnamed Rust fields. -fn fields_from_tuple_fields(fields: &[TupleField]) -> Vec { +fn fields_from_tuple_fields(fields: &[TupleField], mod_level: &ModLevel) -> Vec { fields .iter() .map(|TupleField { name, ty }| { // NOTE: Currently we ignore tuple field names. let _name = name; - let ty = ty_from_pint_ty(ty); + let ty = ty_from_pint_ty(ty, mod_level); syn::parse_quote!(#ty) }) .collect() } /// Convert the given pint tuple to an equivalent Rust tuple type. -fn ty_from_tuple(tuple: &[TupleField]) -> syn::Type { - let fields = fields_from_tuple_fields(tuple); +fn ty_from_tuple(tuple: &[TupleField], mod_level: &ModLevel) -> syn::Type { + let fields = fields_from_tuple_fields(tuple, mod_level); syn::parse_quote! { ( #( #fields ),* ) } } /// Convert the given pint array to an equivalent Rust array type. -fn ty_from_array(ty: &TypeABI, size: i64) -> syn::Type { - let syn_ty = ty_from_pint_ty(ty); +fn ty_from_array(ty: &TypeABI, size: i64, mod_level: &ModLevel) -> syn::Type { + let syn_ty = ty_from_pint_ty(ty, mod_level); let len = usize::try_from(size).expect("array size out of range of `usize`"); syn::parse_quote! { [#syn_ty; #len] @@ -99,16 +116,16 @@ fn ty_from_array(ty: &TypeABI, size: i64) -> syn::Type { } /// Convert the given pint ABI type to an equivalent Rust type. -fn ty_from_pint_ty(ty: &TypeABI) -> syn::Type { +fn ty_from_pint_ty(ty: &TypeABI, mod_level: &ModLevel) -> syn::Type { match ty { TypeABI::Bool => syn::parse_quote!(bool), TypeABI::Int => syn::parse_quote!(i64), TypeABI::Real => syn::parse_quote!(f64), TypeABI::String => syn::parse_quote!(String), TypeABI::B256 => syn::parse_quote!([i64; 4]), - TypeABI::Union { name, .. } => ty_from_union(name), - TypeABI::Tuple(tuple) => ty_from_tuple(tuple), - TypeABI::Array { ty, size } => ty_from_array(ty, *size), + TypeABI::Union { name, .. } => ty_from_union(name, mod_level), + TypeABI::Tuple(tuple) => ty_from_tuple(tuple, mod_level), + TypeABI::Array { ty, size } => ty_from_array(ty, *size, mod_level), TypeABI::Map { .. } => unreachable!("Maps are not allowed as non-storage types"), } } @@ -151,20 +168,12 @@ fn items_from_predicate( fn mod_from_predicate( name: &str, predicate: &PredicateABI, - unions: &BTreeSet<(String, Vec)>, addr: Option, ) -> syn::ItemMod { let doc_str = format!("Items for the `{name}` predicate."); let ident = syn::Ident::new(name, Span::call_site()); let mut items = vec![]; - for (name, _) in unions { - let ident = syn::Ident::new(strip_colons_prefix(name), Span::call_site()); - items.push(syn::parse_quote! { - use super::#ident; - }); - } - items.extend(items_from_predicate(predicate, addr.as_ref())); syn::parse_quote! { #[allow(non_snake_case)] @@ -199,7 +208,6 @@ fn predicates_with_addrs<'a>( /// Generate a module for each named predicate. fn mods_from_named_predicates( predicates: &[PredicateABI], - unions: &BTreeSet<(String, Vec)>, addrs: Option<&Addresses>, ) -> Vec { predicates_with_addrs(predicates, addrs) @@ -207,7 +215,7 @@ fn mods_from_named_predicates( .filter(|(predicate, _)| predicate.name != ROOT_MOD_NAME) .map(|(predicate, addr)| { let name = strip_colons_prefix(&predicate.name); - mod_from_predicate(name, predicate, unions, addr) + mod_from_predicate(name, predicate, addr) }) .collect() } @@ -223,21 +231,17 @@ fn find_root_predicate<'a>( /// Given the set of predicates, generate all items. /// /// This includes a module for each named predicate, and types for the root predicate. -fn items_from_predicates( - predicates: &[PredicateABI], - unions: &BTreeSet<(String, Vec)>, - addrs: Option<&Addresses>, -) -> Vec { +fn items_from_predicates(predicates: &[PredicateABI], addrs: Option<&Addresses>) -> Vec { let mut items = vec![]; // Add the root predicate items. if let Some((root_pred, addr)) = find_root_predicate(predicates, addrs) { // By default, the root predicate has no name. We name its predicate module `root`. let name = "root"; - items.push(mod_from_predicate(name, root_pred, unions, addr).into()); + items.push(mod_from_predicate(name, root_pred, addr).into()); } // Add the named predicate modules. items.extend( - mods_from_named_predicates(predicates, unions, addrs) + mods_from_named_predicates(predicates, addrs) .into_iter() .map(syn::Item::from), ); @@ -334,21 +338,18 @@ fn nesting_key_doc_str(nesting: &[Nesting]) -> String { /// The `mutations` and `keys` items for the given keyed vars. /// /// This is used for both `storage` and `pub_vars` mod generation. -fn items_from_keyed_vars( - vars: &[VarABI], - unions: &BTreeSet<(String, Vec)>, -) -> Vec { +fn items_from_keyed_vars(vars: &[VarABI]) -> Vec { let mut items = vec![]; // The `mutations` module and re-exports. - items.push(mutations::module(vars, unions).into()); + items.push(mutations::module(vars).into()); items.push(syn::parse_quote! { #[doc(inline)] pub use mutations::{mutations, Mutations}; }); // The `keys` module and re-exports. - items.push(keys::module(vars, unions).into()); + items.push(keys::module(vars).into()); items.push(syn::parse_quote! { #[doc(inline)] pub use keys::{keys, Keys}; @@ -360,22 +361,11 @@ fn items_from_keyed_vars( /// Create a module with `mutations` and `keys` fns for the given keyed vars. /// /// This is used for `storage` mod generation. -fn mod_from_keyed_vars( - mod_name: &str, - vars: &[VarABI], - unions: &BTreeSet<(String, Vec)>, -) -> syn::ItemMod { +fn mod_from_keyed_vars(mod_name: &str, vars: &[VarABI]) -> syn::ItemMod { let mut items: Vec = vec![]; let mod_ident = syn::Ident::new(mod_name, Span::call_site()); - for (name, _) in unions { - let ident = syn::Ident::new(strip_colons_prefix(name), Span::call_site()); - items.push(syn::parse_quote! { - use super::#ident; - }); - } - - items.extend(items_from_keyed_vars(vars, unions)); + items.extend(items_from_keyed_vars(vars)); syn::parse_quote! { pub mod #mod_ident { @@ -412,7 +402,7 @@ fn variants_from_union_variants(variants: &[UnionVariant]) -> Vec let split = name.split("::").collect::>(); let variant_name = syn::Ident::new(split[1], Span::call_site()); if let Some(ty) = ty { - let ty = ty_from_pint_ty(ty); + let ty = ty_from_pint_ty(ty, &ModLevel::Top); syn::parse_quote!( #variant_name ( #ty ) ) @@ -541,14 +531,14 @@ fn items_from_abi_and_addrs(abi: &ContractABI, addrs: Option<&Addresses>) -> Vec items.extend(items_from_unions(&unions)); - items.extend(items_from_predicates(&abi.predicates, &unions, addrs)); + items.extend(items_from_predicates(&abi.predicates, addrs)); if let Some(addrs) = addrs { items.push(addr::contract_const(&addrs.contract).into()); } if !abi.storage.is_empty() { - items.push(mod_from_keyed_vars("storage", &abi.storage, &unions).into()); + items.push(mod_from_keyed_vars("storage", &abi.storage).into()); } items diff --git a/pint-abi-gen/src/map.rs b/pint-abi-gen/src/map.rs index 720ed342..516973ed 100644 --- a/pint-abi-gen/src/map.rs +++ b/pint-abi-gen/src/map.rs @@ -2,7 +2,7 @@ use crate::{ array, construct_key_expr, keys, mutations, nesting_expr, nesting_key_doc_str, nesting_ty_str, - tuple, ty_from_pint_ty, SingleKeyTy, + tuple, ty_from_pint_ty, ModLevel, SingleKeyTy, }; use pint_abi_types::TypeABI; use pint_abi_visit::{KeyedVarTree, Nesting, NodeIx}; @@ -36,7 +36,7 @@ fn keys_struct(struct_name: &str, nesting: &[Nesting]) -> syn::ItemStruct { /// A map key builder method for entries with nested array values. fn key_method_for_array(ty_from: &TypeABI, array_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let struct_name = array::struct_name(array_nesting); let struct_ident = syn::Ident::new(&struct_name, Span::call_site()); syn::parse_quote! { @@ -53,7 +53,7 @@ fn key_method_for_array(ty_from: &TypeABI, array_nesting: &[Nesting]) -> syn::Im /// A map key builder method for entries with nested map values. fn key_method_for_map(ty_from: &TypeABI, map_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let struct_name = struct_name(map_nesting); let struct_ident = syn::Ident::new(&struct_name, Span::call_site()); syn::parse_quote! { @@ -70,7 +70,7 @@ fn key_method_for_map(ty_from: &TypeABI, map_nesting: &[Nesting]) -> syn::ImplIt /// A map key builder method for entries with tuple values. fn key_method_for_tuple(ty_from: &TypeABI, tup_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let struct_name = tuple::struct_name(tup_nesting); let struct_ident = syn::Ident::new(&struct_name, Span::call_site()); syn::parse_quote! { @@ -87,7 +87,7 @@ fn key_method_for_tuple(ty_from: &TypeABI, tup_nesting: &[Nesting]) -> syn::Impl /// A map key builder method for an entry with a single-key value. fn key_method_for_single_key(ty_from: &TypeABI, val_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let nesting_expr: syn::ExprArray = nesting_expr(val_nesting); let construct_key_expr: syn::Expr = construct_key_expr(); let abi_key_doc_str = nesting_key_doc_str(val_nesting); @@ -206,7 +206,7 @@ fn mutations_struct(struct_name: &str, nesting: &[Nesting]) -> syn::ItemStruct { /// A map mutation builder method for entries with nested array values. fn mutation_method_for_array(ty_from: &TypeABI, array_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let struct_name = array::struct_name(array_nesting); let struct_ident = syn::Ident::new(&struct_name, Span::call_site()); syn::parse_quote! { @@ -223,7 +223,7 @@ fn mutation_method_for_array(ty_from: &TypeABI, array_nesting: &[Nesting]) -> sy /// A map mutation builder method for entries with nested map values. fn mutation_method_for_map(ty_from: &TypeABI, map_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let struct_name = struct_name(map_nesting); let struct_ident = syn::Ident::new(&struct_name, Span::call_site()); syn::parse_quote! { @@ -240,7 +240,7 @@ fn mutation_method_for_map(ty_from: &TypeABI, map_nesting: &[Nesting]) -> syn::I /// A map mutation builder method for entries with tuple values. fn mutation_method_for_tuple(ty_from: &TypeABI, tup_nesting: &[Nesting]) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); let struct_name = tuple::struct_name(tup_nesting); let struct_ident = syn::Ident::new(&struct_name, Span::call_site()); syn::parse_quote! { @@ -261,8 +261,8 @@ fn mutation_method_for_single_key( val_ty: &SingleKeyTy, val_nesting: &[Nesting], ) -> syn::ImplItemFn { - let key_ty = ty_from_pint_ty(ty_from); - let val_ty = val_ty.syn_ty(); + let key_ty = ty_from_pint_ty(ty_from, &ModLevel::Mutations); + let val_ty = val_ty.syn_ty(&ModLevel::Mutations); let nesting_expr: syn::ExprArray = nesting_expr(val_nesting); let construct_key_expr: syn::Expr = construct_key_expr(); let abi_key_doc_str = nesting_key_doc_str(val_nesting); diff --git a/pint-abi-gen/src/mutations.rs b/pint-abi-gen/src/mutations.rs index f8af1301..59e11a6f 100644 --- a/pint-abi-gen/src/mutations.rs +++ b/pint-abi-gen/src/mutations.rs @@ -2,11 +2,10 @@ //! aimed at making it easier to build a set of //! [`Mutation`][essential_types::solution::Mutation]s for a solution. -use crate::{array, map, tuple, SingleKeyTy}; -use pint_abi_types::{TypeABI, UnionVariant, VarABI}; +use crate::{array, map, tuple, ModLevel, SingleKeyTy}; +use pint_abi_types::{TypeABI, VarABI}; use pint_abi_visit::{KeyedVarTree, Nesting, NodeIx}; use proc_macro2::Span; -use std::collections::BTreeSet; /// Recursively traverse the given keyed var and create a builder struct and /// associated impl for each tuple, map and array. @@ -112,7 +111,7 @@ fn method_for_single_key(name: &str, arg_ty: &SingleKeyTy, nesting: &[Nesting]) let method_ident = syn::Ident::new(name, Span::call_site()); let nesting_key_doc_str = crate::nesting_key_doc_str(nesting); let doc_str = format!("{}\n\nKey: `{nesting_key_doc_str}`", method_doc_str(name)); - let arg_ty = arg_ty.syn_ty(); + let arg_ty = arg_ty.syn_ty(&ModLevel::Mutations); let nesting_expr: syn::ExprArray = crate::nesting_expr(nesting); let construct_key_expr: syn::Expr = crate::construct_key_expr(); syn::parse_quote! { @@ -248,18 +247,8 @@ fn mutations_items(vars: &[VarABI]) -> Vec { } /// A `mutations` module for all `Mutations`-related items. -pub(crate) fn module( - vars: &[VarABI], - unions: &BTreeSet<(String, Vec)>, -) -> syn::ItemMod { +pub(crate) fn module(vars: &[VarABI]) -> syn::ItemMod { let mut items = vec![]; - for (name, _) in unions { - let ident = syn::Ident::new(crate::strip_colons_prefix(name), Span::call_site()); - items.push(syn::parse_quote! { - use super::super::#ident; - }); - } - items.extend(mutations_items(vars)); syn::parse_quote! { diff --git a/pint-abi-gen/src/pub_vars.rs b/pint-abi-gen/src/pub_vars.rs index c5586a11..9e8073e5 100644 --- a/pint-abi-gen/src/pub_vars.rs +++ b/pint-abi-gen/src/pub_vars.rs @@ -1,11 +1,14 @@ //! Struct and impls for the decision variables `PubVars` type. -use crate::utils::{field_idents, fields}; +use crate::{ + utils::{field_idents, fields}, + ModLevel, +}; use pint_abi_types::VarABI; /// Generate a struct for an predicate's decision variables. fn struct_decl(pub_vars: &[VarABI]) -> syn::ItemStruct { - let fields = fields(pub_vars); + let fields = fields(pub_vars, &ModLevel::Predicate); syn::parse_quote! { /// The predicate's decision variables. #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord)] diff --git a/pint-abi-gen/src/utils.rs b/pint-abi-gen/src/utils.rs index ecad2768..45ad5a5e 100644 --- a/pint-abi-gen/src/utils.rs +++ b/pint-abi-gen/src/utils.rs @@ -1,16 +1,16 @@ //! Utility functions for working with `VarABI` -use crate::{field_name_from_var_name, ty_from_pint_ty}; +use crate::{field_name_from_var_name, ty_from_pint_ty, ModLevel}; use pint_abi_types::VarABI; use proc_macro2::Span; /// A named field for each of the decision variables. -pub(super) fn fields(vars: &[VarABI]) -> Vec { +pub(super) fn fields(vars: &[VarABI], mod_level: &ModLevel) -> Vec { vars.iter() .map(|var| { let name = field_name_from_var_name(&var.name); let ident = syn::Ident::new(&name, Span::call_site()); - let ty = ty_from_pint_ty(&var.ty); + let ty = ty_from_pint_ty(&var.ty, mod_level); syn::parse_quote! { pub #ident: #ty } diff --git a/pint-abi-gen/src/vars.rs b/pint-abi-gen/src/vars.rs index e15a953c..fd171bae 100644 --- a/pint-abi-gen/src/vars.rs +++ b/pint-abi-gen/src/vars.rs @@ -1,11 +1,14 @@ //! Struct and impls for the decision variables `Vars` type. -use crate::utils::{field_idents, fields}; +use crate::{ + utils::{field_idents, fields}, + ModLevel, +}; use pint_abi_types::VarABI; /// Generate a struct for an predicate's decision variables. fn struct_decl(vars: &[VarABI]) -> syn::ItemStruct { - let fields = fields(vars); + let fields = fields(vars, &ModLevel::Predicate); syn::parse_quote! { /// The predicate's decision variables. #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord)]