Return the value without a control flow wrapper if a loop only has co…

…ntinue.
cryspen · Feb 6, 2025 · cd76d3d · cd76d3d
1 parent 321a613
commit cd76d3d
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 0 deletions.
diff --git a/engine/lib/phases/phase_drop_return_break_continue.ml b/engine/lib/phases/phase_drop_return_break_continue.ml
@@ -101,6 +101,9 @@ module%inlined_contents Make (F : Features.T) = struct
               Some ({ return_type; break_type }, _) ) ->
               UA.M.expr_Constructor_CF ~return_type ~span ~break_type ~e ~acc
                 `Break
+          | ( Continue { acc = Some (acc, _); _ },
+              Some ({ return_type = None; break_type = None }, _) ) ->
+              acc
           | ( Continue { acc = Some (acc, _); _ },
               Some ({ return_type; break_type }, _) ) ->
               UA.M.expr_Constructor_CF ~return_type ~span ~break_type ~acc

diff --git a/test-harness/src/snapshots/toolchain__loops into-fstar.snap b/test-harness/src/snapshots/toolchain__loops into-fstar.snap
@@ -377,6 +377,57 @@ let nested_return (_: Prims.unit) : i32 =
   with
   | Core.Ops.Control_flow.ControlFlow_Break ret -> ret
   | Core.Ops.Control_flow.ControlFlow_Continue sum -> sum *! mk_i32 2
+
+let continue_only (x: t_Slice i32) : (i32 & Prims.unit) =
+  let product:i32 = mk_i32 1 in
+  Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(t_Slice i32)
+        #FStar.Tactics.Typeclasses.solve
+        x
+      <:
+      Core.Slice.Iter.t_Iter i32)
+    product
+    (fun product i ->
+        let product:i32 = product in
+        let i:i32 = i in
+        if i =. mk_i32 0 <: bool
+        then product
+        else Core.Ops.Arith.f_mul_assign #i32 #i32 #FStar.Tactics.Typeclasses.solve product i <: i32
+    ),
+  ()
+  <:
+  (i32 & Prims.unit)
+
+let continue_and_break (x: t_Slice i32) : (i32 & Prims.unit) =
+  let product:i32 = mk_i32 1 in
+  Rust_primitives.Hax.Folds.fold_cf (Core.Iter.Traits.Collect.f_into_iter #(t_Slice i32)
+        #FStar.Tactics.Typeclasses.solve
+        x
+      <:
+      Core.Slice.Iter.t_Iter i32)
+    product
+    (fun product i ->
+        let product:i32 = product in
+        let i:i32 = i in
+        if i =. mk_i32 0 <: bool
+        then
+          Core.Ops.Control_flow.ControlFlow_Continue product
+          <:
+          Core.Ops.Control_flow.t_ControlFlow (Prims.unit & i32) i32
+        else
+          if i <. mk_i32 0 <: bool
+          then
+            Core.Ops.Control_flow.ControlFlow_Break ((), product <: (Prims.unit & i32))
+            <:
+            Core.Ops.Control_flow.t_ControlFlow (Prims.unit & i32) i32
+          else
+            Core.Ops.Control_flow.ControlFlow_Continue
+            (Core.Ops.Arith.f_mul_assign #i32 #i32 #FStar.Tactics.Typeclasses.solve product i <: i32
+            )
+            <:
+            Core.Ops.Control_flow.t_ControlFlow (Prims.unit & i32) i32),
+  ()
+  <:
+  (i32 & Prims.unit)
 '''
 "Loops.For_loops.fst" = '''
 module Loops.For_loops

diff --git a/tests/loops/src/lib.rs b/tests/loops/src/lib.rs
@@ -243,6 +243,27 @@ mod control_flow {
         sum *= 2;
         sum
     }
+    fn continue_only(x: &[i32]) {
+        let mut product = 1;
+        for i in x {
+            if *i == 0 {
+                continue;
+            }
+            product *= i
+        }
+    }
+    fn continue_and_break(x: &[i32]) {
+        let mut product = 1;
+        for i in x {
+            if *i == 0 {
+                continue;
+            }
+            if *i < 0 {
+                break;
+            }
+            product *= i
+        }
+    }
 }
 
 mod and_mut_side_effect_loop {