Restate quantifier-free argument refinements at lemma calls (fixes #4591)

src/typechecker/FStarC.TypeChecker.TcTerm.fst

@@ -2963,6 +2963,12 @@ and check_application_args env head (chead:comp) ghead args expected_topt : ML (
             name_opt |> Option.map (Env.push_bv env)
                      |> Option.dflt env) in
 
+        (* see the comment on [no_capture] below *)
+        let result_is_refined_unit =
+          Cons? arg_comps_rev
+          && U.is_pure_or_ghost_comp cres
+          && TcUtil.is_refined_unit env (U.comp_result cres) in
+
         //Bind arguments
         let _, comp, g_comp =
           List.fold_left
@@ -3011,7 +3017,23 @@ and check_application_args env head (chead:comp) ghead args expected_topt : ML (
                  its value before the solver ever sees it, so the callee's
                  typing axiom never fires and the refinement is simply gone.
                  [FStar.UInt32.lognot 0xff00ul] reduces to [0xffff00fful], and
-                 with it the only statement that the two are related. *)
+                 with it the only statement that the two are related.
+
+                 A lemma call is a third exception.  Its result type
+                 [unit{post}] is not the type of any value -- it is dissolved
+                 into a hypothesis as soon as the call is sequenced -- so
+                 restating an argument's refinement there costs no type
+                 pollution, and it puts the fact right next to the
+                 postcondition that mentions the argument.  Recovering it from
+                 the callee's typing axiom instead takes two instantiations
+                 (the typing axiom, then the refinement's interpretation), and
+                 nonlinear arithmetic is very sensitive to exactly these sign
+                 facts: [eval s >= 0] and [pow2 n > 0] next to a
+                 [distributivity_add_right] postcondition (issue #4591).  Only
+                 the quantifier-free conjuncts are restated, though: a
+                 quantified one -- [Seq.init]'s [forall i. index s i == f i],
+                 say -- is new instantiation work in every goal that follows
+                 the call. *)
               let arg_head_is_reducible_primop () =
                 let hd, _ = U.head_and_args_full e in
                 match (U.un_uinst hd).n with
@@ -3022,16 +3044,19 @@ and check_application_args env head (chead:comp) ghead args expected_topt : ML (
                   && is_empty (Free.uvars e)
                 | _ -> false
               in
+              let capture_all =
+                head_is_data_constructor || arg_head_is_reducible_primop () in
               let no_capture =
-                (not head_is_data_constructor
-                 && not (arg_head_is_reducible_primop ()))
+                (not capture_all && not result_is_refined_unit)
                 || S.is_aqual_implicit q
                 || not (is_empty (Free.uvars (U.comp_result c))) in
               let e_opt = if U.is_pure_or_ghost_comp c then Some e else None in
               let c_out, g_out =
                 if no_capture
                 then TcUtil.bind_no_capture e.pos false env e_opt (c, Env.trivial_guard) (x, out_c, g_out)
-                else TcUtil.bind e.pos false env e_opt (c, Env.trivial_guard) (x, out_c, g_out) in
+                else if capture_all
+                then TcUtil.bind e.pos false env e_opt (c, Env.trivial_guard) (x, out_c, g_out)
+                else TcUtil.bind_capture_quantifier_free e.pos false env e_opt (c, Env.trivial_guard) (x, out_c, g_out) in
               i+1, c_out, g_out)
           (1, cres, Env.trivial_guard)
           arg_comps_rev in

src/typechecker/FStarC.TypeChecker.Util.fst

@@ -1137,16 +1137,16 @@ let captured_typing
 (* Restating a conjunct that the continuation's result type already carries
    costs a duplicated hypothesis at every enclosing bind, so a long statement
    sequence would accumulate the same facts quadratically.  Drop those. *)
+let rec conjuncts (phi:term) : ML (list term) =
+  let hd, args = U.head_and_args_full phi in
+  match (U.un_uinst hd).n, args with
+  | Tm_fvar fv, [(a, _); (b, _)] when S.fv_eq_lid fv C.and_lid ->
+    conjuncts a @ conjuncts b
+  | _ -> [phi]
+
 let drop_redundant_conjuncts (env:Env.env) (already_says:typ) (phi:term) : ML term =
   if U.is_t_true phi then phi
   else
-    let rec conjuncts (phi:term) : ML (list term) =
-      let hd, args = U.head_and_args_full phi in
-      match (U.un_uinst hd).n, args with
-      | Tm_fvar fv, [(a, _); (b, _)] when S.fv_eq_lid fv C.and_lid ->
-        conjuncts a @ conjuncts b
-      | _ -> [phi]
-    in
     let already =
       match (N.normalize_refinement N.whnf_steps env already_says).n with
       | Tm_refine {b; phi} ->
@@ -1165,6 +1165,22 @@ let drop_redundant_conjuncts (env:Env.env) (already_says:typ) (phi:term) : ML te
     in
     U.mk_conj_l keep
 
+(* A quantifier is a binder -- [l_Forall (fun x -> ...)] -- so a conjunct with
+   no binder anywhere in it is quantifier-free.  Keep only those. *)
+let quantifier_free_conjuncts (phi:term) : ML term =
+  if U.is_t_true phi then phi
+  else
+    let binds (t:term) : ML bool =
+      let found = mk_ref false in
+      let _ = FStarC.Syntax.Visit.visit_term false (fun t ->
+        (match t.n with
+         | Tm_abs _ | Tm_arrow _ | Tm_refine _ | Tm_let _ | Tm_match _ -> found := true
+         | _ -> ());
+        t) t in
+      !found
+    in
+    U.mk_conj_l (List.filter (fun c -> not (binds c)) (conjuncts phi))
+
 (* The result type of the composite.
 
    This is the *sole* authority on a bind's result type.  [simplify_bind] and
@@ -1174,11 +1190,12 @@ let drop_redundant_conjuncts (env:Env.env) (already_says:typ) (phi:term) : ML te
    they produce with what this returns.  When there is nothing to say, this
    returns [lc2]'s result type unchanged, so the overwrite is a no-op. *)
 let composite_result_typ
-      (capture:bool) (is_let_binding:bool)
+      (capture:bool) (quantifier_free:bool) (is_let_binding:bool)
       (env:Env.env) (e1opt:option term) (lc1:comp) (b:option bv) (lc2:comp)
 : ML (typ & guard_t)
 = let subst_x = bind_result_subst env e1opt lc1 b lc2 in
   let phi = captured_typing env capture is_let_binding (Cons? subst_x) lc1 e1opt b in
+  let phi = if quantifier_free then quantifier_free_conjuncts phi else phi in
   (* [g_esc] is [mzero] except on [eliminate_binder_from_typ]'s last resort,
      where [check_no_escape] may equate the type to a fresh uvar. *)
   let res_typ_base, g_esc =
@@ -1470,7 +1487,7 @@ let bind_general (bi:bind_input) : ML (comp & guard_t) =
   else mk_bind c1 b c2 trivial_guard
 
 let bind_maybe_capture
-      (capture:bool)
+      (capture:bool) (quantifier_free:bool)
       (r1:Range.t)
       (is_let_binding:bool)
       (env:Env.env) (e1opt:option term) (lc1_g : comp & guard_t) (binder_lc2:comp_with_binder) : ML (comp & guard_t) =
@@ -1488,7 +1505,7 @@ let bind_maybe_capture
   (* The result type is computed here and nowhere else: the comps returned below
      are derived from [lc2] and may still mention [b], which is out of scope for
      the caller. *)
-  let res_typ, g_esc = composite_result_typ capture is_let_binding env e1opt lc1 b lc2 in
+  let res_typ, g_esc = composite_result_typ capture quantifier_free is_let_binding env e1opt lc1 b lc2 in
   let c, g =
     match simplify_bind bi with
     | Inl (c, g, reason) ->
@@ -1503,10 +1520,16 @@ let bind_maybe_capture
   U.set_result_typ c res_typ, Env.conj_guard g g_esc
 
 let bind r1 is_let_binding env e1opt lc1 binder_lc2 : ML (comp & guard_t) =
-  bind_maybe_capture true r1 is_let_binding env e1opt lc1 binder_lc2
+  bind_maybe_capture true false r1 is_let_binding env e1opt lc1 binder_lc2
 
 let bind_no_capture r1 is_let_binding env e1opt lc1 binder_lc2 : ML (comp & guard_t) =
-  bind_maybe_capture false r1 is_let_binding env e1opt lc1 binder_lc2
+  bind_maybe_capture false false r1 is_let_binding env e1opt lc1 binder_lc2
+
+let bind_capture_quantifier_free r1 is_let_binding env e1opt lc1 binder_lc2 : ML (comp & guard_t) =
+  bind_maybe_capture true true r1 is_let_binding env e1opt lc1 binder_lc2
+
+let is_refined_unit (env:env) (t:typ) : ML bool =
+  Refined_unit? (unit_shape_of env t)
 
 let weaken_guard g1 g2 : ML _ = match g1, g2 with
     | NonTrivial f1, NonTrivial f2 ->
@@ -2510,7 +2533,7 @@ let weaken_result_typ env (e:term) (lc_g : comp & guard_t) (t:typ) (use_eq:bool)
                           with [x == e]. *)
                       let x = {x with sort=(U.comp_result lc)} in
                       //AR: M_M bind
-                      let c, g_lc = bind_maybe_capture false e.pos false env (Some e) (c, Env.trivial_guard) (Some x, eq_ret, g_eq) in
+                      let c, g_lc = bind_maybe_capture false false e.pos false env (Some e) (c, Env.trivial_guard) (Some x, eq_ret, g_eq) in
                       if Debug.extreme ()
                       then Format.print1 "Strengthened to %s\n" (Normalize.comp_to_string env c);
                       c, Env.conj_guards [g_c; gret; g_lc]

src/typechecker/FStarC.TypeChecker.Util.fsti

@@ -79,6 +79,12 @@ val bind: Range.t -> is_let_binding:bool -> env -> option term -> (comp & guard_
    type, above all, which is the image of a precondition and carries an
    obligation discharged at the call rather than a fact about a result. *)
 val bind_no_capture: Range.t -> is_let_binding:bool -> env -> option term -> (comp & guard_t) -> comp_with_binder -> ML (comp & guard_t)
+(* [bind_capture_quantifier_free] is [bind], except that only the
+   quantifier-free conjuncts of [e1]'s refinement are restated. *)
+val bind_capture_quantifier_free: Range.t -> is_let_binding:bool -> env -> option term -> (comp & guard_t) -> comp_with_binder -> ML (comp & guard_t)
+
+(* Is [t] (up to abbreviations and [squash]) a refinement of [unit]? *)
+val is_refined_unit: env -> typ -> ML bool
 
 val weaken_guard: guard_formula -> guard_formula -> ML guard_formula
 

tests/bug-reports/closed/Bug4591.fst

@@ -0,0 +1,37 @@
+module Bug4591
+
+module Seq = FStar.Seq
+module Math = FStar.Math.Lemmas
+
+let rad : pos = 64
+let base : pos = pow2 rad
+
+#push-options "--fuel 2 --ifuel 1 --z3rlimit 60"
+
+let rec eval (s: Seq.seq nat) : Tot nat (decreases Seq.length s) =
+  let n = Seq.length s in
+  if n = 0 then 0
+  else Seq.index s 0 + base * eval (Seq.slice s 1 n)
+
+// An unused lemma above [eval_split] used to be enough to push its final
+// goal past the rlimit (issue #4591): the sign facts [eval _ >= 0] and
+// [pow2 _ > 0] of the lemma arguments were not in the VC.
+let eval_empty (s: Seq.seq nat) : Lemma (requires Seq.length s == 0) (ensures eval s == 0) = ()
+
+let rec eval_split (s: Seq.seq nat) (k: nat { k <= Seq.length s })
+  : Lemma (ensures eval s == eval (Seq.slice s 0 k) +
+                             pow2 (rad * k) * eval (Seq.slice s k (Seq.length s)))
+          (decreases k)
+  = let n = Seq.length s in
+    if k = 0 then Seq.lemma_eq_intro (Seq.slice s 0 n) s
+    else begin
+      let t = Seq.slice s 1 n in
+      eval_split t (k - 1);
+      Seq.lemma_eq_intro (Seq.slice t 0 (k - 1)) (Seq.slice (Seq.slice s 0 k) 1 k);
+      Seq.lemma_eq_intro (Seq.slice t (k - 1) (n - 1)) (Seq.slice s k n);
+      Math.pow2_plus rad (rad * (k - 1));
+      Math.distributivity_add_right base (eval (Seq.slice t 0 (k - 1)))
+        (pow2 (rad * (k - 1)) * eval (Seq.slice s k n));
+      Math.paren_mul_right base (pow2 (rad * (k - 1))) (eval (Seq.slice s k n))
+    end
+#pop-options

tests/tactics/Postprocess.fst.output.expected

@@ -378,9 +378,9 @@ visible let  xx  : t1 = (C1 (fun uu___0 -> (match uu___0@0:(Tm_unknown) with
 [@ ]
 visible let  q_as_lem  : (p:(squash (l_Forall (fun x -> (b@1:(Tm_unknown) x@0:(Tm_unknown))))) -> x:a@2:(Tm_unknown) -> Lemma ((squash (b@2:(Tm_unknown) x@0:(Tm_unknown))))) = (fun p x -> ())
 [@ ]
-visible let  congruence_fun  : (f:(x:a@1:(Tm_unknown) -> Tot (b@1:(Tm_unknown) x@0:(Tm_unknown))) -> g:(x:a@2:(Tm_unknown) -> Tot (b@2:(Tm_unknown) x@0:(Tm_unknown))) -> x:(squash (l_Forall (fun x -> (eq2 (f@2:(Tm_unknown) x@0:(Tm_unknown)) (g@1:(Tm_unknown) x@0:(Tm_unknown)))))) -> Lemma ((squash (eq2 (fun x -> (f@3:(Tm_unknown) x@0:(Tm_unknown))) (fun x -> (g@2:(Tm_unknown) x@0:(Tm_unknown))))))) = (fun f g x -> (assert_by_tactic (eq2 (fun x -> (f@3:(Tm_unknown) x@0:(Tm_unknown))) (fun x -> (g@2:(Tm_unknown) x@0:(Tm_unknown)))) (fun uu___ -> let  [@ (inline_let)]uu___#2548  : unit = ()
+visible let  congruence_fun  : (f:(x:a@1:(Tm_unknown) -> Tot (b@1:(Tm_unknown) x@0:(Tm_unknown))) -> g:(x:a@2:(Tm_unknown) -> Tot (b@2:(Tm_unknown) x@0:(Tm_unknown))) -> x:(squash (l_Forall (fun x -> (eq2 (f@2:(Tm_unknown) x@0:(Tm_unknown)) (g@1:(Tm_unknown) x@0:(Tm_unknown)))))) -> Lemma ((squash (eq2 (fun x -> (f@3:(Tm_unknown) x@0:(Tm_unknown))) (fun x -> (g@2:(Tm_unknown) x@0:(Tm_unknown))))))) = (fun f g x -> (assert_by_tactic (eq2 (fun x -> (f@3:(Tm_unknown) x@0:(Tm_unknown))) (fun x -> (g@2:(Tm_unknown) x@0:(Tm_unknown)))) (fun uu___ -> let  [@ (inline_let)]uu___#2550  : unit = ()
 in
-let  uu___#2549  : unit = let  uu___#2550  : (list term) = let  uu___#2551  : term = quote ((q_as_lem x@2:(Tm_unknown)))
+let  uu___#2551  : unit = let  uu___#2552  : (list term) = let  uu___#2553  : term = quote ((q_as_lem x@2:(Tm_unknown)))
 in
 (Cons uu___@0:(Tm_unknown) (Nil ))
 in

ulib/FStar.UInt128.fst

@@ -847,6 +847,7 @@ let eq_mask (a b: t) : Pure t
   (ensures (fun r -> (v a = v b ==> v r = pow2 128 - 1) /\ (v a <> v b ==> v r = 0))) =
   let mask = U64.logand (U64.eq_mask a.low b.low)
                         (U64.eq_mask a.high b.high) in
+  if v a = v b then v_inj a b;
   { low = mask; high = mask; }
 
 private let gte_characterization (a b: t) :