/-
Copyright (c) 2022 Andrew Yang. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Andrew Yang
-/
import ring_theory.ideal.cotangent
import ring_theory.dedekind_domain.basic
import ring_theory.valuation.valuation_ring
import ring_theory.nakayama
/-!

# Equivalent conditions for DVR

In `discrete_valuation_ring.tfae`, we show that the following are equivalent for a
noetherian local domain `(R, m, k)`:
- `R` is a discrete valuation ring
- `R` is a valuation ring
- `R` is a dedekind domain
- `R` is integrally closed with a unique prime ideal
- `m` is principal
- `dimₖ m/m² = 1`
- Every nonzero ideal is a power of `m`.

-/

variables (R : Type*) [comm_ring R] (K : Type*) [field K] [algebra R K] [is_fraction_ring R K]

open_locale discrete_valuation
open local_ring

open_locale big_operators

lemma exists_maximal_ideal_pow_eq_of_principal [is_noetherian_ring R] [local_ring R] [is_domain R]
  (h : ¬ is_field R) (h' : (maximal_ideal R).is_principal) (I : ideal R) (hI : I ≠ ⊥) :
    ∃ n : ℕ, I = (maximal_ideal R) ^ n :=
begin
  classical,
  unfreezingI { obtain ⟨x, hx : _ = ideal.span _⟩ := h' },
  by_cases hI' : I = ⊤, { use 0, rw [pow_zero, hI', ideal.one_eq_top] },
  have H : ∀ r : R, ¬ (is_unit r) ↔ x ∣ r :=
    λ r, (set_like.ext_iff.mp hx r).trans ideal.mem_span_singleton,
  have : x ≠ 0,
  { rintro rfl,
    apply ring.ne_bot_of_is_maximal_of_not_is_field (maximal_ideal.is_maximal R) h,
    simp [hx] },
  have hx' := discrete_valuation_ring.irreducible_of_span_eq_maximal_ideal x this hx,
  have H' : ∀ r : R, r ≠ 0 → r ∈ nonunits R → ∃ (n : ℕ), associated (x ^ n) r,
  { intros r hr₁ hr₂,
    obtain ⟨f, hf₁, rfl, hf₂⟩ := (wf_dvd_monoid.not_unit_iff_exists_factors_eq r hr₁).mp hr₂,
    have : ∀ b ∈ f, associated x b,
    { intros b hb,
      exact irreducible.associated_of_dvd hx' (hf₁ b hb) ((H b).mp (hf₁ b hb).1) },
    clear hr₁ hr₂ hf₁,
    induction f using multiset.induction with fa fs fh,
    { exact (hf₂ rfl).elim },
    rcases eq_or_ne fs ∅ with rfl|hf',
    { use 1,
      rw [pow_one, multiset.prod_cons, multiset.empty_eq_zero, multiset.prod_zero, mul_one],
      exact this _ (multiset.mem_cons_self _ _) },
    { obtain ⟨n, hn⟩ := fh hf' (λ b hb, this _ (multiset.mem_cons_of_mem hb)),
      use n + 1,
      rw [pow_add, multiset.prod_cons, mul_comm, pow_one],
      exact associated.mul_mul (this _ (multiset.mem_cons_self _ _)) hn } },
  have : ∃ n : ℕ, x ^ n ∈ I,
  { obtain ⟨r, hr₁, hr₂⟩ : ∃ r : R, r ∈ I ∧ r ≠ 0,
    { by_contra h, push_neg at h, apply hI, rw eq_bot_iff, exact h },
    obtain ⟨n, u, rfl⟩ := H' r hr₂ (le_maximal_ideal hI' hr₁),
    use n,
    rwa [← I.unit_mul_mem_iff_mem u.is_unit, mul_comm] },
  use nat.find this,
  apply le_antisymm,
  { change ∀ s ∈ I, s ∈ _,
    by_contra hI'',
    push_neg at hI'',
    obtain ⟨s, hs₁, hs₂⟩ := hI'',
    apply hs₂,
    by_cases hs₃ : s = 0, { rw hs₃, exact zero_mem _ },
    obtain ⟨n, u, rfl⟩ := H' s hs₃ (le_maximal_ideal hI' hs₁),
    rw [mul_comm, ideal.unit_mul_mem_iff_mem _ u.is_unit] at ⊢ hs₁,
    apply ideal.pow_le_pow (nat.find_min' this hs₁),
    apply ideal.pow_mem_pow,
    exact (H _).mpr (dvd_refl _) },
  { rw [hx, ideal.span_singleton_pow, ideal.span_le, set.singleton_subset_iff],
    exact nat.find_spec this }
end

lemma maximal_ideal_is_principal_of_is_dedekind_domain
  [local_ring R] [is_domain R] [is_dedekind_domain R] : (maximal_ideal R).is_principal :=
begin
  classical,
  by_cases ne_bot : maximal_ideal R = ⊥,
  { rw ne_bot, apply_instance },
  obtain ⟨a, ha₁, ha₂⟩ : ∃ a ∈ maximal_ideal R, a ≠ (0 : R),
  { by_contra h', push_neg at h', apply ne_bot, rwa eq_bot_iff },
  have hle : ideal.span {a} ≤ maximal_ideal R,
  { rwa [ideal.span_le, set.singleton_subset_iff] },
  have : (ideal.span {a}).radical = maximal_ideal R,
  { rw ideal.radical_eq_Inf,
    apply le_antisymm,
    { exact Inf_le ⟨hle, infer_instance⟩ },
    { refine le_Inf (λ I hI, (eq_maximal_ideal $
        is_dedekind_domain.dimension_le_one _ (λ e, ha₂ _) hI.2).ge),
      rw [← ideal.span_singleton_eq_bot, eq_bot_iff, ← e], exact hI.1 } },
  have : ∃ n, maximal_ideal R ^ n ≤ ideal.span {a},
  { rw ← this, apply ideal.exists_radical_pow_le_of_fg, exact is_noetherian.noetherian _ },
  cases hn : nat.find this,
  { have := nat.find_spec this,
    rw [hn, pow_zero, ideal.one_eq_top] at this,
    exact (ideal.is_maximal.ne_top infer_instance (eq_top_iff.mpr $ this.trans hle)).elim },
  obtain ⟨b, hb₁, hb₂⟩ : ∃ b ∈ maximal_ideal R ^ n, ¬ b ∈ ideal.span {a},
  { by_contra h', push_neg at h', rw nat.find_eq_iff at hn,
    exact hn.2 n n.lt_succ_self (λ x hx, not_not.mp (h' x hx)) },
  have hb₃ : ∀ m ∈ maximal_ideal R, ∃ k : R, k * a = b * m,
  { intros m hm, rw ← ideal.mem_span_singleton', apply nat.find_spec this,
    rw [hn, pow_succ'], exact ideal.mul_mem_mul hb₁ hm },
  have hb₄ : b ≠ 0,
  { rintro rfl, apply hb₂, exact zero_mem _ },
  let K := fraction_ring R,
  let x : K := algebra_map R K b / algebra_map R K a,
  let M := submodule.map (algebra.of_id R K).to_linear_map (maximal_ideal R),
  have ha₃ : algebra_map R K a ≠ 0 := is_fraction_ring.to_map_eq_zero_iff.not.mpr ha₂,
  by_cases hx : ∀ y ∈ M, x * y ∈ M,
  { have := is_integral_of_smul_mem_submodule M _ _ x hx,
    { obtain ⟨y, e⟩ := is_integrally_closed.algebra_map_eq_of_integral this,
      refine (hb₂ (ideal.mem_span_singleton'.mpr ⟨y, _⟩)).elim,
      apply is_fraction_ring.injective R K,
      rw [map_mul, e, div_mul_cancel _ ha₃] },
    { rw submodule.ne_bot_iff, refine ⟨_, ⟨a, ha₁, rfl⟩, _⟩,
      exact is_fraction_ring.to_map_eq_zero_iff.not.mpr ha₂ },
    { apply submodule.fg.map, exact is_noetherian.noetherian _ } },
  { have : (M.map (distrib_mul_action.to_linear_map R K x)).comap
      (algebra.of_id R K).to_linear_map = ⊤,
    { by_contra h, apply hx,
      rintros m' ⟨m, hm, (rfl : algebra_map R K m = m')⟩,
      obtain ⟨k, hk⟩ := hb₃ m hm,
      have hk' : x * algebra_map R K m = algebra_map R K k,
      { rw [← mul_div_right_comm, ← map_mul, ← hk, map_mul, mul_div_cancel _ ha₃] },
      exact ⟨k, le_maximal_ideal h ⟨_, ⟨_, hm, rfl⟩, hk'⟩, hk'.symm⟩ },
    obtain ⟨y, hy₁, hy₂⟩ : ∃ y ∈ maximal_ideal R, b * y = a,
    { rw [ideal.eq_top_iff_one, submodule.mem_comap] at this,
      obtain ⟨_, ⟨y, hy, rfl⟩, hy' : x * algebra_map R K y = algebra_map R K 1⟩ := this,
      rw [map_one, ← mul_div_right_comm, div_eq_one_iff_eq ha₃, ← map_mul] at hy',
      exact ⟨y, hy, is_fraction_ring.injective R K hy'⟩ },
    refine ⟨⟨y, _⟩⟩,
    apply le_antisymm,
    { intros m hm, obtain ⟨k, hk⟩ := hb₃ m hm, rw [← hy₂, mul_comm, mul_assoc] at hk,
      rw [← mul_left_cancel₀ hb₄ hk, mul_comm], exact ideal.mem_span_singleton'.mpr ⟨_, rfl⟩ },
    { rwa [submodule.span_le, set.singleton_subset_iff] } }
end

lemma discrete_valuation_ring.tfae [is_noetherian_ring R] [local_ring R] [is_domain R]
  (h : ¬ is_field R) :
  tfae [discrete_valuation_ring R,
    valuation_ring R,
    is_dedekind_domain R,
    is_integrally_closed R ∧ ∃! P : ideal R, P ≠ ⊥ ∧ P.is_prime,
    (maximal_ideal R).is_principal,
    finite_dimensional.finrank (residue_field R) (cotangent_space R) = 1,
    ∀ I ≠ ⊥, ∃ n : ℕ, I = (maximal_ideal R) ^ n] :=
begin
  have ne_bot := ring.ne_bot_of_is_maximal_of_not_is_field (maximal_ideal.is_maximal R) h,
  classical,
  rw finrank_eq_one_iff',
  tfae_have : 1 → 2,
  { introI _, apply_instance },
  tfae_have : 2 → 1,
  { introI _,
    haveI := is_bezout.to_gcd_domain R,
    haveI : unique_factorization_monoid R := ufm_of_gcd_of_wf_dvd_monoid,
    apply discrete_valuation_ring.of_ufd_of_unique_irreducible,
    { obtain ⟨x, hx₁, hx₂⟩ := ring.exists_not_is_unit_of_not_is_field h,
      obtain ⟨p, hp₁, hp₂⟩ := wf_dvd_monoid.exists_irreducible_factor hx₂ hx₁,
      exact ⟨p, hp₁⟩ },
    { exact valuation_ring.unique_irreducible } },
  tfae_have : 1 → 4,
  { introI H,
    exact ⟨infer_instance, ((discrete_valuation_ring.iff_pid_with_one_nonzero_prime R).mp H).2⟩ },
  tfae_have : 4 → 3,
  { rintros ⟨h₁, h₂⟩, exact ⟨infer_instance, λ I hI hI', unique_of_exists_unique h₂
      ⟨ne_bot, infer_instance⟩ ⟨hI, hI'⟩ ▸ maximal_ideal.is_maximal R, h₁⟩ },
  tfae_have : 3 → 5,
  { introI h, exact maximal_ideal_is_principal_of_is_dedekind_domain R },
  tfae_have : 5 → 6,
  { rintro ⟨x, hx⟩,
    have : x ∈ maximal_ideal R := by { rw hx, exact submodule.subset_span (set.mem_singleton x) },
    let x' : maximal_ideal R := ⟨x, this⟩,
    use submodule.quotient.mk x',
    split,
    { intro e,
      rw submodule.quotient.mk_eq_zero at e,
      apply ring.ne_bot_of_is_maximal_of_not_is_field (maximal_ideal.is_maximal R) h,
      apply submodule.eq_bot_of_le_smul_of_le_jacobson_bot (maximal_ideal R),
      { exact ⟨{x}, (finset.coe_singleton x).symm ▸ hx.symm⟩ },
      { conv_lhs { rw hx },
        rw submodule.mem_smul_top_iff at e,
        rwa [submodule.span_le, set.singleton_subset_iff] },
      { rw local_ring.jacobson_eq_maximal_ideal (⊥ : ideal R) bot_ne_top, exact le_refl _ } },
    { refine λ w, quotient.induction_on' w $ λ y, _,
      obtain ⟨y, hy⟩ := y,
      rw [hx, submodule.mem_span_singleton] at hy,
      obtain ⟨a, rfl⟩ := hy,
      exact ⟨ideal.quotient.mk _ a, rfl⟩ } },
  tfae_have : 6 → 5,
  { rintro ⟨x, hx, hx'⟩,
    induction x using quotient.induction_on',
    use x,
    apply le_antisymm,
    swap, { rw [submodule.span_le, set.singleton_subset_iff], exact x.prop },
    have h₁ : (ideal.span {x} : ideal R) ⊔ maximal_ideal R ≤
      ideal.span {x} ⊔ (maximal_ideal R) • (maximal_ideal R),
    { refine sup_le le_sup_left _,
      rintros m hm,
      obtain ⟨c, hc⟩ := hx' (submodule.quotient.mk ⟨m, hm⟩),
      induction c using quotient.induction_on',
      rw ← sub_sub_cancel (c * x) m,
      apply sub_mem _ _,
      { apply_instance },
      { refine ideal.mem_sup_left (ideal.mem_span_singleton'.mpr ⟨c, rfl⟩) },
      { have := (submodule.quotient.eq _).mp hc,
        rw [submodule.mem_smul_top_iff] at this,
        exact ideal.mem_sup_right this } },
    have h₂ : maximal_ideal R ≤ (⊥ : ideal R).jacobson,
    { rw local_ring.jacobson_eq_maximal_ideal, exacts [le_refl _, bot_ne_top] },
    have := submodule.smul_sup_eq_smul_sup_of_le_smul_of_le_jacobson
      (is_noetherian.noetherian _) h₂ h₁,
    rw [submodule.bot_smul, sup_bot_eq] at this,
    rw [← sup_eq_left, eq_comm],
    exact le_sup_left.antisymm (h₁.trans $ le_of_eq this) },
  tfae_have : 5 → 7,
  { exact exists_maximal_ideal_pow_eq_of_principal R h },
  tfae_have : 7 → 2,
  { rw valuation_ring.iff_ideal_total,
    intro H,
    constructor,
    intros I J,
    by_cases hI : I = ⊥, { subst hI,  left, exact bot_le },
    by_cases hJ : J = ⊥, { subst hJ, right, exact bot_le },
    obtain ⟨n, rfl⟩ := H I hI,
    obtain ⟨m, rfl⟩ := H J hJ,
    cases le_total m n with h' h',
    {  left, exact ideal.pow_le_pow h' },
    { right, exact ideal.pow_le_pow h' } },
  tfae_finish,
end


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.Transform
import Lean.Meta.Tactic.Injection
import Lean.Meta.Tactic.Apply
import Lean.Meta.Tactic.Refl
import Lean.Meta.Tactic.Cases
import Lean.Meta.Tactic.Subst
import Lean.Meta.Tactic.Simp.Types
import Lean.Meta.Tactic.Assumption

namespace Lean.Meta

private def mkAnd? (args : Array Expr) : Option Expr := Id.run do
  if args.isEmpty then
    return none
  else
    let mut result := args.back
    for arg in args.reverse[1:] do
      result := mkApp2 (mkConst ``And) arg result
    return result

def elimOptParam (type : Expr) : CoreM Expr := do
  Core.transform type fun e =>
    if e.isAppOfArity  ``optParam 2 then
      return TransformStep.visit (e.getArg! 0)
    else
      return .continue

private partial def mkInjectiveTheoremTypeCore? (ctorVal : ConstructorVal) (useEq : Bool) : MetaM (Option Expr) := do
  let us := ctorVal.levelParams.map mkLevelParam
  let type ← elimOptParam ctorVal.type
  forallBoundedTelescope type ctorVal.numParams fun params type =>
  forallTelescope type fun args1 resultType => do
    let jp (args2 args2New : Array Expr) : MetaM (Option Expr) := do
      let lhs := mkAppN (mkAppN (mkConst ctorVal.name us) params) args1
      let rhs := mkAppN (mkAppN (mkConst ctorVal.name us) params) args2
      let eq ← mkEq lhs rhs
      let mut eqs := #[]
      for arg1 in args1, arg2 in args2 do
        let arg1Type ← inferType arg1
        if !(← isProp arg1Type) && arg1 != arg2 then
          eqs := eqs.push (← mkEqHEq arg1 arg2)
      if let some andEqs := mkAnd? eqs then
        let result ← if useEq then
          mkEq eq andEqs
        else
          mkArrow eq andEqs
        mkForallFVars params (← mkForallFVars args1 (← mkForallFVars args2New result))
      else
        return none
    let rec mkArgs2 (i : Nat) (type : Expr) (args2 args2New : Array Expr) : MetaM (Option Expr) := do
      if h : i < args1.size then
        match (← whnf type) with
        | Expr.forallE n d b _ =>
          let arg1 := args1.get ⟨i, h⟩
          if arg1.occurs resultType then
            mkArgs2 (i + 1) (b.instantiate1 arg1) (args2.push arg1) args2New
          else
            withLocalDecl n (if useEq then BinderInfo.default else BinderInfo.implicit) d fun arg2 =>
              mkArgs2 (i + 1) (b.instantiate1 arg2) (args2.push arg2) (args2New.push arg2)
        | _ => throwError "unexpected constructor type for '{ctorVal.name}'"
      else
        jp args2 args2New
    if useEq then
      mkArgs2 0 type #[] #[]
    else
      withNewBinderInfos (params.map fun param => (param.fvarId!, BinderInfo.implicit)) <|
      withNewBinderInfos (args1.map fun arg1 => (arg1.fvarId!, BinderInfo.implicit)) <|
        mkArgs2 0 type #[] #[]

private def mkInjectiveTheoremType? (ctorVal : ConstructorVal) : MetaM (Option Expr) :=
  mkInjectiveTheoremTypeCore? ctorVal false

private def injTheoremFailureHeader (ctorName : Name) : MessageData :=
  m!"failed to prove injectivity theorem for constructor '{ctorName}', use 'set_option genInjectivity false' to disable the generation"

private def throwInjectiveTheoremFailure {α} (ctorName : Name) (mvarId : MVarId) : MetaM α :=
  throwError "{injTheoremFailureHeader ctorName}{indentD <| MessageData.ofGoal mvarId}"

private def solveEqOfCtorEq (ctorName : Name) (mvarId : MVarId) (h : FVarId) : MetaM Unit := do
  match (← injection mvarId h) with
  | InjectionResult.solved => unreachable!
  | InjectionResult.subgoal mvarId .. =>
    (←  mvarId.splitAnd).forM fun mvarId =>
      unless (← mvarId.assumptionCore) do
        throwInjectiveTheoremFailure ctorName mvarId

private def mkInjectiveTheoremValue (ctorName : Name) (targetType : Expr) : MetaM Expr :=
  forallTelescopeReducing targetType fun xs type => do
    let mvar ← mkFreshExprSyntheticOpaqueMVar type
    solveEqOfCtorEq ctorName mvar.mvarId! xs.back.fvarId!
    mkLambdaFVars xs mvar

def mkInjectiveTheoremNameFor (ctorName : Name) : Name :=
  ctorName ++ `inj

private def mkInjectiveTheorem (ctorVal : ConstructorVal) : MetaM Unit := do
  let some type ← mkInjectiveTheoremType? ctorVal
    | return ()
  let value ← mkInjectiveTheoremValue ctorVal.name type
  let name := mkInjectiveTheoremNameFor ctorVal.name
  addDecl <| Declaration.thmDecl {
    name
    levelParams := ctorVal.levelParams
    type        := (← instantiateMVars type)
    value       := (← instantiateMVars value)
  }

def mkInjectiveEqTheoremNameFor (ctorName : Name) : Name :=
  ctorName ++ `injEq

private def mkInjectiveEqTheoremType? (ctorVal : ConstructorVal) : MetaM (Option Expr) :=
  mkInjectiveTheoremTypeCore? ctorVal true

private def mkInjectiveEqTheoremValue (ctorName : Name) (targetType : Expr) : MetaM Expr := do
  forallTelescopeReducing targetType fun xs type => do
    let mvar ← mkFreshExprSyntheticOpaqueMVar type
    let [mvarId₁, mvarId₂] ← mvar.mvarId!.apply (mkConst ``Eq.propIntro)
      | throwError "unexpected number of subgoals when proving injective theorem for constructor '{ctorName}'"
    let (h, mvarId₁) ← mvarId₁.intro1
    let (_, mvarId₂) ← mvarId₂.intro1
    solveEqOfCtorEq ctorName mvarId₁ h
    let mvarId₂ ← mvarId₂.casesAnd
    if let some mvarId₂ ← mvarId₂.substEqs then
      try mvarId₂.refl catch _ => throwError (injTheoremFailureHeader ctorName)
    mkLambdaFVars xs mvar

private def mkInjectiveEqTheorem (ctorVal : ConstructorVal) : MetaM Unit := do
  let some type ← mkInjectiveEqTheoremType? ctorVal
    | return ()
  let value ← mkInjectiveEqTheoremValue ctorVal.name type
  let name := mkInjectiveEqTheoremNameFor ctorVal.name
  addDecl <| Declaration.thmDecl {
    name
    levelParams := ctorVal.levelParams
    type        := (← instantiateMVars type)
    value       := (← instantiateMVars value)
  }
  addSimpTheorem (ext := simpExtension) name (post := true) (inv := false) AttributeKind.global (prio := eval_prio default)

register_builtin_option genInjectivity : Bool := {
  defValue := true
  descr    := "generate injectivity theorems for inductive datatype constructors"
}

def mkInjectiveTheorems (declName : Name) : MetaM Unit := do
  if (← getEnv).contains ``Eq.propIntro && genInjectivity.get (← getOptions) &&  !(← isInductivePredicate declName) then
    let info ← getConstInfoInduct declName
    unless info.isUnsafe do
      for ctor in info.ctors do
        let ctorVal ← getConstInfoCtor ctor
        if ctorVal.numFields > 0 then
          mkInjectiveTheorem ctorVal
          mkInjectiveEqTheorem ctorVal

end Lean.Meta


import tactic.basic data.list.defs data.prod data.sum tactic.rcases
universes u₁ u₂

open interactive interactive.types
open lean.parser nat tactic

meta def get_ext_subject : expr → tactic name
| (expr.pi n bi d b) :=
  do v  ← mk_local' n bi d,
     b' ← whnf $ b.instantiate_var v,
     get_ext_subject b'
| (expr.app _ e) :=
  do t ← infer_type e >>= instantiate_mvars >>= head_beta,
     if t.get_app_fn.is_constant then
       pure $ t.get_app_fn.const_name
     else if t.is_pi then
       pure $ name.mk_numeral 0 name.anonymous
     else if t.is_sort then
       pure $ name.mk_numeral 1 name.anonymous
     else do
       t ← pp t,
       fail format!"only constants and Pi types are supported: {t}"
| e := fail format!"Only expressions of the form `_ → _ → ... → R ... e are supported: {e}"

open native

@[reducible] def ext_param_type := option name ⊕ option name

meta def opt_minus : lean.parser (option name → ext_param_type) :=
sum.inl <$ tk "-" <|> pure sum.inr

meta def ext_param :=
opt_minus <*> ( name.mk_numeral 0 name.anonymous <$ brackets "(" ")" (tk "→" <|> tk "->") <|>
                none <$  tk "*" <|>
                some <$> ident )

meta def saturate_fun : name → tactic expr
| (name.mk_numeral 0 name.anonymous) :=
do v₀ ← mk_mvar,
   v₁ ← mk_mvar,
   return $ v₀.imp v₁
| (name.mk_numeral 1 name.anonymous) :=
do u ← mk_meta_univ,
   pure $ expr.sort u
| n :=
do e ← resolve_constant n >>= mk_const,
   a ← get_arity e,
   e.mk_app <$> (list.iota a).mmap (λ _, mk_mvar)

meta def equiv_type_constr (n n' : name) : tactic unit :=
do e  ← saturate_fun n,
   e' ← saturate_fun n',
   unify e e' <|> fail format!"{n} and {n'} are not definitionally equal types"

/--
 Tag lemmas of the form:

 ```
 @[extensionality]
 lemma my_collection.ext (a b : my_collection)
   (h : ∀ x, a.lookup x = b.lookup y) :
   a = b := ...
 ```

 The attribute indexes extensionality lemma using the type of the
 objects (i.e. `my_collection`) which it gets from the statement of
 the lemma.  In some cases, the same lemma can be used to state the
 extensionality of multiple types that are definitionally equivalent.

 ```
 attribute [extensionality [(→),thunk,stream]] funext
 ```

 Those parameters are cumulative. The following are equivalent:

 ```
 attribute [extensionality [(→),thunk]] funext
 attribute [extensionality [stream]] funext
 ```
 and
 ```
 attribute [extensionality [(→),thunk,stream]] funext
 ```

 One removes type names from the list for one lemma with:
 ```
 attribute [extensionality [-stream,-thunk]] funext
  ```

 Finally, the following:

 ```
 @[extensionality]
 lemma my_collection.ext (a b : my_collection)
   (h : ∀ x, a.lookup x = b.lookup y) :
   a = b := ...
 ```

 is equivalent to

 ```
 @[extensionality *]
 lemma my_collection.ext (a b : my_collection)
   (h : ∀ x, a.lookup x = b.lookup y) :
   a = b := ...
 ```

 This allows us specify type synonyms along with the type
 that referred to in the lemma statement.

 ```
 @[extensionality [*,my_type_synonym]]
 lemma my_collection.ext (a b : my_collection)
   (h : ∀ x, a.lookup x = b.lookup y) :
   a = b := ...
 ```
 -/
@[user_attribute]
meta def extensional_attribute : user_attribute (name_map name) (bool × list ext_param_type × list name × list (name × name)) :=
{ name := `extensionality,
  descr := "lemmas usable by `ext` tactic",
  cache_cfg := { mk_cache := λ ls,
                          do { attrs ← ls.mmap $ λ l,
                                     do { ⟨_,_,ls,_⟩ ← extensional_attribute.get_param l,
                                          pure $ prod.mk <$> ls <*> pure l },
                               pure $ rb_map.of_list $ attrs.join },
                 dependencies := [] },
  parser :=
    do { ls ← pure <$> ext_param <|> list_of ext_param <|> pure [],
         m ← extensional_attribute.get_cache,
         pure $ (ff,ls,[],m.to_list)  },
  after_set := some $ λ n _ b,
    do (ff,ls,_,ls') ← extensional_attribute.get_param n | pure (),
       s ← mk_const n >>= infer_type >>= get_ext_subject,
       let (rs,ls'') := if ls.empty
                           then ([],[s])
                           else ls.partition_map (sum.map (flip option.get_or_else s) (flip option.get_or_else s)),
       ls''.mmap' (equiv_type_constr s),
       let l := ls'' ∪ (ls'.filter $ λ l, prod.snd l = n).map prod.fst \ rs,
       extensional_attribute.set n (tt,[],l,[]) b }

attribute [extensionality] array.ext propext prod.ext
attribute [extensionality [(→),thunk]] _root_.funext

namespace ulift
@[extensionality] lemma ext {α : Type u₁} (X Y : ulift.{u₂} α) (w : X.down = Y.down) : X = Y :=
begin
  cases X, cases Y, dsimp at w, rw w,
end
end ulift

namespace plift
@[extensionality] lemma ext {P : Prop} (a b : plift P) : a = b :=
begin
  cases a, cases b, refl
end
end plift

namespace tactic

meta def try_intros : ext_patt → tactic ext_patt
| [] := try intros $> []
| (x::xs) :=
do tgt ← target >>= whnf,
   if tgt.is_pi
     then rintro [x] >> try_intros xs
     else pure (x :: xs)

meta def ext1 (xs : ext_patt) (cfg : apply_cfg := {}): tactic ext_patt :=
do subject ← target >>= get_ext_subject,
   m ← extensional_attribute.get_cache,
   do { rule ← m.find subject,
        applyc rule cfg } <|>
     do { ls ← attribute.get_instances `extensionality,
          ls.any_of (λ n, applyc n cfg) } <|>
     fail format!"no applicable extensionality rule found for {subject}",
   try_intros xs

meta def ext : ext_patt → option ℕ → tactic unit
| _  (some 0) := skip
| xs n        := focus1 $ do
  ys ← ext1 xs, try (ext ys (nat.pred <$> n))

local postfix `?`:9001 := optional
local postfix *:9001 := many

/--
  `ext1 id` selects and apply one extensionality lemma (with attribute
  `extensionality`), using `id`, if provided, to name a local constant
  introduced by the lemma. If `id` is omitted, the local constant is
  named automatically, as per `intro`.
 -/
meta def interactive.ext1 (xs : parse ext_parse) : tactic unit :=
ext1 xs $> ()

/--
  - `ext` applies as many extensionality lemmas as possible;
  - `ext ids`, with `ids` a list of identifiers, finds extentionality and applies them
    until it runs out of identifiers in `ids` to name the local constants.

  When trying to prove:

  ```
  α β : Type,
  f g : α → set β
  ⊢ f = g
  ```

  applying `ext x y` yields:

  ```
  α β : Type,
  f g : α → set β,
  x : α,
  y : β
  ⊢ y ∈ f x ↔ y ∈ f x
  ```

  by applying functional extensionality and set extensionality.

  A maximum depth can be provided with `ext x y z : 3`.
  -/
meta def interactive.ext : parse ext_parse → parse (tk ":" *> small_nat)? → tactic unit
 | [] (some n) := iterate_range 1 n (ext1 [] $> ())
 | [] none     := repeat1 (ext1 [] $> ())
 | xs n        := tactic.ext xs n

end tactic