Native decide performance on bounded forall problems

Stream: lean4

[Josh Clune]

I'm currently working on some proofs that should, in principle, be easy to automate with brute force computation. For example, something like this:

import Lean

example : ∀ n : Nat, n < 2^15 → n >>> 15 = 0 := by
  decide +native

I was surprised to find that proving such goals with decide +native took longer than I expected, and resulted in stack overflow errors sooner than I expected.

import Lean

-- This yields a stack overflow error when run in VS Code
example : ∀ n : Nat, n < 2^20 → n >>> 20 = 0 := by
  decide +native

In principle, the actual computation required by these goals is not very expensive.

def main : IO Unit := do
  for n in [0:2^20] do
    if (n >>> 20) != 0 then
      IO.println "Found a counterexample!"
      return
  IO.println "Success!"
  return

#eval main -- This succeeds very quickly

Currently, I think decide's comparably worse performance is due to the instance Nat.decidableBallLT, which dictates how decide computes goals of the form ∀ n : Nat, n < d → .... As a sanity check, I implemented this simple tactic which computes goals of the same form via a fold function:

import Lean

open Lean Meta Parser Elab Tactic

syntax (name := brute) "brute" : tactic

/-- If `goalBinders` has the form `#[(x : Nat), (h : x < d)]` then returns `(x, d)` -/
def validateGoalBinders (goalBinders : Array FVarId) : TacticM (Array (FVarId × Expr)) := do
  match goalBinders with
  | #[b1, b2] =>
    let lctx ← getLCtx
    let some b1LocalDecl := lctx.find? b1
      | throwError "{decl_name%} :: Unable to find type of goal binder {Expr.fvar b1}"
    let some b2LocalDecl := lctx.find? b2
      | throwError "{decl_name%} :: Unable to find type of goal binder {Expr.fvar b2}"
    let b1Type := b1LocalDecl.type
    let b2Type := b2LocalDecl.type
    if b1Type != mkConst ``Nat then throwError "{decl_name%} :: Wrong b1Type: {b1Type}"
    let b1UpperBound ←
      match b2Type with
      | .app (.app (.app (.app (.const ``LT.lt _) _) _) x) y =>
        if x != Expr.fvar b1 then
          throwError "{decl_name%} :: Second binder type does not have the form {Expr.fvar b1} < _: {b2Type}"
        else
          pure y
      | _ => throwError "{decl_name%} :: Wrong b2Type: {b2Type}"
    return #[(b1, b1UpperBound)]
  | _ => throwError "Wrong number of goal binders"

def mkFold (f : Nat → Bool) (b : Nat) (acc : Bool) : Bool :=
  (List.range' 0 b).foldr
    (fun (x : Nat) (acc : Bool) => acc && f x) acc

theorem ofMkFoldEqTrueAux (f : Nat → Bool) (b : Nat) (acc : Bool) :
  mkFold f b acc = (acc ∧ ∀ x < b, f x) := by
  simp only [mkFold, eq_iff_iff]
  induction b generalizing acc
  . simp
  . next b ih =>
    simp only [List.range'_1_concat, Nat.zero_add, List.foldr_append, List.foldr_cons,
      List.foldr_nil, ih (acc && f b), Bool.and_eq_true]
    constructor
    . simp only [and_imp]
      intro hacc hfb hb
      constructor
      . exact hacc
      . intro x hx
        by_cases hxb : x = b
        . rw [hxb, hfb]
        . exact hb x (by omega)
    . simp only [and_imp]
      intro hacc h
      constructor
      . rw [hacc, h b (by omega), and_self]
      . intro x hxb
        exact h x (by omega)

theorem ofMkFoldEqTrue (f : Nat → Bool) (b : Nat) :
  mkFold f b true → ∀ x < b, f x := by simp only [ofMkFoldEqTrueAux f b true, true_and, imp_self]

@[tactic brute]
def evalBrute : Tactic
| `(tactic| brute) => withMainContext do
  let g ← getMainGoal
  let (goalBinders, g') ← g.intros
  g'.withContext do
    let #[(x, xBound)] ← validateGoalBinders goalBinders
      | throwError "Unexpected output from validateGoalBinders"
    let f ← mkLambdaFVars #[Expr.fvar x] (← mkDecide (← g'.getType))
    let res ← mkAppM ``mkFold #[f, xBound, mkConst ``true]

    let levels := (collectLevelParams {} res).params.toList
    let auxDeclName ← Term.mkAuxName `_brute
    let decl := Declaration.defnDecl {
      name := auxDeclName
      levelParams := levels
      type := mkConst ``Bool
      value := res
      hints := .abbrev
      safety := .safe
    }
    addAndCompile decl

    let rflPrf ← mkEqRefl (toExpr true)
    let levelParams := levels.map .param
    let pf := mkApp3 (mkConst ``ofMkFoldEqTrue) f xBound <|
      mkApp3 (mkConst ``Lean.ofReduceBool) (mkConst auxDeclName levelParams) (toExpr true) rflPrf
    g'.assign $ ← mkAppOptM ``of_decide_eq_true #[none, none, ← mkAppM' pf (goalBinders.map Expr.fvar)]
| _ => throwUnsupportedSyntax

-- Both of these work with minimal delay and no stack overflow
example : ∀ n : Nat, n < 2^15 → n >>> 15 = 0 := by
  brute

example : ∀ n : Nat, n < 2^20 → n >>> 20 = 0 := by
  brute

My questions are:

• Is the root issue of decide's performance related to Nat.decidableBallLT, or have I misdiagnosed the problem?
• If the issue is related to Nat.decidableBallLT, is there a way to override this instance with a more efficient one to improve decide's performance? Or would it be better to just use/write a different tactic that can better handle these sorts of goals?

[Aaron Liu]

oh no it's not tail recursive

[Lawrence Wu (llllvvuu)]

@Josh Clune you could take a look at https://github.com/leanprover/lean4/pull/6651 . My PR is slightly unoptimized for native_decide, as it also needs to handle non-native decide where tail-recursion has no effect. But you could modify the instance for your use case.

To disable the old one you can do attribute [-instance] Nat.decidableBallLT