/-
Copyright (c) 2021 Gabriel Ebner. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Gabriel Ebner, Mario Carneiro
-/
import Std.Tactic.RCases
import Std.Lean.Command

namespace Std.Tactic.Ext
open Lean Meta

/-- `declare_ext_theorems_for A` declares the extensionality theorems for the structure `A`. -/
syntax "declare_ext_theorems_for" ident: Parser.Parser
ident prio: Parser.Category
prio ? : command: Parser.Category
command

/-- Information about an extensionality theorem, stored in the environment extension. -/
structure ExtTheorem: Type
ExtTheorem where
  /-- Declaration name of the extensionality theorem. -/
  declName: ExtTheorem → Name
declName : Name: Type
Name
  /-- Priority of the extensionality theorem. -/
  priority: ExtTheorem → Nat
priority : Nat: Type
Nat
  /-- Key in the discrimination tree. -/
  keys: ExtTheorem → Array (DiscrTree.Key true)
keys : Array: Type ?u.68 → Type ?u.68
Array (DiscrTree.Key: Bool → Type
DiscrTree.Key true: Bool
true)
  deriving Inhabited: Sort u → Sort (max1u)
Inhabited, Repr: Type u → Type u
Repr, BEq: Type u → Type u
BEq, Hashable: Sort u → Sort (max1u)
Hashable

/-- The environment extension to track `@[ext]` lemmas. -/
initialize extExtension: SimpleScopedEnvExtension ExtTheorem (DiscrTree ExtTheorem true)
extExtension :
    SimpleScopedEnvExtension: Type → Type → Type
SimpleScopedEnvExtension ExtTheorem: Type
ExtTheorem (DiscrTree: Type → Bool → Type
DiscrTree ExtTheorem: Type
ExtTheorem true: Bool
true) ←
  registerSimpleScopedEnvExtension: {α σ : Type} → SimpleScopedEnvExtension.Descr α σ → IO (SimpleScopedEnvExtension α σ)
registerSimpleScopedEnvExtension {
    addEntry := fun dt: ?m.3411
dt thm: ?m.3414
thm => dt: ?m.3411
dt.insertCore: {α : Type} → {s : Bool} → [inst : BEq α] → DiscrTree α s → Array (DiscrTree.Key s) → α → DiscrTree α s
insertCore thm: ?m.3414
thm.keys: ExtTheorem → Array (DiscrTree.Key true)
keys thm: ?m.3414
thm
    initial := {}: DiscrTree ?m.3458 ?m.3459
{}
  }

/-- Get the list of `@[ext]` lemmas corresponding to the key `ty`. -/
@[inline] def getExtLemmas: Expr → MetaM (Array ExtTheorem)
getExtLemmas (ty: Expr
ty : Expr: Type
Expr) : MetaM: Type → Type
MetaM (Array: Type ?u.6695 → Type ?u.6695
Array ExtTheorem: Type
ExtTheorem) :=
  return (← (extExtension.getState (← getEnv)).getMatch ty): ?m.6848
(← (extExtension: SimpleScopedEnvExtension ExtTheorem (DiscrTree ExtTheorem true)
extExtension(← (extExtension.getState (← getEnv)).getMatch ty): ?m.6848
.getState: {σ α β : Type} → [inst : Inhabited σ] → ScopedEnvExtension α β σ → Environment → σ
getState(← (extExtension.getState (← getEnv)).getMatch ty): ?m.6848
 (← getEnv): ?m.6764
(← getEnv: {m : Type → Type} → [self : MonadEnv m] → m Environment
getEnv(← getEnv): ?m.6764
)(← (extExtension.getState (← getEnv)).getMatch ty): ?m.6848
).getMatch: {α : Type} → {s : Bool} → DiscrTree α s → Expr → MetaM (Array α)
getMatch(← (extExtension.getState (← getEnv)).getMatch ty): ?m.6848
 ty: Expr
ty(← (extExtension.getState (← getEnv)).getMatch ty): ?m.6848
)
    |>.insertionSort: {α : Type ?u.6875} → Array α → (α → α → Bool) → Array α
insertionSort fun a: ?m.6882
a b: ?m.6885
b => a: ?m.6882
a.priority: ExtTheorem → Nat
priority > b: ?m.6885
b.priority: ExtTheorem → Nat
priority

/-- Registers an extensionality lemma.

When `@[ext]` is applied to a structure,
it generates `.ext` and `.ext_iff` theorems
and registers them for the `ext` tactic.

When `@[ext]` is applied to a theorem,
the theorem is registered for the `ext` tactic.

You can use `@[ext 9000]` to specify a priority for the attribute. -/
syntax (name := ext: ParserDescr
ext) "ext" prio: Parser.Category
prio ? : attr: Parser.Category
attr

initialize registerBuiltinAttribute: AttributeImpl → IO Unit
registerBuiltinAttribute {
  name := `ext: Name
`ext
  descr := "Marks a lemma as extensionality lemma": String
"Marks a lemma as extensionality lemma"
  add := fun declName: ?m.8013
declName stx: ?m.8016
stx kind: ?m.8019
kind => do
    let `(attr: Parser.Category
attr| ext $[$prio: ?m.8073 x✝
prio]?) := stx: ?m.8016
stx | throwError "unexpected @[ext] attribute {stx: ?m.8016
stx}"
    if isStructure: Environment → Name → Bool
isStructure (← getEnv): ?m.8487
(← getEnv: {m : Type → Type} → [self : MonadEnv m] → m Environment
getEnv(← getEnv): ?m.8487
) declName: ?m.8013
declName then
      liftCommandElabM: {α : Type} → Elab.Command.CommandElabM α → CoreM α
liftCommandElabM <| Elab.Command.elabCommand: Syntax → Elab.Command.CommandElabM Unit
Elab.Command.elabCommand <|
        ← `(declare_ext_theorems_for $(mkCIdentFrom: Syntax → Name → optParam Bool false → Ident
mkCIdentFrom stx: ?m.8016
stx declName: ?m.8013
declName) $[$prio: ?m.8073 x✝
prio]?)
    else MetaM.run': {α : Type} →
  MetaM α →
    optParam Context
        {
          config :=
            { foApprox := false, ctxApprox := false, quasiPatternApprox := false, constApprox := false,
              isDefEqStuckEx := false, transparency := TransparencyMode.default, zetaNonDep := true, trackZeta := false,
              unificationHints := true, proofIrrelevance := true, assignSyntheticOpaque := false, offsetCnstrs := true,
              etaStruct := EtaStructMode.all },
          lctx :=
            {
              fvarIdToDecl :=
                { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
              decls :=
                { root := PersistentArrayNode.node (Array.mkEmpty (USize.toNat PersistentArray.branching)),
                  tail := Array.mkEmpty (USize.toNat PersistentArray.branching), size := 0,
                  shift := PersistentArray.initShift, tailOff := 0 } },
          localInstances := #[], defEqCtx? := none, synthPendingDepth := 0, canUnfold? := none } →
      optParam State
          {
            mctx :=
              { depth := 0, levelAssignDepth := 0, mvarCounter := 0,
                lDepth := { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                decls := { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                userNames :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                lAssignment :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                eAssignment :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                dAssignment :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 } },
            cache :=
              {
                inferType :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                funInfo := { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                synthInstance :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                whnfDefault :=
                  { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                whnfAll := { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 },
                defEq := { root := PersistentHashMap.Node.entries PersistentHashMap.mkEmptyEntriesArray, size := 0 } },
            zetaFVarIds := ∅,
            postponed :=
              { root := PersistentArrayNode.node (Array.mkEmpty (USize.toNat PersistentArray.branching)),
                tail := Array.mkEmpty (USize.toNat PersistentArray.branching), size := 0,
                shift := PersistentArray.initShift, tailOff := 0 } } →
        CoreM α
MetaM.run' do
      let declTy: ?m.9438
declTy := (← getConstInfo declName): ?m.9435
(← getConstInfo: {m : Type → Type} → [inst : Monad m] → [inst : MonadEnv m] → [inst : MonadError m] → Name → m ConstantInfo
getConstInfo(← getConstInfo declName): ?m.9435
 declName: ?m.8013
declName(← getConstInfo declName): ?m.9435
).type: ConstantInfo → Expr
type
      let (_, _, declTy: Expr
declTy) ← withDefault: {n : Type → Type ?u.9477} → [inst : MonadControlT MetaM n] → [inst : Monad n] → {α : Type} → n α → n α
withDefault <| forallMetaTelescopeReducing: Expr →
  optParam (Option Nat) none → optParam MetavarKind MetavarKind.natural → MetaM (Array Expr × Array BinderInfo × Expr)
forallMetaTelescopeReducing declTy: ?m.9438
declTy
      let failNotEq: ?m.9620
failNotEq := throwError
        "@[ext] attribute only applies to structures or lemmas proving x = y, got {declTy: Expr
declTy}"
      let some: {α : Type ?u.9833} → α → Option α
some (ty: Expr
ty, lhs: Expr
lhs, rhs: Expr
rhs) := declTy: Expr
declTy.eq?: Expr → Option (Expr × Expr × Expr)
eq? | failNotEq: ?m.9620
failNotEq
      unless lhs: Expr
lhs.isMVar: Expr → Bool
isMVar && rhs: Expr
rhs.isMVar: Expr → Bool
isMVar do failNotEq: MetaM (?m.9883 y✝)
failNotEq
      let keys: ?m.10223
keys ← withReducible: {n : Type → Type ?u.10166} → [inst : MonadControlT MetaM n] → [inst : Monad n] → {α : Type} → n α → n α
withReducible <| DiscrTree.mkPath: {s : Bool} → Expr → MetaM (Array (DiscrTree.Key s))
DiscrTree.mkPath ty: Expr
ty
      let priority: ?m.10959
priority ← liftCommandElabM: {α : Type} → Elab.Command.CommandElabM α → CoreM α
liftCommandElabM do Elab.liftMacroM: {m : Type → Type} →
  {α : Type} →
    [inst : Monad m] →
      [inst : Elab.MonadMacroAdapter m] →
        [inst : MonadEnv m] →
          [inst : MonadRecDepth m] →
            [inst : MonadError m] →
              [inst : MonadResolveName m] →
                [inst : MonadTrace m] →
                  [inst : MonadOptions m] → [inst : AddMessageContext m] → [inst : MonadLiftT IO m] → MacroM α → m α
Elab.liftMacroM do
        evalPrio: Syntax → MacroM Nat
evalPrio (prio: ?m.8073 x✝
prio.getD: {α : Type ?u.10712} → Option α → α → α
getD (← `(prio| default)): ?m.10708
(← `(prio: Parser.Category
prio(← `(prio| default)): ?m.10708
| (← `(prio| default)): ?m.10708
default(← `(prio| default)): ?m.10708
)))
      extExtension: SimpleScopedEnvExtension ExtTheorem (DiscrTree ExtTheorem true)
extExtension.add: {m : Type → Type} →
  {α β σ : Type} →
    [inst : Monad m] →
      [inst : MonadResolveName m] →
        [inst : MonadEnv m] → ScopedEnvExtension α β σ → β → optParam AttributeKind AttributeKind.global → m Unit
add {declName: ?m.8013
declName, keys, priority} kind: ?m.8019
kind
}