Fix remaining dag_in_context tests

egraphs-good · Jun 4, 2024 · c976ac5 · c976ac5
1 parent b348242
commit c976ac5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/dag_in_context/src/optimizations/switch_rewrites.rs b/dag_in_context/src/optimizations/switch_rewrites.rs
@@ -4,10 +4,11 @@ use crate::egglog_test;
 #[test]
 fn switch_rewrite_three_quarters_and() -> crate::Result {
     use crate::ast::*;
+    use crate::schema::Assumption;
 
     let (build, build_cache) = tif(and(tfalse(), ttrue()), empty(), int(1), int(2))
         .with_arg_types(emptyt(), base(intt()))
-        .add_dummy_ctx();
+        .add_ctx(Assumption::dummy());
 
     let (check, check_cache) = tif(
         tfalse(),
@@ -16,7 +17,7 @@ fn switch_rewrite_three_quarters_and() -> crate::Result {
         int(2),
     )
     .with_arg_types(emptyt(), base(intt()))
-    .add_dummy_ctx();
+    .add_ctx(Assumption::dummy());
 
     egglog_test(
         &format!("(let build_ {build})\n{}", build_cache.get_unions()),
@@ -34,10 +35,11 @@ fn switch_rewrite_three_quarters_and() -> crate::Result {
 #[test]
 fn switch_rewrite_three_quarters_or() -> crate::Result {
     use crate::ast::*;
+    use crate::schema::Assumption;
 
     let (build, build_cache) = tif(or(tfalse(), ttrue()), empty(), int(1), int(2))
         .with_arg_types(emptyt(), base(intt()))
-        .add_dummy_ctx();
+        .add_ctx(Assumption::dummy());
 
     let (check, check_cache) = tif(
         tfalse(),
@@ -46,7 +48,7 @@ fn switch_rewrite_three_quarters_or() -> crate::Result {
         tif(get(arg(), 0), empty(), int(1), int(2)),
     )
     .with_arg_types(emptyt(), base(intt()))
-    .add_dummy_ctx();
+    .add_ctx(Assumption::dummy());
 
     egglog_test(
         &format!("(let build_ {build})\n{}", build_cache.get_unions()),

diff --git a/dag_in_context/src/utility/context_of.rs b/dag_in_context/src/utility/context_of.rs
@@ -18,12 +18,12 @@ fn test_context_of() -> crate::Result {
     // If statement should have the context of its predicate
     let check = format!(
         "
-        (let pred-ctx {ctx})
-        (let pred (Bop (Eq) (Arg (Base (IntT)) {ctx}) (Const (Int 5) (Base (IntT)) {ctx})))
-        (check (ContextOf pred pred-ctx))
-        (let if {})
-        {body_with_context}
-        (check (ContextOf if pred-ctx))",
+(let pred-ctx {ctx})
+(let pred (Bop (Eq) (Arg (Base (IntT)) {ctx}) (Const (Int 5) (Base (IntT)) {ctx})))
+(check (ContextOf pred pred-ctx))
+(let if {body_with_context})
+{}
+(check (ContextOf if pred-ctx))",
         body_cache.get_unions(),
     );