Fix loop-peel contexts

dag_in_context/src/optimizations/loop_unroll.egg

@@ -73,6 +73,7 @@
 
 ;; loop peeling rule
 ;; Only peel loops that we know iterate < 3 times
+(function LoopPeeledPlaceholder (Expr) Assumption :unextractable)
 (rule
  ((= lhs (DoWhile inputs outputs))
   (ContextOf lhs ctx)
@@ -90,15 +91,22 @@
     (InIf true (Get executed-once 0) executed-once-body))
   (let else-ctx
     (InIf false (Get executed-once 0) executed-once-body))
+
+  (let new-loop-arg
+    (Arg inputs-ty then-ctx))
+  (let new-loop-body
+    (Subst (LoopPeeledPlaceholder lhs) new-loop-arg outputs))
+  (union (InLoop new-loop-arg new-loop-body) (LoopPeeledPlaceholder lhs))
+
   (union lhs
     ;; check if we need to continue executing the loop
     (If (Get executed-once 0)
-        executed-once-body ;; inputs are the body executed once
-        (DoWhile (Arg inputs-ty then-ctx)
-          outputs) ;; right now, loop unrolling shares the same outputs, but we could add more context here
-        (Arg inputs-ty else-ctx)))
-   (set (LoopNumItersGuess (Arg inputs-ty then-ctx) outputs) (- old_cost 1))
-    )
+      executed-once-body ;; inputs are the body executed once
+      (DoWhile new-loop-arg new-loop-body)
+      (Arg inputs-ty else-ctx)))
+
+  (set (LoopNumItersGuess new-loop-arg new-loop-body) (- old_cost 1))
+  )
  :ruleset loop-peel)
 
 ;; unroll a loop with constant bounds and initial value

tests/snapshots/files__block-diamond-optimize.snap

@@ -4,30 +4,30 @@ expression: visualization.result
 ---
 @main(v0: int) {
 .b1_:
-  c2_: int = const 1;
-  c3_: int = const 2;
-  v4_: bool = lt v0 c3_;
-  v5_: bool = not v4_;
-  c6_: int = const 0;
+  c2_: int = const 2;
+  v3_: bool = lt v0 c2_;
+  v4_: bool = not v3_;
+  c5_: int = const 0;
+  c6_: int = const 1;
   c7_: int = const 5;
-  v8_: int = id c2_;
-  v9_: int = id c2_;
-  v10_: int = id c3_;
-  br v4_ .b11_ .b12_;
+  v8_: int = id c6_;
+  v9_: int = id c6_;
+  v10_: int = id c2_;
+  br v3_ .b11_ .b12_;
 .b11_:
   c13_: int = const 4;
   v8_: int = id c13_;
-  v9_: int = id c2_;
-  v10_: int = id c3_;
+  v9_: int = id c6_;
+  v10_: int = id c2_;
 .b12_:
   v14_: int = id v8_;
   v15_: int = id v9_;
-  br v5_ .b16_ .b17_;
+  br v4_ .b16_ .b17_;
 .b16_:
   v18_: int = add v10_ v8_;
   v14_: int = id v18_;
   v15_: int = id v9_;
 .b17_:
-  v19_: int = add c2_ v14_;
+  v19_: int = add c6_ v14_;
   print v19_;
 }

tests/snapshots/files__gamma_condition_and-optimize.snap

@@ -5,8 +5,16 @@ expression: visualization.result
 @main(v0: int) {
 .b1_:
   c2_: int = const 0;
+<<<<<<< HEAD
   v3_: bool = lt v0 c2_;
   v4_: bool = lt c2_ v0;
+||||||| parent of 96126495 (Fix loop-peel contexts)
+  v3_: bool = lt v0 c2_;
+  v4_: bool = gt v0 c2_;
+=======
+  v3_: bool = gt v0 c2_;
+  v4_: bool = lt v0 c2_;
+>>>>>>> 96126495 (Fix loop-peel contexts)
   c5_: int = const 1;
   c6_: int = const 3;
   v7_: int = id c6_;

tests/snapshots/files__if_dead_code_nested-optimize.snap

@@ -31,7 +31,7 @@ expression: visualization.result
   print v18_;
   ret;
 .b8_:
-  v20_: bool = lt c6_ v0;
+  v20_: bool = gt v0 c6_;
   c21_: bool = const false;
   c22_: int = const 2;
   v23_: int = id c22_;