Change assumption related to st.balance

lib/gen_mlw.ml

@@ -126,6 +126,7 @@ module T = struct
           (* Hack: && => /\ *)
           Tbinnop (of_expr e1, Dterm.DTand, of_expr e2)
       | Eapply (e1, e2) -> Tapply (of_expr e1, of_expr e2)
+      | Eif (e1, e2, e3) -> Tif (of_expr e1, of_expr e2, of_expr e3)
       | _ ->
           Format.eprintf "T.of_expr: unsupported: %a@."
             (Mlw_printer.pp_expr ~attr:true).closed e;
@@ -474,14 +475,7 @@ module Generator (D : Desc) = struct
       (fun _ c e ->
         E.mk_if
           (is_contract_of c @@ Step_constant.self st)
-          (wrap_assume
-             ~assumption:
-               (T.mk_and
-                  (T.of_expr @@ call_param_wf_of c gp)
-                  (T.of_expr
-                     (E.mk_bin (Step_constant.balance st) "="
-                        (E.mk_bin (balance_of c ctx) "+"
-                           (Step_constant.amount st)))))
+          (wrap_assume ~assumption:(T.of_expr @@ call_param_wf_of c gp)
           @@ call_func_of c st gp ctx)
           e)
       contracts (call_unknown ctx)
@@ -672,6 +666,19 @@ module Generator (D : Desc) = struct
          in
          let+ p = E.mk_any gparam_pty in
          wrap_assume ~assumption:(wf st)
+         @@ wrap_assume
+              ~assumption:
+                (T.of_expr
+                @@ M.fold
+                     (fun _ c e ->
+                       E.mk_if
+                         (is_contract_of c @@ Step_constant.self st)
+                         (E.mk_bin (Step_constant.balance st) "="
+                            (E.mk_bin
+                               (balance_of c (E.var_of_binder ctx))
+                               "+" (Step_constant.amount st)))
+                         e)
+                     contracts (expr Etrue))
          @@ let+ ctx = dispatch_transfer (E.var_of_binder ctx) st p in
             call_unknown ctx
     in