Fix various issues when using a state

2 files changed, 12 insertions, 30 deletions

diff --git a/src/ExtractToFStar.ml b/src/ExtractToFStar.ml
index 4baa1fd6..1c93c9da 100644
--- a/src/ExtractToFStar.ml
+++ b/src/ExtractToFStar.ml
@@ -1310,8 +1310,8 @@ let extract_template_decreases_clause (ctx : extraction_ctx) (fmt : F.formatter)
     it is useful for the decrease clause.
  *)
 let extract_fun_decl (ctx : extraction_ctx) (fmt : F.formatter)
-    (qualif : fun_decl_qualif) (has_decreases_clause : bool)
-    (has_state_param : bool) (fwd_def : fun_decl) (def : fun_decl) : unit =
+    (qualif : fun_decl_qualif) (has_decreases_clause : bool) (def : fun_decl) :
+    unit =
   (* Retrieve the function name *)
   let def_name = ctx_get_local_function def.def_id def.back_id ctx in
   (* (* Add the type parameters - note that we need those bindings only for the
@@ -1397,31 +1397,19 @@ let extract_fun_decl (ctx : extraction_ctx) (fmt : F.formatter)
        * backward functions have no influence on termination: we thus
        * share the decrease clauses between the forward and the backward
        * functions).
-       * Something annoying is that there may be a state parameter, for
-       * the state-error monad, in which case we need to forget the input
-       * parameter before last:
-       * ```
-       * val f_fwd (x : u32) (st : state) :
-       *   Tot (result (state & u32)) (decreases (f_decreases x st))
-       *
-       *              We ignore this parameter
-       *                       VVV
-       * val f_back (x : u32) (ret : u32) (st : state) :
-       *   Tot (result (state & u32)) (decreases (f_decreases x st))
-       * ```
-       * Rk.: if a function has a decreases clause, it is necessarily
-       * a transparent function
        *)
       let inputs_lvs =
-        let num_fwd_inputs = List.length (Option.get fwd_def.body).inputs_lvs in
+        let all_inputs = (Option.get def.body).inputs_lvs in
+        (* We have to count:
+         * - the forward inputs
+         * - the state
+         *)
+        let num_fwd_inputs = def.signature.info.num_fwd_inputs in
         let num_fwd_inputs =
-          if has_state_param then num_fwd_inputs - 1 else num_fwd_inputs
+          if def.signature.info.input_state then 1 + num_fwd_inputs
+          else num_fwd_inputs
         in
-        let all_inputs = (Option.get def.body).inputs_lvs in
-        let inputs = Collections.List.prefix num_fwd_inputs all_inputs in
-        if has_state_param then
-          inputs @ [ List.nth all_inputs (List.length all_inputs - 1) ]
-        else inputs
+        Collections.List.prefix num_fwd_inputs all_inputs
       in
       let _ =
         List.fold_left
diff --git a/src/Translate.ml b/src/Translate.ml
index a264a121..d5505894 100644
--- a/src/Translate.ml
+++ b/src/Translate.ml
@@ -481,19 +481,13 @@ let extract_definitions (fmt : Format.formatter) (config : gen_config)
          let has_decr_clause =
            has_decreases_clause def && config.extract_decreases_clauses
          in
-         (* Is there an input parameter "visible" for the state used in
-          * the state error monad (if we use a state error monad)? *)
-         let has_state_param =
-           config.use_state_monad
-           && config.mp_config.unfold_monadic_let_bindings
-         in
          (* Check if the definition needs to be filtered or not *)
          if
            ((not is_opaque) && config.extract_transparent)
            || (is_opaque && config.extract_opaque)
          then
            ExtractToFStar.extract_fun_decl ctx.extract_ctx fmt qualif
-             has_decr_clause has_state_param fwd_def def)
+             has_decr_clause def)
        fls);
     (* Insert unit tests if necessary *)
     if config.test_unit_functions then