Fix an issue with the type of the values given back by loops

3 files changed, 237 insertions, 157 deletions

diff --git a/compiler/InterpreterLoops.ml b/compiler/InterpreterLoops.ml
index 7714f5bb..90161196 100644
--- a/compiler/InterpreterLoops.ml
+++ b/compiler/InterpreterLoops.ml
@@ -70,6 +70,202 @@ let eval_loop_concrete (span : Meta.span) (eval_loop_body : stl_cm_fun) :
   in
   (ctx_resl, cf)
 
+(** Auxiliary function for {!eval_loop_symbolic}.
+
+    Match the context upon entering the loop with the loop fixed-point to
+    compute how to "apply" the fixed-point. Compute the correspondance from
+    the borrow ids in the current context to the loans which appear in the
+    loop context (we need this in order to know how to introduce the region
+    abstractions of the loop).
+
+    We check the fixed-point at the same time to make sure the loans and borrows
+    inside the region abstractions are properly ordered (this is necessary for the
+    synthesis).
+    Ex.: if in the fixed point we have:
+    {[
+      abs { MB l0, MB l1, ML l2, ML l3 }
+    ]}
+    we want to make sure that borrow l0 actually corresponds to loan l2, and
+    borrow l1 to loan l3.
+ *)
+let eval_loop_symbolic_synthesize_fun_end (config : config) (span : span)
+    (loop_id : LoopId.id) (init_ctx : eval_ctx) (fixed_ids : ids_sets)
+    (fp_ctx : eval_ctx) (fp_input_svalues : SymbolicValueId.id list)
+    (rg_to_abs : AbstractionId.id RegionGroupId.Map.t) :
+    ((eval_ctx * statement_eval_res) * (eval_result -> eval_result))
+    * borrow_loan_corresp =
+  (* First, preemptively end borrows/move values by matching the current
+     context with the target context *)
+  let ctx, cf_prepare =
+    log#ldebug
+      (lazy
+        ("eval_loop_symbolic_synthesize_fun_end: about to reorganize the \
+          original context to match the fixed-point ctx with it:\n\
+          - src ctx (fixed-point ctx):\n" ^ eval_ctx_to_string fp_ctx
+       ^ "\n\n-tgt ctx (original context):\n"
+        ^ eval_ctx_to_string init_ctx));
+
+    prepare_match_ctx_with_target config span loop_id fixed_ids fp_ctx init_ctx
+  in
+
+  (* Actually match *)
+  log#ldebug
+    (lazy
+      ("eval_loop_symbolic_synthesize_fun_end: about to compute the id \
+        correspondance between the fixed-point ctx and the original ctx:\n\
+        - src ctx (fixed-point ctx)\n" ^ eval_ctx_to_string fp_ctx
+     ^ "\n\n-tgt ctx (original context):\n" ^ eval_ctx_to_string ctx));
+
+  (* Compute the id correspondance between the contexts *)
+  let fp_bl_corresp =
+    compute_fixed_point_id_correspondance span fixed_ids ctx fp_ctx
+  in
+  log#ldebug
+    (lazy
+      ("eval_loop_symbolic_synthesize_fun_end: about to match the fixed-point \
+        context with the original context:\n\
+        - src ctx (fixed-point ctx)"
+      ^ eval_ctx_to_string ~span:(Some span) fp_ctx
+      ^ "\n\n-tgt ctx (original context):\n"
+      ^ eval_ctx_to_string ~span:(Some span) ctx
+      ^ "\n\n- fp_bl_corresp:\n"
+      ^ show_borrow_loan_corresp fp_bl_corresp
+      ^ "\n"));
+
+  (* Compute the end expression, that is the expresion corresponding to the
+     end of the function where we call the loop (for now, when calling a loop
+     we never get out) *)
+  let res_fun_end =
+    comp cf_prepare
+      (match_ctx_with_target config span loop_id true fp_bl_corresp
+         fp_input_svalues fixed_ids fp_ctx ctx)
+  in
+
+  (* Sanity check: the mutable borrows/loans are properly ordered.
+     TODO: it seems that the way the fixed points are computed makes this check
+     always succeed. If it happens to fail we can reorder the borrows/loans
+     inside the region abstractions. *)
+  let check_abs (abs_id : AbstractionId.id) =
+    let abs = ctx_lookup_abs fp_ctx abs_id in
+    let is_borrow (av : typed_avalue) : bool =
+      match av.value with
+      | ABorrow _ -> true
+      | ALoan _ -> false
+      | _ -> craise __FILE__ __LINE__ span "Unreachable"
+    in
+    let borrows, loans = List.partition is_borrow abs.avalues in
+
+    let mut_borrows =
+      List.filter_map
+        (fun (av : typed_avalue) ->
+          match av.value with
+          | ABorrow (AMutBorrow (pm, bid, child_av)) ->
+              sanity_check __FILE__ __LINE__ (pm = PNone) span;
+              sanity_check __FILE__ __LINE__ (is_aignored child_av.value) span;
+              Some bid
+          | ABorrow (ASharedBorrow (pm, _)) ->
+              sanity_check __FILE__ __LINE__ (pm = PNone) span;
+              None
+          | _ -> craise __FILE__ __LINE__ span "Unreachable")
+        borrows
+    in
+
+    let mut_loans =
+      List.filter_map
+        (fun (av : typed_avalue) ->
+          match av.value with
+          | ALoan (AMutLoan (pm, bid, child_av)) ->
+              sanity_check __FILE__ __LINE__ (pm = PNone) span;
+              sanity_check __FILE__ __LINE__ (is_aignored child_av.value) span;
+              Some bid
+          | ALoan (ASharedLoan (pm, _, _, _)) ->
+              sanity_check __FILE__ __LINE__ (pm = PNone) span;
+              None
+          | _ -> craise __FILE__ __LINE__ span "Unreachable")
+        loans
+    in
+
+    sanity_check __FILE__ __LINE__
+      (List.length mut_borrows = List.length mut_loans)
+      span;
+
+    let borrows_loans = List.combine mut_borrows mut_loans in
+    List.iter
+      (fun (bid, lid) ->
+        let lid_of_bid =
+          BorrowId.InjSubst.find bid fp_bl_corresp.borrow_to_loan_id_map
+        in
+        sanity_check __FILE__ __LINE__ (lid_of_bid = lid) span)
+      borrows_loans
+  in
+  List.iter check_abs (RegionGroupId.Map.values rg_to_abs);
+
+  (* Return *)
+  (res_fun_end, fp_bl_corresp)
+
+(** Auxiliary function for {!eval_loop_symbolic}.
+
+    Synthesize the body of the loop.
+ *)
+let eval_loop_symbolic_synthesize_loop_body (config : config) (span : span)
+    (eval_loop_body : stl_cm_fun) (loop_id : LoopId.id) (fixed_ids : ids_sets)
+    (fp_ctx : eval_ctx) (fp_input_svalues : SymbolicValueId.id list)
+    (fp_bl_corresp : borrow_loan_corresp) :
+    (eval_ctx * statement_eval_res) list
+    * (SymbolicAst.expression list option -> eval_result) =
+  (* First, evaluate the loop body starting from the **fixed-point** context *)
+  let ctx_resl, cf_loop = eval_loop_body fp_ctx in
+
+  (* Then, do a special treatment of the break and continue cases.
+     For now, we forbid having breaks in loops (and eliminate breaks
+     in the prepasses) *)
+  let eval_after_loop_iter (ctx, res) =
+    log#ldebug (lazy "eval_loop_symbolic: eval_after_loop_iter");
+    match res with
+    | Return ->
+        (* We replace the [Return] with a [LoopReturn] *)
+        ((ctx, LoopReturn loop_id), fun e -> e)
+    | Panic -> ((ctx, res), fun e -> e)
+    | Break _ ->
+        (* Breaks should have been eliminated in the prepasses *)
+        craise __FILE__ __LINE__ span "Unexpected break"
+    | Continue i ->
+        (* We don't support nested loops for now *)
+        cassert __FILE__ __LINE__ (i = 0) span
+          "Nested loops are not supported yet";
+        log#ldebug
+          (lazy
+            ("eval_loop_symbolic: about to match the fixed-point context with \
+              the context at a continue:\n\
+              - src ctx (fixed-point ctx)"
+            ^ eval_ctx_to_string ~span:(Some span) fp_ctx
+            ^ "\n\n-tgt ctx (ctx at continue):\n"
+            ^ eval_ctx_to_string ~span:(Some span) ctx));
+        match_ctx_with_target config span loop_id false fp_bl_corresp
+          fp_input_svalues fixed_ids fp_ctx ctx
+    | Unit | LoopReturn _ | EndEnterLoop _ | EndContinue _ ->
+        (* For why we can't get [Unit], see the comments inside {!eval_loop_concrete}.
+           For [EndEnterLoop] and [EndContinue]: we don't support nested loops for now.
+        *)
+        craise __FILE__ __LINE__ span "Unreachable"
+  in
+
+  (* Apply and compose *)
+  let ctx_resl, cfl = List.split (List.map eval_after_loop_iter ctx_resl) in
+  let cc (el : SymbolicAst.expression list option) : eval_result =
+    match el with
+    | None -> None
+    | Some el ->
+        let el =
+          List.map
+            (fun (cf, e) -> Option.get (cf (Some e)))
+            (List.combine cfl el)
+        in
+        cf_loop (Some el)
+  in
+
+  (ctx_resl, cc)
+
 (** Evaluate a loop in symbolic mode *)
 let eval_loop_symbolic (config : config) (span : span)
     (eval_loop_body : stl_cm_fun) : stl_cm_fun =
@@ -105,115 +301,25 @@ let eval_loop_symbolic (config : config) (span : span)
 
   (* Synthesize the end of the function - we simply match the context at the
      loop entry with the fixed point: in the synthesized code, the function
-     will end with a call to the loop translation
-  *)
-  let ((res_fun_end, cf_fun_end), fp_bl_corresp) :
-      ((eval_ctx * statement_eval_res) * (eval_result -> eval_result)) * _ =
-    (* First, preemptively end borrows/move values by matching the current
-       context with the target context *)
-    let ctx, cf_prepare =
-      log#ldebug
-        (lazy
-          ("eval_loop_symbolic: about to reorganize the original context to \
-            match the fixed-point ctx with it:\n\
-            - src ctx (fixed-point ctx):\n" ^ eval_ctx_to_string fp_ctx
-         ^ "\n\n-tgt ctx (original context):\n" ^ eval_ctx_to_string ctx));
-
-      prepare_match_ctx_with_target config span loop_id fixed_ids fp_ctx ctx
-    in
+     will end with a call to the loop translation.
 
-    (* Actually match *)
-    log#ldebug
-      (lazy
-        ("eval_loop_symbolic: about to compute the id correspondance between \
-          the fixed-point ctx and the original ctx:\n\
-          - src ctx (fixed-point ctx)" ^ eval_ctx_to_string fp_ctx
-       ^ "\n\n-tgt ctx (original context):\n" ^ eval_ctx_to_string ctx));
-
-    (* Compute the id correspondance between the contexts *)
-    let fp_bl_corresp =
-      compute_fixed_point_id_correspondance span fixed_ids ctx fp_ctx
-    in
-    log#ldebug
-      (lazy
-        ("eval_loop_symbolic: about to match the fixed-point context with the \
-          original context:\n\
-          - src ctx (fixed-point ctx)"
-        ^ eval_ctx_to_string ~span:(Some span) fp_ctx
-        ^ "\n\n-tgt ctx (original context):\n"
-        ^ eval_ctx_to_string ~span:(Some span) ctx));
-
-    (* Compute the end expression, that is the expresion corresponding to the
-       end of the function where we call the loop (for now, when calling a loop
-       we never get out) *)
-    let res_fun_end =
-      comp cf_prepare
-        (match_ctx_with_target config span loop_id true fp_bl_corresp
-           fp_input_svalues fixed_ids fp_ctx ctx)
-    in
-    (res_fun_end, fp_bl_corresp)
+     We update the loop fixed point at the same time by reordering the borrows/
+     loans which appear inside it.
+  *)
+  let (res_fun_end, cf_fun_end), fp_bl_corresp =
+    eval_loop_symbolic_synthesize_fun_end config span loop_id ctx fixed_ids
+      fp_ctx fp_input_svalues rg_to_abs
   in
+
   log#ldebug
     (lazy
       "eval_loop_symbolic: matched the fixed-point context with the original \
-       context");
+       context.");
 
   (* Synthesize the loop body *)
-  let (resl_loop_body, cf_loop_body) :
-      (eval_ctx * statement_eval_res) list
-      * (SymbolicAst.expression list option -> eval_result) =
-    (* First, evaluate the loop body starting from the **fixed-point** context *)
-    let ctx_resl, cf_loop = eval_loop_body fp_ctx in
-
-    (* Then, do a special treatment of the break and continue cases.
-       For now, we forbid having breaks in loops (and eliminate breaks
-       in the prepasses) *)
-    let eval_after_loop_iter (ctx, res) =
-      log#ldebug (lazy "eval_loop_symbolic: eval_after_loop_iter");
-      match res with
-      | Return ->
-          (* We replace the [Return] with a [LoopReturn] *)
-          ((ctx, LoopReturn loop_id), fun e -> e)
-      | Panic -> ((ctx, res), fun e -> e)
-      | Break _ ->
-          (* Breaks should have been eliminated in the prepasses *)
-          craise __FILE__ __LINE__ span "Unexpected break"
-      | Continue i ->
-          (* We don't support nested loops for now *)
-          cassert __FILE__ __LINE__ (i = 0) span
-            "Nested loops are not supported yet";
-          log#ldebug
-            (lazy
-              ("eval_loop_symbolic: about to match the fixed-point context \
-                with the context at a continue:\n\
-                - src ctx (fixed-point ctx)"
-              ^ eval_ctx_to_string ~span:(Some span) fp_ctx
-              ^ "\n\n-tgt ctx (ctx at continue):\n"
-              ^ eval_ctx_to_string ~span:(Some span) ctx));
-          match_ctx_with_target config span loop_id false fp_bl_corresp
-            fp_input_svalues fixed_ids fp_ctx ctx
-      | Unit | LoopReturn _ | EndEnterLoop _ | EndContinue _ ->
-          (* For why we can't get [Unit], see the comments inside {!eval_loop_concrete}.
-             For [EndEnterLoop] and [EndContinue]: we don't support nested loops for now.
-          *)
-          craise __FILE__ __LINE__ span "Unreachable"
-    in
-
-    (* Apply and compose *)
-    let ctx_resl, cfl = List.split (List.map eval_after_loop_iter ctx_resl) in
-    let cc (el : SymbolicAst.expression list option) : eval_result =
-      match el with
-      | None -> None
-      | Some el ->
-          let el =
-            List.map
-              (fun (cf, e) -> Option.get (cf (Some e)))
-              (List.combine cfl el)
-          in
-          cf_loop (Some el)
-    in
-
-    (ctx_resl, cc)
+  let resl_loop_body, cf_loop_body =
+    eval_loop_symbolic_synthesize_loop_body config span eval_loop_body loop_id
+      fixed_ids fp_ctx fp_input_svalues fp_bl_corresp
   in
 
   log#ldebug
@@ -242,61 +348,33 @@ let eval_loop_symbolic (config : config) (span : span)
      return nothing.
   *)
   let rg_to_given_back =
-    let compute_abs_given_back_tys (abs : abs) : rty list =
+    let compute_abs_given_back_tys (abs_id : AbstractionId.id) : rty list =
+      let abs = ctx_lookup_abs fp_ctx abs_id in
+      log#ldebug
+        (lazy
+          ("eval_loop_symbolic: compute_abs_given_back_tys:\n- abs:\n"
+         ^ abs_to_string span ctx abs ^ "\n"));
+
       let is_borrow (av : typed_avalue) : bool =
         match av.value with
         | ABorrow _ -> true
         | ALoan _ -> false
         | _ -> craise __FILE__ __LINE__ span "Unreachable"
       in
-      let borrows, loans = List.partition is_borrow abs.avalues in
-
-      let borrows =
-        List.filter_map
-          (fun (av : typed_avalue) ->
-            match av.value with
-            | ABorrow (AMutBorrow (pm, bid, child_av)) ->
-                sanity_check __FILE__ __LINE__ (pm = PNone) span;
-                sanity_check __FILE__ __LINE__ (is_aignored child_av.value) span;
-                Some (bid, child_av.ty)
-            | ABorrow (ASharedBorrow (pm, _)) ->
-                sanity_check __FILE__ __LINE__ (pm = PNone) span;
-                None
-            | _ -> craise __FILE__ __LINE__ span "Unreachable")
-          borrows
-      in
-      let borrows = ref (BorrowId.Map.of_list borrows) in
-
-      let loan_ids =
-        List.filter_map
-          (fun (av : typed_avalue) ->
-            match av.value with
-            | ALoan (AMutLoan (pm, bid, child_av)) ->
-                sanity_check __FILE__ __LINE__ (pm = PNone) span;
-                sanity_check __FILE__ __LINE__ (is_aignored child_av.value) span;
-                Some bid
-            | ALoan (ASharedLoan (pm, _, _, _)) ->
-                sanity_check __FILE__ __LINE__ (pm = PNone) span;
-                None
-            | _ -> craise __FILE__ __LINE__ span "Unreachable")
-          loans
-      in
-
-      (* List the given back types, in the order given by the loans *)
-      let given_back_tys =
-        List.map
-          (fun lid ->
-            let bid =
-              BorrowId.InjSubst.find lid fp_bl_corresp.loan_to_borrow_id_map
-            in
-            let ty = BorrowId.Map.find bid !borrows in
-            borrows := BorrowId.Map.remove bid !borrows;
-            ty)
-          loan_ids
-      in
-      sanity_check __FILE__ __LINE__ (BorrowId.Map.is_empty !borrows) span;
-
-      given_back_tys
+      let borrows, _ = List.partition is_borrow abs.avalues in
+
+      List.filter_map
+        (fun (av : typed_avalue) ->
+          match av.value with
+          | ABorrow (AMutBorrow (pm, _, child_av)) ->
+              sanity_check __FILE__ __LINE__ (pm = PNone) span;
+              sanity_check __FILE__ __LINE__ (is_aignored child_av.value) span;
+              Some child_av.ty
+          | ABorrow (ASharedBorrow (pm, _)) ->
+              sanity_check __FILE__ __LINE__ (pm = PNone) span;
+              None
+          | _ -> craise __FILE__ __LINE__ span "Unreachable")
+        borrows
     in
     RegionGroupId.Map.map compute_abs_given_back_tys rg_to_abs
   in
diff --git a/compiler/InterpreterLoopsFixedPoint.ml b/compiler/InterpreterLoopsFixedPoint.ml
index 79beb761..b68f2a4d 100644
--- a/compiler/InterpreterLoopsFixedPoint.ml
+++ b/compiler/InterpreterLoopsFixedPoint.ml
@@ -393,7 +393,7 @@ let prepare_ashared_loans_no_synth (span : Meta.span) (loop_id : LoopId.id)
 
 let compute_loop_entry_fixed_point (config : config) (span : Meta.span)
     (loop_id : LoopId.id) (eval_loop_body : stl_cm_fun) (ctx0 : eval_ctx) :
-    eval_ctx * ids_sets * abs RegionGroupId.Map.t =
+    eval_ctx * ids_sets * AbstractionId.id RegionGroupId.Map.t =
   (* Introduce "reborrows" for the shared values in the abstractions, so that
      the shared values in the fixed abstractions never get modified (technically,
      they are immutable, but in practice we can introduce more shared loans, or
@@ -702,7 +702,9 @@ let compute_loop_entry_fixed_point (config : config) (span : Meta.span)
     let all_abs_ids =
       List.filter_map
         (function EAbs abs -> Some abs.abs_id | _ -> None)
-        !fp.env
+        (* TODO: we may want to use a different order, for instance the order
+           in which the regions were ended. *)
+        (List.rev !fp.env)
     in
     let _ =
       RegionGroupId.Map.iter
@@ -764,8 +766,7 @@ let compute_loop_entry_fixed_point (config : config) (span : Meta.span)
                     craise __FILE__ __LINE__ span "Unexpected")
                 ids;
               (* Register the mapping *)
-              let abs = ctx_lookup_abs !fp !id0 in
-              rg_to_abs := RegionGroupId.Map.add_strict rg_id abs !rg_to_abs)
+              rg_to_abs := RegionGroupId.Map.add_strict rg_id !id0 !rg_to_abs)
         !fp_ended_aids
     in
     let rg_to_abs = !rg_to_abs in
diff --git a/compiler/InterpreterLoopsFixedPoint.mli b/compiler/InterpreterLoopsFixedPoint.mli
index 59d42812..8db7b973 100644
--- a/compiler/InterpreterLoopsFixedPoint.mli
+++ b/compiler/InterpreterLoopsFixedPoint.mli
@@ -81,10 +81,11 @@ val compute_loop_entry_fixed_point :
   Meta.span ->
   loop_id ->
   (* This function is the function to evaluate the loop body (eval_statement applied
-     to the proper arguments) *)
+     to the proper arguments). We need to take it as input because [compute_loop_entry_fixed_point]
+     is mutually recursive with [eval_statement], but doesn't live in the same module. *)
   Cps.stl_cm_fun ->
   eval_ctx ->
-  eval_ctx * ids_sets * abs SymbolicAst.region_group_id_map
+  eval_ctx * ids_sets * AbstractionId.id SymbolicAst.region_group_id_map
 
 (** For the abstractions in the fixed point, compute the correspondance between
     the borrows ids and the loans ids, if we want to introduce equivalent