Fix a bug when composing the continuations in eval_statement

4 files changed, 131 insertions, 122 deletions

diff --git a/compiler/Cps.ml b/compiler/Cps.ml
index 142c2b08..f7694ba2 100644
--- a/compiler/Cps.ml
+++ b/compiler/Cps.ml
@@ -106,6 +106,12 @@ let cc_singleton file line span cf el =
   | Some _ -> internal_error file line span
   | _ -> None
 
+let cf_singleton file line span el =
+  match el with
+  | Some [ e ] -> Some e
+  | Some _ -> internal_error file line span
+  | _ -> None
+
 (** It happens that we need to concatenate lists of results, for
     instance when evaluating the branches of a match. When applying
     the continuations to build the symbolic expressions, we need
diff --git a/compiler/InterpreterExpansion.ml b/compiler/InterpreterExpansion.ml
index 388d7382..4393e89f 100644
--- a/compiler/InterpreterExpansion.ml
+++ b/compiler/InterpreterExpansion.ml
@@ -478,7 +478,7 @@ let expand_symbolic_value_no_branching (config : config) (span : Meta.span)
   (* Debug *)
   log#ldebug
     (lazy
-      ("expand_symbolic_value_no_branching:" ^ symbolic_value_to_string ctx sv));
+      ("expand_symbolic_value_no_branching: " ^ symbolic_value_to_string ctx sv));
   (* Remember the initial context for printing purposes *)
   let ctx0 = ctx in
   (* Compute the expanded value - note that when doing so, we may introduce
@@ -631,6 +631,8 @@ let greedy_expand_symbolics_with_borrows (config : config) (span : Meta.span) :
       (* We reverse the environment before exploring it - this way the values
          get expanded in a more "logical" order (this is only for convenience) *)
       obj#visit_env () (List.rev ctx.env);
+      log#ldebug
+        (lazy "greedy_expand_symbolics_with_borrows: no value to expand\n");
       (* Nothing to expand: continue *)
       (ctx, fun e -> e)
     with FoundSymbolicValue sv ->
@@ -674,6 +676,12 @@ let greedy_expand_symbolics_with_borrows (config : config) (span : Meta.span) :
         | TVar _ | TLiteral _ | TNever | TTraitType _ | TArrow _ | TRawPtr _ ->
             craise __FILE__ __LINE__ span "Unreachable"
       in
+      (* *)
+      log#ldebug
+        (lazy
+          ("\ngreedy_expand_symbolics_with_borrows: after expansion:\n"
+          ^ eval_ctx_to_string ~span:(Some span) ctx
+          ^ "\n\n"));
       (* Compose and continue *)
       comp cc (expand ctx)
   in
diff --git a/compiler/InterpreterStatements.ml b/compiler/InterpreterStatements.ml
index c6a65757..19510c2e 100644
--- a/compiler/InterpreterStatements.ml
+++ b/compiler/InterpreterStatements.ml
@@ -914,116 +914,111 @@ let rec eval_statement (config : config) (st : statement) : stl_cm_fun =
   (* Sanity check *)
   Invariants.check_invariants st.span ctx;
 
-  (* Evaluate *)
-  let stl, cf_eval_st =
-    log#ldebug
-      (lazy
-        ("\neval_statement: cf_eval_st: statement:\n"
-        ^ statement_to_string_with_tab ctx st
-        ^ "\n\n"));
-    match st.content with
-    | Assign (p, rvalue) -> (
-        (* We handle global assignments separately *)
-        match rvalue with
-        | Global (gid, generics) ->
-            (* Evaluate the global *)
-            eval_global config st.span p gid generics ctx
-        | _ ->
-            (* Evaluate the rvalue *)
-            let res, ctx, cc =
-              eval_rvalue_not_global config st.span rvalue ctx
-            in
-            (* Assign *)
-            log#ldebug
-              (lazy
-                ("about to assign to place: " ^ place_to_string ctx p
-               ^ "\n- Context:\n"
-                ^ eval_ctx_to_string ~span:(Some st.span) ctx));
-            let (ctx, res), cf_assign =
-              match res with
-              | Error EPanic -> ((ctx, Panic), fun e -> e)
-              | Ok rv ->
-                  (* Update the synthesized AST - here we store additional span-information.
-                   * We do it only in specific cases (it is not always useful, and
-                   * also it can lead to issues - for instance, if we borrow a
-                   * reserved borrow, we later can't translate it to pure values...) *)
-                  let cc =
-                    match rvalue with
-                    | Global _ -> craise __FILE__ __LINE__ st.span "Unreachable"
-                    | Use _
-                    | RvRef (_, (BShared | BMut | BTwoPhaseMut | BShallow))
-                    | UnaryOp _ | BinaryOp _ | Discriminant _ | Aggregate _ ->
-                        let rp = rvalue_get_place rvalue in
-                        let rp =
-                          match rp with
-                          | Some rp -> Some (S.mk_mplace st.span rp ctx)
-                          | None -> None
-                        in
-                        S.synthesize_assignment ctx
-                          (S.mk_mplace st.span p ctx)
-                          rv rp
-                  in
-                  let ctx, cc =
-                    comp cc (assign_to_place config st.span rv p ctx)
-                  in
-                  ((ctx, Unit), cc)
-            in
-            let cc = cc_comp cc cf_assign in
-            (* Compose and apply *)
-            ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc))
-    | FakeRead p ->
-        let ctx, cc = eval_fake_read config st.span p ctx in
-        ([ (ctx, Unit) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | SetDiscriminant (p, variant_id) ->
-        let (ctx, res), cc = set_discriminant config st.span p variant_id ctx in
-        ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Drop p ->
-        let ctx, cc = drop_value config st.span p ctx in
-        ([ (ctx, Unit) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Assert assertion ->
-        let (ctx, res), cc = eval_assertion config st.span assertion ctx in
-        ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Call call -> eval_function_call config st.span call ctx
-    | Panic -> ([ (ctx, Panic) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Return -> ([ (ctx, Return) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Break i -> ([ (ctx, Break i) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Continue i ->
-        ([ (ctx, Continue i) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Nop -> ([ (ctx, Unit) ], cc_singleton __FILE__ __LINE__ st.span cc)
-    | Sequence (st1, st2) ->
-        (* Evaluate the first statement *)
-        let ctx_resl, cf_st1 = eval_statement config st1 ctx in
-        (* Evaluate the sequence (evaluate the second statement if the first
-           statement successfully evaluated, otherwise transfmit the control-flow
-           break) *)
-        let ctx_res_cfl =
-          List.map
-            (fun (ctx, res) ->
-              match res with
-              (* Evaluation successful: evaluate the second statement *)
-              | Unit -> eval_statement config st2 ctx
-              (* Control-flow break: transmit. We enumerate the cases on purpose *)
-              | Panic | Break _ | Continue _ | Return | LoopReturn _
-              | EndEnterLoop _ | EndContinue _ ->
-                  ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc))
-            ctx_resl
-        in
-        (* Put everything together:
-           - we return the flattened list of contexts and results
-           - we need to build the continuation which will build the whole
-             expression from the continuations for the individual branches
-        *)
-        let ctx_resl, cf_st2 =
-          comp_seqs __FILE__ __LINE__ st.span ctx_res_cfl
-        in
-        (ctx_resl, cc_comp cf_st1 cf_st2)
-    | Loop loop_body ->
-        let eval_loop_body = eval_statement config loop_body in
-        InterpreterLoops.eval_loop config st.span eval_loop_body ctx
-    | Switch switch -> eval_switch config st.span switch ctx
-  in
-  (* Compose and apply *)
-  (stl, cc_comp cc cf_eval_st)
+  (* Evaluate the statement *)
+  comp cc (eval_statement_raw config st ctx)
+
+and eval_statement_raw (config : config) (st : statement) : stl_cm_fun =
+ fun ctx ->
+  log#ldebug
+    (lazy
+      ("\neval_statement_raw: statement:\n"
+      ^ statement_to_string_with_tab ctx st
+      ^ "\n\n"));
+  match st.content with
+  | Assign (p, rvalue) -> (
+      (* We handle global assignments separately *)
+      match rvalue with
+      | Global (gid, generics) ->
+          (* Evaluate the global *)
+          eval_global config st.span p gid generics ctx
+      | _ ->
+          (* Evaluate the rvalue *)
+          let res, ctx, cc = eval_rvalue_not_global config st.span rvalue ctx in
+          (* Assign *)
+          log#ldebug
+            (lazy
+              ("about to assign to place: " ^ place_to_string ctx p
+             ^ "\n- Context:\n"
+              ^ eval_ctx_to_string ~span:(Some st.span) ctx));
+          let (ctx, res), cf_assign =
+            match res with
+            | Error EPanic -> ((ctx, Panic), fun e -> e)
+            | Ok rv ->
+                (* Update the synthesized AST - here we store additional span-information.
+                 * We do it only in specific cases (it is not always useful, and
+                 * also it can lead to issues - for instance, if we borrow a
+                 * reserved borrow, we later can't translate it to pure values...) *)
+                let cc =
+                  match rvalue with
+                  | Global _ -> craise __FILE__ __LINE__ st.span "Unreachable"
+                  | Use _
+                  | RvRef (_, (BShared | BMut | BTwoPhaseMut | BShallow))
+                  | UnaryOp _ | BinaryOp _ | Discriminant _ | Aggregate _ ->
+                      let rp = rvalue_get_place rvalue in
+                      let rp =
+                        match rp with
+                        | Some rp -> Some (S.mk_mplace st.span rp ctx)
+                        | None -> None
+                      in
+                      S.synthesize_assignment ctx
+                        (S.mk_mplace st.span p ctx)
+                        rv rp
+                in
+                let ctx, cc =
+                  comp cc (assign_to_place config st.span rv p ctx)
+                in
+                ((ctx, Unit), cc)
+          in
+          let cc = cc_comp cc cf_assign in
+          (* Compose and apply *)
+          ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc))
+  | FakeRead p ->
+      let ctx, cc = eval_fake_read config st.span p ctx in
+      ([ (ctx, Unit) ], cc_singleton __FILE__ __LINE__ st.span cc)
+  | SetDiscriminant (p, variant_id) ->
+      let (ctx, res), cc = set_discriminant config st.span p variant_id ctx in
+      ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc)
+  | Drop p ->
+      let ctx, cc = drop_value config st.span p ctx in
+      ([ (ctx, Unit) ], cc_singleton __FILE__ __LINE__ st.span cc)
+  | Assert assertion ->
+      let (ctx, res), cc = eval_assertion config st.span assertion ctx in
+      ([ (ctx, res) ], cc_singleton __FILE__ __LINE__ st.span cc)
+  | Call call -> eval_function_call config st.span call ctx
+  | Panic -> ([ (ctx, Panic) ], cf_singleton __FILE__ __LINE__ st.span)
+  | Return -> ([ (ctx, Return) ], cf_singleton __FILE__ __LINE__ st.span)
+  | Break i -> ([ (ctx, Break i) ], cf_singleton __FILE__ __LINE__ st.span)
+  | Continue i -> ([ (ctx, Continue i) ], cf_singleton __FILE__ __LINE__ st.span)
+  | Nop -> ([ (ctx, Unit) ], cf_singleton __FILE__ __LINE__ st.span)
+  | Sequence (st1, st2) ->
+      (* Evaluate the first statement *)
+      let ctx_resl, cf_st1 = eval_statement config st1 ctx in
+      (* Evaluate the sequence (evaluate the second statement if the first
+         statement successfully evaluated, otherwise transfmit the control-flow
+         break) *)
+      let ctx_res_cfl =
+        List.map
+          (fun (ctx, res) ->
+            match res with
+            (* Evaluation successful: evaluate the second statement *)
+            | Unit -> eval_statement config st2 ctx
+            (* Control-flow break: transmit. We enumerate the cases on purpose *)
+            | Panic | Break _ | Continue _ | Return | LoopReturn _
+            | EndEnterLoop _ | EndContinue _ ->
+                ([ (ctx, res) ], cf_singleton __FILE__ __LINE__ st.span))
+          ctx_resl
+      in
+      (* Put everything together:
+         - we return the flattened list of contexts and results
+         - we need to build the continuation which will build the whole
+           expression from the continuations for the individual branches
+      *)
+      let ctx_resl, cf_st2 = comp_seqs __FILE__ __LINE__ st.span ctx_res_cfl in
+      (ctx_resl, cc_comp cf_st1 cf_st2)
+  | Loop loop_body ->
+      let eval_loop_body = eval_statement config loop_body in
+      InterpreterLoops.eval_loop config st.span eval_loop_body ctx
+  | Switch switch -> eval_switch config st.span switch ctx
 
 and eval_global (config : config) (span : Meta.span) (dest : place)
     (gid : GlobalDeclId.id) (generics : generic_args) : stl_cm_fun =
@@ -1609,11 +1604,11 @@ and eval_function_body (config : config) (body : statement) : stl_cm_fun =
   let ctx_res_cfl =
     List.map
       (fun (ctx, res) ->
-        log#ldebug (lazy "eval_function_body: cf_finish");
         (* Note that we *don't* check the result ({!Panic}, {!Return}, etc.): we
-           * delegate the check to the caller. *)
+           delegate the check to the caller. *)
+        log#ldebug (lazy "eval_function_body: cf_finish");
         (* Expand the symbolic values if necessary - we need to do that before
-           * checking the invariants *)
+           checking the invariants *)
         let ctx, cf = greedy_expand_symbolic_values config body.span ctx in
         (* Sanity check *)
         Invariants.check_invariants body.span ctx;
diff --git a/compiler/SymbolicAst.ml b/compiler/SymbolicAst.ml
index e9143ab5..750297e4 100644
--- a/compiler/SymbolicAst.ml
+++ b/compiler/SymbolicAst.ml
@@ -37,7 +37,7 @@ type call_id =
 
 type call = {
   call_id : call_id;
-  ctx : Contexts.eval_ctx;
+  ctx : (Contexts.eval_ctx[@opaque]);
       (** The context upon calling the function (after the operands have been
           evaluated). We need it to compute the translated values for shared
           borrows (we need to perform lookups).
@@ -123,9 +123,9 @@ class ['self] iter_expression_base =
 (** **Rem.:** here, {!expression} is not at all equivalent to the expressions
     used in LLBC or in lambda-calculus: they are simply a first step towards
     lambda-calculus expressions.
- *)
+*)
 type expression =
-  | Return of Contexts.eval_ctx * typed_value option
+  | Return of (Contexts.eval_ctx[@opaque]) * typed_value option
       (** There are two cases:
           - the AST is for a forward function: the typed value should contain
             the value which was in the return variable
@@ -137,7 +137,7 @@ type expression =
        *)
   | Panic
   | FunCall of call * expression
-  | EndAbstraction of Contexts.eval_ctx * abs * expression
+  | EndAbstraction of (Contexts.eval_ctx[@opaque]) * abs * expression
       (** The context is the evaluation context upon ending the abstraction,
           just after we removed the abstraction from the context.
 
@@ -146,7 +146,7 @@ type expression =
        *)
   | EvalGlobal of global_decl_id * generic_args * symbolic_value * expression
       (** Evaluate a global to a fresh symbolic value *)
-  | Assertion of Contexts.eval_ctx * typed_value * expression
+  | Assertion of (Contexts.eval_ctx[@opaque]) * typed_value * expression
       (** An assertion.
 
           The context is the evaluation context from after evaluating the asserted
@@ -162,7 +162,7 @@ type expression =
           to prettify the generated code.
        *)
   | IntroSymbolic of
-      Contexts.eval_ctx
+      (Contexts.eval_ctx[@opaque])
       * mplace option
       * symbolic_value
       * value_aggregate
@@ -179,7 +179,7 @@ type expression =
           value. It has the same purpose as for the {!Return} case.
        *)
   | ForwardEnd of
-      Contexts.eval_ctx
+      (Contexts.eval_ctx[@opaque])
       * typed_value symbolic_value_id_map option
       * expression
       * expression region_group_id_map
@@ -211,7 +211,7 @@ type expression =
           The boolean is [true].
           TODO: merge this with Return.
        *)
-  | Meta of espan * expression  (** Meta information *)
+  | Meta of (espan[@opaque]) * expression  (** Meta information *)
   | Error of Meta.span option * string
 
 and loop = {