Skip to content

Faster validation with tree-based constraint checking #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laurmaedje merged 3 commits into from
Jul 31, 2025
Merged

Faster validation with tree-based constraint checking #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
017b4d6
Faster validation with tree-based constraint checking
laurmaedje Jul 30, 2025
6c41328
Test for nested mutable tracking
laurmaedje Jul 30, 2025
cda3b63
Move struct
laurmaedje Jul 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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]]
Expand Down
43 changes: 36 additions & 7 deletions macros/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,7 @@ use syn::{Error, Result, parse_quote};
/// - _Mutably tracked:_ The argument is of the form `TrackedMut<T>`. Through
/// this type, you can safely mutate an argument from within a memoized
/// function. If there is a cache hit, comemo will replay all mutations.
/// Mutable tracked methods can also have return values that are tracked just
/// like immutable methods.
/// Mutable tracked methods cannot have return values.
///
/// # Restrictions
/// The following restrictions apply to memoized functions:
Expand All @@ -50,11 +49,40 @@ use syn::{Error, Result, parse_quote};
/// expose to the hasher**. Otherwise, memoized results might get reused
/// invalidly.
///
/// - The **only obversable impurity memoized functions may exhibit are
/// - The **only observable impurity memoized functions may exhibit are
/// mutations through `TrackedMut<T>` arguments.** Comemo stops you from using
/// basic mutable arguments, but it cannot determine all sources of impurity,
/// so this is your responsibility.
///
/// - Memoized functions must **call tracked methods in _reorderably
/// deterministic_ fashion.** Consider two executions A and B of a memoized
/// function. We define the following two properties:
///
/// - _In-order deterministic:_ If the first N tracked calls and their results
/// are the same in A and B, then the N+1th call must also be the same. This
/// is a fairly natural property as far as deterministic functions go, as,
/// if the first N calls and their results were the same across two
/// execution, the available information for choosing the N+1th call is the
/// same. However, this property is a bit too restrictive in practice. For
/// instance, a function that internally uses multi-threading may call
/// tracked methods out-of-order while still producing a deterministic
/// result.
///
/// - _Reorderably deterministic:_ If, for the first N calls in A, B has
/// matching calls (same arguments, same return value) somewhere in its call
/// sequence, then the N+1th call invoked by A must also occur _somewhere_
/// in the call sequence of B. This is a somewhat relaxed version of
/// in-order determinism that still allows comemo to perform internal
/// optimizations while permitting memoization of many more functions (e.g.
/// ones that use internal multi-threading in an outwardly deterministic
/// fashion).
///
/// Reorderable determinism is necessary for efficient cache lookups. If a
/// memoized function is not reorderably determinstic, comemo may panic in
/// debug mode to bring your attention to this. Meanwhile, in release mode,
/// memoized functions will still yield correct results, but caching may prove
/// ineffective.
///
/// - The output of a memoized function must be `Send` and `Sync` because it is
/// stored in the global cache.
///
Expand Down Expand Up @@ -126,10 +154,6 @@ pub fn memoize(args: BoundaryStream, stream: BoundaryStream) -> BoundaryStream {
/// arguments, tracking is the only option, so that comemo can replay the side
/// effects when there is a cache hit.
///
/// If you attempt to track any mutable methods, your type must implement
/// [`Clone`] so that comemo can roll back attempted mutations which did not
/// result in a cache hit.
///
/// # Restrictions
/// Tracked impl blocks or traits may not be generic and may only contain
/// methods. Just like with memoized functions, certain restrictions apply to
Expand All @@ -147,6 +171,11 @@ pub fn memoize(args: BoundaryStream, stream: BoundaryStream) -> BoundaryStream {
/// [`Hash`](std::hash::Hash) and **must feed all the information they expose
/// to the hasher**. Otherwise, memoized results might get reused invalidly.
///
/// - Mutable tracked methods must not have a return value.
///
/// - A tracked implementation cannot have a mix of mutable and immutable
/// methods.
///
/// - The arguments to a tracked method must be `Send` and `Sync` because they
/// are stored in the global cache.
///
Expand Down
4 changes: 2 additions & 2 deletions macros/src/memoize.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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,
)
Expand Down
124 changes: 51 additions & 73 deletions macros/src/track.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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();
Expand All @@ -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
Expand All @@ -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,)* }
}
}

Expand Down Expand Up @@ -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.
Expand All @@ -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
}
}
}
Loading