mod_nesting

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;

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<Key> = simple::storage::keys()
  .s0()

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);

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<syn::Item> {
 }
 
 /// A `keys` module for all `Keys`-related items.
-pub(crate) fn module(
- vars: &[VarABI],
- unions: &BTreeSet<(String, Vec<UnionVariant>)>,
-) -> 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! {

pint-abi-gen/src/lib.rs

@@ -51,64 +51,81 @@ 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<syn::Field> {
+fn fields_from_tuple_fields(fields: &[TupleField], mod_level: &ModLevel) -> Vec<syn::Field> {
  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]
  }
 }
 
 /// 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<UnionVariant>)>,
  addr: Option<PredicateAddress>,
 ) -> 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,15 +208,14 @@ fn predicates_with_addrs<'a>(
 /// Generate a module for each named predicate.
 fn mods_from_named_predicates(
  predicates: &[PredicateABI],
- unions: &BTreeSet<(String, Vec<UnionVariant>)>,
  addrs: Option<&Addresses>,
 ) -> Vec<syn::ItemMod> {
  predicates_with_addrs(predicates, addrs)
  .filter(|(predicate, addr)| !is_predicate_empty(predicate) || addr.is_some())
  .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<UnionVariant>)>,
- addrs: Option<&Addresses>,
-) -> Vec<syn::Item> {
+fn items_from_predicates(predicates: &[PredicateABI], addrs: Option<&Addresses>) -> Vec<syn::Item> {
  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<UnionVariant>)>,
-) -> Vec<syn::Item> {
+fn items_from_keyed_vars(vars: &[VarABI]) -> Vec<syn::Item> {
  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<UnionVariant>)>,
-) -> syn::ItemMod {
+fn mod_from_keyed_vars(mod_name: &str, vars: &[VarABI]) -> syn::ItemMod {
  let mut items: Vec<syn::Item> = 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<syn::Variant>
  let split = name.split("::").collect::<Vec<_>>();
  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