Resolve review comments

dag_in_context/src/add_context.rs

@@ -16,8 +16,10 @@ use crate::{
 pub struct ContextCache {
     with_ctx: HashMap<(*const Expr, AssumptionRef), RcExpr>,
     symbol_gen: HashMap<(*const Expr, AssumptionRef), String>,
-    /// Information for generating the unions for loop contexts
-    placeholder: usize,
+    /// 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
@@ -48,7 +50,7 @@ impl ContextCache {
         ContextCache {
             with_ctx: HashMap::new(),
             symbol_gen: HashMap::new(),
-            placeholder: 0,
+            loop_context_placeholder_counter: 0,
             loop_context_unions: Vec::new(),
             symbolic_ctx: false,
             dummy_ctx: false,
@@ -59,7 +61,7 @@ impl ContextCache {
         ContextCache {
             with_ctx: HashMap::new(),
             symbol_gen: HashMap::new(),
-            placeholder: 0,
+            loop_context_placeholder_counter: 0,
             loop_context_unions: Vec::new(),
             symbolic_ctx: true,
             dummy_ctx: false,
@@ -70,7 +72,7 @@ impl ContextCache {
         ContextCache {
             with_ctx: HashMap::new(),
             symbol_gen: HashMap::new(),
-            placeholder: 0,
+            loop_context_placeholder_counter: 0,
             loop_context_unions: Vec::new(),
             symbolic_ctx: false,
             dummy_ctx: true,
@@ -120,8 +122,14 @@ impl ContextCache {
     }
 
     pub fn new_placeholder(&mut self) -> Assumption {
-        let placeholder = Assumption::InFunc(format!(" loop_ctx_{}", self.placeholder));
-        self.placeholder += 1;
+        // 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
     }
 
@@ -148,14 +156,16 @@ impl TreeProgram {
 
     fn add_context_internal(
         &self,
-        func: impl Fn(&RcExpr) -> Assumption,
+        get_assumption: impl Fn(&RcExpr) -> Assumption,
         mut cache: ContextCache,
     ) -> (TreeProgram, ContextCache) {
-        let entry = self.entry.add_ctx_with_cache(func(&self.entry), &mut cache);
+        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(func(f), &mut cache))
+            .map(|f| f.add_ctx_with_cache(get_assumption(f), &mut cache))
             .collect();
         (TreeProgram { functions, entry }, cache)
     }
@@ -228,6 +238,16 @@ 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);

dag_in_context/src/lib.rs

@@ -200,8 +200,8 @@ pub fn check_roundtrip_egraph(program: &TreeProgram) {
         DefaultCostModel,
     );
 
-    let original_with_ctx = program.add_dummy_ctx().0;
-    let res_with_ctx = res.add_dummy_ctx().0;
+    let (original_with_ctx, _) = program.add_dummy_ctx();
+    let (res_with_ctx, _) = res.add_dummy_ctx();
 
     if !are_progs_eq(original_with_ctx.clone(), res_with_ctx.clone()) {
         eprintln!("Original program: {}", tree_to_svg(&original_with_ctx));