Update Pure.fun_sig_info

7 files changed, 129 insertions, 68 deletions

diff --git a/compiler/Config.ml b/compiler/Config.ml
index b09544ba..9cd1ebc2 100644
--- a/compiler/Config.ml
+++ b/compiler/Config.ml
@@ -92,6 +92,10 @@ let loop_fixed_point_max_num_iters = 2
 
 (** {1 Translation} *)
 
+(** If true, do not define separate forward/backward functions, but make the
+    forward functions return the backward function. *)
+let return_back_funs = ref false
+
 (** Forbids using field projectors for structures.
 
     If we don't use field projectors, whenever we symbolically expand a structure
diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index 20cdb20b..93fcf416 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -1469,8 +1469,10 @@ let extract_fun_decl_gen (ctx : extraction_ctx) (fmt : F.formatter)
        *)
       let inputs_lvs =
         let all_inputs = (Option.get def.body).inputs_lvs in
+        (* TODO: *)
+        assert (not !Config.return_back_funs);
         let num_fwd_inputs =
-          def.signature.info.num_fwd_inputs_with_fuel_with_state
+          def.signature.info.fwd_info.num_inputs_with_fuel_with_state
         in
         Collections.List.prefix num_fwd_inputs all_inputs
       in
@@ -1515,8 +1517,10 @@ let extract_fun_decl_gen (ctx : extraction_ctx) (fmt : F.formatter)
   if has_decreases_clause && !backend = Lean then (
     let def_body = Option.get def.body in
     let all_vars = List.map (fun (v : var) -> v.id) def_body.inputs in
+    (* TODO: *)
+    assert (not !Config.return_back_funs);
     let num_fwd_inputs =
-      def.signature.info.num_fwd_inputs_with_fuel_with_state
+      def.signature.info.fwd_info.num_inputs_with_fuel_with_state
     in
     let vars = Collections.List.prefix num_fwd_inputs all_vars in
 
diff --git a/compiler/Pure.ml b/compiler/Pure.ml
index d7aea0f7..80d8782b 100644
--- a/compiler/Pure.ml
+++ b/compiler/Pure.ml
@@ -873,23 +873,25 @@ type fun_effect_info = {
 }
 [@@deriving show]
 
-(** Meta information about a function signature *)
-type fun_sig_info = {
+type inputs_info = {
   has_fuel : bool;
-  num_fwd_inputs_no_fuel_no_state : int;
-      (** The number of input types for forward computation, ignoring the fuel (if used)
+  num_inputs_no_fuel_no_state : int;
+      (** The number of input types ignoring the fuel (if used)
           and ignoring the state (if used) *)
-  num_fwd_inputs_with_fuel_no_state : int;
-      (** The number of input types for forward computation, with the fuel (if used)
+  num_inputs_with_fuel_no_state : int;
+      (** The number of input types, with the fuel (if used)
           and ignoring the state (if used) *)
-  num_fwd_inputs_with_fuel_with_state : int;
-      (** The number of input types for forward computation, with fuel and state (if used) *)
-  num_back_inputs_no_state : int option;
-      (** The number of additional inputs for the backward computation (if pertinent),
-          ignoring the state (if there is one) *)
-  num_back_inputs_with_state : int option;
-      (** The number of additional inputs for the backward computation (if pertinent),
-          with the state (if there is one) *)
+  num_inputs_with_fuel_with_state : int;
+      (** The number of input types, with fuel and state (if used) *)
+}
+[@@deriving show]
+
+(** Meta information about a function signature *)
+type fun_sig_info = {
+  fwd_info : inputs_info;
+      (** Information about the inputs of the forward function *)
+  back_info : inputs_info option;
+      (** Information about the inputs of the backward function, if pertinent *)
   effect_info : fun_effect_info;
 }
 [@@deriving show]
diff --git a/compiler/PureMicroPasses.ml b/compiler/PureMicroPasses.ml
index 34578750..7f122f15 100644
--- a/compiler/PureMicroPasses.ml
+++ b/compiler/PureMicroPasses.ml
@@ -1326,6 +1326,9 @@ let decompose_loops (def : fun_decl) : fun_decl * fun_decl list =
             let fun_sig_info = fun_sig.info in
             let fun_effect_info = fun_sig_info.effect_info in
 
+            (* TODO: *)
+            assert (not !Config.return_back_funs);
+
             (* Generate the loop definition *)
             let loop_effect_info =
               {
@@ -1340,38 +1343,44 @@ let decompose_loops (def : fun_decl) : fun_decl * fun_decl list =
             let loop_sig_info =
               let fuel = if !Config.use_fuel then 1 else 0 in
               let num_inputs = List.length loop.inputs in
-              let num_fwd_inputs_no_fuel_no_state = num_inputs in
-              let num_fwd_inputs_with_fuel_no_state = fuel + num_inputs in
-              let fwd_state =
-                fun_sig_info.num_fwd_inputs_with_fuel_with_state
-                - fun_sig_info.num_fwd_inputs_with_fuel_no_state
-              in
-              let num_fwd_inputs_with_fuel_with_state =
-                num_fwd_inputs_with_fuel_no_state + fwd_state
+              let fwd_info : inputs_info =
+                let info = fun_sig_info.fwd_info in
+                let fwd_state =
+                  info.num_inputs_with_fuel_with_state
+                  - info.num_inputs_with_fuel_no_state
+                in
+                {
+                  has_fuel = !Config.use_fuel;
+                  num_inputs_no_fuel_no_state = num_inputs;
+                  num_inputs_with_fuel_no_state = num_inputs + fuel;
+                  num_inputs_with_fuel_with_state =
+                    num_inputs + fuel + fwd_state;
+                }
               in
+
               {
-                has_fuel = !Config.use_fuel;
-                num_fwd_inputs_no_fuel_no_state;
-                num_fwd_inputs_with_fuel_no_state;
-                num_fwd_inputs_with_fuel_with_state;
-                num_back_inputs_no_state = fun_sig_info.num_back_inputs_no_state;
-                num_back_inputs_with_state =
-                  fun_sig_info.num_back_inputs_with_state;
+                fwd_info;
+                back_info = fun_sig_info.back_info;
                 effect_info = loop_effect_info;
               }
             in
+            assert (fun_sig_info_is_wf loop_sig_info);
 
             let inputs_tys =
+              (* TODO: *)
+              assert (not !Config.return_back_funs);
+
               let fuel = if !Config.use_fuel then [ mk_fuel_ty ] else [] in
               let fwd_inputs = List.map (fun (v : var) -> v.ty) loop.inputs in
+              let info = fun_sig_info.fwd_info in
               let state =
                 Collections.List.subslice fun_sig.inputs
-                  fun_sig_info.num_fwd_inputs_with_fuel_no_state
-                  fun_sig_info.num_fwd_inputs_with_fuel_with_state
+                  info.num_inputs_with_fuel_no_state
+                  info.num_inputs_with_fuel_with_state
               in
               let _, back_inputs =
                 Collections.List.split_at fun_sig.inputs
-                  fun_sig_info.num_fwd_inputs_with_fuel_with_state
+                  info.num_inputs_with_fuel_with_state
               in
               List.concat [ fuel; fwd_inputs; state; back_inputs ]
             in
@@ -1432,14 +1441,17 @@ let decompose_loops (def : fun_decl) : fun_decl * fun_decl list =
               in
 
               (* Introduce the additional backward inputs *)
+              (* TODO: *)
+              assert (not !Config.return_back_funs);
               let fun_body = Option.get def.body in
+              let info = fun_sig_info.fwd_info in
               let _, back_inputs =
                 Collections.List.split_at fun_body.inputs
-                  fun_sig_info.num_fwd_inputs_with_fuel_with_state
+                  info.num_inputs_with_fuel_with_state
               in
               let _, back_inputs_lvs =
                 Collections.List.split_at fun_body.inputs_lvs
-                  fun_sig_info.num_fwd_inputs_with_fuel_with_state
+                  info.num_inputs_with_fuel_with_state
               in
 
               let inputs =
@@ -2055,12 +2067,14 @@ let filter_loop_inputs (transl : pure_fun_translation list) :
 
   (* We start by computing the filtering information, for each function *)
   let compute_one_filter_info (decl : fun_decl) =
+    (* TODO: *)
+    assert (not !Config.return_back_funs);
     (* There should be a body *)
     let body = Option.get decl.body in
     (* We only look at the forward inputs, without the state *)
     let inputs_prefix, _ =
       Collections.List.split_at body.inputs
-        decl.signature.info.num_fwd_inputs_with_fuel_no_state
+        decl.signature.info.fwd_info.num_inputs_with_fuel_no_state
     in
     let used = ref (List.map (fun v -> (var_get_id v, false)) inputs_prefix) in
     let inputs_prefix_length = List.length inputs_prefix in
@@ -2080,7 +2094,10 @@ let filter_loop_inputs (transl : pure_fun_translation list) :
 
     (* Set the fuel as used *)
     let sg_info = decl.signature.info in
-    if sg_info.has_fuel then set_used (fst (Collections.List.nth inputs 0));
+    (* TODO: *)
+    assert (not !Config.return_back_funs);
+    if sg_info.fwd_info.has_fuel then
+      set_used (fst (Collections.List.nth inputs 0));
 
     let visitor =
       object (self : 'self)
@@ -2168,34 +2185,35 @@ let filter_loop_inputs (transl : pure_fun_translation list) :
               =
             decl.signature
           in
+          (* TODO: *)
+          assert (not !Config.return_back_funs);
+          let { fwd_info; back_info; effect_info } = info in
+
           let {
             has_fuel;
-            num_fwd_inputs_no_fuel_no_state;
-            num_fwd_inputs_with_fuel_no_state;
-            num_fwd_inputs_with_fuel_with_state;
-            num_back_inputs_no_state;
-            num_back_inputs_with_state;
-            effect_info;
+            num_inputs_no_fuel_no_state;
+            num_inputs_with_fuel_no_state;
+            num_inputs_with_fuel_with_state;
           } =
-            info
+            fwd_info
           in
 
           let inputs = filter_prefix used_info inputs in
 
-          let info =
+          let fwd_info =
             {
               has_fuel;
-              num_fwd_inputs_no_fuel_no_state =
-                num_fwd_inputs_no_fuel_no_state - num_filtered;
-              num_fwd_inputs_with_fuel_no_state =
-                num_fwd_inputs_with_fuel_no_state - num_filtered;
-              num_fwd_inputs_with_fuel_with_state =
-                num_fwd_inputs_with_fuel_with_state - num_filtered;
-              num_back_inputs_no_state;
-              num_back_inputs_with_state;
-              effect_info;
+              num_inputs_no_fuel_no_state =
+                num_inputs_no_fuel_no_state - num_filtered;
+              num_inputs_with_fuel_no_state =
+                num_inputs_with_fuel_no_state - num_filtered;
+              num_inputs_with_fuel_with_state =
+                num_inputs_with_fuel_with_state - num_filtered;
             }
           in
+
+          let info = { fwd_info; back_info; effect_info } in
+          assert (fun_sig_info_is_wf info);
           let signature =
             { generics; llbc_generics; preds; inputs; output; doutputs; info }
           in
diff --git a/compiler/PureUtils.ml b/compiler/PureUtils.ml
index 39dcd52d..23a41f0e 100644
--- a/compiler/PureUtils.ml
+++ b/compiler/PureUtils.ml
@@ -57,6 +57,25 @@ end
 module FunLoopIdMap = Collections.MakeMap (FunLoopIdOrderedType)
 module FunLoopIdSet = Collections.MakeSet (FunLoopIdOrderedType)
 
+let inputs_info_is_wf (info : inputs_info) : bool =
+  let {
+    has_fuel;
+    num_inputs_no_fuel_no_state;
+    num_inputs_with_fuel_no_state;
+    num_inputs_with_fuel_with_state;
+  } =
+    info
+  in
+  let fuel = if has_fuel then 1 else 0 in
+  num_inputs_no_fuel_no_state >= 0
+  && num_inputs_with_fuel_no_state = num_inputs_no_fuel_no_state + fuel
+  && num_inputs_with_fuel_with_state >= num_inputs_with_fuel_no_state
+
+let fun_sig_info_is_wf (info : fun_sig_info) : bool =
+  inputs_info_is_wf info.fwd_info
+  &&
+  match info.back_info with None -> true | Some info -> inputs_info_is_wf info
+
 let dest_arrow_ty (ty : ty) : ty * ty =
   match ty with
   | TArrow (arg_ty, ret_ty) -> (arg_ty, ret_ty)
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 2ef313e6..971a8cbd 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -1034,6 +1034,8 @@ let translate_fun_sig (decls_ctx : C.decls_ctx) (fun_id : A.fun_id)
   (* Generic parameters *)
   let generics = translate_generic_params sg.generics in
   (* Return *)
+  (* TODO: *)
+  assert (not !Config.return_back_funs);
   let has_fuel = fuel <> [] in
   let num_fwd_inputs_no_fuel_no_state = List.length fwd_inputs in
   let num_fwd_inputs_with_fuel_no_state =
@@ -1043,24 +1045,32 @@ let translate_fun_sig (decls_ctx : C.decls_ctx) (fun_id : A.fun_id)
   let num_back_inputs_no_state =
     if bid = None then None else Some (List.length back_inputs)
   in
-  let info =
+  let fwd_info : inputs_info =
     {
       has_fuel;
-      num_fwd_inputs_no_fuel_no_state;
-      num_fwd_inputs_with_fuel_no_state;
-      num_fwd_inputs_with_fuel_with_state =
+      num_inputs_no_fuel_no_state = num_fwd_inputs_no_fuel_no_state;
+      num_inputs_with_fuel_no_state = num_fwd_inputs_with_fuel_no_state;
+      num_inputs_with_fuel_with_state =
         (* We use the fact that [fwd_state_ty] has length 1 if there is a state,
            and 0 otherwise *)
         num_fwd_inputs_with_fuel_no_state + List.length fwd_state_ty;
-      num_back_inputs_no_state;
-      num_back_inputs_with_state =
-        (* Length of [back_state_ty]: similar trick as for [fwd_state_ty] *)
-        Option.map
-          (fun n -> n + List.length back_state_ty)
-          num_back_inputs_no_state;
-      effect_info;
     }
   in
+  let back_info : inputs_info option =
+    Option.map
+      (fun n ->
+        (* Note that backward functions never use fuel *)
+        {
+          has_fuel = false;
+          num_inputs_no_fuel_no_state = n;
+          num_inputs_with_fuel_no_state = n;
+          (* Length of [back_state_ty]: similar trick as for [fwd_state_ty] *)
+          num_inputs_with_fuel_with_state = n + List.length back_state_ty;
+        })
+      num_back_inputs_no_state
+  in
+  let info = { fwd_info; back_info; effect_info } in
+  assert (fun_sig_info_is_wf info);
   let preds = translate_predicates sg.preds in
   let sg =
     {
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 221d4e73..54e24066 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -216,11 +216,15 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
         (* We need to ignore the forward inputs, and the state input (if there is) *)
         let backward_inputs =
           let sg = backward_sg.sg in
+          (* TODO: *)
+          assert (not !Config.return_back_funs);
           (* We need to ignore the forward state and the backward state *)
           let num_forward_inputs =
-            sg.info.num_fwd_inputs_with_fuel_with_state
+            sg.info.fwd_info.num_inputs_with_fuel_with_state
+          in
+          let num_back_inputs =
+            (Option.get sg.info.back_info).num_inputs_no_fuel_no_state
           in
-          let num_back_inputs = Option.get sg.info.num_back_inputs_no_state in
           Collections.List.subslice sg.inputs num_forward_inputs
             (num_forward_inputs + num_back_inputs)
         in