Fix another bug

7 files changed, 127 insertions, 39 deletions

diff --git a/compiler/Contexts.ml b/compiler/Contexts.ml
index 3d09b58a..f2fa36c5 100644
--- a/compiler/Contexts.ml
+++ b/compiler/Contexts.ml
@@ -491,6 +491,12 @@ let ctx_subst_abs (ctx : eval_ctx) (abs_id : V.AbstractionId.id) (nabs : V.abs)
   let env, abs_opt = env_subst_abs ctx.env abs_id nabs in
   ({ ctx with env }, abs_opt)
 
+let ctx_set_abs_can_end (ctx : eval_ctx) (abs_id : V.AbstractionId.id)
+    (can_end : bool) : eval_ctx =
+  let abs = ctx_lookup_abs ctx abs_id in
+  let abs = { abs with can_end } in
+  fst (ctx_subst_abs ctx abs_id abs)
+
 let ctx_type_decl_is_rec (ctx : eval_ctx) (id : TypeDeclId.id) : bool =
   let decl_group = TypeDeclId.Map.find id ctx.type_context.type_decls_groups in
   match decl_group with Rec _ -> true | NonRec _ -> false
diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index d5032756..85da809e 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -96,7 +96,7 @@ let initialize_symbolic_context_for_fun (type_context : C.type_context)
     in
     (ctx, avalues)
   in
-  let region_can_end _ = true in
+  let region_can_end _ = false in
   let ctx =
     create_push_abstractions_from_abs_region_groups
       (fun rg_id -> V.SynthInput rg_id)
@@ -134,6 +134,17 @@ let evaluate_function_symbolic_synthesize_backward_from_return
     (config : C.config) (fdef : A.fun_decl) (inst_sg : A.inst_fun_sig)
     (back_id : T.RegionGroupId.id) (loop_id : V.LoopId.id option)
     (inside_loop : bool) (ctx : C.eval_ctx) : SA.expression option =
+  log#ldebug
+    (lazy
+      ("evaluate_function_symbolic_synthesize_backward_from_return:"
+     ^ "\n- back_id: "
+      ^ T.RegionGroupId.to_string back_id
+      ^ "\n- loop_id: "
+      ^ Print.option_to_string V.LoopId.to_string loop_id
+      ^ "\n- inside_loop: "
+      ^ Print.bool_to_string inside_loop
+      ^ "\n- ctx:\n"
+      ^ Print.Contexts.eval_ctx_to_string_gen true true ctx));
   let entered_loop = Option.is_some loop_id in
   (* We need to instantiate the function signature - to retrieve
    * the return type. Note that it is important to re-generate
@@ -200,7 +211,7 @@ let evaluate_function_symbolic_synthesize_backward_from_return
       else ctx
     in
 
-    (* Set the proper loop abstractions as endable *)
+    (* Set the proper abstractions as endable *)
     let ctx =
       let visit_loop_abs =
         object
@@ -208,10 +219,28 @@ let evaluate_function_symbolic_synthesize_backward_from_return
 
           method! visit_abs _ abs =
             match abs.kind with
-            | V.Loop (_, rg_id', _) ->
+            | V.Loop (loop_id', rg_id', V.LoopSynthInput) ->
+                (* We only allow to end the loop synth input abs for the region
+                   group [rg_id] *)
+                assert (
+                  if Option.is_some loop_id then loop_id = Some loop_id'
+                  else true);
+                (* Loop abstractions *)
                 let rg_id' = Option.get rg_id' in
-                if rg_id' = back_id then { abs with can_end = true } else abs
-            | _ -> abs
+                if rg_id' = back_id && inside_loop then
+                  { abs with can_end = true }
+                else abs
+            | V.Loop (loop_id', _, V.LoopCall) ->
+                (* We can end all the loop call abstractions *)
+                assert (loop_id = Some loop_id');
+                { abs with can_end = true }
+            | V.SynthInput rg_id' ->
+                if rg_id' = back_id && not inside_loop then
+                  { abs with can_end = true }
+                else abs
+            | _ ->
+                (* Other abstractions *)
+                abs
         end
       in
       visit_loop_abs#visit_eval_ctx () ctx
@@ -247,6 +276,11 @@ let evaluate_function_symbolic_synthesize_backward_from_return
         let abs = Option.get (C.ctx_find_abs ctx pred) in
         abs.abs_id
     in
+    log#ldebug
+      (lazy
+        ("evaluate_function_symbolic_synthesize_backward_from_return: ending \
+          input abstraction: "
+        ^ V.AbstractionId.to_string current_abs_id));
     let target_abs_ids = List.append parent_input_abs_ids [ current_abs_id ] in
     let cf_end_target_abs cf =
       List.fold_left
@@ -296,7 +330,7 @@ let evaluate_function_symbolic (synthesize : bool)
         ^ Cps.show_statement_eval_res res));
 
     match res with
-    | Return | LoopReturn ->
+    | Return | LoopReturn _ ->
         if synthesize then
           (* We have to "play different endings":
            * - one execution for the forward function
@@ -326,16 +360,17 @@ let evaluate_function_symbolic (synthesize : bool)
              abstractions to consume the return value, then end all the
              abstractions up to the one in which we are interested.
           *)
-          let inside_loop =
+          let loop_id =
             match res with
-            | Return -> false
-            | LoopReturn -> true
+            | Return -> None
+            | LoopReturn loop_id -> Some loop_id
             | _ -> raise (Failure "Unreachable")
           in
+          let inside_loop = Option.is_some loop_id in
           let finish_back_eval back_id =
             Option.get
               (evaluate_function_symbolic_synthesize_backward_from_return config
-                 fdef inst_sg back_id None inside_loop ctx)
+                 fdef inst_sg back_id loop_id inside_loop ctx)
           in
           let back_el =
             T.RegionGroupId.mapi
diff --git a/compiler/InterpreterBorrows.ml b/compiler/InterpreterBorrows.ml
index 208918af..0b7b4be5 100644
--- a/compiler/InterpreterBorrows.ml
+++ b/compiler/InterpreterBorrows.ml
@@ -1005,7 +1005,13 @@ and end_abstraction_aux (config : C.config) (chain : borrow_or_abs_ids)
   let abs = C.ctx_lookup_abs ctx abs_id in
 
   (* Check that we can end the abstraction *)
-  assert abs.can_end;
+  if abs.can_end then ()
+  else
+    raise
+      (Failure
+         ("Can't end abstraction "
+         ^ V.AbstractionId.to_string abs.abs_id
+         ^ " as it is set as non-endable"));
 
   (* End the parent abstractions first *)
   let cc = end_abstractions_aux config chain abs.parents in
diff --git a/compiler/InterpreterLoops.ml b/compiler/InterpreterLoops.ml
index 7a72f8e3..96f6e7b5 100644
--- a/compiler/InterpreterLoops.ml
+++ b/compiler/InterpreterLoops.ml
@@ -2323,8 +2323,11 @@ let compute_loop_entry_fixed_point (config : C.config) (loop_id : V.LoopId.id)
                 let abs_id =
                   V.AbstractionId.of_int (T.RegionGroupId.to_int rg_id)
                 in
-                let abs = C.ctx_lookup_abs ctx abs_id in
-                assert (abs.kind = V.SynthInput rg_id);
+                (* By default, the [SynthInput] abs can't end *)
+                let ctx = C.ctx_set_abs_can_end ctx abs_id true in
+                assert (
+                  let abs = C.ctx_lookup_abs ctx abs_id in
+                  abs.kind = V.SynthInput rg_id);
                 (* End this abstraction *)
                 let ctx =
                   InterpreterBorrows.end_abstraction_no_synth config abs_id ctx
@@ -3171,7 +3174,10 @@ let eval_loop_symbolic (config : C.config) (eval_loop_body : st_cm_fun) :
   let cf_loop : st_m_fun =
    fun res ctx ->
     match res with
-    | Return | Panic -> cf res ctx
+    | Return ->
+        (* We replace the [Return] with a [LoopReturn] *)
+        cf (LoopReturn loop_id) ctx
+    | Panic -> cf res ctx
     | Break i ->
         (* Break out of the loop by calling the continuation *)
         let res = if i = 0 then Unit else Break (i - 1) in
diff --git a/compiler/Print.ml b/compiler/Print.ml
index e2bc3777..d040680f 100644
--- a/compiler/Print.ml
+++ b/compiler/Print.ml
@@ -10,6 +10,8 @@ module Expressions = Charon.PrintExpressions
 let list_to_string (to_string : 'a -> string) (ls : 'a list) : string =
   "[" ^ String.concat "; " (List.map to_string ls) ^ "]"
 
+let bool_to_string (b : bool) : string = if b then "true" else "false"
+
 (** Pretty-printing for values *)
 module Values = struct
   type value_formatter = {
@@ -292,16 +294,41 @@ module Values = struct
         ^ abstract_shared_borrows_to_string fmt sb
         ^ ")"
 
-  let abs_to_string (fmt : value_formatter) (indent : string)
+  let loop_abs_kind_to_string (kind : V.loop_abs_kind) : string =
+    match kind with
+    | LoopSynthInput -> "LoopSynthInput"
+    | LoopCall -> "LoopCall"
+
+  let abs_kind_to_string (kind : V.abs_kind) : string =
+    match kind with
+    | V.FunCall (fid, rg_id) ->
+        "FunCall(fid:" ^ V.FunCallId.to_string fid ^ ", rg_id:"
+        ^ T.RegionGroupId.to_string rg_id
+        ^ ")"
+    | SynthInput rg_id ->
+        "SynthInput(rg_id:" ^ T.RegionGroupId.to_string rg_id ^ ")"
+    | SynthRet rg_id ->
+        "SynthRet(rg_id:" ^ T.RegionGroupId.to_string rg_id ^ ")"
+    | Loop (lp_id, rg_id, abs_kind) ->
+        "Loop(loop_id:" ^ V.LoopId.to_string lp_id ^ ", rg_id:"
+        ^ option_to_string T.RegionGroupId.to_string rg_id
+        ^ ", loop abs kind: "
+        ^ loop_abs_kind_to_string abs_kind
+        ^ ")"
+
+  let abs_to_string (fmt : value_formatter) (verbose : bool) (indent : string)
       (indent_incr : string) (abs : V.abs) : string =
     let indent2 = indent ^ indent_incr in
     let avs =
       List.map (fun av -> indent2 ^ typed_avalue_to_string fmt av) abs.avalues
     in
     let avs = String.concat ",\n" avs in
+    let kind =
+      if verbose then "[kind:" ^ abs_kind_to_string abs.kind ^ "]" else ""
+    in
     indent ^ "abs@"
     ^ V.AbstractionId.to_string abs.abs_id
-    ^ "{parents="
+    ^ kind ^ "{parents="
     ^ V.AbstractionId.Set.to_string None abs.parents
     ^ "}" ^ "{regions="
     ^ T.RegionId.Set.to_string None abs.regions
@@ -325,20 +352,21 @@ module Contexts = struct
     | VarBinder b -> var_binder_to_string b
     | DummyBinder bid -> dummy_var_id_to_string bid
 
-  let env_elem_to_string (fmt : PV.value_formatter) (indent : string)
-      (indent_incr : string) (ev : C.env_elem) : string =
+  let env_elem_to_string (fmt : PV.value_formatter) (verbose : bool)
+      (indent : string) (indent_incr : string) (ev : C.env_elem) : string =
     match ev with
     | Var (var, tv) ->
         let bv = binder_to_string var in
         indent ^ bv ^ " -> " ^ PV.typed_value_to_string fmt tv ^ " ;"
-    | Abs abs -> PV.abs_to_string fmt indent indent_incr abs
+    | Abs abs -> PV.abs_to_string fmt verbose indent indent_incr abs
     | Frame -> raise (Failure "Can't print a Frame element")
 
-  let opt_env_elem_to_string (fmt : PV.value_formatter) (indent : string)
-      (indent_incr : string) (ev : C.env_elem option) : string =
+  let opt_env_elem_to_string (fmt : PV.value_formatter) (verbose : bool)
+      (indent : string) (indent_incr : string) (ev : C.env_elem option) : string
+      =
     match ev with
     | None -> indent ^ "..."
-    | Some ev -> env_elem_to_string fmt indent indent_incr ev
+    | Some ev -> env_elem_to_string fmt verbose indent indent_incr ev
 
   (** Filters "dummy" bindings from an environment, to gain space and clarity/
       See [env_to_string]. *)
@@ -376,14 +404,16 @@ module Contexts = struct
       allows to filter them when printing, replacing groups of such bindings with
       "..." to gain space and clarity.
    *)
-  let env_to_string (filter : bool) (fmt : PV.value_formatter) (env : C.env) :
-      string =
+  let env_to_string (filter : bool) (fmt : PV.value_formatter) (verbose : bool)
+      (env : C.env) : string =
     let env =
       if filter then filter_env env else List.map (fun ev -> Some ev) env
     in
     "{\n"
     ^ String.concat "\n"
-        (List.map (fun ev -> opt_env_elem_to_string fmt "  " "  " ev) env)
+        (List.map
+           (fun ev -> opt_env_elem_to_string fmt verbose "  " "  " ev)
+           env)
     ^ "\n}"
 
   type ctx_formatter = PV.value_formatter
@@ -486,8 +516,8 @@ module Contexts = struct
     let frames = split_aux [] [] env in
     frames
 
-  let fmt_eval_ctx_to_string_gen (fmt : ctx_formatter) (filter : bool)
-      (ctx : C.eval_ctx) : string =
+  let fmt_eval_ctx_to_string_gen (fmt : ctx_formatter) (verbose : bool)
+      (filter : bool) (ctx : C.eval_ctx) : string =
     let ended_regions = T.RegionId.Set.to_string None ctx.ended_regions in
     let frames = split_env_according_to_frames ctx.env in
     let num_frames = List.length frames in
@@ -509,21 +539,23 @@ module Contexts = struct
           ^ string_of_int !num_bindings
           ^ "\n- dummy bindings: " ^ string_of_int !num_dummies
           ^ "\n- abstractions: " ^ string_of_int !num_abs ^ "\n"
-          ^ env_to_string filter fmt f ^ "\n")
+          ^ env_to_string filter fmt verbose f
+          ^ "\n")
         frames
     in
     "# Ended regions: " ^ ended_regions ^ "\n" ^ "# " ^ string_of_int num_frames
     ^ " frame(s)\n" ^ String.concat "" frames
 
-  let eval_ctx_to_string_gen (filter : bool) (ctx : C.eval_ctx) : string =
+  let eval_ctx_to_string_gen (verbose : bool) (filter : bool) (ctx : C.eval_ctx)
+      : string =
     let fmt = eval_ctx_to_ctx_formatter ctx in
-    fmt_eval_ctx_to_string_gen fmt filter ctx
+    fmt_eval_ctx_to_string_gen fmt verbose filter ctx
 
   let eval_ctx_to_string (ctx : C.eval_ctx) : string =
-    eval_ctx_to_string_gen true ctx
+    eval_ctx_to_string_gen false true ctx
 
   let eval_ctx_to_string_no_filter (ctx : C.eval_ctx) : string =
-    eval_ctx_to_string_gen false ctx
+    eval_ctx_to_string_gen false false ctx
 end
 
 module PC = Contexts (* local module *)
@@ -593,10 +625,10 @@ module EvalCtxLlbcAst = struct
   let env_elem_to_string (ctx : C.eval_ctx) (indent : string)
       (indent_incr : string) (ev : C.env_elem) : string =
     let fmt = PC.eval_ctx_to_ctx_formatter ctx in
-    PC.env_elem_to_string fmt indent indent_incr ev
+    PC.env_elem_to_string fmt false indent indent_incr ev
 
   let abs_to_string (ctx : C.eval_ctx) (indent : string) (indent_incr : string)
       (abs : V.abs) : string =
     let fmt = PC.eval_ctx_to_ctx_formatter ctx in
-    PV.abs_to_string fmt indent indent_incr abs
+    PV.abs_to_string fmt false indent indent_incr abs
 end
diff --git a/compiler/PrintPure.ml b/compiler/PrintPure.ml
index c83858b3..b19e0be6 100644
--- a/compiler/PrintPure.ml
+++ b/compiler/PrintPure.ml
@@ -619,14 +619,16 @@ and switch_to_string (fmt : ast_formatter) (indent : string)
 and loop_to_string (fmt : ast_formatter) (indent : string)
     (indent_incr : string) (loop : loop) : string =
   let indent1 = indent ^ indent_incr in
+  let indent2 = indent1 ^ indent_incr in
   let fun_end =
-    texpression_to_string fmt false indent1 indent_incr loop.fun_end
+    texpression_to_string fmt false indent2 indent_incr loop.fun_end
   in
   let loop_body =
-    texpression_to_string fmt false indent1 indent_incr loop.loop_body
+    texpression_to_string fmt false indent2 indent_incr loop.loop_body
   in
-  "loop {\n" ^ indent1 ^ "fun_end: " ^ fun_end ^ "\n" ^ indent1 ^ "loop_body:"
-  ^ loop_body ^ "\n" ^ indent ^ "}"
+  "loop {\n" ^ indent1 ^ "fun_end: {\n" ^ indent2 ^ fun_end ^ "\n" ^ indent1
+  ^ "}\n" ^ indent1 ^ "loop_body: {\n" ^ indent2 ^ loop_body ^ "\n" ^ indent1
+  ^ "}\n" ^ indent ^ "}"
 
 and meta_to_string (fmt : ast_formatter) (meta : meta) : string =
   let meta =
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index f06e6542..3488b334 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -284,9 +284,10 @@ let fun_decl_to_string (ctx : bs_ctx) (def : Pure.fun_decl) : string =
 let abs_to_string (ctx : bs_ctx) (abs : V.abs) : string =
   let fmt = bs_ctx_to_ast_formatter ctx in
   let fmt = Print.Contexts.ast_to_value_formatter fmt in
+  let verbose = false in
   let indent = "" in
   let indent_incr = "  " in
-  Print.Values.abs_to_string fmt indent indent_incr abs
+  Print.Values.abs_to_string fmt verbose indent indent_incr abs
 
 let get_instantiated_fun_sig (fun_id : A.fun_id)
     (back_id : T.RegionGroupId.id option) (tys : ty list) (ctx : bs_ctx) :