Fix some issues

2 files changed, 10 insertions, 9 deletions

diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index 0634eed7..fd3e334b 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -358,6 +358,7 @@ let evaluate_function_symbolic_synthesize_backward_from_return (config : config)
       in
       if not inside_loop then (Some fun_abs_id, true)
       else
+        (* We are inside a loop *)
         let pred (abs : abs) =
           match abs.kind with
           | Loop (_, rg_id', kind) ->
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 8b4362a2..922f0375 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -3044,16 +3044,19 @@ and translate_forward_end (ectx : C.eval_ctx)
     let ctx, pure_fwd_var = fresh_var None ctx.sg.fwd_output ctx in
     let fwd_e = translate_one_end ctx None in
 
-    (* If we are translating a loop, we need to ignore the backward functions
-       which are not associated to region abstractions of the loop. *)
+    (* If we reached a loop: if we are *inside* a loop, we need to ignore the
+       backward functions which are not associated to region abstractions.
+    *)
     let keep_rg_ids =
       match ctx.loop_id with
       | None -> None
       | Some loop_id ->
-          let loop_info = LoopId.Map.find loop_id ctx.loops in
-          Some
-            (RegionGroupId.Set.of_list
-               (RegionGroupId.Map.keys loop_info.back_outputs))
+          if ctx.inside_loop then
+            let loop_info = LoopId.Map.find loop_id ctx.loops in
+            Some
+              (RegionGroupId.Set.of_list
+                 (RegionGroupId.Map.keys loop_info.back_outputs))
+          else None
     in
     let keep_rg_id =
       match keep_rg_ids with
@@ -3517,9 +3520,6 @@ and translate_loop (loop : S.loop) (ctx : bs_ctx) : texpression =
         loop_body;
       }
   in
-  (* If we translate forward functions: the return type of a loop body is the
-     same as the parent function *)
-  assert (Option.is_some ctx.bid || fun_end.ty = loop_body.ty);
   let ty = fun_end.ty in
   { e = loop; ty }