Remove (crate), add new_ methods to the cache, disentangle entry points

dag_in_context/src/add_context.rs

@@ -22,7 +22,7 @@ struct ContextCache {
 }
 
 impl ContextCache {
-    pub(crate) fn get_symbolic_ctx(&mut self, expr: &RcExpr, ctx: &Assumption) -> Assumption {
+    pub fn get_symbolic_ctx(&mut self, expr: &RcExpr, ctx: &Assumption) -> Assumption {
         let ctx_ref = ctx.to_ref();
         let key = (expr.as_ref() as *const Expr, ctx_ref.clone());
         if let Some(sym) = self.symbol_gen.get(&key) {
@@ -32,6 +32,36 @@ impl ContextCache {
         self.symbol_gen.insert(key, sym.clone());
         Assumption::WildCard(sym)
     }
+
+    fn new() -> ContextCache {
+        ContextCache {
+            with_ctx: HashMap::new(),
+            symbol_gen: HashMap::new(),
+            unions: Vec::new(),
+            symbolic_ctx: false,
+            dummy_ctx: false,
+        }
+    }
+
+    fn new_symbolic_ctx() -> ContextCache {
+        ContextCache {
+            with_ctx: HashMap::new(),
+            symbol_gen: HashMap::new(),
+            unions: Vec::new(),
+            symbolic_ctx: true,
+            dummy_ctx: false,
+        }
+    }
+
+    fn new_dummy_ctx() -> ContextCache {
+        ContextCache {
+            with_ctx: HashMap::new(),
+            symbol_gen: HashMap::new(),
+            unions: Vec::new(),
+            symbolic_ctx: false,
+            dummy_ctx: true,
+        }
+    }
 }
 
 pub struct UnionsAnd<T> {
@@ -54,73 +84,68 @@ impl<T> UnionsAnd<T> {
 
 impl TreeProgram {
     pub fn add_context(&self) -> TreeProgram {
-        self.add_context_internal(Expr::func_add_ctx)
+        self.add_context_internal(Expr::func_get_ctx, &mut ContextCache::new())
     }
 
     /// add stand-in variables for all the contexts in the program
     /// useful for testing if you don't care about context in the test
-    #[allow(dead_code)]
-    pub(crate) fn add_symbolic_ctx(&self) -> TreeProgram {
-        self.add_context_internal(Expr::add_symbolic_ctx)
+    pub fn add_symbolic_ctx(&self) -> TreeProgram {
+        self.add_context_internal(
+            |_| Assumption::dummy(),
+            &mut ContextCache::new_symbolic_ctx(),
+        )
     }
 
-    pub(crate) fn add_dummy_ctx(&self) -> TreeProgram {
-        self.add_context_internal(Expr::add_dummy_ctx)
+    pub fn add_dummy_ctx(&self) -> TreeProgram {
+        self.add_context_internal(|_| Assumption::dummy(), &mut ContextCache::new_dummy_ctx())
     }
 
-    fn add_context_internal(&self, func: impl Fn(&RcExpr) -> RcExpr) -> TreeProgram {
+    fn add_context_internal(
+        &self,
+        func: impl Fn(&RcExpr) -> Assumption,
+        cache: &mut ContextCache,
+    ) -> TreeProgram {
         TreeProgram {
-            functions: self.functions.iter().map(&func).collect(),
-            entry: func(&self.entry),
+            functions: self
+                .functions
+                .iter()
+                .map(|f| f.add_ctx_with_cache(func(f), cache))
+                .collect(),
+            entry: self.entry.add_ctx_with_cache(func(&self.entry), cache),
         }
     }
 }
 
 impl Expr {
-    pub(crate) fn func_add_ctx(self: &RcExpr) -> RcExpr {
-        let Expr::Function(name, arg_ty, ret_ty, body) = &self.as_ref() else {
+    fn func_get_ctx(self: &RcExpr) -> Assumption {
+        let Expr::Function(name, _arg_ty, _ret_ty, _body) = &self.as_ref() else {
+            panic!("Expected Function, got {:?}", self);
+        };
+        Assumption::InFunc(name.clone())
+    }
+
+    pub fn func_add_ctx(self: &RcExpr) -> RcExpr {
+        let Expr::Function(name, arg_ty, ret_ty, body) = self.as_ref() else {
             panic!("Expected Function, got {:?}", self);
         };
-        let current_ctx = Assumption::InFunc(name.clone());
         RcExpr::new(Expr::Function(
             name.clone(),
             arg_ty.clone(),
             ret_ty.clone(),
-            body.add_ctx(current_ctx),
+            body.add_ctx_with_cache(self.func_get_ctx(), &mut ContextCache::new()),
         ))
     }
 
-    pub(crate) fn add_dummy_ctx(self: &RcExpr) -> RcExpr {
-        let mut cache = ContextCache {
-            with_ctx: HashMap::new(),
-            symbol_gen: HashMap::new(),
-            unions: Vec::new(),
-            symbolic_ctx: false,
-            dummy_ctx: true,
-        };
-        self.add_ctx_with_cache(Assumption::dummy(), &mut cache)
+    pub fn add_dummy_ctx(self: &RcExpr) -> RcExpr {
+        self.add_ctx_with_cache(Assumption::dummy(), &mut ContextCache::new_dummy_ctx())
     }
 
-    pub(crate) fn add_symbolic_ctx(self: &RcExpr) -> RcExpr {
-        let mut cache = ContextCache {
-            with_ctx: HashMap::new(),
-            symbol_gen: HashMap::new(),
-            unions: Vec::new(),
-            symbolic_ctx: true,
-            dummy_ctx: false,
-        };
-        self.add_ctx_with_cache(Assumption::dummy(), &mut cache)
+    pub fn add_symbolic_ctx(self: &RcExpr) -> RcExpr {
+        self.add_ctx_with_cache(Assumption::dummy(), &mut ContextCache::new_symbolic_ctx())
     }
 
-    pub(crate) fn add_ctx(self: &RcExpr, current_ctx: Assumption) -> RcExpr {
-        let mut cache = ContextCache {
-            with_ctx: HashMap::new(),
-            symbol_gen: HashMap::new(),
-            unions: Vec::new(),
-            symbolic_ctx: false,
-            dummy_ctx: false,
-        };
-        self.add_ctx_with_cache(current_ctx, &mut cache)
+    pub fn add_ctx(self: &RcExpr, current_ctx: Assumption) -> RcExpr {
+        self.add_ctx_with_cache(current_ctx, &mut ContextCache::new())
     }
 
     fn add_ctx_with_cache(