Add an option to eagerly end abstractions if a function has return type

unit and the abstractions don't contain loans

4 files changed, 61 insertions, 2 deletions

diff --git a/src/Contexts.ml b/src/Contexts.ml
index 1fb8cac1..a4551420 100644
--- a/src/Contexts.ml
+++ b/src/Contexts.ml
@@ -179,6 +179,15 @@ type config = {
           meaning at intermediate steps of the assignment where the invariants
           might actually be broken.
        *)
+  return_unit_end_abs_with_no_loans : bool;
+      (** If a function doesn't return any borrows, we can immediately call
+          its backward functions. If this option is on, whenever we call a
+          function *and* this function returns unit, we immediately end all the
+          abstractions which are introduced and don't contain loans. This can be
+          useful to make the code cleaner (the backward function is introduced
+          where the function call happened) and make sure all forward functions
+          with no return value are followed by a backward function.
+       *)
 }
 [@@deriving show]
 
@@ -186,6 +195,7 @@ type partial_config = {
   check_invariants : bool;
   greedy_expand_symbolics_with_borrows : bool;
   allow_bottom_below_borrow : bool;
+  return_unit_end_abs_with_no_loans : bool;
 }
 (** See [config] *)
 
@@ -197,6 +207,7 @@ let config_of_partial (mode : interpreter_mode) (config : partial_config) :
     greedy_expand_symbolics_with_borrows =
       config.greedy_expand_symbolics_with_borrows;
     allow_bottom_below_borrow = config.allow_bottom_below_borrow;
+    return_unit_end_abs_with_no_loans = config.return_unit_end_abs_with_no_loans;
   }
 
 type type_context = {
diff --git a/src/InterpreterStatements.ml b/src/InterpreterStatements.ml
index e8e9c6c4..1083d643 100644
--- a/src/InterpreterStatements.ml
+++ b/src/InterpreterStatements.ml
@@ -1120,6 +1120,7 @@ and eval_function_call_symbolic_from_inst_sig (config : C.config)
   let cc = eval_operands config args in
 
   (* Generate the abstractions and insert them in the context *)
+  let abs_ids = List.map (fun rg -> rg.T.id) inst_sg.regions_hierarchy in
   let cf_call cf (args : V.typed_value list) : m_fun =
    fun ctx ->
     let args_with_rtypes = List.combine args inst_sg.A.inputs in
@@ -1168,7 +1169,6 @@ and eval_function_call_symbolic_from_inst_sig (config : C.config)
     let expr = cf ctx in
 
     (* Synthesize the symbolic AST *)
-    let abs_ids = List.map (fun rg -> rg.T.id) inst_sg.regions_hierarchy in
     S.synthesize_regular_function_call fid call_id abs_ids type_args args
       args_places ret_spc dest_place expr
   in
@@ -1177,6 +1177,53 @@ and eval_function_call_symbolic_from_inst_sig (config : C.config)
   (* Move the return value to its destination *)
   let cc = comp cc (assign_to_place config ret_value dest) in
 
+  (* End the abstractions which don't contain loans and don't have parent
+   * abstractions.
+   * We do the general, nested borrows case here: we end abstractions, then
+   * retry (because then we might end their children abstractions)
+   *)
+  let abs_ids = ref abs_ids in
+  let rec end_abs_with_no_loans cf : m_fun =
+   fun ctx ->
+    (* Find the abstractions which don't contain loans *)
+    let no_loans_abs, with_loans_abs =
+      List.partition
+        (fun abs_id ->
+          (* Lookup the abstraction *)
+          let abs = C.ctx_lookup_abs ctx abs_id in
+          (* Check if it has parents *)
+          V.AbstractionId.Set.is_empty abs.parents
+          (* Check if it contains non-ignored loans *)
+          && Option.is_none
+               (InterpreterBorrowsCore
+                .get_first_non_ignored_aloan_in_abstraction abs))
+        !abs_ids
+    in
+    (* Check if there are abstractions to end *)
+    if no_loans_abs <> [] then (
+      (* Update the reference to the list of asbtraction ids, for the recursive calls *)
+      abs_ids := with_loans_abs;
+      (* End the abstractions which can be ended *)
+      let no_loans_abs = V.AbstractionId.Set.of_list no_loans_abs in
+      let cc = InterpreterBorrows.end_abstractions config [] no_loans_abs in
+      (* Recursive call *)
+      let cc = comp cc end_abs_with_no_loans in
+      (* Continue *)
+      cc cf ctx)
+    else (* No abstractions to end: continue *)
+      cf ctx
+  in
+  (* Try to end the abstractions with no loans if:
+   * - the option is enabled
+   * - the function returns unit
+   * (see the documentation of [config] for more information)
+   *)
+  let cc =
+    if config.return_unit_end_abs_with_no_loans && ty_is_unit inst_sg.output
+    then comp cc end_abs_with_no_loans
+    else cc
+  in
+
   (* Continue - note that we do as if the function call has been successful,
    * by giving [Unit] to the continuation, because we place us in the case
    * where we haven't panicked. Of course, the translation needs to take the
diff --git a/src/Values.ml b/src/Values.ml
index 4799f3b9..4e45db03 100644
--- a/src/Values.ml
+++ b/src/Values.ml
@@ -545,7 +545,7 @@ and aloan_content =
        *)
   | AIgnoredSharedLoan of typed_avalue
       (** An ignored shared loan.
-      
+            
           Example:
           ========
           ```
diff --git a/src/main.ml b/src/main.ml
index 3fb162f3..6b1083f5 100644
--- a/src/main.ml
+++ b/src/main.ml
@@ -159,6 +159,7 @@ let () =
           C.check_invariants = not !no_check_inv;
           greedy_expand_symbolics_with_borrows = true;
           allow_bottom_below_borrow = true;
+          return_unit_end_abs_with_no_loans = true;
         }
       in