Faster validation with tree-based constraint validation

Author: laurmaedje

Cargo.toml

@@ -17,10 +17,13 @@ comemo-macros = { version = "0.4.0", path = "macros" }
 once_cell = "1.18"
 parking_lot = "0.12"
 proc-macro2 = "1"
+quickcheck = "1"
+quickcheck_macros = "1"
 quote = "1"
 rustc-hash = "2.1"
 serial_test = "3"
 siphasher = "1"
+slab = "0.4"
 syn = { version = "2", features = ["full"] }
 
 [package]
@@ -45,8 +48,11 @@ comemo-macros = { workspace = true, optional = true }
 parking_lot = { workspace = true }
 rustc-hash = { workspace = true }
 siphasher = { workspace = true }
+slab = { workspace = true }
 
 [dev-dependencies]
+quickcheck = { workspace = true }
+quickcheck_macros = { workspace = true }
 serial_test = { workspace = true }
 
 [[test]]

macros/src/memoize.rs

@@ -148,7 +148,7 @@ fn process(function: &Function) -> Result<TokenStream> {
 
     wrapped.block = parse_quote! { {
         static __CACHE: ::comemo::internal::Cache<
-            <::comemo::internal::Multi<#arg_ty_tuple> as ::comemo::internal::Input>::Constraint,
+            <::comemo::internal::Multi<#arg_ty_tuple> as ::comemo::internal::Input>::Call,
             #output,
         > = ::comemo::internal::Cache::new(|| {
             ::comemo::internal::register_evictor(|max_age| __CACHE.evict(max_age));
@@ -160,7 +160,7 @@ fn process(function: &Function) -> Result<TokenStream> {
         ::comemo::internal::memoize(
             &__CACHE,
             ::comemo::internal::Multi(#arg_tuple),
-            &::core::default::Default::default(),
+            &mut ::core::default::Default::default(),
             #enabled,
             #closure,
         )

macros/src/track.rs

@@ -42,6 +42,13 @@ pub fn expand(item: &syn::Item) -> Result<TokenStream> {
         _ => bail!(item, "`track` can only be applied to impl blocks and traits"),
     };
 
+    if methods.iter().any(|m| m.mutable) && methods.iter().any(|m| !m.mutable) {
+        bail!(
+            item,
+            "`track` cannot be applied to a mix of mutable and immutable methods"
+        );
+    }
+
     // Produce the necessary items for the type to become trackable.
     let variants = create_variants(&methods);
     let scope = create(&ty, generics, trait_, &methods)?;
@@ -168,6 +175,12 @@ fn prepare_method(vis: syn::Visibility, sig: &syn::Signature) -> Result<Method>
         bail!(ty, "tracked methods cannot return mutable references");
     }
 
+    if let syn::ReturnType::Type(_, ty) = &sig.output
+        && receiver.mutability.is_some()
+    {
+        bail!(ty, "mutable tracked methods cannot have a return value");
+    }
+
     Ok(Method {
         vis,
         sig: sig.clone(),
@@ -225,11 +238,6 @@ fn create(
     let t: syn::GenericParam = parse_quote! { '__comemo_tracked };
     let r: syn::GenericParam = parse_quote! { '__comemo_retrack };
     let d: syn::GenericParam = parse_quote! { '__comemo_dynamic };
-    let maybe_cloned = if methods.iter().any(|it| it.mutable) {
-        quote! { ::core::clone::Clone::clone(self) }
-    } else {
-        quote! { self }
-    };
 
     // Prepare generics.
     let (impl_gen, type_gen, where_clause) = generics.split_for_impl();
@@ -245,37 +253,9 @@ fn create(
     impl_params_t.params.push(t.clone());
     type_params_t.params.push(t.clone());
 
-    // Prepare validations.
     let prefix = trait_.as_ref().map(|name| quote! { #name for });
-    let validations: Vec<_> = methods.iter().map(create_validation).collect();
-    let validate = if !methods.is_empty() {
-        quote! {
-            let mut this = #maybe_cloned;
-            constraint.validate(|call| match &call.0 { #(#validations,)* })
-        }
-    } else {
-        quote! { true }
-    };
-    let validate_with_id = if !methods.is_empty() {
-        quote! {
-            let mut this = #maybe_cloned;
-            constraint.validate_with_id(
-                |call| match &call.0 { #(#validations,)* },
-                id,
-            )
-        }
-    } else {
-        quote! { true }
-    };
-
-    // Prepare replying.
-    let immutable = methods.iter().all(|m| !m.mutable);
-    let replays = methods.iter().map(create_replay);
-    let replay = (!immutable).then(|| {
-        quote! {
-            constraint.replay(|call| match &call.0 { #(#replays,)* });
-        }
-    });
+    let calls: Vec<_> = methods.iter().map(create_call).collect();
+    let calls_mut: Vec<_> = methods.iter().map(create_call_mut).collect();
 
     // Prepare variants and wrapper methods.
     let wrapper_methods = methods
@@ -284,32 +264,18 @@ fn create(
         .map(|m| create_wrapper(m, false));
     let wrapper_methods_mut = methods.iter().map(|m| create_wrapper(m, true));
 
-    let constraint = if immutable {
-        quote! { ImmutableConstraint }
-    } else {
-        quote! { MutableConstraint }
-    };
-
     Ok(quote! {
-        impl #impl_params ::comemo::Track for #ty #where_clause {}
-
-        impl #impl_params ::comemo::Validate for #ty #where_clause {
-            type Constraint = ::comemo::internal::#constraint<__ComemoCall>;
+        impl #impl_params ::comemo::Track for #ty #where_clause {
+            type Call = __ComemoCall;
 
             #[inline]
-            fn validate(&self, constraint: &Self::Constraint) -> bool {
-                #validate
+            fn call(&self, call: &Self::Call) -> u128 {
+                match call.0 { #(#calls,)* }
             }
 
             #[inline]
-            fn validate_with_id(&self, constraint: &Self::Constraint, id: usize) -> bool {
-                #validate_with_id
-            }
-
-            #[inline]
-            #[allow(unused_variables)]
-            fn replay(&mut self, constraint: &Self::Constraint) {
-                #replay
+            fn call_mut(&mut self, call: &Self::Call) {
+                match call.0 { #(#calls_mut,)* }
             }
         }
 
@@ -363,41 +329,50 @@ fn create(
     })
 }
 
-/// Produce a constraint validation for a method.
+/// Produce a call enum variant for a method.
 fn create_variant(method: &Method) -> TokenStream {
     let name = &method.sig.ident;
     let types = &method.types;
     quote! { #name(#(<#types as ::std::borrow::ToOwned>::Owned),*) }
 }
 
-/// Produce a constraint validation for a method.
-fn create_validation(method: &Method) -> TokenStream {
+/// Produce a call branch for a method.
+fn create_call(method: &Method) -> TokenStream {
     let name = &method.sig.ident;
     let args = &method.args;
     let prepared = method.args.iter().zip(&method.kinds).map(|(arg, kind)| match kind {
         Kind::Normal => quote! { #arg.to_owned() },
         Kind::Reference => quote! { #arg },
     });
-    quote! {
-        __ComemoVariant::#name(#(#args),*)
-            => ::comemo::internal::hash(&this.#name(#(#prepared),*))
+    if method.mutable {
+        quote! {
+            __ComemoVariant::#name(..) => 0
+        }
+    } else {
+        quote! {
+            __ComemoVariant::#name(#(ref #args),*)
+                => ::comemo::internal::hash(&self.#name(#(#prepared),*))
+        }
     }
 }
 
-/// Produce a constraint validation for a method.
-fn create_replay(method: &Method) -> TokenStream {
+/// Produce a mutable call branch for a method.
+fn create_call_mut(method: &Method) -> TokenStream {
     let name = &method.sig.ident;
     let args = &method.args;
     let prepared = method.args.iter().zip(&method.kinds).map(|(arg, kind)| match kind {
         Kind::Normal => quote! { #arg.to_owned() },
         Kind::Reference => quote! { #arg },
     });
-    let body = method.mutable.then(|| {
+    if method.mutable {
         quote! {
-            self.#name(#(#prepared),*);
+            __ComemoVariant::#name(#(ref #args),*)  => self.#name(#(#prepared),*)
         }
-    });
-    quote! { __ComemoVariant::#name(#(#args),*) => { #body } }
+    } else {
+        quote! {
+            __ComemoVariant::#name(..) => {}
+        }
+    }
 }
 
 /// Produce a wrapped surface method.
@@ -417,16 +392,19 @@ fn create_wrapper(method: &Method, tracked_mut: bool) -> TokenStream {
         #[track_caller]
         #[inline]
         #vis #sig {
-            let __comemo_variant = __ComemoVariant::#name(#(#args.to_owned()),*);
-            let (__comemo_value, __comemo_constraint) = ::comemo::internal::#to_parts;
-            let output = __comemo_value.#name(#(#args,)*);
-            if let Some(constraint) = __comemo_constraint {
-                constraint.push(
+            let (__comemo_value, __comemo_sink) = ::comemo::internal::#to_parts;
+            if let Some(__comemo_sink) = __comemo_sink {
+                let __comemo_variant = __ComemoVariant::#name(#(#args.to_owned()),*);
+                let output = __comemo_value.#name(#(#args,)*);
+                ::comemo::internal::Sink::emit(
+                    __comemo_sink,
                     __ComemoCall(__comemo_variant),
                     ::comemo::internal::hash(&output),
                 );
+                output
+            } else {
+                __comemo_value.#name(#(#args,)*)
             }
-            output
         }
     }
 }