From a76f4652095683ffddd5014cf3e1fbd2220c2b3a Mon Sep 17 00:00:00 2001 From: Mohammad Fawaz Date: Mon, 7 Oct 2024 11:43:25 -0400 Subject: [PATCH 1/2] unions abi wip --- .../test-pkgs/simple/src/contract.pnt | 17 ++ pint-abi-gen-tests/tests/simple.rs | 6 + pint-abi-gen/src/array.rs | 10 +- pint-abi-gen/src/keys.rs | 36 +++- pint-abi-gen/src/lib.rs | 204 ++++++++++++++++-- pint-abi-gen/src/map.rs | 12 +- pint-abi-gen/src/mutations.rs | 29 ++- pint-abi-types/src/lib.rs | 14 +- pint-abi-visit/src/lib.rs | 17 +- pintc/src/asm_gen/asm_builder.rs | 2 +- pintc/src/types.rs | 26 ++- 11 files changed, 337 insertions(+), 36 deletions(-) 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 3776adf8..25f5d4c6 100644 --- a/pint-abi-gen-tests/test-pkgs/simple/src/contract.pnt +++ b/pint-abi-gen-tests/test-pkgs/simple/src/contract.pnt @@ -1,3 +1,6 @@ +union U = A(int) | B | C(b256); +union WW = D | E(U); + storage { s0: bool, s1: int, @@ -9,6 +12,10 @@ 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, + w1: WW, } predicate Foo { @@ -73,6 +80,16 @@ predicate Foo { state my_nested_map1_key2_3: { int, { b256, int } } = mut storage::my_nested_map1[2][0x3333333333333333333333333333333333333333333333333333333333333333]; state my_array: int[5] = mut storage::my_array; + state u1 = mut storage::u1; + state u2 = mut storage::u2; + 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)); + // Check the state. constraint s0' == true; constraint s1' == 42; diff --git a/pint-abi-gen-tests/tests/simple.rs b/pint-abi-gen-tests/tests/simple.rs index 9d9fdab8..8ce30086 100644 --- a/pint-abi-gen-tests/tests/simple.rs +++ b/pint-abi-gen-tests/tests/simple.rs @@ -96,8 +96,14 @@ 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))) .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 d26503eb..bbad323c 100644 --- a/pint-abi-gen/src/array.rs +++ b/pint-abi-gen/src/array.rs @@ -111,7 +111,12 @@ fn key_method_for_single_key(elem_nesting: &[Nesting]) -> syn::ImplItemFn { fn key_method(tree: &KeyedVarTree, elem: NodeIx, elem_ty: &TypeABI) -> syn::ImplItemFn { let elem_nesting = tree.nesting(elem); match elem_ty { - TypeABI::Bool | TypeABI::Int | TypeABI::Real | TypeABI::String | TypeABI::B256 => (), + TypeABI::Bool + | TypeABI::Int + | TypeABI::Real + | TypeABI::String + | TypeABI::B256 + | TypeABI::Union { .. } => (), TypeABI::Array { .. } => return key_method_for_array(&elem_nesting), TypeABI::Tuple { .. } => { return key_method_for_tuple(&elem_nesting); @@ -259,12 +264,13 @@ fn mutation_method(tree: &KeyedVarTree, elem: NodeIx, elem_ty: &TypeABI) -> syn: TypeABI::Bool => SingleKeyTy::Bool, TypeABI::Int => SingleKeyTy::Int, TypeABI::Real => SingleKeyTy::Real, - TypeABI::Array { .. } => return mutation_method_for_array(&elem_nesting), TypeABI::String => SingleKeyTy::String, TypeABI::B256 => SingleKeyTy::B256, + TypeABI::Array { .. } => return mutation_method_for_array(&elem_nesting), TypeABI::Tuple { .. } => { return mutation_method_for_tuple(&elem_nesting); } + TypeABI::Union { name, .. } => SingleKeyTy::Union(name.clone()), TypeABI::Map { .. } => { return mutation_method_for_map(&elem_nesting); } diff --git a/pint-abi-gen/src/keys.rs b/pint-abi-gen/src/keys.rs index f524531a..5596c79d 100644 --- a/pint-abi-gen/src/keys.rs +++ b/pint-abi-gen/src/keys.rs @@ -3,9 +3,10 @@ //! a solution. use crate::{array, map, tuple}; -use pint_abi_types::{TypeABI, VarABI}; +use pint_abi_types::{TypeABI, UnionVariant, 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. @@ -21,7 +22,12 @@ fn nested_items_from_node(tree: &KeyedVarTree, n: NodeIx) -> Vec { TypeABI::Tuple(_fields) => { items.extend(tuple::keys_items(tree, n)); } - TypeABI::Bool | TypeABI::Int | TypeABI::Real | TypeABI::String | TypeABI::B256 => (), + TypeABI::Bool + | TypeABI::Int + | TypeABI::Real + | TypeABI::String + | TypeABI::B256 + | TypeABI::Union { .. } => (), } items } @@ -128,7 +134,12 @@ fn method_for_single_key(name: &str, nesting: &[Nesting]) -> syn::ImplItemFn { pub(crate) fn method_from_node(tree: &KeyedVarTree, n: NodeIx, name: &str) -> syn::ImplItemFn { let nesting = tree.nesting(n); match &tree[n].ty { - TypeABI::Bool | TypeABI::Int | TypeABI::Real | TypeABI::String | TypeABI::B256 => (), + TypeABI::Bool + | TypeABI::Int + | TypeABI::Real + | TypeABI::String + | TypeABI::B256 + | TypeABI::Union { .. } => (), TypeABI::Array { ty: _, size: _ } => { return method_for_array(name, &nesting); } @@ -222,7 +233,7 @@ pub(crate) fn impl_deref_for_nested(struct_name: &str) -> Vec { } /// All items for the `Keys` type, nested keys builder types and their impls. -fn items(vars: &[VarABI]) -> Vec { +fn keys_items(vars: &[VarABI]) -> Vec { let mut items = vec![ keys_struct().into(), keys_fn().into(), @@ -234,8 +245,21 @@ fn items(vars: &[VarABI]) -> Vec { } /// A `keys` module for all `Keys`-related items. -pub(crate) fn module(vars: &[VarABI]) -> syn::ItemMod { - let items = items(vars); +pub(crate) fn module( + vars: &[VarABI], + unions: &BTreeSet<(String, Vec)>, +) -> 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! { pub mod keys { //! All items related to building a set of [`Keys`]. diff --git a/pint-abi-gen/src/lib.rs b/pint-abi-gen/src/lib.rs index 1facb1ef..6e75afce 100644 --- a/pint-abi-gen/src/lib.rs +++ b/pint-abi-gen/src/lib.rs @@ -17,11 +17,12 @@ use addr::Addresses; use essential_types::{contract::Contract, PredicateAddress}; -use pint_abi_types::{ContractABI, PredicateABI, TupleField, TypeABI, VarABI}; +use pint_abi_types::{ContractABI, PredicateABI, TupleField, TypeABI, UnionVariant, VarABI}; use pint_abi_visit::Nesting; use proc_macro::TokenStream; use proc_macro2::Span; use quote::ToTokens; +use std::collections::BTreeSet; use syn::parse_macro_input; mod addr; @@ -45,6 +46,7 @@ enum SingleKeyTy { Real, String, B256, + Union(String), } impl SingleKeyTy { @@ -56,10 +58,16 @@ impl SingleKeyTy { 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), } } } +fn ty_from_union(name: &str) -> syn::Type { + let ident = syn::Ident::new(strip_colons_prefix(name), Span::call_site()); + syn::parse_quote!(#ident) +} + /// Convert the given pint tuple fields to unnamed Rust fields. fn fields_from_tuple_fields(fields: &[TupleField]) -> Vec { fields @@ -98,6 +106,7 @@ fn ty_from_pint_ty(ty: &TypeABI) -> syn::Type { 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::Map { .. } => unreachable!("Maps are not allowed as non-storage types"), @@ -142,11 +151,21 @@ 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 items = items_from_predicate(predicate, addr.as_ref()); + 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)] #[doc = #doc_str] @@ -180,6 +199,7 @@ 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) @@ -187,7 +207,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, addr) + mod_from_predicate(name, predicate, unions, addr) }) .collect() } @@ -203,17 +223,21 @@ 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], addrs: Option<&Addresses>) -> Vec { +fn items_from_predicates( + predicates: &[PredicateABI], + unions: &BTreeSet<(String, Vec)>, + 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, addr).into()); + items.push(mod_from_predicate(name, root_pred, unions, addr).into()); } // Add the named predicate modules. items.extend( - mods_from_named_predicates(predicates, addrs) + mods_from_named_predicates(predicates, unions, addrs) .into_iter() .map(syn::Item::from), ); @@ -310,18 +334,21 @@ 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]) -> Vec { +fn items_from_keyed_vars( + vars: &[VarABI], + unions: &BTreeSet<(String, Vec)>, +) -> Vec { let mut items = vec![]; // The `mutations` module and re-exports. - items.push(mutations::module(vars).into()); + items.push(mutations::module(vars, unions).into()); items.push(syn::parse_quote! { #[doc(inline)] pub use mutations::{mutations, Mutations}; }); // The `keys` module and re-exports. - items.push(keys::module(vars).into()); + items.push(keys::module(vars, unions).into()); items.push(syn::parse_quote! { #[doc(inline)] pub use keys::{keys, Keys}; @@ -333,9 +360,23 @@ fn items_from_keyed_vars(vars: &[VarABI]) -> Vec { /// 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]) -> syn::ItemMod { - let items = items_from_keyed_vars(vars); +fn mod_from_keyed_vars( + mod_name: &str, + vars: &[VarABI], + unions: &BTreeSet<(String, Vec)>, +) -> 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)); + syn::parse_quote! { pub mod #mod_ident { //! Items related to simplifying the process of building sets of @@ -363,16 +404,153 @@ fn mod_from_keyed_vars(mod_name: &str, vars: &[VarABI]) -> syn::ItemMod { } } +/// Convert the given pint union variant to Rust enum variants. +fn variants_from_union_variants(variants: &[UnionVariant]) -> Vec { + variants + .iter() + .map(|UnionVariant { name, ty }| { + 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); + syn::parse_quote!( + #variant_name ( #ty ) + ) + } else { + syn::parse_quote!( + #variant_name + ) + } + }) + .collect() +} + +fn enum_decl_from_union(name: &str, variants: &[UnionVariant]) -> syn::Item { + // Most nest mods here if needed + let enum_name = syn::Ident::new(strip_colons_prefix(name), Span::call_site()); + let enum_variants = variants_from_union_variants(variants); + + syn::parse_quote! { + #[derive(Debug)] + pub enum #enum_name { + #( #enum_variants ),* + } + } +} + +fn impl_encode_from_union(name: &str, variants: &[UnionVariant]) -> syn::Item { + let enum_name = syn::Ident::new(strip_colons_prefix(name), Span::call_site()); + + let union_size = 1 + variants + .iter() + .filter_map(|variant| variant.ty.as_ref()) + .map(pint_abi_visit::ty_size) + .max() + .unwrap_or_default(); + + let variant_matches: Vec = variants + .iter() + .enumerate() + .map(|(tag, variant)| { + let split = variant.name.split("::").collect::>(); + let variant_name = syn::Ident::new(split[1], Span::call_site()); + let variant_ty = &variant.ty; + let variant_size = variants[tag] + .ty + .as_ref() + .map(pint_abi_visit::ty_size) + .unwrap_or_default(); + + let padding = union_size - 1 - variant_size; + + let padding_code: Option = if padding > 0 { + // Generate the loop code only if padding is greater than 0 + Some(syn::parse_quote! { + for _ in 0..#padding { + pint_abi::Encode::encode(&0, w)?; + } + }) + } else { + None + }; + + if variant_ty.is_none() { + syn::parse_quote! { + #enum_name::#variant_name => { + let tag = #tag; + pint_abi::Encode::encode(&(tag as i64), w)?; + #padding_code + } + } + } else { + syn::parse_quote! { + #enum_name::#variant_name (ty) => { + let tag = #tag; + pint_abi::Encode::encode(&(tag as i64), w)?; + pint_abi::Encode::encode(ty, w)?; + #padding_code + } + } + } + }) + .collect::>(); + + syn::parse_quote! { + impl pint_abi::Encode for #enum_name { + fn encode(&self, w: &mut W) -> Result<(), W::Error> { + match self { + #(#variant_matches),* // Insert generated matches + } + Ok(()) + } + } + } +} + +fn items_from_union(name: &str, variants: &[UnionVariant]) -> Vec { + vec![ + enum_decl_from_union(name, variants), + impl_encode_from_union(name, variants), + ] +} + +fn items_from_unions(unions: &BTreeSet<(String, Vec)>) -> Vec { + unions + .iter() + .flat_map(|(name, variants)| items_from_union(name, variants)) + .collect() +} + /// Given an ABI, generate all items. fn items_from_abi_and_addrs(abi: &ContractABI, addrs: Option<&Addresses>) -> Vec { let mut items = vec![]; - items.extend(items_from_predicates(&abi.predicates, addrs)); + + let unions: BTreeSet<(String, Vec)> = abi + .storage + .iter() + .chain( + abi.predicates + .iter() + .flat_map(|predicate| predicate.vars.iter().chain(predicate.pub_vars.iter())), + ) + .filter_map(|var| match &var.ty { + TypeABI::Union { name, variants } => Some((name.clone(), variants.clone())), + _ => None, + }) + .collect(); + + items.extend(items_from_unions(&unions)); + + items.extend(items_from_predicates(&abi.predicates, &unions, 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).into()); + items.push(mod_from_keyed_vars("storage", &abi.storage, &unions).into()); } + items } diff --git a/pint-abi-gen/src/map.rs b/pint-abi-gen/src/map.rs index 1d7d220d..720ed342 100644 --- a/pint-abi-gen/src/map.rs +++ b/pint-abi-gen/src/map.rs @@ -125,7 +125,12 @@ fn key_method( ) -> syn::ImplItemFn { let nesting = tree.nesting(entry); match ty_to { - TypeABI::Bool | TypeABI::Int | TypeABI::Real | TypeABI::String | TypeABI::B256 => (), + TypeABI::Bool + | TypeABI::Int + | TypeABI::Real + | TypeABI::String + | TypeABI::B256 + | TypeABI::Union { .. } => (), TypeABI::Array { ty: _, size: _ } => { return key_method_for_array(ty_from, &nesting); } @@ -303,14 +308,15 @@ fn mutation_method( TypeABI::Bool => SingleKeyTy::Bool, TypeABI::Int => SingleKeyTy::Int, TypeABI::Real => SingleKeyTy::Real, + TypeABI::String => SingleKeyTy::String, + TypeABI::B256 => SingleKeyTy::B256, TypeABI::Array { ty: _, size: _ } => { return mutation_method_for_array(ty_from, &nesting); } - TypeABI::String => SingleKeyTy::String, - TypeABI::B256 => SingleKeyTy::B256, TypeABI::Tuple(_) => { return mutation_method_for_tuple(ty_from, &nesting); } + TypeABI::Union { name, .. } => SingleKeyTy::Union(name.clone()), TypeABI::Map { .. } => { return mutation_method_for_map(ty_from, &nesting); } diff --git a/pint-abi-gen/src/mutations.rs b/pint-abi-gen/src/mutations.rs index 5a374666..f8af1301 100644 --- a/pint-abi-gen/src/mutations.rs +++ b/pint-abi-gen/src/mutations.rs @@ -3,9 +3,10 @@ //! [`Mutation`][essential_types::solution::Mutation]s for a solution. use crate::{array, map, tuple, SingleKeyTy}; -use pint_abi_types::{TypeABI, VarABI}; +use pint_abi_types::{TypeABI, UnionVariant, 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. @@ -21,7 +22,12 @@ fn nested_items_from_node(tree: &KeyedVarTree, n: NodeIx) -> Vec { TypeABI::Tuple(_fields) => { items.extend(tuple::mutations_items(tree, n)); } - TypeABI::Bool | TypeABI::Int | TypeABI::Real | TypeABI::String | TypeABI::B256 => (), + TypeABI::Bool + | TypeABI::Int + | TypeABI::Real + | TypeABI::String + | TypeABI::B256 + | TypeABI::Union { .. } => (), } items } @@ -142,6 +148,7 @@ pub(crate) fn method_from_node(tree: &KeyedVarTree, n: NodeIx, name: &str) -> sy TypeABI::Tuple(_) => { return method_for_tuple(name, &nesting); } + TypeABI::Union { name, .. } => SingleKeyTy::Union(name.clone()), TypeABI::Map { .. } => { return method_for_map(name, &nesting); } @@ -229,7 +236,7 @@ pub(crate) fn impl_deref_for_nested(struct_name: &str) -> Vec { } /// All items for the `Mutations` type, nested mutations builder types and their impls. -fn items(vars: &[VarABI]) -> Vec { +fn mutations_items(vars: &[VarABI]) -> Vec { let mut items = vec![ mutations_struct().into(), mutations_fn().into(), @@ -241,8 +248,20 @@ fn items(vars: &[VarABI]) -> Vec { } /// A `mutations` module for all `Mutations`-related items. -pub(crate) fn module(vars: &[VarABI]) -> syn::ItemMod { - let items = items(vars); +pub(crate) fn module( + vars: &[VarABI], + unions: &BTreeSet<(String, Vec)>, +) -> 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! { pub mod mutations { //! All items related to building a set of [`Mutations`]. diff --git a/pint-abi-types/src/lib.rs b/pint-abi-types/src/lib.rs index e6ba6bbb..c69a2790 100644 --- a/pint-abi-types/src/lib.rs +++ b/pint-abi-types/src/lib.rs @@ -23,13 +23,19 @@ pub struct VarABI { pub ty: TypeABI, } -#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, Serialize, Deserialize)] pub struct TupleField { pub name: Option, pub ty: TypeABI, } -#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, Serialize, Deserialize)] +pub struct UnionVariant { + pub name: String, + pub ty: Option, +} + +#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, Serialize, Deserialize)] pub enum TypeABI { Bool, Int, @@ -41,6 +47,10 @@ pub enum TypeABI { ty: Box, size: i64, }, + Union { + name: String, + variants: Vec, + }, Map { ty_from: Box, ty_to: Box, diff --git a/pint-abi-visit/src/lib.rs b/pint-abi-visit/src/lib.rs index 6fc6535a..d67edb08 100644 --- a/pint-abi-visit/src/lib.rs +++ b/pint-abi-visit/src/lib.rs @@ -219,6 +219,15 @@ pub fn ty_size(ty: &TypeABI) -> usize { TypeABI::Array { ty, size } => { ty_size(ty) * usize::try_from(*size).expect("size out of range") } + TypeABI::Union { variants, .. } => { + // 1 word for the tag + max variant type size + 1 + variants + .iter() + .filter_map(|variant| variant.ty.as_ref()) + .map(ty_size) + .max() + .unwrap_or_default() + } TypeABI::Map { .. } => 1, } } @@ -228,7 +237,12 @@ pub fn ty_size(ty: &TypeABI) -> usize { fn add_children<'a>(graph: &mut KeyedVarGraph<'a>, a: NodeIx, ty: &'a TypeABI) { match ty { // Leaf types have no further nesting. - TypeABI::Bool | TypeABI::Int | TypeABI::Real | TypeABI::String | TypeABI::B256 => {} + TypeABI::Bool + | TypeABI::Int + | TypeABI::Real + | TypeABI::String + | TypeABI::B256 + | TypeABI::Union { .. } => {} // Recurse for nested tuple types. TypeABI::Tuple(fields) => { @@ -313,6 +327,7 @@ fn flattened_key_count(ty: &TypeABI) -> usize { | TypeABI::Int | TypeABI::String | TypeABI::B256 + | TypeABI::Union { .. } | TypeABI::Map { .. } => 1, } } diff --git a/pintc/src/asm_gen/asm_builder.rs b/pintc/src/asm_gen/asm_builder.rs index 3049f5f4..f04c623b 100644 --- a/pintc/src/asm_gen/asm_builder.rs +++ b/pintc/src/asm_gen/asm_builder.rs @@ -398,7 +398,7 @@ impl<'a> AsmBuilder<'a> { asm.push(Stack::Push(0).into()); value_size += 1; } - value_size + 1 + value_size } Immediate::Error | Immediate::Nil | Immediate::Real(_) | Immediate::String(_) => { unreachable!("Unexpected literal") diff --git a/pintc/src/types.rs b/pintc/src/types.rs index cb90e54e..798d022e 100644 --- a/pintc/src/types.rs +++ b/pintc/src/types.rs @@ -661,6 +661,7 @@ impl Type { PrimitiveKind::B256 => TypeABI::B256, _ => unimplemented!(), }), + Type::Tuple { fields, .. } => Ok(TypeABI::Tuple( fields .iter() @@ -672,6 +673,7 @@ impl Type { }) .collect::, _>>()?, )), + Type::Array { ty, range, size, .. } => Ok(TypeABI::Array { @@ -685,15 +687,33 @@ impl Type { contract, )?), }), + + Type::Union { path, .. } => Ok(TypeABI::Union { + name: path.clone(), + variants: self + .get_union_variant_names(&contract.unions) + .into_iter() + .zip(self.get_union_variant_types(&contract.unions)) + .map(|(name, ty)| { + Ok(pint_abi_types::UnionVariant { + name: name.to_string(), + ty: ty + .as_ref() + .map(|ty| ty.abi(handler, contract)) + .transpose()?, + }) + }) + .collect::, _>>()?, + }), + Type::Map { ty_from, ty_to, .. } => Ok(TypeABI::Map { ty_from: Box::new((*ty_from).abi(handler, contract)?), ty_to: Box::new((*ty_to).abi(handler, contract)?), }), - // These, of course, are incorrect. It's just a placeholder until we can support ABI gen - // for them, which is non-trivial. + // This, of course, is incorrect. It's just a placeholder until we can support ABI gen + // for vectors, which is non-trivial. Type::Vector { .. } => Ok(TypeABI::Int), - Type::Union { .. } => Ok(TypeABI::Int), _ => unimplemented!("other types are not yet supported"), } From 6b21ddc854bc0e9c54764669d6105e27732e86dc Mon Sep 17 00:00:00 2001 From: Mohammad Fawaz Date: Wed, 9 Oct 2024 10:59:36 -0400 Subject: [PATCH 2/2] mod_nesting --- .../test-pkgs/simple/src/contract.pnt | 18 ++-- pint-abi-gen-tests/tests/simple.rs | 10 +- pint-abi-gen/src/array.rs | 4 +- pint-abi-gen/src/keys.rs | 21 +--- pint-abi-gen/src/lib.rs | 102 ++++++++---------- pint-abi-gen/src/map.rs | 20 ++-- pint-abi-gen/src/mutations.rs | 24 ++--- pint-abi-gen/src/pub_vars.rs | 7 +- pint-abi-gen/src/utils.rs | 6 +- pint-abi-gen/src/vars.rs | 7 +- 10 files changed, 92 insertions(+), 127 deletions(-) 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..e73d6a18 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. @@ -233,7 +232,7 @@ pub(crate) fn impl_deref_for_nested(struct_name: &str) -> Vec { } /// All items for the `Keys` type, nested keys builder types and their impls. -fn keys_items(vars: &[VarABI]) -> Vec { +fn items(vars: &[VarABI]) -> Vec { let mut items = vec![ keys_struct().into(), keys_fn().into(), @@ -245,20 +244,8 @@ 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 { - 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)); +pub(crate) fn module(vars: &[VarABI]) -> syn::ItemMod { + let items = items(vars); syn::parse_quote! { pub mod keys { diff --git a/pint-abi-gen/src/lib.rs b/pint-abi-gen/src/lib.rs index 6e75afce..53ca9bec 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,21 +168,11 @@ 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())); + let items = items_from_predicate(predicate, addr.as_ref()); syn::parse_quote! { #[allow(non_snake_case)] #[doc = #doc_str] @@ -199,7 +206,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 +213,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 +229,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 +336,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,23 +359,10 @@ 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 { - let mut items: Vec = vec![]; +fn mod_from_keyed_vars(mod_name: &str, vars: &[VarABI]) -> syn::ItemMod { + let items = items_from_keyed_vars(vars); 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)); - syn::parse_quote! { pub mod #mod_ident { //! Items related to simplifying the process of building sets of @@ -412,7 +398,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 +527,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..e6db968d 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! { @@ -236,7 +235,7 @@ pub(crate) fn impl_deref_for_nested(struct_name: &str) -> Vec { } /// All items for the `Mutations` type, nested mutations builder types and their impls. -fn mutations_items(vars: &[VarABI]) -> Vec { +fn items(vars: &[VarABI]) -> Vec { let mut items = vec![ mutations_struct().into(), mutations_fn().into(), @@ -248,19 +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 { - 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)); +pub(crate) fn module(vars: &[VarABI]) -> syn::ItemMod { + let items = items(vars); syn::parse_quote! { pub mod mutations { 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)]