Propagate constraints

Author: laurmaedje

macros/src/track.rs

@@ -226,7 +226,7 @@ fn create(
             #(#wrapper_methods)*
         }
 
-        #[derive(Default)]
+        #[derive(Debug, Default)]
         pub struct #constraint {
             #(#constraint_fields)*
         }
@@ -260,6 +260,7 @@ fn create_wrapper(method: &Method) -> TokenStream {
     let name = &method.sig.ident;
     let args = &method.args;
     quote! {
+        #[track_caller]
         #vis #sig {
             let input = (#(#args.to_owned(),)*);
             let (value, constraint) = ::comemo::internal::to_parts(self.0);

src/cache.rs

@@ -14,52 +14,6 @@ thread_local! {
     static CACHE: RefCell<Cache> = RefCell::new(Cache::default());
 }
 
-/// Configure the caching and eviction behaviour.
-pub fn config(config: Config) {
-    CACHE.with(|cache| cache.borrow_mut().config = config);
-}
-
-/// Configuration for caching and eviction behaviour.
-pub struct Config {
-    max_age: u32,
-}
-
-impl Config {
-    /// The maximum number of evictions an entry can survive without having been
-    /// used in between.
-    ///
-    /// Default: 5
-    pub fn max_age(mut self, age: u32) -> Self {
-        self.max_age = age;
-        self
-    }
-}
-
-impl Default for Config {
-    fn default() -> Self {
-        Self { max_age: 5 }
-    }
-}
-
-/// Evict cache entries that haven't been used in a while.
-///
-/// The eviction behaviour can be customized with the [`config`] function.
-/// Currently, comemo does not evict the cache automatically (this might
-/// change in the future).
-pub fn evict() {
-    CACHE.with(|cache| {
-        let mut cache = cache.borrow_mut();
-        let max = cache.config.max_age;
-        cache.map.retain(|_, entries| {
-            entries.retain_mut(|entry| {
-                entry.age += 1;
-                entry.age <= max
-            });
-            !entries.is_empty()
-        });
-    });
-}
-
 /// Execute a function or use a cached result for it.
 pub fn memoized<In, Out, F>(id: TypeId, input: In, func: F) -> Out
 where
@@ -76,55 +30,57 @@ where
             (id, hash)
         };
 
-        // Check if there is a cached output.
-        let mut borrow = cache.borrow_mut();
-        if let Some(output) = borrow.lookup::<In, Out>(key, &input) {
-            return output;
-        }
-
-        borrow.depth += 1;
-        drop(borrow);
-
         // Point all tracked parts of the input to these constraints.
         let constraint = In::Constraint::default();
-        let (tracked, outer) = input.retrack(&constraint);
+
+        // Check if there is a cached output.
+        if let Some(constrained) = cache.borrow().lookup::<In, Out>(key, &input) {
+            // Add the cached constraints to the outer ones.
+            let (_, outer) = input.retrack(&constraint);
+            outer.join(&constrained.constraint);
+            return constrained.output.clone();
+        }
 
         // Execute the function with the new constraints hooked in.
-        let output = func(tracked);
+        let (input, outer) = input.retrack(&constraint);
+        let output = func(input);
 
-        // Add the new constraints to the previous outer ones.
+        // Add the new constraints to the outer ones.
         outer.join(&constraint);
 
         // Insert the result into the cache.
-        borrow = cache.borrow_mut();
-        borrow.insert::<In, Out>(key, constraint, output.clone());
-        borrow.depth -= 1;
+        cache.borrow_mut().insert::<In, Out>(key, constraint, output.clone());
 
         output
     })
 }
 
+/// Completely clear the cache.
+pub fn clear() {
+    CACHE.with(|cache| cache.borrow_mut().map.clear());
+}
+
 /// The global cache.
 #[derive(Default)]
 struct Cache {
     /// Maps from function IDs + hashes to memoized results.
     map: HashMap<(TypeId, u128), Vec<Entry>>,
-    /// The current depth of the memoized call stack.
-    depth: usize,
-    /// The current configuration.
-    config: Config,
 }
 
 impl Cache {
     /// Look for a matching entry in the cache.
-    fn lookup<In, Out>(&mut self, key: (TypeId, u128), input: &In) -> Option<Out>
+    fn lookup<In, Out>(
+        &self,
+        key: (TypeId, u128),
+        input: &In,
+    ) -> Option<&Constrained<In::Constraint, Out>>
     where
         In: Input,
         Out: Clone + 'static,
     {
         self.map
-            .get_mut(&key)?
-            .iter_mut()
+            .get(&key)?
+            .iter()
             .find_map(|entry| entry.lookup::<In, Out>(input))
     }
 
@@ -151,8 +107,6 @@ struct Entry {
     ///
     /// This is of type `Constrained<In::Constraint, Out>`.
     constrained: Box<dyn Any>,
-    /// How many evictions have passed since the entry has been last used.
-    age: u32,
 }
 
 /// A value with a constraint.
@@ -172,22 +126,18 @@ impl Entry {
     {
         Self {
             constrained: Box::new(Constrained { constraint, output }),
-            age: 0,
         }
     }
 
     /// Return the entry's output if it is valid for the given input.
-    fn lookup<In, Out>(&mut self, input: &In) -> Option<Out>
+    fn lookup<In, Out>(&self, input: &In) -> Option<&Constrained<In::Constraint, Out>>
     where
         In: Input,
         Out: Clone + 'static,
     {
-        let Constrained::<In::Constraint, Out> { constraint, output } =
+        let constrained: &Constrained<In::Constraint, Out> =
             self.constrained.downcast_ref().expect("wrong entry type");
 
-        input.valid(constraint).then(|| {
-            self.age = 0;
-            output.clone()
-        })
+        input.valid(&constrained.constraint).then(|| constrained)
     }
 }

src/constraint.rs

@@ -1,4 +1,5 @@
 use std::cell::{Cell, RefCell};
+use std::fmt::Debug;
 use std::hash::Hash;
 
 use siphasher::sip128::{Hasher128, SipHasher};
@@ -10,6 +11,7 @@ pub trait Join<T = Self> {
 }
 
 impl<T: Join> Join<T> for Option<&T> {
+    #[inline]
     fn join(&self, inner: &T) {
         if let Some(outer) = self {
             outer.join(inner);
@@ -18,18 +20,28 @@ impl<T: Join> Join<T> for Option<&T> {
 }
 
 /// Defines a constraint for a tracked method without arguments.
-#[derive(Default)]
+#[derive(Debug, Default)]
 pub struct SoloConstraint {
     cell: Cell<Option<u128>>,
 }
 
 impl SoloConstraint {
     /// Set the constraint for the value.
+    #[inline]
+    #[track_caller]
     pub fn set(&self, _: (), hash: u128) {
-        self.cell.set(Some(hash));
+        // If there's already a constraint, it must match.
+        // This assertion can fail if a tracked function isn't pure
+        // (which violates comemo's contract).
+        if let Some(existing) = self.cell.get() {
+            check(hash, existing);
+        } else {
+            self.cell.set(Some(hash));
+        }
     }
 
     /// Whether the value fulfills the constraint.
+    #[inline]
     pub fn valid<F>(&self, f: F) -> bool
     where
         F: Fn(()) -> u128,
@@ -39,15 +51,16 @@ impl SoloConstraint {
 }
 
 impl Join for SoloConstraint {
+    #[inline]
     fn join(&self, inner: &Self) {
-        let inner = inner.cell.get();
-        if inner.is_some() {
-            self.cell.set(inner);
+        if let Some(hash) = inner.cell.get() {
+            self.set((), hash);
         }
     }
 }
 
 /// Defines a constraint for a tracked method with arguments.
+#[derive(Debug)]
 pub struct MultiConstraint<In> {
     calls: RefCell<Vec<(In, u128)>>,
 }
@@ -57,45 +70,65 @@ where
     In: Clone + PartialEq,
 {
     /// Enter a constraint for a pair of inputs and output.
+    #[inline]
+    #[track_caller]
     pub fn set(&self, input: In, hash: u128) {
         let mut calls = self.calls.borrow_mut();
-        if calls.iter().all(|item| item.0 != input) {
+        if let Some(item) = calls.iter().find(|item| item.0 == input) {
+            check(item.1, hash);
+        } else {
             calls.push((input, hash));
         }
     }
 
     /// Whether the method satisfies as all input-output pairs.
+    #[inline]
     pub fn valid<F>(&self, f: F) -> bool
     where
         F: Fn(&In) -> u128,
     {
-        let calls = self.calls.borrow();
-        calls.iter().all(|(input, hash)| *hash == f(input))
+        self.calls.borrow().iter().all(|(input, hash)| *hash == f(input))
     }
 }
 
 impl<In> Join for MultiConstraint<In>
 where
-    In: Clone + PartialEq,
+    In: Debug + Clone + PartialEq,
 {
+    #[inline]
     fn join(&self, inner: &Self) {
         let mut calls = self.calls.borrow_mut();
-        let inner = inner.calls.borrow();
-        for (input, hash) in inner.iter() {
-            if calls.iter().all(|item| &item.0 != input) {
+        for (input, hash) in inner.calls.borrow().iter() {
+            if let Some(item) = calls.iter().find(|item| &item.0 == input) {
+                check(item.1, *hash);
+            } else {
                 calls.push((input.clone(), *hash));
             }
         }
     }
 }
 
 impl<In> Default for MultiConstraint<In> {
+    #[inline]
     fn default() -> Self {
         Self { calls: RefCell::new(vec![]) }
     }
 }
 
+/// Check for a constraint violation.
+#[inline]
+#[track_caller]
+fn check(left: u128, right: u128) {
+    if left != right {
+        panic!(
+            "comemo: found conflicting constraints. \
+             is this tracked function pure?"
+        )
+    }
+}
+
 /// Produce a 128-bit hash of a value.
+#[inline]
 pub fn hash<T: Hash>(value: &T) -> u128 {
     let mut state = SipHasher::new();
     value.hash(&mut state);