diff --git a/dag_in_context/src/add_context.rs b/dag_in_context/src/add_context.rs index 1ad4fd8e..f316c4db 100644 --- a/dag_in_context/src/add_context.rs +++ b/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); diff --git a/dag_in_context/src/lib.rs b/dag_in_context/src/lib.rs index 0acebae4..83d21460 100644 --- a/dag_in_context/src/lib.rs +++ b/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));