Do not introduce match on the fuel for non-recursive functions

3 files changed, 50 insertions, 11 deletions

diff --git a/compiler/FunsAnalysis.ml b/compiler/FunsAnalysis.ml
index 75a6c0ce..b72fa078 100644
--- a/compiler/FunsAnalysis.ml
+++ b/compiler/FunsAnalysis.ml
@@ -18,7 +18,14 @@ type fun_info = {
   can_fail : bool;
       (* Not used yet: all the extracted functions use an error monad *)
   stateful : bool;
-  divergent : bool; (* Not used yet *)
+  can_diverge : bool;
+  (* The function can diverge if:
+     - it is recursive
+     - it contains a loop
+     - it calls a functions which can diverge
+  *)
+  is_rec : bool;
+      (* [true] if the function is recursive (or in a mutually recursive group) *)
 }
 [@@deriving show]
 
@@ -48,7 +55,8 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
       fun_info =
     let can_fail = ref false in
     let stateful = ref false in
-    let divergent = ref false in
+    let can_diverge = ref false in
+    let is_rec = ref false in
 
     let visit_fun (f : fun_decl) : unit =
       let obj =
@@ -70,14 +78,16 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
           method! visit_Call env call =
             (match call.func with
             | Regular id ->
-                if FunDeclId.Set.mem id fun_ids then divergent := true
+                if FunDeclId.Set.mem id fun_ids then (
+                  can_diverge := true;
+                  is_rec := true)
                 else
                   let info = FunDeclId.Map.find id !infos in
                   self#may_fail info.can_fail;
                   stateful := !stateful || info.stateful;
-                  divergent := !divergent || info.divergent
+                  can_diverge := !can_diverge || info.can_diverge
             | Assumed id ->
-                (* None of the assumed functions is stateful for now *)
+                (* None of the assumed functions can diverge nor are considered stateful *)
                 can_fail := !can_fail || Assumed.assumed_can_fail id);
             super#visit_Call env call
 
@@ -86,7 +96,7 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
             super#visit_Panic env
 
           method! visit_Loop env loop =
-            divergent := true;
+            can_diverge := true;
             super#visit_Loop env loop
         end
       in
@@ -110,7 +120,12 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
      * However, we do keep the result of the analysis for global bodies.
      * *)
     can_fail := (not is_global_decl_body) || !can_fail;
-    { can_fail = !can_fail; stateful = !stateful; divergent = !divergent }
+    {
+      can_fail = !can_fail;
+      stateful = !stateful;
+      can_diverge = !can_diverge;
+      is_rec = !is_rec;
+    }
   in
 
   let analyze_fun_decl_group (d : fun_declaration_group) : unit =
diff --git a/compiler/Pure.ml b/compiler/Pure.ml
index 11f627d7..d9d3a404 100644
--- a/compiler/Pure.ml
+++ b/compiler/Pure.ml
@@ -525,6 +525,8 @@ type fun_effect_info = {
   can_fail : bool;  (** [true] if the return type is a [result] *)
   can_diverge : bool;
       (** [true] if the function can diverge (i.e., not terminate) *)
+  is_rec : bool;
+      (** [true] if the function is recursive (or in a mutually recursive group) *)
 }
 
 (** Meta information about a function signature *)
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 9d8724ab..3bd6c5b3 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -495,7 +495,17 @@ let list_ancestor_abstractions (ctx : bs_ctx) (abs : V.abs) :
   let abs_ids = list_ancestor_abstractions_ids ctx abs in
   List.map (fun id -> V.AbstractionId.Map.find id ctx.abstractions) abs_ids
 
-(** Small utility *)
+(** Small utility.
+
+    Does the function *decrease* the fuel? [true] if recursive.
+ *)
+let function_decreases_fuel (info : fun_effect_info) : bool =
+  !Config.use_fuel && info.is_rec
+
+(** Small utility.
+
+    Does the function *use* the fuel? [true] if can diverge.
+ *)
 let function_uses_fuel (info : fun_effect_info) : bool =
   !Config.use_fuel && info.can_diverge
 
@@ -522,7 +532,8 @@ let get_fun_effect_info (fun_infos : FA.fun_info A.FunDeclId.Map.t)
         can_fail = info.can_fail;
         stateful_group;
         stateful;
-        can_diverge = info.divergent;
+        can_diverge = info.can_diverge;
+        is_rec = info.is_rec;
       }
   | A.Assumed aid ->
       {
@@ -530,6 +541,7 @@ let get_fun_effect_info (fun_infos : FA.fun_info A.FunDeclId.Map.t)
         stateful_group = false;
         stateful = false;
         can_diverge = false;
+        is_rec = false;
       }
 
 (** Translate a function signature.
@@ -1268,6 +1280,7 @@ and translate_function_call (call : S.call) (e : S.expression) (ctx : bs_ctx) :
             stateful_group = false;
             stateful = false;
             can_diverge = false;
+            is_rec = false;
           }
         in
         (ctx, Unop Not, effect_info, args, None)
@@ -1283,6 +1296,7 @@ and translate_function_call (call : S.call) (e : S.expression) (ctx : bs_ctx) :
                 stateful_group = false;
                 stateful = false;
                 can_diverge = false;
+                is_rec = false;
               }
             in
             (ctx, Unop (Neg int_ty), effect_info, args, None)
@@ -1295,6 +1309,7 @@ and translate_function_call (call : S.call) (e : S.expression) (ctx : bs_ctx) :
             stateful_group = false;
             stateful = false;
             can_diverge = false;
+            is_rec = false;
           }
         in
         (ctx, Unop (Cast (src_ty, tgt_ty)), effect_info, args, None)
@@ -1310,6 +1325,7 @@ and translate_function_call (call : S.call) (e : S.expression) (ctx : bs_ctx) :
                 stateful_group = false;
                 stateful = false;
                 can_diverge = false;
+                is_rec = false;
               }
             in
             (ctx, Binop (binop, int_ty0), effect_info, args, None)
@@ -1913,14 +1929,20 @@ let translate_fun_decl (ctx : bs_ctx) (body : S.expression option) : fun_decl =
         let body = translate_expression body ctx in
         (* Add a match over the fuel, if necessary *)
         let body =
-          if function_uses_fuel effect_info then wrap_in_match_fuel body ctx
+          if function_decreases_fuel effect_info then
+            wrap_in_match_fuel body ctx
           else body
         in
         (* Sanity check *)
         type_check_texpression ctx body;
         (* Introduce the fuel parameter, if necessary *)
         let fuel =
-          if function_uses_fuel effect_info then [ mk_fuel_var ctx.fuel0 ]
+          if function_uses_fuel effect_info then
+            let fuel_var =
+              if function_decreases_fuel effect_info then ctx.fuel0
+              else ctx.fuel
+            in
+            [ mk_fuel_var fuel_var ]
           else []
         in
         (* Introduce the forward input state (the state at call site of the