Fix state passthrough (banish the shadow realm)

dag_in_context/src/add_context.rs

@@ -3,25 +3,39 @@
 //! by remembering the most recent context (ex. DoWhile or If).
 //! Mantains the sharing invariant (see restore_sharing_invariant) by using a cache.
 
+use egglog::Term;
 use std::collections::HashMap;
 
 use crate::{
+    print_with_intermediate_helper,
     schema::{Assumption, Expr, RcExpr, TreeProgram},
     schema_helpers::AssumptionRef,
+    to_egglog::TreeToEgglog,
 };
 
-struct ContextCache {
+pub struct ContextCache {
     with_ctx: HashMap<(*const Expr, AssumptionRef), RcExpr>,
     symbol_gen: HashMap<(*const Expr, AssumptionRef), String>,
+    /// How many placeholder contexts we've created (used to make a unique name each time)
+    loop_context_placeholder_counter: usize,
+    /// The unions that we need to make between assumptions
+    /// One assumption will be a placeholder and the other will be an `InLoop`
+    loop_context_unions: Vec<(Assumption, Assumption)>,
     /// When true, don't add context- instead, make fresh query variables
     /// and put these in place of context
     symbolic_ctx: bool,
     /// Replace all context with (InFunc "dummy")
     dummy_ctx: bool,
 }
 
+impl Default for ContextCache {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 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) {
@@ -31,91 +45,172 @@ impl ContextCache {
         self.symbol_gen.insert(key, sym.clone());
         Assumption::WildCard(sym)
     }
+
+    pub fn new() -> ContextCache {
+        ContextCache {
+            with_ctx: HashMap::new(),
+            symbol_gen: HashMap::new(),
+            loop_context_placeholder_counter: 0,
+            loop_context_unions: Vec::new(),
+            symbolic_ctx: false,
+            dummy_ctx: false,
+        }
+    }
+
+    pub fn new_symbolic_ctx() -> ContextCache {
+        ContextCache {
+            with_ctx: HashMap::new(),
+            symbol_gen: HashMap::new(),
+            loop_context_placeholder_counter: 0,
+            loop_context_unions: Vec::new(),
+            symbolic_ctx: true,
+            dummy_ctx: false,
+        }
+    }
+
+    pub fn new_dummy_ctx() -> ContextCache {
+        ContextCache {
+            with_ctx: HashMap::new(),
+            symbol_gen: HashMap::new(),
+            loop_context_placeholder_counter: 0,
+            loop_context_unions: Vec::new(),
+            symbolic_ctx: false,
+            dummy_ctx: true,
+        }
+    }
+
+    pub fn get_unions(&self) -> String {
+        use std::fmt::Write;
+
+        self.loop_context_unions
+            .iter()
+            .fold(String::new(), |mut output, (a, b)| {
+                let _ = writeln!(output, "(union {a} {b})");
+                output
+            })
+    }
+
+    pub(crate) fn get_unions_with_sharing(
+        &self,
+        printed: &mut String,
+        tree_state: &mut TreeToEgglog,
+        term_cache: &mut HashMap<Term, String>,
+    ) -> String {
+        self.loop_context_unions
+            .iter()
+            .map(|(a, b)| {
+                let internal_a = a.to_egglog_internal(tree_state);
+                let internal_b = b.to_egglog_internal(tree_state);
+
+                let shared_a = print_with_intermediate_helper(
+                    &tree_state.termdag,
+                    internal_a,
+                    term_cache,
+                    printed,
+                );
+                let shared_b = print_with_intermediate_helper(
+                    &tree_state.termdag,
+                    internal_b,
+                    term_cache,
+                    printed,
+                );
+
+                format!("(union {shared_a} {shared_b})")
+            })
+            .collect::<Vec<_>>()
+            .join("\n")
+    }
+
+    pub fn new_placeholder(&mut self) -> Assumption {
+        // Using `InFunc` is a hack to avoid changing the schema
+        // We just need a unique `Assumption` e-node
+        // The string contains a space to avoid conflicts
+        let placeholder = Assumption::InFunc(format!(
+            " loop_ctx_{}",
+            self.loop_context_placeholder_counter
+        ));
+        self.loop_context_placeholder_counter += 1;
+        placeholder
+    }
+
+    pub fn push_loop_context_union(&mut self, a: Assumption, b: Assumption) {
+        self.loop_context_unions.push((a, b))
+    }
 }
 
 impl TreeProgram {
-    pub fn add_context(&self) -> TreeProgram {
-        TreeProgram {
-            functions: self
-                .functions
-                .iter()
-                .map(|f| f.clone().func_add_ctx())
-                .collect(),
-            entry: self.entry.clone().func_add_ctx(),
-        }
+    pub fn add_context(&self) -> (TreeProgram, ContextCache) {
+        self.add_context_internal(Expr::func_get_ctx, 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 {
-        TreeProgram {
-            functions: self
-                .functions
-                .iter()
-                .map(|f| f.clone().add_symbolic_ctx())
-                .collect(),
-            entry: self.entry.clone().add_symbolic_ctx(),
-        }
+    pub fn add_symbolic_ctx(&self) -> TreeProgram {
+        self.add_context_internal(|_| Assumption::dummy(), ContextCache::new_symbolic_ctx())
+            .0
     }
 
-    pub(crate) fn add_dummy_ctx(&self) -> TreeProgram {
-        TreeProgram {
-            functions: self
-                .functions
-                .iter()
-                .map(|f| f.clone().add_dummy_ctx())
-                .collect(),
-            entry: self.entry.clone().add_dummy_ctx(),
-        }
+    pub fn add_dummy_ctx(&self) -> (TreeProgram, ContextCache) {
+        self.add_context_internal(|_| Assumption::dummy(), ContextCache::new_dummy_ctx())
+    }
+
+    fn add_context_internal(
+        &self,
+        get_assumption: impl Fn(&RcExpr) -> Assumption,
+        mut cache: ContextCache,
+    ) -> (TreeProgram, ContextCache) {
+        let entry = self
+            .entry
+            .add_ctx_with_cache(get_assumption(&self.entry), &mut cache);
+        let functions = self
+            .functions
+            .iter()
+            .map(|f| f.add_ctx_with_cache(get_assumption(f), &mut cache))
+            .collect();
+        (TreeProgram { functions, 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);
         };
-        let current_ctx = Assumption::InFunc(name.clone());
-        RcExpr::new(Expr::Function(
+        Assumption::InFunc(name.clone())
+    }
+
+    pub fn func_add_ctx(self: &RcExpr) -> (RcExpr, ContextCache) {
+        let Expr::Function(name, arg_ty, ret_ty, body) = self.as_ref() else {
+            panic!("Expected Function, got {:?}", self);
+        };
+        let mut cache = ContextCache::new();
+        let value = 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 cache),
+        ));
+        (value, cache)
     }
 
-    pub(crate) fn add_dummy_ctx(self: RcExpr) -> RcExpr {
-        let mut cache = ContextCache {
-            with_ctx: HashMap::new(),
-            symbol_gen: HashMap::new(),
-            symbolic_ctx: false,
-            dummy_ctx: true,
-        };
-        self.add_ctx_with_cache(Assumption::dummy(), &mut cache)
+    pub fn add_dummy_ctx(self: &RcExpr) -> (RcExpr, ContextCache) {
+        let mut cache = ContextCache::new_dummy_ctx();
+        let value = self.add_ctx_with_cache(Assumption::dummy(), &mut cache);
+        (value, cache)
     }
 
-    pub(crate) fn add_symbolic_ctx(self: RcExpr) -> RcExpr {
-        let mut cache = ContextCache {
-            with_ctx: HashMap::new(),
-            symbol_gen: HashMap::new(),
-            symbolic_ctx: true,
-            dummy_ctx: false,
-        };
+    pub fn add_symbolic_ctx(self: &RcExpr) -> RcExpr {
+        let mut cache = ContextCache::new_symbolic_ctx();
         self.add_ctx_with_cache(Assumption::dummy(), &mut cache)
     }
 
-    pub(crate) fn add_ctx(self: &RcExpr, current_ctx: Assumption) -> RcExpr {
-        let mut cache = ContextCache {
-            with_ctx: HashMap::new(),
-            symbol_gen: HashMap::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, ContextCache) {
+        let mut cache = ContextCache::new();
+        let value = self.add_ctx_with_cache(current_ctx, &mut cache);
+        (value, cache)
     }
 
-    fn add_ctx_with_cache(
+    pub fn add_ctx_with_cache(
         self: &RcExpr,
         current_ctx: Assumption,
         cache: &mut ContextCache,
@@ -143,12 +238,26 @@ impl Expr {
             Expr::Arg(ty, _oldctx) => RcExpr::new(Expr::Arg(ty.clone(), context_to_add)),
             // create new contexts for let, loop, and if
             Expr::DoWhile(inputs, pred_and_body) => {
+                // We need to make loop contexts point to their enclosing loops in the e-graph.
+                // (This creates a cycle in the e-graph, but we can still extract because
+                // our extractor ignores context.)
+
+                // To do this we have to first make a placeholder context
+                // that we can use to build the body. Once we've built the body, we can
+                // build the real `InLoop` context. Finally, we'll have to union the
+                // placeholder context with the `InLoop` one later, which we track using
+                // the `push_loop_context_union` function.
+
+                let placeholder = cache.new_placeholder();
+
                 let new_inputs = inputs.add_ctx_with_cache(current_ctx.clone(), cache);
-                let new_ctx = Assumption::InLoop(new_inputs.clone(), pred_and_body.clone());
-                RcExpr::new(Expr::DoWhile(
-                    new_inputs,
-                    pred_and_body.add_ctx_with_cache(new_ctx, cache),
-                ))
+                let new_pred_and_body =
+                    pred_and_body.add_ctx_with_cache(placeholder.clone(), cache);
+
+                let new_ctx = Assumption::InLoop(new_inputs.clone(), new_pred_and_body.clone());
+                cache.push_loop_context_union(placeholder, new_ctx);
+
+                RcExpr::new(Expr::DoWhile(new_inputs, new_pred_and_body))
             }
             Expr::If(pred, input, then_case, else_calse) => {
                 let new_pred = pred.add_ctx_with_cache(current_ctx.clone(), cache);