Factor out the symbolic execution for the forward/backward translations

5 files changed, 77 insertions, 92 deletions

diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index f776c7ff..dc203105 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -210,16 +210,10 @@ let evaluate_function_symbolic_synthesize_backward_from_return
  *)
 let evaluate_function_symbolic (synthesize : bool)
     (type_context : C.type_context) (fun_context : C.fun_context)
-    (global_context : C.global_context) (fdef : A.fun_decl)
-    (back_id : T.RegionGroupId.id option) :
+    (global_context : C.global_context) (fdef : A.fun_decl) :
     V.symbolic_value list * SA.expression option =
   (* Debug *)
-  let name_to_string () =
-    Print.fun_name_to_string fdef.A.name
-    ^ " ("
-    ^ Print.option_to_string T.RegionGroupId.to_string back_id
-    ^ ")"
-  in
+  let name_to_string () = Print.fun_name_to_string fdef.A.name in
   log#ldebug (lazy ("evaluate_function_symbolic: " ^ name_to_string ()));
 
   (* Create the evaluation context *)
@@ -234,34 +228,46 @@ let evaluate_function_symbolic (synthesize : bool)
     match res with
     | Return ->
         if synthesize then
+          (* We have to "play different endings":
+           * - one execution for the forward function
+           * - one execution per backward function
+           * We then group everything together.
+           *)
           (* There are two cases:
            * - if this is a forward translation, we retrieve the returned value.
            * - if this is a backward translation, we introduce "return"
            *   abstractions to consume the return value, then end all the
            *   abstractions up to the one in which we are interested.
            *)
-          match back_id with
-          | None ->
-              (* Forward translation *)
-              (* Pop the frame and retrieve the returned value at the same time*)
-              let cf_pop = pop_frame config in
-              (* Generate the Return node *)
-              let cf_return ret_value : m_fun =
-               fun _ -> Some (SA.Return (Some ret_value))
-              in
-              (* Apply *)
-              cf_pop cf_return ctx
-          | Some back_id ->
-              (* Backward translation *)
-              let e =
-                evaluate_function_symbolic_synthesize_backward_from_return
-                  config fdef inst_sg back_id ctx
-              in
-              (* We insert a delimiter to indicate the point where we switch
-               * from the part which is common to all the functions (forwards
-               * and backwards) and the part specific to this backward function.
-               *)
-              S.synthesize_forward_end e
+          (* Forward translation: retrieve the returned value *)
+          let fwd_e =
+            (* Pop the frame and retrieve the returned value at the same time*)
+            let cf_pop = pop_frame config in
+            (* Generate the Return node *)
+            let cf_return ret_value : m_fun =
+             fun _ -> Some (SA.Return (Some ret_value))
+            in
+            (* Apply *)
+            cf_pop cf_return ctx
+          in
+          let fwd_e = Option.get fwd_e in
+          (* Backward translation: introduce "return"
+             abstractions to consume the return value, then end all the
+             abstractions up to the one in which we are interested.
+          *)
+          let finish_back_eval back_id =
+            Option.get
+              (evaluate_function_symbolic_synthesize_backward_from_return config
+                 fdef inst_sg back_id ctx)
+          in
+          let back_el =
+            T.RegionGroupId.mapi
+              (fun gid _ -> (gid, finish_back_eval gid))
+              fdef.signature.regions_hierarchy
+          in
+          let back_el = T.RegionGroupId.Map.of_list back_el in
+          (* Put everything together *)
+          S.synthesize_forward_end fwd_e back_el
         else None
     | Panic ->
         (* Note that as we explore all the execution branches, one of
@@ -355,18 +361,10 @@ module Test = struct
       (lazy ("test_function_symbolic: " ^ Print.fun_name_to_string fdef.A.name));
 
     (* Evaluate *)
-    let evaluate =
+    let _ =
       evaluate_function_symbolic synthesize type_context fun_context
         global_context fdef
     in
-    (* Execute the forward function *)
-    let _ = evaluate None in
-    (* Execute the backward functions *)
-    let _ =
-      T.RegionGroupId.mapi
-        (fun gid _ -> evaluate (Some gid))
-        fdef.signature.regions_hierarchy
-    in
 
     ()
 
diff --git a/compiler/SymbolicAst.ml b/compiler/SymbolicAst.ml
index 9d9adf4f..235e33e4 100644
--- a/compiler/SymbolicAst.ml
+++ b/compiler/SymbolicAst.ml
@@ -77,13 +77,16 @@ type expression =
           We use it to compute meaningful names for the variables we introduce,
           to prettify the generated code.
        *)
-  | ForwardEnd of expression
+  | ForwardEnd of expression * expression T.RegionGroupId.Map.t
       (** We use this delimiter to indicate at which point we switch to the
-          generation of code specific to the backward function(s).
+          generation of code specific to the backward function(s). This allows
+          us in particular to factor the work out: we don't need to replay the
+          symbolic execution up to this point, and can reuse it for the forward
+          function and all the backward functions.
 
-          TODO: use this to factorize the generation of the forward and backward
-          functions (today we replay the *whole* symbolic execution once per
-          generated function).
+          The first expression gives the end of the translation for the forward
+          function, the map from region group ids to expressions gives the end
+          of the translation for the backward functions.
        *)
   | Meta of meta * expression  (** Meta information *)
 
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 6b44a69c..fafcb348 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -1113,13 +1113,16 @@ let rec translate_expression (e : S.expression) (ctx : bs_ctx) : texpression =
   | Assertion (v, e) -> translate_assertion v e ctx
   | Expansion (p, sv, exp) -> translate_expansion p sv exp ctx
   | Meta (meta, e) -> translate_meta meta e ctx
-  | ForwardEnd e ->
+  | ForwardEnd (e, back_e) ->
       (* Update the current state with the additional state received by the backward
-         function, if needs be *)
-      let ctx =
+         function, if needs be, and lookup the proper expression *)
+      let ctx, e =
         match ctx.bid with
-        | None -> ctx
-        | Some _ -> { ctx with state_var = ctx.back_state_var }
+        | None -> (ctx, e)
+        | Some bid ->
+            let ctx = { ctx with state_var = ctx.back_state_var } in
+            let e = T.RegionGroupId.Map.find bid back_e in
+            (ctx, e)
       in
       translate_expression e ctx
 
diff --git a/compiler/SynthesizeSymbolic.ml b/compiler/SynthesizeSymbolic.ml
index 8d4dac82..5bcbd482 100644
--- a/compiler/SynthesizeSymbolic.ml
+++ b/compiler/SynthesizeSymbolic.ml
@@ -152,5 +152,6 @@ let synthesize_assignment (lplace : mplace) (rvalue : V.typed_value)
 let synthesize_assertion (v : V.typed_value) (e : expression option) =
   Option.map (fun e -> Assertion (v, e)) e
 
-let synthesize_forward_end (e : expression option) =
-  Option.map (fun e -> ForwardEnd e) e
+let synthesize_forward_end (e : expression)
+    (el : expression T.RegionGroupId.Map.t) =
+  Some (ForwardEnd (e, el))
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index e943c78a..c4185f41 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -21,7 +21,7 @@ type symbolic_fun_translation = V.symbolic_value list * SA.expression
     for the forward function and the backward functions.
 *)
 let translate_function_to_symbolics (trans_ctx : trans_ctx) (fdef : A.fun_decl)
-    : (symbolic_fun_translation * symbolic_fun_translation list) option =
+    : symbolic_fun_translation option =
   (* Debug *)
   log#ldebug
     (lazy
@@ -36,24 +36,11 @@ let translate_function_to_symbolics (trans_ctx : trans_ctx) (fdef : A.fun_decl)
   | Some _ ->
       (* Evaluate *)
       let synthesize = true in
-      let evaluate gid =
-        let inputs, symb =
-          evaluate_function_symbolic synthesize type_context fun_context
-            global_context fdef gid
-        in
-        (inputs, Option.get symb)
-      in
-      (* Execute the forward function *)
-      let forward = evaluate None in
-      (* Execute the backward functions *)
-      let backwards =
-        T.RegionGroupId.mapi
-          (fun gid _ -> evaluate (Some gid))
-          fdef.signature.regions_hierarchy
+      let inputs, symb =
+        evaluate_function_symbolic synthesize type_context fun_context
+          global_context fdef
       in
-
-      (* Return *)
-      Some (forward, backwards)
+      Some (inputs, Option.get symb)
 
 (** Translate a function, by generating its forward and backward translations.
 
@@ -75,11 +62,6 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
 
   (* Compute the symbolic ASTs, if the function is transparent *)
   let symbolic_trans = translate_function_to_symbolics trans_ctx fdef in
-  let symbolic_forward, symbolic_backwards =
-    match symbolic_trans with
-    | None -> (None, None)
-    | Some (fwd, backs) -> (Some fwd, Some backs)
-  in
 
   (* Convert the symbolic ASTs to pure ASTs: *)
 
@@ -133,11 +115,13 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
     }
   in
 
-  (* We need to initialize the input/output variables *)
-  let add_forward_inputs input_svs ctx =
-    match fdef.body with
-    | None -> ctx
-    | Some body ->
+  (* Add the forward inputs (the initialized input variables for the forward
+     function)
+  *)
+  let ctx =
+    match (fdef.body, symbolic_trans) with
+    | None, None -> ctx
+    | Some body, Some (input_svs, _) ->
         let forward_input_vars = LlbcAstUtils.fun_body_get_input_vars body in
         let forward_input_varnames =
           List.map (fun (v : A.var) -> v.name) forward_input_vars
@@ -147,16 +131,14 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
           SymbolicToPure.fresh_named_vars_for_symbolic_values input_svs ctx
         in
         { ctx with forward_inputs }
+    | _ -> raise (Failure "Unreachable")
   in
 
   (* Translate the forward function *)
   let pure_forward =
-    match symbolic_forward with
+    match symbolic_trans with
     | None -> SymbolicToPure.translate_fun_decl ctx None
-    | Some (fwd_svs, fwd_ast) ->
-        SymbolicToPure.translate_fun_decl
-          (add_forward_inputs fwd_svs ctx)
-          (Some fwd_ast)
+    | Some (_, ast) -> SymbolicToPure.translate_fun_decl ctx (Some ast)
   in
 
   (* Translate the backward functions *)
@@ -167,7 +149,7 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
     assert (rg.parents = []);
     let back_id = rg.id in
 
-    match symbolic_backwards with
+    match symbolic_trans with
     | None ->
         (* Initialize the context - note that the ret_ty is not really
          * useful as we don't translate a body *)
@@ -178,12 +160,10 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
 
         (* Translate *)
         SymbolicToPure.translate_fun_decl ctx None
-    | Some symbolic_backwards ->
-        let input_svs, symbolic =
-          T.RegionGroupId.nth symbolic_backwards back_id
-        in
-        let ctx = add_forward_inputs input_svs ctx in
-        (* TODO: the computation of the backward inputs is a bit awckward... *)
+    | Some (_, symbolic) ->
+        (* Finish initializing the context by adding the additional input
+           variables required by the backward function.
+        *)
         let backward_sg =
           RegularFunIdMap.find (A.Regular def_id, Some back_id) fun_sigs
         in
@@ -206,7 +186,7 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
         in
         (* The outputs for the backward functions, however, come from borrows
          * present in the input values of the rust function: for those we reuse
-         * the names of the  input values. *)
+         * the names of the input values. *)
         let backward_outputs =
           List.combine backward_sg.output_names backward_sg.sg.doutputs
         in