feat: linear-size noConfusionType construction (#8037)

This PR introduces a `noConfusionType` construction that’s sub-quadratic
in size, and reduces faster.

The previous `noConfusion` construction with two nested `match`
statements is quadratic in size and reduction behavior. Using some
helper definitions, a linear size construction is possible.

With this, processing the RISC-V-AST definition from
https://github.com/opencompl/sail-riscv-lean takes 6s instead of 60s.

The previous construction is still used when processing the early
prelude, and can be enabled elsewhere using `set_option
backwards.linearNoConfusionType false`.

Author: nomeata

src/Lean/Meta/Constructions/NoConfusion.lean

@@ -7,6 +7,13 @@ prelude
 import Lean.AddDecl
 import Lean.Meta.AppBuilder
 import Lean.Meta.CompletionName
+import Lean.Meta.Constructions.NoConfusionLinear
+
+
+register_builtin_option backwards.linearNoConfusionType : Bool := {
+  defValue := true
+  descr    := "use the linear-size construction for the `noConfusionType` declaration of an inductive type. Set to false to use the previous, simpler but quadratic-size construction. "
+}
 
 namespace Lean
 
@@ -21,12 +28,23 @@ def mkNoConfusionCore (declName : Name) : MetaM Unit := do
   let recInfo ← getConstInfo (mkRecName declName)
   unless recInfo.levelParams.length > indVal.levelParams.length do return
 
-  let name := Name.mkStr declName "noConfusionType"
-  let decl ← ofExceptKernelException (mkNoConfusionTypeCoreImp (← getEnv) declName)
-  addDecl decl
-  setReducibleAttribute name
-  modifyEnv fun env => addToCompletionBlackList env name
-  modifyEnv fun env => addProtected env name
+  let useLinear ←
+    if backwards.linearNoConfusionType.get (← getOptions) then
+      NoConfusionLinear.deps.allM (hasConst · (skipRealize := true))
+    else
+      pure false
+
+  if useLinear then
+    NoConfusionLinear.mkWithCtorType declName
+    NoConfusionLinear.mkWithCtor declName
+    NoConfusionLinear.mkNoConfusionTypeLinear declName
+  else
+    let name := Name.mkStr declName "noConfusionType"
+    let decl ← ofExceptKernelException (mkNoConfusionTypeCoreImp (← getEnv) declName)
+    addDecl decl
+    setReducibleAttribute name
+    modifyEnv fun env => addToCompletionBlackList env name
+    modifyEnv fun env => addProtected env name
 
   let name := Name.mkStr declName "noConfusion"
   let decl ← ofExceptKernelException (mkNoConfusionCoreImp (← getEnv) declName)

src/Lean/Meta/Constructions/NoConfusionLinear.lean

@@ -0,0 +1,226 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+prelude
+import Lean.AddDecl
+import Lean.Meta.AppBuilder
+import Lean.Meta.CompletionName
+
+/-!
+This module produces a construction for the `noConfusionType` that is linear in size in the number of
+constructors of the inductive type. This is in contrast to the previous construction (definde in
+`no_confusion.cpp`), that is quadratic in size due to nested `.brecOn` applications.
+
+We still use the old construction when processing the prelude, for the few inductives that we need
+until below (`Nat`, `Bool`, `Decidable`).
+
+The main trick is to use a `withCtor` helper that is like a match with one constructor pattern and
+one catch-all pattern, and thus linear in size. And the helper itself is a single function
+definition, rather than one for each constructor, using a `withCtorType` helper in the function.
+
+See the `linearNoConfusion.lean` test for exemplary output of this translation (checked to be
+up-to-date).
+
+The `withCtor` functions could be generally useful, but for that they should probably eliminate into
+`Sort _` rather than `Type _`, and then writing the `withCtorType` function runs into universe level
+confusion, which may be solvable if the kernel had more complete univere level normalization.
+Until then we put these helper in the `noConfusionType` namespace to indicate that they are not
+general purpose.
+
+This module is written in a rather manual style, constructing the `Expr` directly. It's best
+read with the expected output to the side.
+-/
+
+namespace Lean.NoConfusionLinear
+
+open Meta
+
+/--
+List of constants that the linear `noConfusionType` construction depends on.
+-/
+def deps : Array Lean.Name :=
+  #[ ``Nat.lt, ``cond, ``Nat, ``PUnit, ``Eq, ``Not, ``dite, ``Nat.decEq, ``Nat.blt ]
+
+def mkNatLookupTable (n : Expr) (es : Array Expr) (default : Expr) : MetaM Expr := do
+  let type ← inferType default
+  let u ← getLevel type
+  let rec go (start stop : Nat) (hstart : start < stop := by omega) (hstop : stop ≤ es.size := by omega) : MetaM Expr := do
+    if h : start + 1 = stop then
+      return es[start]
+    else
+      let mid := (start + stop) / 2
+      let low ← go start mid
+      let high ← go mid stop
+      return mkApp4 (mkConst ``cond [u]) type (mkApp2 (mkConst ``Nat.blt) n (mkRawNatLit mid)) low high
+  if h : es.size = 0 then
+    pure default
+  else
+    go 0 es.size
+
+def mkWithCtorTypeName (indName : Name) : Name :=
+  Name.str indName "noConfusionType" |>.str "withCtorType"
+
+def mkWithCtorName (indName : Name) : Name :=
+  Name.str indName "noConfusionType" |>.str "withCtor"
+
+def mkNoConfusionTypeName (indName : Name) : Name :=
+  Name.str indName "noConfusionType"
+
+def mkWithCtorType (indName : Name) : MetaM Unit := do
+  let ConstantInfo.inductInfo info ← getConstInfo indName | unreachable!
+  let casesOnName := mkCasesOnName indName
+  let casesOnInfo ← getConstVal casesOnName
+  let v::us := casesOnInfo.levelParams.map mkLevelParam | panic! "unexpected universe levels on `casesOn`"
+  let indTyCon := mkConst indName us
+  let indTyKind ← inferType indTyCon
+  let indLevel ← getLevel indTyKind
+  let e ← forallBoundedTelescope indTyKind info.numParams fun xs  _ => do
+    withLocalDeclD `P (mkSort v.succ) fun P => do
+    withLocalDeclD `ctorIdx (mkConst ``Nat) fun ctorIdx => do
+      let default ← mkArrow (mkConst ``PUnit [indLevel]) P
+      let es ← info.ctors.toArray.mapM fun ctorName => do
+        let ctor := mkAppN (mkConst ctorName us) xs
+        let ctorType ← inferType ctor
+        let argType ← forallTelescope ctorType fun ys _ =>
+          mkForallFVars ys P
+        mkArrow (mkConst ``PUnit [indLevel]) argType
+      let e ← mkNatLookupTable ctorIdx es default
+      mkLambdaFVars ((xs.push P).push ctorIdx) e
+
+  let declName := mkWithCtorTypeName indName
+  addAndCompile (.defnDecl (← mkDefinitionValInferrringUnsafe
+    (name        := declName)
+    (levelParams := casesOnInfo.levelParams)
+    (type        := (← inferType e))
+    (value       := e)
+    (hints       := ReducibilityHints.abbrev)
+  ))
+  modifyEnv fun env => addToCompletionBlackList env declName
+  modifyEnv fun env => addProtected env declName
+  setReducibleAttribute declName
+
+def mkWithCtor (indName : Name) : MetaM Unit := do
+  let ConstantInfo.inductInfo info ← getConstInfo indName | unreachable!
+  let withCtorTypeName := mkWithCtorTypeName indName
+  let casesOnName := mkCasesOnName indName
+  let casesOnInfo ← getConstVal casesOnName
+  let v::us := casesOnInfo.levelParams.map mkLevelParam | panic! "unexpected universe levels on `casesOn`"
+  let indTyCon := mkConst indName us
+  let indTyKind ← inferType indTyCon
+  let indLevel ← getLevel indTyKind
+  let e ← forallBoundedTelescope indTyKind info.numParams fun xs t => do
+    withLocalDeclD `P (mkSort v.succ) fun P => do
+    withLocalDeclD `ctorIdx (mkConst ``Nat) fun ctorIdx => do
+      let withCtorTypeNameApp := mkAppN (mkConst withCtorTypeName (v :: us)) (xs.push P)
+      let kType := mkApp withCtorTypeNameApp  ctorIdx
+      withLocalDeclD `k kType fun k =>
+      withLocalDeclD `k' P fun k' =>
+      forallBoundedTelescope t info.numIndices fun ys t' => do
+        let t' ← whnfD t'
+        assert! t'.isSort
+        withLocalDeclD `x (mkAppN indTyCon (xs ++ ys)) fun x => do
+          let e := mkConst (mkCasesOnName indName) (v.succ :: us)
+          let e := mkAppN e xs
+          let motive ← mkLambdaFVars (ys.push x) P
+          let e := mkApp e motive
+          let e := mkAppN e ys
+          let e := mkApp e x
+          let alts ← info.ctors.toArray.mapIdxM fun i ctorName => do
+            let ctor := mkAppN (mkConst ctorName us) xs
+            let ctorType ← inferType ctor
+            forallTelescope ctorType fun zs _ => do
+              let heq := mkApp3 (mkConst ``Eq [1]) (mkConst ``Nat) ctorIdx (mkRawNatLit i)
+              let «then» ← withLocalDeclD `h heq fun h => do
+                let e ← mkEqNDRec (motive := withCtorTypeNameApp) k h
+                let e := mkApp e (mkConst ``PUnit.unit [indLevel])
+                let e := mkAppN e zs
+                -- ``Eq.ndrec
+                mkLambdaFVars #[h] e
+              let «else» ← withLocalDeclD `h (mkNot heq) fun h =>
+                mkLambdaFVars #[h] k'
+              let alt := mkApp5 (mkConst ``dite [v.succ])
+                  P heq (mkApp2 (mkConst ``Nat.decEq) ctorIdx (mkRawNatLit i))
+                  «then» «else»
+              mkLambdaFVars zs alt
+          let e := mkAppN e alts
+          mkLambdaFVars (xs ++ #[P, ctorIdx, k, k'] ++ ys ++ #[x]) e
+
+  let declName := mkWithCtorName indName
+  -- not compiled to avoid old code generator bug #1774
+  addDecl (.defnDecl (← mkDefinitionValInferrringUnsafe
+    (name        := declName)
+    (levelParams := casesOnInfo.levelParams)
+    (type        := (← inferType e))
+    (value       := e)
+    (hints       := ReducibilityHints.abbrev)
+  ))
+  modifyEnv fun env => addToCompletionBlackList env declName
+  modifyEnv fun env => addProtected env declName
+  setReducibleAttribute declName
+
+def mkNoConfusionTypeLinear (indName : Name) : MetaM Unit := do
+  let declName := mkNoConfusionTypeName indName
+  let ConstantInfo.inductInfo info ← getConstInfo indName | unreachable!
+  let casesOnName := mkCasesOnName indName
+  let casesOnInfo ← getConstVal casesOnName
+  let v::us := casesOnInfo.levelParams.map mkLevelParam | panic! "unexpected universe levels on `casesOn`"
+  let e := mkConst casesOnName (v.succ::us)
+  let t ← inferType e
+  let e ← forallBoundedTelescope t info.numParams fun xs t => do
+    let e := mkAppN e xs
+    let PType := mkSort v
+    withLocalDeclD `P PType fun P => do
+      let motive ← forallTelescope (← whnfD t).bindingDomain! fun ys _ =>
+        mkLambdaFVars ys PType
+      let t ← instantiateForall t #[motive]
+      let e := mkApp e motive
+      forallBoundedTelescope t info.numIndices fun ys t => do
+        let e := mkAppN e ys
+        let xType := mkAppN (mkConst indName us) (xs ++ ys)
+        withLocalDeclD `x1 xType fun x1 => do
+        withLocalDeclD `x2 xType fun x2 => do
+          let t ← instantiateForall t #[x1]
+          let e := mkApp e x1
+          forallBoundedTelescope t info.numCtors fun alts _ => do
+            let alts' ← alts.mapIdxM fun i alt => do
+              let altType ← inferType alt
+              forallTelescope altType fun zs1 _ => do
+                let alt := mkConst (mkWithCtorName indName) (v :: us)
+                let alt := mkAppN alt xs
+                let alt := mkApp alt PType
+                let alt := mkApp alt (mkRawNatLit i)
+                let k ← forallTelescopeReducing (← inferType alt).bindingDomain! fun zs2 _ => do
+                  let eqs ← (Array.zip zs1 zs2[1:]).filterMapM fun (z1,z2) => do
+                    if (← isProof z1) then
+                      return none
+                    else
+                      return some (← mkEqHEq z1 z2)
+                  let k ← mkArrowN eqs P
+                  let k ← mkArrow k P
+                  mkLambdaFVars zs2 k
+                let alt := mkApp alt k
+                let alt := mkApp alt P
+                let alt := mkAppN alt ys
+                let alt := mkApp alt x2
+                mkLambdaFVars zs1 alt
+            let e := mkAppN e alts'
+            let e ← mkLambdaFVars #[x1, x2] e
+            let e ← mkLambdaFVars #[P] e
+            let e ← mkLambdaFVars ys e
+            let e ← mkLambdaFVars xs e
+            pure e
+
+  addDecl (.defnDecl (← mkDefinitionValInferrringUnsafe
+    (name        := declName)
+    (levelParams := casesOnInfo.levelParams)
+    (type        := (← inferType e))
+    (value       := e)
+    (hints       := ReducibilityHints.abbrev)
+  ))
+  modifyEnv fun env => addToCompletionBlackList env declName
+  modifyEnv fun env => addProtected env declName
+  setReducibleAttribute declName
+
+end Lean.NoConfusionLinear

tests/lean/run/815.lean

@@ -1,14 +1,16 @@
-def is_smooth {α β} (f : α → β) : Prop := sorry
+axiom testSorry : α
+
+def is_smooth {α β} (f : α → β) : Prop := testSorry
 
 class IsSmooth {α β} (f : α → β) : Prop where
   (proof : is_smooth f)
 
-instance identity : IsSmooth fun a : α => a := sorry
-instance const (b : β) : IsSmooth fun a : α => b := sorry
-instance swap (f : α → β → γ) [∀ a, IsSmooth (f a)] : IsSmooth (λ b a => f a b) := sorry
-instance parm (f : α → β → γ) [IsSmooth f] (b : β) : IsSmooth (λ a => f a b) := sorry
-instance comp (f : β → γ) (g : α → β) [IsSmooth f] [IsSmooth g] : IsSmooth (fun a => f (g a)) := sorry
-instance diag (f : β → δ → γ) (g : α → β) (h : α → δ) [IsSmooth f] [∀ b, IsSmooth (f b)] [IsSmooth g] [IsSmooth h] : IsSmooth (λ a => f (g a) (h a)) := sorry
+instance identity : IsSmooth fun a : α => a := testSorry
+instance const (b : β) : IsSmooth fun a : α => b := testSorry
+instance swap (f : α → β → γ) [∀ a, IsSmooth (f a)] : IsSmooth (λ b a => f a b) := testSorry
+instance parm (f : α → β → γ) [IsSmooth f] (b : β) : IsSmooth (λ a => f a b) := testSorry
+instance comp (f : β → γ) (g : α → β) [IsSmooth f] [IsSmooth g] : IsSmooth (fun a => f (g a)) := testSorry
+instance diag (f : β → δ → γ) (g : α → β) (h : α → δ) [IsSmooth f] [∀ b, IsSmooth (f b)] [IsSmooth g] [IsSmooth h] : IsSmooth (λ a => f (g a) (h a)) := testSorry
 
 example (f : β → δ → γ) [IsSmooth f] (g : α → β) [IsSmooth g] (d : δ) : IsSmooth (λ a => f (g a) d) := by infer_instance
 example (f : β → δ → γ) [IsSmooth f] (g : α → β) [IsSmooth g] : IsSmooth (λ a d => f (g a) d) := by infer_instance

tests/lean/run/grind_pre.lean

@@ -108,8 +108,8 @@ bs : List Point
 b₂ : Nat
 b₃ : Int
 head_eq : a₁ = b₁
-x_eq : a₂ = b₂
-y_eq : a₃ = b₃
+h_1 : a₂ = b₂
+h_2 : a₃ = b₃
 tail_eq_1 : as = bs
 ⊢ False
 [grind] Goal diagnostics

tests/lean/run/linearNoConfusion.lean

@@ -0,0 +1,86 @@
+/-!
+This test tests and also explains the noConfusionType construction.
+
+It contains copies of the definitions of the constructions, for manual experimentation with
+the code, and uses `#guard_msgs` and `rfl` to compare them to the generated ones.
+
+This also serves as documentation to the `NoConfusionLinear` module, so do not delete or remove
+from this file without updating that module's docstring.
+-/
+
+inductive Vec.{u} (α : Type) : Nat → Type u where
+  | nil : Vec α 0
+  | cons {n} : α → Vec α n → Vec α (n + 1)
+
+@[reducible] protected def Vec.noConfusionType.withCtorType'.{u_1, u} :
+    Type → Type u_1 → Nat → Type (max (u + 1) u_1) := fun α P ctorIdx =>
+  bif Nat.blt ctorIdx 1
+  then PUnit.{u + 2} → P
+  else PUnit.{u + 2} → {n : Nat} → α → Vec.{u} α n → P
+
+/--
+info: @[reducible] protected def Vec.noConfusionType.withCtorType.{u_1, u} : Type → Type u_1 → Nat → Type (max (u + 1) u_1) :=
+fun α P ctorIdx => bif ctorIdx.blt 1 then PUnit → P else PUnit → {n : Nat} → α → Vec α n → P
+-/
+#guard_msgs in
+#print Vec.noConfusionType.withCtorType
+
+example : @Vec.noConfusionType.withCtorType.{u_1,u} = @Vec.noConfusionType.withCtorType'.{u_1,u} := rfl
+
+@[reducible] protected noncomputable def Vec.noConfusionType.withCtor'.{u_1, u} : (α : Type) →
+  (P : Type u_1) → (ctorIdx : Nat) → Vec.noConfusionType.withCtorType' α P ctorIdx → P → (a : Nat) → Vec.{u} α a → P :=
+fun _α _P ctorIdx k k' _a x =>
+  Vec.casesOn x
+    (if h : ctorIdx = 0 then Eq.ndrec k h PUnit.unit else k')
+    (fun a a_1 => if h : ctorIdx = 1 then Eq.ndrec k h PUnit.unit a a_1 else k')
+
+/--
+info: @[reducible] protected def Vec.noConfusionType.withCtor.{u_1, u} : (α : Type) →
+  (P : Type u_1) → (ctorIdx : Nat) → Vec.noConfusionType.withCtorType α P ctorIdx → P → (a : Nat) → Vec α a → P :=
+fun α P ctorIdx k k' a x =>
+  Vec.casesOn x (if h : ctorIdx = 0 then (h ▸ k) PUnit.unit else k') fun {n} a a_1 =>
+    if h : ctorIdx = 1 then (h ▸ k) PUnit.unit a a_1 else k'
+-/
+#guard_msgs in
+#print Vec.noConfusionType.withCtor
+
+example : @Vec.noConfusionType.withCtor.{u_1,u} = @Vec.noConfusionType.withCtor'.{u_1,u} := rfl
+
+@[reducible] protected def Vec.noConfusionType'.{u_1, u} : {α : Type} →
+  {a : Nat} → Sort u_1 → Vec.{u} α a → Vec α a → Sort u_1 :=
+fun {α} {a} P x1 x2 =>
+  Vec.casesOn x1
+    (Vec.noConfusionType.withCtor' α (Sort u_1) 0 (fun _x => P → P) P a x2)
+    (fun {n} a_1 a_2 => Vec.noConfusionType.withCtor' α (Sort u_1) 1 (fun _x {n_1} a a_3 => (n = n_1 → a_1 = a → HEq a_2 a_3 → P) → P) P a x2)
+
+/--
+info: @[reducible] protected def Vec.noConfusionType.{u_1, u} : {α : Type} →
+  {a : Nat} → Sort u_1 → Vec α a → Vec α a → Sort u_1 :=
+fun {α} {a} P x1 x2 =>
+  Vec.casesOn x1 (Vec.noConfusionType.withCtor α (Sort u_1) 0 (fun x => P → P) P a x2) fun {n} a_1 a_2 =>
+    Vec.noConfusionType.withCtor α (Sort u_1) 1 (fun x {n_1} a a_3 => (n = n_1 → a_1 = a → HEq a_2 a_3 → P) → P) P a x2
+-/
+#guard_msgs in
+#print Vec.noConfusionType
+
+example : @Vec.noConfusionType.{u_1,u} = @Vec.noConfusionType'.{u_1,u} := rfl
+
+/-
+run_meta do
+  let mut i := 0
+  for (n, _c) in (← getEnv).constants do
+    if let .str indName "noConfusion" := n then
+      let ConstantInfo.inductInfo _ ← getConstInfo indName | continue
+      logInfo m!"Looking at {.ofConstName indName}"
+      mkToCtorIdx' indName
+      mkWithCtorType indName
+      mkWithCtor indName
+      mkNoConfusionType' indName
+      i := i + 1
+      if i > 10 then
+        return
+-/
+
+-- inductive Enum.{u} : Type u where | a | b
+-- set_option pp.universes true in
+-- #print noConfusionTypeEnum