missing whnf in NormNum.Core?

Stream: mathlib4

[David Renshaw]

Currently, norm_num can sometimes get stuck on sums:

import Mathlib.Tactic.NormNum
import Mathlib.Tactic.NormNum.BigOperators

open BigOperators

-- works
example : (∑ i in Finset.range 10, (i^2 : ℕ)) = 285 := by norm_num1

-- also works (via the evalOfNat extension)
example : (∑ i in Finset.range 10, (OfNat.ofNat i)) = 45 := by norm_num1

-- does not work
example : (∑ i in Finset.range 10, i) = 45 := by norm_num1

The problem is that when it tries to evaluate the inner body of the third example, derive is passed an expression of the form (fun x : Nat => x) 9 where 9 is a nat literal. This expression is not itself a nat literal, and no norm num extensions apply to it.

All three examples work if I make the following change:

diff --git a/Mathlib/Tactic/NormNum/Core.lean b/Mathlib/Tactic/NormNum/Core.lean
index 6ebf7ed63e..ced2559f41 100644
--- a/Mathlib/Tactic/NormNum/Core.lean
+++ b/Mathlib/Tactic/NormNum/Core.lean
@@ -80,8 +80,9 @@ initialize normNumExt : ScopedEnvExtension Entry (Entry × NormNumExt) NormNums

 /-- Run each registered `norm_num` extension on an expression, returning a `NormNum.Result`. -/
 def derive {α : Q(Type u)} (e : Q($α)) (post := false) : MetaM (Result e) := do
-  if e.isNatLit then
-    let lit : Q(ℕ) := e
+  let e' ← whnfR e
+  if e'.isNatLit then
+    let lit : Q(ℕ) := e'
     return .isNat (q(instAddMonoidWithOneNat) : Q(AddMonoidWithOne ℕ))
       lit (q(IsNat.raw_refl $lit) : Expr)
   profileitM Exception "norm_num" (← getOptions) do

[David Renshaw]

Is that a reasonable change to make? I can prepare a PR if so.

[Mario Carneiro]

no, I think we don't want to whnf. This should probably happen in the norm_num extension that has to handle the lambda

[David Renshaw]

Is your concern that this would be too expensive?

The normNums.tree.getMatch e discrTreeConfig call later down already does some kind of whnf reduction.

For reference, this is the current derive function:

/-- Run each registered `norm_num` extension on an expression, returning a `NormNum.Result`. -/
def derive {α : Q(Type u)} (e : Q($α)) (post := false) : MetaM (Result e) := do
  if e.isNatLit then
    let lit : Q(ℕ) := e
    return .isNat (q(instAddMonoidWithOneNat) : Q(AddMonoidWithOne ℕ))
      lit (q(IsNat.raw_refl $lit) : Expr)
  profileitM Exception "norm_num" (← getOptions) do
    let s ← saveState
    let normNums := normNumExt.getState (← getEnv)
    let arr ← normNums.tree.getMatch e discrTreeConfig
    for ext in arr do
      if (bif post then ext.post else ext.pre) && ! normNums.erased.contains ext.name then
        try
          let new ← withReducibleAndInstances <| ext.eval e
          trace[Tactic.norm_num] "{ext.name}:\n{e} ==> {new}"
          return new
        catch err =>
          trace[Tactic.norm_num] "{e} failed: {err.toMessageData}"
          s.restore
    throwError "{e}: no norm_nums apply"

[David Renshaw]

This should probably happen in the norm_num extension that has to handle the lambda

I don't understand what you mean. By the time norm_num extensions are seeing this expression, the beta redex has been reduced and the lambda is gone.

[David Renshaw]

Ah, I was wrong -- apparently the beta redex does not always go away. I tried writing an extension like this:

@[norm_num _] def evalNatLiteral : NormNumExt where eval {u α} e := do
  guard e.isNatLit
  let sα : Q(AddMonoidWithOne ℕ) := q(inferInstance)
  let n : Q(ℕ) := e
  let pa' ← mkAppM ``IsNat.mk #[q(ℕ), sα, e, e, q(@rfl Type ℕ)]
  return .isNat sα n pa'

and this fails to solve the problem because it sees e as the un-reduced (fun (i : Nat) => i) 9.

[David Renshaw]

aha! This seems to work:

@[norm_num _] def evalBetaRedex : NormNumExt where eval {u α} e := do
  let .app (.lam _ _ _ _) _ := e | throwError "not beta redex"
  let e' ← Core.betaReduce e
  derive (u := u) (α := α) e'

I would prefer to something like @[norm_num (fun _ => _ ) _] , but that still seems to match everything, and so would still need the let .app (.lam ...) ... line.

[David Renshaw]

Another way to achieve basically the same thing:

diff --git a/Mathlib/Tactic/NormNum/Core.lean b/Mathlib/Tactic/NormNum/Core.lean
index 6ebf7ed63e..adf7f96b2c 100644
--- a/Mathlib/Tactic/NormNum/Core.lean
+++ b/Mathlib/Tactic/NormNum/Core.lean
@@ -80,6 +80,12 @@ initialize normNumExt : ScopedEnvExtension Entry (Entry × NormNumExt) NormNums

 /-- Run each registered `norm_num` extension on an expression, returning a `NormNum.Result`. -/
 def derive {α : Q(Type u)} (e : Q($α)) (post := false) : MetaM (Result e) := do
+  let mut e := e
+  while true do
+    if let .app (.lam _ _ _ _) _ := e
+    then e ← Core.betaReduce e
+    else break
+
   if e.isNatLit then

[Mario Carneiro]

I don't want to have reduction there both because it is slow and also because norm_num extensions might want to spot the pattern even though it normalizes. The whole idea is that norm_num tries to use more efficient reduction strategies than kernel computation because it has more domain knowledge.

[David Renshaw]

yeah, I understand now why we don't want to whnf at the beginning of NormNum.derive

[Mario Carneiro]

In this case, I'm saying that you should back up one step to find out who put a beta redex in your term and deal with the issue there

[Mario Carneiro]

in this case I think the culprit is the norm_num extension for sums

[David Renshaw]

ok, I'll look at that

[David Renshaw]

#9446