/-
Copyright (c) 2022 Andrew Yang. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Andrew Yang
-/
import morphisms.quasi_separated
import morphisms.isomorphism

/-!

# Affine morphisms

A morphism of schemes is affine if the preimages of affine open sets are affine.

-/

noncomputable theory

open category_theory category_theory.limits opposite topological_space

universe u

namespace algebraic_geometry

variables {X Y : Scheme.{u}} (f : X ⟶ Y)

/-- A morphism is `affine` if the preimages of affine open sets are affine. -/
@[mk_iff]
class affine (f : X ⟶ Y) : Prop :=
(is_affine_preimage : ∀ U : opens Y.carrier,
  is_affine_open U → is_affine_open ((opens.map f.1.base).obj U))

/-- The `affine_target_morphism_property` corresponding to affine morphisms. -/
def affine.affine_property : affine_target_morphism_property :=
λ X Y f hf, is_affine X

@[simp] lemma affine_affine_property_to_property {X Y : Scheme} (f : X ⟶ Y) :
  affine_target_morphism_property.to_property affine.affine_property f ↔
    is_affine Y ∧ is_affine X :=
by { delta affine_target_morphism_property.to_property affine.affine_property, simp }

@[priority 900]
instance affine_of_is_iso {X Y : Scheme} (f : X ⟶ Y) [is_iso f] : affine f :=
⟨λ U hU, hU.map_is_iso f⟩

@[priority 100]
instance affine.to_quasi_compact [affine f] : quasi_compact f :=
(quasi_compact_iff_forall_affine f).mpr (λ U hU, (affine.is_affine_preimage U hU).is_compact)

instance affine_comp {X Y Z : Scheme} (f : X ⟶ Y) (g : Y ⟶ Z)
  [affine f] [affine g] : affine (f ≫ g) :=
begin
  constructor,
  intros U hU,
  rw [Scheme.comp_val_base, opens.map_comp_obj],
  apply affine.is_affine_preimage,
  apply affine.is_affine_preimage,
  exact hU
end

lemma affine_iff_affine_property :
  affine f ↔ target_affine_locally affine.affine_property f :=
(affine_iff f).trans ⟨λ H U, H U U.prop, λ H U hU, H ⟨U, hU⟩⟩

lemma affine_eq_affine_property :
  @affine = target_affine_locally affine.affine_property :=
by { ext, exact affine_iff_affine_property _ }

instance {X : Scheme} (r : X.presheaf.obj (op ⊤)) :
  affine (X.of_restrict (X.basic_open r).open_embedding) :=
begin
  constructor,
  intros U hU,
  fapply (is_affine_open_iff_of_is_open_immersion (X.of_restrict _) _).mp,
  swap,
  { apply_instance },
  convert hU.basic_open_is_affine (X.presheaf.map (hom_of_le le_top).op r),
  rw X.basic_open_res,
  ext1,
  refine set.image_preimage_eq_inter_range.trans _,
  erw subtype.range_coe,
  refl
end

/-- The preimage of an affine open as an `Scheme.affine_opens`. -/
@[simps]
def affine_preimage {X Y : Scheme} (f : X ⟶ Y) [affine f] (U : Y.affine_opens) :
  X.affine_opens :=
⟨(opens.map f.1.base).obj (U : opens Y.carrier), affine.is_affine_preimage _ U.prop⟩

lemma is_affine_of_span_top_of_is_affine_open (X : Scheme) (s : set (X.presheaf.obj $ op ⊤))
  (h₁ : ideal.span s = ⊤) (h₂ : ∀ r : s, is_affine_open (X.basic_open r.1)) : is_affine X :=
begin
  haveI hX' : quasi_separated_space X.carrier,
  { obtain ⟨s', hs', e⟩ := (ideal.span_eq_top_iff_finite _).mp h₁,
    rw quasi_separated_space_iff_affine,
    intros U V,
    rw [← set.inter_univ (U ∩ V : set X.carrier), ← (show _ = set.univ, from
      (congr_arg subtype.val $ supr_basic_open_eq_top_of_span_eq_top _ _ e : _)), opens.supr_def,
      subtype.val_eq_coe, subtype.coe_mk, set.inter_Union],
    apply is_compact_Union,
    intro i,
    convert_to is_compact ((U.1 ⊓ (X.basic_open i.val)) ⊓ (V.1 ⊓ (X.basic_open i.val))).1,
    { conv_rhs { rw [inf_assoc, ← @inf_assoc _ _ (X.basic_open i.1),
        @inf_comm _ _ (X.basic_open i.1), inf_assoc, inf_idem, ← inf_assoc] },
      refl },
    have : ∀ (S : opens X.carrier), S ⊓ X.basic_open i.1 = X.basic_open
      (X.presheaf.map (hom_of_le le_top : S ⟶ _).op i.1) := λ S, (X.basic_open_res _ _).symm,
    apply (h₂ ⟨i.1, hs' i.2⟩).is_quasi_separated,
    { exact @inf_le_right _ _ U.1 _ },
    { exact (U.val ⊓ X.basic_open i.val).2 },
    { rw this, exact (U.prop.basic_open_is_affine _).is_compact },
    { exact @inf_le_right _ _ V.1 _ },
    { exact (V.val ⊓ X.basic_open i.val).2 },
    { rw this, exact (V.prop.basic_open_is_affine _).is_compact } },
  have hX : compact_space X.carrier,
  { obtain ⟨s', hs', e⟩ := (ideal.span_eq_top_iff_finite _).mp h₁,
    rw [← is_compact_univ_iff, ← (show _ = set.univ, from
      (congr_arg subtype.val $ supr_basic_open_eq_top_of_span_eq_top _ _ e : _)), opens.supr_def],
    apply is_compact_Union,
    intro i, exact (h₂ ⟨i.1, hs' i.2⟩).is_compact },
  constructor,
  rw (is_iso.open_cover_tfae (Γ_Spec.adjunction.unit.app X)).out 0 5,
  refine ⟨s, λ i, prime_spectrum.basic_open i.1, _, _⟩,
  { rw prime_spectrum.Union_basic_open_eq_top_iff, convert h₁, simp },
  { intro r,
    apply_with is_iso_of_is_affine_is_iso { instances := ff },
    { change is_affine_open _,
      rw preimage_adjunction_unit_basic_open,
      exact h₂ r },
    { change is_affine_open _,
      rw ← basic_open_eq_of_affine,
      apply is_affine_open.basic_open_is_affine,
      apply_with top_is_affine_open { instances := ff },
      exact algebraic_geometry.Spec_is_affine _ },
    { suffices : ∀ (U = prime_spectrum.basic_open r.val),
        is_iso ((Γ_Spec.adjunction.unit.app X).val.c.app (op $ U)),
      { rw morphism_restrict_c_app,
        apply_with is_iso.comp_is_iso { instances := ff },
        { apply this, rw opens.open_embedding_obj_top },
        { apply_instance } },
      rintros _ rfl,
      rw CommRing.is_iso_iff_bijective,
      fapply bijective_of_is_localization (submonoid.powers r.1),
      rotate,
      { apply structure_sheaf.open_algebra },
      { apply ring_hom.to_algebra,
        exact X.presheaf.map (hom_of_le le_top :
          (opens.map (Γ_Spec.adjunction.unit.app X).val.base).obj _ ⟶ _).op },
      { apply structure_sheaf.is_localization.to_basic_open },
      { dsimp,
        rw ← is_compact_univ_iff at hX,
        rw ← is_quasi_separated_univ_iff at hX',
        convert is_localization_basic_open_of_qcqs
          (show is_compact (⊤ : opens _).1, from hX) hX' r.1;
          apply Γ_Spec.adjunction.unit_app_map_basic_open },
      { rw [ring_hom.algebra_map_to_algebra, ring_hom.algebra_map_to_algebra,
          Γ_Spec.adjunction_unit_app],
        exact X.1.to_Γ_Spec_SheafedSpace_app_spec r.1 } } }
end

lemma affine_affine_property_is_local :
  affine_target_morphism_property.is_local affine.affine_property :=
begin
  split,
  { apply affine_target_morphism_property.respects_iso_mk,
    all_goals { rintros X Y Z _ _ _ (H : is_affine _), resetI },
    exacts [is_affine_of_iso e.hom, H] },
  { introv H,
    change is_affine_open _,
    rw Scheme.preimage_basic_open f r,
    exact (@@top_is_affine_open X H).basic_open_is_affine _ },
  { rintros X Y H f S hS hS',
    resetI,
    apply_fun ideal.map (f.1.c.app (op ⊤)) at hS,
    rw [ideal.map_span, ideal.map_top] at hS,
    delta affine.affine_property,
    unfreezingI { change ∀ (i : S), is_affine_open _ at hS',
      simp_rw Scheme.preimage_basic_open at hS' },
    apply is_affine_of_span_top_of_is_affine_open X _ hS,
    rintro ⟨_, r, hr, rfl⟩,
    exact hS' ⟨r, hr⟩ }
end

lemma affine_is_local_at_target :
  property_is_local_at_target @affine :=
affine_eq_affine_property.symm ▸
  affine_affine_property_is_local.target_affine_locally_is_local

lemma affine_affine_property_stable_under_base_change :
  affine_target_morphism_property.stable_under_base_change affine.affine_property :=
begin
  introv X H,
  delta affine.affine_property at H ⊢,
  resetI,
  apply_instance
end

lemma affine_stable_under_base_change :
  morphism_property.stable_under_base_change @affine :=
affine_eq_affine_property.symm ▸
  affine_affine_property_is_local.stable_under_base_change
    affine_affine_property_stable_under_base_change

lemma affine_stable_under_composition :
  morphism_property.stable_under_composition @affine :=
λ X Y Z f g hf hg, by exactI infer_instance

lemma affine_over_affine_iff {X Y : Scheme} (f : X ⟶ Y) [is_affine Y] :
  affine f ↔ is_affine X :=
affine_eq_affine_property.symm ▸
  affine_affine_property_is_local.affine_target_iff f

@[priority 100]
instance affine_of_is_affine [is_affine X] [is_affine Y] : affine f :=
begin
  rw affine_over_affine_iff,
  apply_instance
end

lemma is_affine_of_affine [affine f] [is_affine Y] : is_affine X :=
begin
  rw ← affine_over_affine_iff f,
  apply_instance
end

end algebraic_geometry


import Lean.Parser.Command

open Lean Parser Command

namespace Papyrus.Internal

def mkSealedEnumReprInst (ty : Syntax) (ctors : Array Syntax) : MacroM Syntax := do
  let currNamespace ← Macro.getCurrNamespace
  let ctorFmts ← ctors.mapM fun ctor =>
    `(Std.format $(quote <| toString (currNamespace ++ ctor[2].getId)))
  `(def reprFormats : Array Std.Format := #[$[$ctorFmts],*]
    instance : Repr $ty := ⟨fun e _ => reprFormats[e.val.val]⟩)

def unpackOptDeriving : (stx : Syntax) → (Array Syntax × Array (Option Syntax))
| `(optDeriving| deriving $[$clss $[with $argss?]?],*) => (clss, argss?)
| _ => (#[], #[])

set_option hygiene false

--------------------------------------------------------------------------------
-- # Open Enums
--------------------------------------------------------------------------------

syntax enumCtor := "\n| " declModifiers ident " := " term

scoped macro (name := enumDecl)
  mods:declModifiers
  "enum " id:ident " : " type:term optional(" := " <|> " where ")
  ctors:many(enumCtor)
  deriv?:optDeriving
: command => do
  let mut defs : Array Syntax := #[]
  -- structure
  defs := defs.push <| ←
    `($mods:declModifiers
      structure $id where
        val : $type
        $deriv?:optDeriving)
  -- constructors
  for ctor in ctors do
    let ctorId := ctor[2]
    let ctorQualId := mkIdentFrom ctorId <|
      id.getId.modifyBase (· ++ ctorId.getId)
    let ctorVal := ctor[4]
    let ctorMods := ctor[1]
    defs := defs.push <| ←
      `($ctorMods:declModifiers def $ctorQualId:ident : $id := mk $ctorVal)
  mkNullNode defs

--------------------------------------------------------------------------------
-- # Sealed Enums
--------------------------------------------------------------------------------

syntax identCtor := "\n| " declModifiers ident

scoped macro (name := sealedEnumDecl)
  mods:declModifiers
  "sealed-enum " id:ident " : " type:term optional(" := " <|> " where ")
  ctors:many(identCtor)
  deriv?:optDeriving
: command => do
  let mut innerDefs := #[]
  let numCtors := ctors.size
  let maxValLit := quote (numCtors - 1)
  -- filter out special deriving instances
  let mut derivBEq := false
  let mut derivDecEq := false
  let mut derivInhabited := false
  let mut derivRepr := false
  let mut remClasses := #[]
  let mut remArgss? := #[]
  let (classes, argss?) := unpackOptDeriving deriv?
  for cls in classes, args? in argss? do
    if cls.matchesIdent ``BEq then
      derivBEq := true
    else if cls.matchesIdent ``DecidableEq then
      derivDecEq := true
    else if cls.matchesIdent ``Inhabited then
      derivInhabited := true
    else if cls.matchesIdent ``Repr then
      derivRepr := true
    else
      remClasses := remClasses.push cls
      remArgss? := remArgss?.push args?
  -- structure
  let structDecl ←
    `($mods:declModifiers
      structure $id where
        val : $type
        h : val ≤ $maxValLit
        deriving $[$remClasses $[with $remArgss?]?],*)
  -- maximum
  innerDefs := innerDefs.push <| ←
    `(def maxVal : $type := $maxValLit)
  -- theorems
  innerDefs := innerDefs.push <| ←
    `(theorem eq_of_val_eq : {a b : $id} → a.val = b.val → a = b
        | ⟨v, h⟩, ⟨_, _⟩, rfl => rfl
      theorem val_eq_of_eq {a b : $id} (h : a = b) : a.val = b.val :=
        h ▸ rfl
      theorem ne_of_val_ne {a b : $id} (h : a.val ≠ b.val) : a ≠ b :=
        fun h' => absurd (val_eq_of_eq h') h)
  -- constructors
  for ctor in ctors, i in [:numCtors] do
    let ctorVal := quote i
    let ctorMods := ctor[1]
    let ctorId := ctor[2]
    innerDefs := innerDefs.push <| ←
      `($ctorMods:declModifiers def $ctorId:ident : $id := mk $ctorVal (by decide))
  -- derive special instance
  if derivBEq then
    innerDefs := innerDefs.push <| ←
      `(instance : BEq $id := ⟨fun a b => a.val == b.val⟩)
  if derivDecEq then
    innerDefs := innerDefs.push <| ←
      `(instance : DecidableEq $id := fun a b =>
          if h : a.val = b.val
            then isTrue (eq_of_val_eq h)
            else isFalse (ne_of_val_ne h))
  if derivInhabited then
    innerDefs := innerDefs.push <| ←
      `(instance : Inhabited $id := ⟨mk maxVal (Nat.le_refl _)⟩)
  if derivRepr then
    innerDefs := innerDefs.push <| ← mkSealedEnumReprInst id ctors
  -- syntax
  `($structDecl:command
    namespace $id:ident
    $(mkNullNode innerDefs)
    end $id:ident)


/-
Copyright (c) 2017 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Mario Carneiro
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.data.nat.default
import Mathlib.Lean3Lib.init.data.bool.default
import Mathlib.Lean3Lib.init.ite_simp

universes u l u_1 w v

namespace Mathlib

/-- In the VM, d_array is implemented as a persistent array. -/
structure d_array (n : ℕ) (α : fin n → Type u) where
  data : (i : fin n) → α i

namespace d_array

/-- The empty array. -/
def nil {α : fin 0 → Type u_1} : d_array 0 α := mk fun (_x : fin 0) => sorry

/-- `read a i` reads the `i`th member of `a`. Has builtin VM implementation. -/
def read {n : ℕ} {α : fin n → Type u} (a : d_array n α) (i : fin n) : α i := data a i

/-- `write a i v` sets the `i`th member of `a` to be `v`. Has builtin VM implementation. -/
def write {n : ℕ} {α : fin n → Type u} (a : d_array n α) (i : fin n) (v : α i) : d_array n α :=
  mk fun (j : fin n) => dite (i = j) (fun (h : i = j) => eq.rec_on h v) fun (h : ¬i = j) => read a j

def iterate_aux {n : ℕ} {α : fin n → Type u} {β : Type w} (a : d_array n α)
    (f : (i : fin n) → α i → β → β) (i : ℕ) : i ≤ n → β → β :=
  sorry

/-- Fold over the elements of the given array in ascending order. Has builtin VM implementation. -/
def iterate {n : ℕ} {α : fin n → Type u} {β : Type w} (a : d_array n α) (b : β)
    (f : (i : fin n) → α i → β → β) : β :=
  iterate_aux a f n sorry b

/-- Map the array. Has builtin VM implementation. -/
def foreach {n : ℕ} {α : fin n → Type u} {α' : fin n → Type v} (a : d_array n α)
    (f : (i : fin n) → α i → α' i) : d_array n α' :=
  mk fun (i : fin n) => f i (read a i)

def map {n : ℕ} {α : fin n → Type u} {α' : fin n → Type v} (f : (i : fin n) → α i → α' i)
    (a : d_array n α) : d_array n α' :=
  foreach a f

def map₂ {n : ℕ} {α : fin n → Type u} {α' : fin n → Type v} {α'' : fin n → Type w}
    (f : (i : fin n) → α i → α' i → α'' i) (a : d_array n α) (b : d_array n α') : d_array n α'' :=
  foreach b fun (i : fin n) => f i (read a i)

def foldl {n : ℕ} {α : fin n → Type u} {β : Type w} (a : d_array n α) (b : β)
    (f : (i : fin n) → α i → β → β) : β :=
  iterate a b f

def rev_iterate_aux {n : ℕ} {α : fin n → Type u} {β : Type w} (a : d_array n α)
    (f : (i : fin n) → α i → β → β) (i : ℕ) : i ≤ n → β → β :=
  sorry

def rev_iterate {n : ℕ} {α : fin n → Type u} {β : Type w} (a : d_array n α) (b : β)
    (f : (i : fin n) → α i → β → β) : β :=
  rev_iterate_aux a f n sorry b

@[simp] theorem read_write {n : ℕ} {α : fin n → Type u} (a : d_array n α) (i : fin n) (v : α i) :
    read (write a i v) i = v :=
  sorry

@[simp] theorem read_write_of_ne {n : ℕ} {α : fin n → Type u} (a : d_array n α) {i : fin n}
    {j : fin n} (v : α i) : i ≠ j → read (write a i v) j = read a j :=
  sorry

protected theorem ext {n : ℕ} {α : fin n → Type u} {a : d_array n α} {b : d_array n α}
    (h : ∀ (i : fin n), read a i = read b i) : a = b :=
  sorry

protected theorem ext' {n : ℕ} {α : fin n → Type u} {a : d_array n α} {b : d_array n α}
    (h :
      ∀ (i : ℕ) (h : i < n),
        read a { val := i, property := h } = read b { val := i, property := h }) :
    a = b :=
  sorry

protected def beq_aux {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)]
    (a : d_array n α) (b : d_array n α) (i : ℕ) : i ≤ n → Bool :=
  sorry

/-- Boolean element-wise equality check. -/
protected def beq {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)] (a : d_array n α)
    (b : d_array n α) : Bool :=
  d_array.beq_aux a b n sorry

theorem of_beq_aux_eq_tt {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)]
    {a : d_array n α} {b : d_array n α} (i : ℕ) (h : i ≤ n) :
    d_array.beq_aux a b i h = tt →
        ∀ (j : ℕ) (h' : j < i),
          read a { val := j, property := lt_of_lt_of_le h' h } =
            read b { val := j, property := lt_of_lt_of_le h' h } :=
  sorry

theorem of_beq_eq_tt {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)]
    {a : d_array n α} {b : d_array n α} : d_array.beq a b = tt → a = b :=
  sorry

theorem of_beq_aux_eq_ff {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)]
    {a : d_array n α} {b : d_array n α} (i : ℕ) (h : i ≤ n) :
    d_array.beq_aux a b i h = false →
        ∃ (j : ℕ),
          ∃ (h' : j < i),
            read a { val := j, property := lt_of_lt_of_le h' h } ≠
              read b { val := j, property := lt_of_lt_of_le h' h } :=
  sorry

theorem of_beq_eq_ff {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)]
    {a : d_array n α} {b : d_array n α} : d_array.beq a b = false → a ≠ b :=
  sorry

protected instance decidable_eq {n : ℕ} {α : fin n → Type u} [(i : fin n) → DecidableEq (α i)] :
    DecidableEq (d_array n α) :=
  fun (a b : d_array n α) =>
    dite (d_array.beq a b = tt) (fun (h : d_array.beq a b = tt) => is_true sorry)
      fun (h : ¬d_array.beq a b = tt) => isFalse sorry

end d_array

/-- A non-dependent array (see `d_array`). Implemented in the VM as a persistent array.  -/
def array (n : ℕ) (α : Type u) := d_array n fun (_x : fin n) => α

/-- `mk_array n v` creates a new array of length `n` where each element is `v`. Has builtin VM implementation. -/
def mk_array {α : Type u_1} (n : ℕ) (v : α) : array n α := d_array.mk fun (_x : fin n) => v

namespace array

def nil {α : Type u_1} : array 0 α := d_array.nil

def read {n : ℕ} {α : Type u} (a : array n α) (i : fin n) : α := d_array.read a i

def write {n : ℕ} {α : Type u} (a : array n α) (i : fin n) (v : α) : array n α :=
  d_array.write a i v

/-- Fold array starting from 0, folder function includes an index argument. -/
def iterate {n : ℕ} {α : Type u} {β : Type v} (a : array n α) (b : β) (f : fin n → α → β → β) : β :=
  d_array.iterate a b f

/-- Map each element of the given array with an index argument. -/
def foreach {n : ℕ} {α : Type u} {β : Type v} (a : array n α) (f : fin n → α → β) : array n β :=
  d_array.foreach a f

def map₂ {n : ℕ} {α : Type u} (f : α → α → α) (a : array n α) (b : array n α) : array n α :=
  foreach b fun (i : fin n) => f (read a i)

def foldl {n : ℕ} {α : Type u} {β : Type v} (a : array n α) (b : β) (f : α → β → β) : β :=
  iterate a b fun (_x : fin n) => f

def rev_list {n : ℕ} {α : Type u} (a : array n α) : List α :=
  foldl a [] fun (_x : α) (_y : List α) => _x :: _y

def rev_iterate {n : ℕ} {α : Type u} {β : Type v} (a : array n α) (b : β) (f : fin n → α → β → β) :
    β :=
  d_array.rev_iterate a b f

def rev_foldl {n : ℕ} {α : Type u} {β : Type v} (a : array n α) (b : β) (f : α → β → β) : β :=
  rev_iterate a b fun (_x : fin n) => f

def to_list {n : ℕ} {α : Type u} (a : array n α) : List α :=
  rev_foldl a [] fun (_x : α) (_y : List α) => _x :: _y

theorem push_back_idx {j : ℕ} {n : ℕ} (h₁ : j < n + 1) (h₂ : j ≠ n) : j < n :=
  nat.lt_of_le_and_ne (nat.le_of_lt_succ h₁) h₂

/-- `push_back a v` pushes value `v` to the end of the array. Has builtin VM implementation. -/
def push_back {n : ℕ} {α : Type u} (a : array n α) (v : α) : array (n + 1) α :=
  d_array.mk fun (_x : fin (n + 1)) => sorry

theorem pop_back_idx {j : ℕ} {n : ℕ} (h : j < n) : j < n + 1 := nat.lt.step h

/-- Discard _last_ element in the array. Has builtin VM implementation. -/
def pop_back {n : ℕ} {α : Type u} (a : array (n + 1) α) : array n α :=
  d_array.mk fun (_x : fin n) => sorry

/-- Auxilliary function for monadically mapping a function over an array. -/
def mmap_core {n : ℕ} {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (a : array n α)
    (f : α → m β) (i : ℕ) (H : i ≤ n) : m (array i β) :=
  sorry

/-- Monadically map a function over the array. -/
def mmap {n : ℕ} {α : Type u} {β : Type v} {m : Type v → Type u_1} [Monad m] (a : array n α)
    (f : α → m β) : m (array n β) :=
  mmap_core a f n sorry

/-- Map a function over the array. -/
def map {n : ℕ} {α : Type u} {β : Type v} (a : array n α) (f : α → β) : array n β :=
  d_array.map (fun (_x : fin n) => f) a

protected def mem {n : ℕ} {α : Type u} (v : α) (a : array n α) := ∃ (i : fin n), read a i = v

protected instance has_mem {n : ℕ} {α : Type u} : has_mem α (array n α) := has_mem.mk array.mem

theorem read_mem {n : ℕ} {α : Type u} (a : array n α) (i : fin n) : read a i ∈ a :=
  exists.intro i rfl

protected instance has_repr {n : ℕ} {α : Type u} [has_repr α] : has_repr (array n α) :=
  has_repr.mk (repr ∘ to_list)

@[simp] theorem read_write {n : ℕ} {α : Type u} (a : array n α) (i : fin n) (v : α) :
    read (write a i v) i = v :=
  d_array.read_write a i v

@[simp] theorem read_write_of_ne {n : ℕ} {α : Type u} (a : array n α) {i : fin n} {j : fin n}
    (v : α) : i ≠ j → read (write a i v) j = read a j :=
  d_array.read_write_of_ne a v

def read' {n : ℕ} {β : Type v} [Inhabited β] (a : array n β) (i : ℕ) : β :=
  dite (i < n) (fun (h : i < n) => read a { val := i, property := h })
    fun (h : ¬i < n) => Inhabited.default

def write' {n : ℕ} {α : Type u} (a : array n α) (i : ℕ) (v : α) : array n α :=
  dite (i < n) (fun (h : i < n) => write a { val := i, property := h } v) fun (h : ¬i < n) => a

theorem read_eq_read' {n : ℕ} {α : Type u} [Inhabited α] (a : array n α) {i : ℕ} (h : i < n) :
    read a { val := i, property := h } = read' a i :=
  sorry

theorem write_eq_write' {n : ℕ} {α : Type u} (a : array n α) {i : ℕ} (h : i < n) (v : α) :
    write a { val := i, property := h } v = write' a i v :=
  sorry

protected theorem ext {n : ℕ} {α : Type u} {a : array n α} {b : array n α}
    (h : ∀ (i : fin n), read a i = read b i) : a = b :=
  d_array.ext h

protected theorem ext' {n : ℕ} {α : Type u} {a : array n α} {b : array n α}
    (h :
      ∀ (i : ℕ) (h : i < n),
        read a { val := i, property := h } = read b { val := i, property := h }) :
    a = b :=
  d_array.ext' h

protected instance decidable_eq {n : ℕ} {α : Type u} [DecidableEq α] : DecidableEq (array n α) :=
  eq.mpr sorry fun (a b : d_array n fun (_x : fin n) => α) => d_array.decidable_eq a b

end Mathlib


/-
Copyright (c) 2021 Floris van Doorn. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Floris van Doorn

! This file was ported from Lean 3 source module tactic.project_dir
! leanprover-community/mathlib commit 29079ba8bc53a8465448a577eb2fc41932b02301
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/

/-!

# Project directory locator

We use the dummy declaration in this file to locate the project directory of mathlib.

-/

/-- This is a dummy declaration that is used to determine the project folder of mathlib, using the
  tactic `tactic.decl_olean`. This is used in `tactic.get_mathlib_dir`. -/
theorem mathlib_dir_locator : True :=
  trivial
#align mathlib_dir_locator mathlib_dir_locator


example (h : x ≠ true) : (x && y) = false := by
  simp [h]

example (h : ¬ (x = true)) : (x && y) = false := by
  simp [h]

example (h : x ≠ false) : (x && y) = y := by
  simp [h]

example (h : ¬ (x = false)) : (x && y) = y := by
  simp [h]


example : Nat → Nat := by
  refine' (fun x => _)
  trace_state
  sorry

example : Nat → Nat := by
  refine' (fun x => ?_)
  trace_state
  sorry


import Lean.Meta
import Lean.Elab
open Lean.Core
open Lean.Meta
open Lean.Elab.Term
open Lean
open Nat

def lower (n : MetaM Nat) : CoreM Nat :=
  n.run' {}

def raise (n: CoreM Nat) : MetaM Nat :=
  do
    let i <- n
    return i

-- copied from source

syntax (name := fooKind) "foo!" term : term

@[termElab fooKind] def elabFoo : TermElab :=
fun stx expectedType? => elabTerm (stx.getArg 1) expectedType?

#check foo! 10

def eg1 : Bool := foo! true

-- example of how to use the elaborator

syntax (name := addone) "addone!" (term)? : term

def addoneMeta (expr0 : Expr) : MetaM Expr :=
    do
      let name2 : Name := `Nat.succ
      let expr : Expr ←
        mkAppM name2  #[expr0]
            -- the expression returned by a function
      return expr

@[termElab addone] def addoneImpl : TermElab :=
fun stx expectedType? =>
  match stx with
  | `(addone! $s) =>
    do
      let s0 ← s
      let expr0 ←  elabTerm s0 (some (Lean.mkConst `Nat))
      let name2 : Name := `Nat.succ
      let expr : Expr ←
        mkAppM name2  #[expr0]
            -- the expression returned by a function
      return expr
  | _ =>
    do
      let name : Name := `Nat.zero
      let name2 : Name := `Nat.succ
      let expr : Expr :=
        mkApp (Lean.mkConst name2)  (Lean.mkConst name)
            -- the expression returned by a function
      return expr

def eg2 := addone! 10

#eval eg2
#eval 3 + addone!

def metaAddOne (n: MetaM Expr) : MetaM Expr :=
  do
    let i <- n
    let env ← getEnv
    IO.println s!"{env.contains ``succ}"
    IO.println s!"{env.contains `succ}"
    let decls ← getOpenDecls
    IO.println s!"{decls}"
    return mkApp (Lean.mkConst ``succ) i

def addOne(n: Nat) : Nat := addone! n

#print addOne
#eval addOne 7
#eval metaAddOne (mkConst `zero)
#eval ``succ

#check Elab.resolveGlobalConstNoOverloadWithInfo

def egNameM : TermElabM Name := do
  let stx : Syntax ← `(succ)
  IO.println s!"syntax : {stx}"
  match stx with
  | stx@(Syntax.ident _ _ n pre) => IO.println (s!"n: {n};  pre :{pre}")
  | _ => IO.println "Error"
  let ns : List Name ← Lean.resolveGlobalConst stx
  return ns.head!

#eval egNameM

#check open List Nat in fun n => cons n

syntax (name := tryapp) term ">>>>>" term : term

@[termElab tryapp] def tryappImpl : TermElab :=
fun stx expectedType? =>
  match stx with
  | `($s >>>>> $t) =>
    do
      let f <- elabTerm s none
      let x <- elabTerm t none
      let expr : Expr := mkApp f x
      let c  ← isTypeCorrect expr
      -- let cc ← hasType expr
      if c  then
        return expr
      else
        return (Lean.mkConst `Nat.zero)
      -- return (Lean.mkConst `Nat.zero)
  | _ =>
    do
      return (Lean.mkConst `Nat)

#check Nat.succ >>>>> Nat.zero
def one := Nat.succ >>>>> Nat.zero
#eval one

#check Eq >>>>> 2
#check (@Eq Nat) >>>>> 2

def shiftnat (n: Nat)(e : Expr) : MetaM Expr :=
  match n with
  | Nat.zero => return e
  | Nat.succ n => do
    let prev ← shiftnat n e
    return ← mkApp (mkConst `Nat.succ) prev

syntax (name := natshift) term ">>+>>" term : term

@[termElab natshift] def natshiftImpl : TermElab :=
fun stx expectedType? =>
  match stx with
  | `($s >>+>> $t) =>
    do
      let e <- elabTerm s (some (Lean.mkConst `Nat))
      let n : Nat <- t.isNatLit?.getD 0
      return ←  shiftnat n e
  | _ =>
    do
      return (Lean.mkConst `Nat)

#check 3 >>+>> 2
#eval 3 >>+>> 12

#check Lean.Syntax.mkScientificLit (Float.toString 3.14)

inductive Someterm  where
  | something  : {α : Type} → (a: α ) → Someterm
  | nothing : Someterm

def Someterm.isEmpty : Someterm → Bool
  | Someterm.something  _ => false
  | Someterm.nothing => true

syntax (name := tryapp2) term " >>>> " term : term

@[termElab tryapp2] def tryappImpl2 : TermElab :=
  let nt := Lean.mkConst `Someterm.nothing
  let st := Lean.mkConst `Someterm.something

  fun stx expectedType? =>
    match stx with
    | `($s >>>> $t) =>
      do
        let f <- elabTerm s none
        let x <- elabTerm t none
        let expr : Expr ←  mkApp f x
        let c ←  isTypeCorrect expr
        -- let cc := !expr.hasExprMVar
        if c then
          return Lean.mkApp st expr
        else
          return nt
    | _ =>
      do
        return nt

#check Nat.succ >>>> 3
#check Nat.succ >>>> true

#check (Eq.trans >>>> (rfl : Nat.zero = Nat.zero))

def optApp {α β γ : Type} (f : α → β) (x : γ)  :=
  f >>>> x

#print optApp

def eg3 := optApp Nat.succ 3

def eg4 := optApp Nat.succ true

#eval eg3.isEmpty -- this fails, the lambda body does not type check
#eval eg4.isEmpty

#print eg3

def exprApp (e1 e1t e2 : Expr) : MetaM Expr :=
  let n := Name.mkSimple "unsafe-name"
  withLetDecl n e1t e1 fun x => do
    let b ←  (mkAppM n #[e2])
    return ← (mkLetFVars #[x] b)

syntax (name := unapp) term " :: " term " |< " term : term

@[termElab unapp] def unappImpl : TermElab :=
  let nt := Lean.mkConst `Someterm.nothing
  let st := Lean.mkConst `Someterm.something
  let n := Name.mkSimple "unsafe-name"
  fun stx expectedType? =>
    match stx with
    | `($s :: $t |< $u) =>
      do
        let f <- elabTerm s none
        let type ← elabTerm t none
        let z <- elabTerm u none
        let expr : Expr ←  withLetDecl n type f fun x => do
                              let b ←  (mkAppM n #[z])
                              return ← (mkLetFVars #[x] b)
        let c <- isTypeCorrect expr
        if c then
          return Lean.mkApp st expr
        else
          return nt
    | _ =>
      do
        return nt

-- #check Nat.succ :: (Nat → Nat) |< 3

#print unappImpl
#print exprApp

syntax (name := minlet) "minlet!" : term

@[termElab minlet] def minletImpl : TermElab :=
  fun stx expectedType? =>
  let n := Name.mkSimple "n"
  let z := Lean.mkConst `Nat.zero
  let ty := Lean.mkConst `Nat
  withLetDecl n ty z fun x => do
    let e <- mkLetFVars #[x] x
    return e

#print minletImpl
#check minlet!

def eglit := minletImpl (Syntax.mkStrLit "minlet!") none

#check eglit
#eval eglit

def blahh := Meta.isExprDefEqAux

def nameLess (name: Name) := 1

syntax (name := ignorename) "ignore!" ident : term

macro_rules
  | `(ignore! $s) => `(Nat.zero)

#check nameLess ``Nat.succ

#check ignore! Nat.succ

inductive WrapTerm where
  | wrap : {α : Type} → (a: α ) → WrapTerm
  | wrapName : Name → WrapTerm
  | wrapExpr : Array Expr → WrapTerm

#check WrapTerm

def makeTypeFamily := Eq 1

#check makeTypeFamily
#check Eq

def makeProp : Prop := by
  apply Eq
  focus
    exact 1
  exact 2

def makeType : Type := by
  apply Option
  exact Nat

def makeIndFam : Nat → Type := by
  intro n
  induction n with
  | zero =>
    exact Nat
  | succ k ih =>
    exact Nat × ih

#check Eq.trans

def eqStatement: Prop := by
  apply Eq
  focus
    apply 1
  focus
    apply 2

def asFunc {α β : Type} (a: α) : (α → β) → β  :=
    fun f => f a

def asPi {α : Type}{motive : α → Type} (a: α) :
      ((x : α) → motive x) → motive a :=
        fun f => f a

def natGen : Nat := by
    apply (asFunc 3)
    exact Nat.succ

def piFactorizer (α : Type)(motive : α → Type) : Type :=
    (a: α ) → motive a

def island : Type := by
  apply piFactorizer
  focus
    exact fun _ => Nat
  focus
    exact Nat

def island2: Type := by
  apply piFactorizer Nat
  intro n
  induction n with
  | zero => exact Nat
  | succ k ih => exact Nat × ih

def egT := island2

open Nat

def egEgt : island2 := fun n =>
  match n with
  | zero =>  Nat.zero
  | succ k  => (k, egEgt k)

def WithType := Σ A : Type, A

def WithType.mk (α : Type) (a : α) : WithType := ⟨α , a⟩

def WithType.getType (w : WithType) : Type := w.1
def WithType.getVal (w : WithType) : w.1 := w.2

def metaWithType(e: Expr) : MetaM Expr := do
  let tp <- inferType e
  let pair  ← mkAppM ``WithType.mk #[tp, e]
  return pair

syntax (name := withtype) "withType! " term : term

@[termElab withtype] def metaWithTypeStx : TermElab :=
  fun stx expectedType? =>
  match stx with
  | `(withType! $s) =>
    do
      let e <- elabTerm s none
      let pair ← metaWithType e
      return pair
  | _ => Elab.throwIllFormedSyntax

def egName := ``Nat

#check egName

def egTyped  := withType! 3

#check egTyped
#check egTyped.getType

def elem : Nat := egTyped.getVal

#eval elem

def infEg : Inhabited Nat := inferInstance

#print infEg

#check @inferInstance (ToString Nat)

def viewExp : ToString Expr := inferInstance

def explicitToString (α : Type)(a: α)(ts: ToString α) : String :=
  ts.toString a

def exprToString (e: Expr) : String := viewExp.toString e

def exprView(e: Expr) : MetaM Expr :=
  do
    let tp ←  inferType e
    let tst ← mkAppM ``ToString #[tp]
    let ts ← synthInstance? tst
    match ts with
    | none =>
      do
        let viewExp : ToString Expr := inferInstance
        let litStr := Literal.strVal (viewExp.toString e)
        let strExp := mkLit litStr
        return  strExp
    | some t => do
      -- let tts ← mkAppM ``ToString.toString #[t]
      let v ← -- mkAppN tts #[e]
         mkAppM ``explicitToString #[tp, e, t]
      return ← whnf v

syntax (name := showexpr) "show! " term : term

@[termElab showexpr] def showexprImpl : TermElab :=
  fun stx expectedType? =>
  match stx with
  | `(show! $s) =>
    do
      let e <- elabTerm s none
      let s ← exprView e
      return s
  | _ => Elab.throwIllFormedSyntax

def egShow : String := show! (3 : Nat)

#eval egShow

def egShow2 : String  := show! (fun n : Nat => 2 + n)

#eval egShow2

def hashExpr : Hashable Expr := inferInstance

#check hashExpr

theorem constfunc{α : Type}{f: Nat → α}:
        (∀ n: Nat, f n = f (succ n)) →  (∀ n: Nat, f n = f zero) := by
          intro hyp
          intro n
          induction n with
          | zero => rfl
          | succ k ih =>
             rw [← ih]
             apply Eq.symm
             apply hyp

theorem constfuncGen{α : Type}{f: Nat → α}:
        (∀ n: Nat, f n = f (succ n)) →
          (∃ c : α, ∀ n: Nat, f n = c) := by
          intro hyp
          apply Exists.intro
          intro n
          induction n with
          | zero => rfl
          | succ k ih =>
             rw [← ih]
             apply Eq.symm
             apply hyp

def mvarMeta : MetaM Expr := do
  let mvar ← mkFreshExprMVar (some (mkConst ``Nat))
  let mvarId := mvar.mvarId!
  let mvar2 ← mkFreshExprMVar (some (mkConst ``Nat)) -- none works too
  let mvarId2 := mvar2.mvarId!
  assignExprMVar mvarId2 mvar
  assignExprMVar mvarId (mkConst ``Nat.zero)
  let mvarUnused := mkFreshExprMVar (some (mkConst ``Nat))
  return mvar2

syntax (name := minass) "minass!" : term

@[termElab minass] def minAssImpl : TermElab :=
  fun stx expectedType? =>
    do
      let e ← mvarMeta
      return mkApp (mkConst ``Nat.succ) e

def chkMinAss  := minass!

#eval chkMinAss

variable {m n: Nat}

constant k : Nat

constant l: Nat

def kl := k * l

def cname := `k

#print kl

#print cname

def mn := m * n

def names := [`n, ``mn]

def nPlusOne := addone! n

#check nPlusOne
#print nPlusOne
#print mn

#eval @nPlusOne 3

#check ignore! n
#check ignore! pqrst

def useName (fn: Nat → Nat) (arg: Name) : MetaM Expr :=
  do
    let lctx ← getLCtx
    match (← getLCtx).findFromUserName? arg with
  | some d => return d.value
  | none   => return mkConst `Nat.zero

def useFVar (z: Bool) (arg: Expr) : MetaM Expr :=
  do
    let lctx ← getLCtx
    let e ← lctx.getFVar! arg
    return mkApp (mkConst `Nat.succ) e.value

syntax (name := minname) "minname! " term : term

@[termElab minname] def minNameImpl : TermElab :=
  fun stx expectedType? =>
   match stx with
  | `(minname! $s) =>
    do
      let n := Name.mkSimple "n"
      let z := Lean.mkConst `Nat.zero
      let ty := Lean.mkConst `Nat
      withLetDecl n ty z fun x =>
        let e := useFVar true x
        return ← e
  | _ => Elab.throwIllFormedSyntax

#eval minname! true


/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Compiler.IR.Basic
import Lean.Compiler.IR.LiveVars
import Lean.Compiler.IR.Format

namespace Lean.IR.ResetReuse
/-! Remark: the insertResetReuse transformation is applied before we have
   inserted `inc/dec` instructions, and perfomed lower level optimizations
   that introduce the instructions `release` and `set`. -/

/-! Remark: the functions `S`, `D` and `R` defined here implement the
  corresponding functions in the paper "Counting Immutable Beans"

  Here are the main differences:
  - We use the State monad to manage the generation of fresh variable names.
  - Support for join points, and `uset` and `sset` instructions for unboxed data.
  - `D` uses the auxiliary function `Dmain`.
  - `Dmain` returns a pair `(b, found)` to avoid quadratic behavior when checking
    the last occurrence of the variable `x`.
  - Because we have join points in the actual implementation, a variable may be live even if it
    does not occur in a function body. See example at `livevars.lean`.
-/

private def mayReuse (c₁ c₂ : CtorInfo) : Bool :=
  c₁.size == c₂.size && c₁.usize == c₂.usize && c₁.ssize == c₂.ssize &&
  /- The following condition is a heuristic.
     We don't want to reuse cells from different types even when they are compatible
     because it produces counterintuitive behavior. -/
  c₁.name.getPrefix == c₂.name.getPrefix

private partial def S (w : VarId) (c : CtorInfo) : FnBody → FnBody
  | FnBody.vdecl x t v@(Expr.ctor c' ys) b   =>
    if mayReuse c c' then
      let updtCidx := c.cidx != c'.cidx
      FnBody.vdecl x t (Expr.reuse w c' updtCidx ys) b
    else
      FnBody.vdecl x t v (S w c b)
  | FnBody.jdecl j ys v b   =>
    let v' := S w c v
    if v == v' then FnBody.jdecl j ys v (S w c b)
    else FnBody.jdecl j ys v' b
  | FnBody.case tid x xType alts    => FnBody.case tid x xType <| alts.map fun alt => alt.modifyBody (S w c)
  | b =>
    if b.isTerminal then b
    else let
      (instr, b) := b.split
      instr.setBody (S w c b)

/-- We use `Context` to track join points in scope. -/
abbrev M := ReaderT LocalContext (StateT Index Id)

private def mkFresh : M VarId := do
  let idx ← getModify (fun n => n + 1)
  pure { idx := idx }

private def tryS (x : VarId) (c : CtorInfo) (b : FnBody) : M FnBody := do
  let w ← mkFresh
  let b' := S w c b
  if b == b' then pure b
  else pure $ FnBody.vdecl w IRType.object (Expr.reset c.size x) b'

private def Dfinalize (x : VarId) (c : CtorInfo) : FnBody × Bool → M FnBody
  | (b, true)  => pure b
  | (b, false) => tryS x c b

private def argsContainsVar (ys : Array Arg) (x : VarId) : Bool :=
  ys.any fun arg => match arg with
    | Arg.var y => x == y
    | _         => false

private def isCtorUsing (b : FnBody) (x : VarId) : Bool :=
  match b with
  | (FnBody.vdecl _ _ (Expr.ctor _ ys) _) => argsContainsVar ys x
  | _ => false

/-- Given `Dmain b`, the resulting pair `(new_b, flag)` contains the new body `new_b`,
   and `flag == true` if `x` is live in `b`.

   Note that, in the function `D` defined in the paper, for each `let x := e; F`,
   `D` checks whether `x` is live in `F` or not. This is great for clarity but it
   is expensive: `O(n^2)` where `n` is the size of the function body. -/
private partial def Dmain (x : VarId) (c : CtorInfo) : FnBody → M (FnBody × Bool)
  | e@(FnBody.case tid y yType alts) => do
    let ctx ← read
    if e.hasLiveVar ctx x then do
      /- If `x` is live in `e`, we recursively process each branch. -/
      let alts ← alts.mapM fun alt => alt.mmodifyBody fun b => Dmain x c b >>= Dfinalize x c
      pure (FnBody.case tid y yType alts, true)
    else pure (e, false)
  | FnBody.jdecl j ys v b   => do
    let (b, found) ← withReader (fun ctx => ctx.addJP j ys v) (Dmain x c b)
    let (v, _ /- found' -/) ← Dmain x c v
    /- If `found' == true`, then `Dmain b` must also have returned `(b, true)` since
       we assume the IR does not have dead join points. So, if `x` is live in `j` (i.e., `v`),
       then it must also live in `b` since `j` is reachable from `b` with a `jmp`.
       On the other hand, `x` may be live in `b` but dead in `j` (i.e., `v`). -/
    pure (FnBody.jdecl j ys v b, found)
  | e => do
    let ctx ← read
    if e.isTerminal then
      pure (e, e.hasLiveVar ctx x)
    else do
      let (instr, b) := e.split
      if isCtorUsing instr x then
        /- If the scrutinee `x` (the one that is providing memory) is being
           stored in a constructor, then reuse will probably not be able to reuse memory at runtime.
           It may work only if the new cell is consumed, but we ignore this case. -/
        pure (e, true)
      else
        let (b, found) ← Dmain x c b
        /- Remark: it is fine to use `hasFreeVar` instead of `hasLiveVar`
           since `instr` is not a `FnBody.jmp` (it is not a terminal) nor it is a `FnBody.jdecl`. -/
        if found || !instr.hasFreeVar x then
          pure (instr.setBody b, found)
        else
          let b ← tryS x c b
          pure (instr.setBody b, true)

private def D (x : VarId) (c : CtorInfo) (b : FnBody) : M FnBody :=
  Dmain x c b >>= Dfinalize x c

partial def R : FnBody → M FnBody
  | FnBody.case tid x xType alts   => do
      let alts ← alts.mapM fun alt => do
        let alt ← alt.mmodifyBody R
        match alt with
        | Alt.ctor c b =>
          if c.isScalar then pure alt
          else Alt.ctor c <$> D x c b
        | _            => pure alt
      pure $ FnBody.case tid x xType alts
  | FnBody.jdecl j ys v b   => do
    let v ← R v
    let b ← withReader (fun ctx => ctx.addJP j ys v) (R b)
    pure $ FnBody.jdecl j ys v b
  | e => do
    if e.isTerminal then pure e
    else do
      let (instr, b) := e.split
      let b ← R b
      pure (instr.setBody b)

end ResetReuse

open ResetReuse

def Decl.insertResetReuse (d : Decl) : Decl :=
  match d with
  | .fdecl (body := b) ..=>
    let nextIndex := d.maxIndex + 1
    let bNew      := (R b {}).run' nextIndex
    d.updateBody! bNew
  | other => other

end Lean.IR


/-
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.ScopedEnvExtension
import Lean.Util.Recognizers
import Lean.Util.CollectMVars
import Lean.Meta.Basic

namespace Lean.Meta

/--
  Data for user-defined theorems marked with the `congr` attribute.

  This type should be confused with `CongrTheorem` which reprents different kinds of automatically
  generated congruence theorems. The `simp` tactic also uses some of them.
-/
structure SimpCongrTheorem where
  theoremName   : Name
  funName       : Name
  hypothesesPos : Array Nat
  priority      : Nat
deriving Inhabited, Repr

structure SimpCongrTheorems where
  lemmas : SMap Name (List SimpCongrTheorem) := {}
  deriving Inhabited, Repr

def SimpCongrTheorems.get (d : SimpCongrTheorems) (declName : Name) : List SimpCongrTheorem :=
  match d.lemmas.find? declName with
  | none    => []
  | some cs => cs

def addSimpCongrTheoremEntry (d : SimpCongrTheorems) (e : SimpCongrTheorem) : SimpCongrTheorems :=
  { d with lemmas :=
      match d.lemmas.find? e.funName with
      | none    => d.lemmas.insert e.funName [e]
      | some es => d.lemmas.insert e.funName <| insert es }
where
  insert : List SimpCongrTheorem → List SimpCongrTheorem
    | []     => [e]
    | e'::es => if e.priority ≥ e'.priority then e::e'::es else e' :: insert es

builtin_initialize congrExtension : SimpleScopedEnvExtension SimpCongrTheorem SimpCongrTheorems ←
  registerSimpleScopedEnvExtension {
    initial        := {}
    addEntry       := addSimpCongrTheoremEntry
    finalizeImport := fun s => { s with lemmas := s.lemmas.switch }
  }

def mkSimpCongrTheorem (declName : Name) (prio : Nat) : MetaM SimpCongrTheorem := withReducible do
  let c ← mkConstWithLevelParams declName
  let (xs, bis, type) ← forallMetaTelescopeReducing (← inferType c)
  match type.eqOrIff? with
  | none => throwError "invalid 'congr' theorem, equality expected{indentExpr type}"
  | some (lhs, rhs) =>
    lhs.withApp fun lhsFn lhsArgs => rhs.withApp fun rhsFn rhsArgs => do
      unless lhsFn.isConst && rhsFn.isConst && lhsFn.constName! == rhsFn.constName! && lhsArgs.size == rhsArgs.size do
        throwError "invalid 'congr' theorem, equality left/right-hand sides must be applications of the same function{indentExpr type}"
      let mut foundMVars : MVarIdSet := {}
      for lhsArg in lhsArgs do
        for mvarId in (lhsArg.collectMVars {}).result do
          foundMVars := foundMVars.insert mvarId
      let mut i := 0
      let mut hypothesesPos := #[]
      for x in xs, bi in bis do
        if bi.isExplicit && !foundMVars.contains x.mvarId! then
          let rhsFn? ← forallTelescopeReducing (← inferType x) fun ys xType => do
            match xType.eqOrIff? with
            | none => pure none -- skip
            | some (xLhs, xRhs) =>
              let mut j := 0
              for y in ys do
                let yType ← inferType y
                unless onlyMVarsAt yType foundMVars do
                  throwError "invalid 'congr' theorem, argument #{j+1} of parameter #{i+1} contains unresolved parameter{indentExpr yType}"
                j := j + 1
              unless onlyMVarsAt xLhs foundMVars do
                throwError "invalid 'congr' theorem, parameter #{i+1} is not a valid hypothesis, the left-hand-side contains unresolved parameters{indentExpr xLhs}"
              let xRhsFn := xRhs.getAppFn
              unless xRhsFn.isMVar do
                throwError "invalid 'congr' theorem, parameter #{i+1} is not a valid hypothesis, the right-hand-side head is not a metavariable{indentExpr xRhs}"
              unless !foundMVars.contains xRhsFn.mvarId! do
                throwError "invalid 'congr' theorem, parameter #{i+1} is not a valid hypothesis, the right-hand-side head was already resolved{indentExpr xRhs}"
              for arg in xRhs.getAppArgs do
                unless arg.isFVar do
                  throwError "invalid 'congr' theorem, parameter #{i+1} is not a valid hypothesis, the right-hand-side argument is not local variable{indentExpr xRhs}"
              pure (some xRhsFn)
          match rhsFn? with
          | none       => pure ()
          | some rhsFn =>
            foundMVars    := foundMVars.insert x.mvarId! |>.insert rhsFn.mvarId!
            hypothesesPos := hypothesesPos.push i
        i := i + 1
      return {
        theoremName   := declName
        funName       := lhsFn.constName!
        hypothesesPos := hypothesesPos
        priority      := prio
      }
where
  /-- Return `true` if `t` contains a metavariable that is not in `mvarSet` -/
  onlyMVarsAt (t : Expr) (mvarSet : MVarIdSet) : Bool :=
    Option.isNone <| t.find? fun e => e.isMVar && !mvarSet.contains e.mvarId!

def addSimpCongrTheorem (declName : Name) (attrKind : AttributeKind) (prio : Nat) : MetaM Unit := do
  let lemma ← mkSimpCongrTheorem declName prio
  congrExtension.add lemma attrKind

builtin_initialize
  registerBuiltinAttribute {
    name  := `congr
    descr := "congruence theorem"
    add   := fun declName stx attrKind => do
      let prio ← getAttrParamOptPrio stx[1]
      discard <| addSimpCongrTheorem declName attrKind prio |>.run {} {}
  }

def getSimpCongrTheorems : MetaM SimpCongrTheorems :=
  return congrExtension.getState (← getEnv)

end Lean.Meta


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.ScopedEnvExtension
import Lean.Util.Recognizers
import Lean.Meta.LevelDefEq
import Lean.Meta.DiscrTree
import Lean.Meta.AppBuilder
import Lean.Meta.Tactic.AuxLemma
namespace Lean.Meta

/--
  The fields `levelParams` and `proof` are used to encode the proof of the simp lemma.
  If the `proof` is a global declaration `c`, we store `Expr.const c []` at `proof` without the universe levels, and `levelParams` is set to `#[]`
  When using the lemma, we create fresh universe metavariables.
  Motivation: most simp lemmas are global declarations, and this approach is faster and saves memory.

  The field `levelParams` is not empty only when we elaborate an expression provided by the user, and it contains universe metavariables.
  Then, we use `abstractMVars` to abstract the universe metavariables and create new fresh universe parameters that are stored at the field `levelParams`.
-/
structure SimpLemma where
  keys        : Array DiscrTree.Key
  levelParams : Array Name -- non empty for local universe polymorhic proofs.
  proof       : Expr
  priority    : Nat
  post        : Bool
  perm        : Bool -- true is lhs and rhs are identical modulo permutation of variables
  name?       : Option Name := none -- for debugging and tracing purposes
  deriving Inhabited

def SimpLemma.getName (s : SimpLemma) : Name :=
  match s.name? with
  | some n => n
  | none   => "<unknown>"

instance : ToFormat SimpLemma where
  format s :=
    let perm := if s.perm then ":perm" else ""
    let name := fmt s.getName
    let prio := f!":{s.priority}"
    name ++ prio ++ perm

instance : ToMessageData SimpLemma where
  toMessageData s := fmt s

instance : BEq SimpLemma where
  beq e₁ e₂ := e₁.proof == e₂.proof

structure SimpLemmas where
  pre         : DiscrTree SimpLemma := DiscrTree.empty
  post        : DiscrTree SimpLemma := DiscrTree.empty
  lemmaNames  : Std.PHashSet Name := {}
  toUnfold    : Std.PHashSet Name := {}
  erased      : Std.PHashSet Name := {}
  deriving Inhabited

def addSimpLemmaEntry (d : SimpLemmas) (e : SimpLemma) : SimpLemmas :=
  if e.post then
    { d with post := d.post.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
  else
    { d with pre := d.pre.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
where
  updateLemmaNames (s : Std.PHashSet Name) : Std.PHashSet Name :=
    match e.name? with
    | none => s
    | some name => s.insert name

def SimpLemmas.addDeclToUnfold (d : SimpLemmas) (declName : Name) : SimpLemmas :=
  { d with toUnfold := d.toUnfold.insert declName }

def SimpLemmas.isDeclToUnfold (d : SimpLemmas) (declName : Name) : Bool :=
  d.toUnfold.contains declName

def SimpLemmas.isLemma (d : SimpLemmas) (declName : Name) : Bool :=
  d.lemmaNames.contains declName

def SimpLemmas.eraseCore [Monad m] [MonadError m] (d : SimpLemmas) (declName : Name) : m SimpLemmas := do
  return { d with erased := d.erased.insert declName, lemmaNames := d.lemmaNames.erase declName, toUnfold := d.toUnfold.erase declName }

def SimpLemmas.erase [Monad m] [MonadError m] (d : SimpLemmas) (declName : Name) : m SimpLemmas := do
  unless d.isLemma declName || d.isDeclToUnfold declName do
    throwError "'{declName}' does not have [simp] attribute"
  d.eraseCore declName

inductive SimpEntry where
  | lemma    : SimpLemma → SimpEntry
  | toUnfold : Name → SimpEntry
  deriving Inhabited

builtin_initialize simpExtension : SimpleScopedEnvExtension SimpEntry SimpLemmas ←
  registerSimpleScopedEnvExtension {
    name     := `simpExt
    initial  := {}
    addEntry := fun d e =>
      match e with
      | SimpEntry.lemma e => addSimpLemmaEntry d e
      | SimpEntry.toUnfold n => d.addDeclToUnfold n
  }

private partial def isPerm : Expr → Expr → MetaM Bool
  | Expr.app f₁ a₁ _, Expr.app f₂ a₂ _ => isPerm f₁ f₂ <&&> isPerm a₁ a₂
  | Expr.mdata _ s _, t => isPerm s t
  | s, Expr.mdata _ t _ => isPerm s t
  | s@(Expr.mvar ..), t@(Expr.mvar ..) => isDefEq s t
  | Expr.forallE n₁ d₁ b₁ _, Expr.forallE n₂ d₂ b₂ _ => isPerm d₁ d₂ <&&> withLocalDeclD n₁ d₁ fun x => isPerm (b₁.instantiate1 x) (b₂.instantiate1 x)
  | Expr.lam n₁ d₁ b₁ _, Expr.lam n₂ d₂ b₂ _ => isPerm d₁ d₂ <&&> withLocalDeclD n₁ d₁ fun x => isPerm (b₁.instantiate1 x) (b₂.instantiate1 x)
  | Expr.letE n₁ t₁ v₁ b₁ _, Expr.letE n₂ t₂ v₂ b₂ _ =>
    isPerm t₁ t₂ <&&> isPerm v₁ v₂ <&&> withLetDecl n₁ t₁ v₁ fun x => isPerm (b₁.instantiate1 x) (b₂.instantiate1 x)
  | Expr.proj _ i₁ b₁ _, Expr.proj _ i₂ b₂ _ => i₁ == i₂ <&&> isPerm b₁ b₂
  | s, t => s == t

private partial def shouldPreprocess (type : Expr) : MetaM Bool :=
  forallTelescopeReducing type fun xs result => return !result.isEq

private partial def preprocess (e type : Expr) : MetaM (List (Expr × Expr)) := do
  let type ← whnf type
  if type.isForall then
    forallTelescopeReducing type fun xs type => do
      let e := mkAppN e xs
      let ps ← preprocess e type
      ps.mapM fun (e, type) =>
        return (← mkLambdaFVars xs e, ← mkForallFVars xs type)
  else if type.isEq then
    return [(e, type)]
  else if let some (lhs, rhs) := type.iff? then
    let type ← mkEq lhs rhs
    let e    ← mkPropExt e
    return [(e, type)]
  else if let some (_, lhs, rhs) := type.ne? then
    let type ← mkEq (← mkEq lhs rhs) (mkConst ``False)
    let e    ← mkEqFalse e
    return [(e, type)]
  else if let some p := type.not? then
    let type ← mkEq p (mkConst ``False)
    let e    ← mkEqFalse e
    return [(e, type)]
  else if let some (type₁, type₂) := type.and? then
    let e₁ := mkProj ``And 0 e
    let e₂ := mkProj ``And 1 e
    return (← preprocess e₁ type₁) ++ (← preprocess e₂ type₂)
  else
    let type ← mkEq type (mkConst ``True)
    let e    ← mkEqTrue e
    return [(e, type)]

private def checkTypeIsProp (type : Expr) : MetaM Unit :=
  unless (← isProp type) do
    throwError "invalid 'simp', proposition expected{indentExpr type}"

private def mkSimpLemmaCore (e : Expr) (levelParams : Array Name) (proof : Expr) (post : Bool) (prio : Nat) (name? : Option Name) : MetaM SimpLemma := do
  let type ← instantiateMVars (← inferType e)
  withNewMCtxDepth do
    let (xs, _, type) ← withReducible <| forallMetaTelescopeReducing type
    let type ← whnfR type
    let (keys, perm) ←
      match type.eq? with
      | some (_, lhs, rhs) => pure (← DiscrTree.mkPath lhs, ← isPerm lhs rhs)
      | none => throwError "unexpected kind of 'simp' theorem{indentExpr type}"
    return { keys := keys, perm := perm, post := post, levelParams := levelParams, proof := proof, name? := name?, priority := prio }

private def mkSimpLemmasFromConst (declName : Name) (post : Bool) (prio : Nat) : MetaM (Array SimpLemma) := do
  let cinfo ← getConstInfo declName
  let val := mkConst declName (cinfo.levelParams.map mkLevelParam)
  withReducible do
    let type ← inferType val
    checkTypeIsProp type
    if (← shouldPreprocess type) then
      let mut r := #[]
      for (val, type) in (← preprocess val type) do
        let auxName ← mkAuxLemma cinfo.levelParams type val
        r := r.push <| (← mkSimpLemmaCore (mkConst auxName (cinfo.levelParams.map mkLevelParam)) #[] (mkConst auxName) post prio declName)
      return r
    else
      #[← mkSimpLemmaCore (mkConst declName (cinfo.levelParams.map mkLevelParam)) #[] (mkConst declName) post prio declName]

def addSimpLemma (declName : Name) (post : Bool) (attrKind : AttributeKind) (prio : Nat) : MetaM Unit := do
  let simpLemmas ← mkSimpLemmasFromConst declName post prio
  for simpLemma in simpLemmas do
    simpExtension.add (SimpEntry.lemma simpLemma) attrKind

builtin_initialize
  registerBuiltinAttribute {
    name  := `simp
    descr := "simplification theorem"
    add   := fun declName stx attrKind =>
      let go : MetaM Unit := do
        let info ← getConstInfo declName
        if (← isProp info.type) then
          let post :=
            if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost
          let prio ← getAttrParamOptPrio stx[2]
          addSimpLemma declName post attrKind prio
        else if info.hasValue then
          simpExtension.add (SimpEntry.toUnfold declName) attrKind
        else
          throwError "invalid 'simp', it is not a proposition nor a definition (to unfold)"
      discard <| go.run {} {}
    erase := fun declName => do
      let s ← simpExtension.getState (← getEnv)
      let s ← s.erase declName
      modifyEnv fun env => simpExtension.modifyState env fun _ => s
  }

def getSimpLemmas : MetaM SimpLemmas :=
  return simpExtension.getState (← getEnv)

/- Auxiliary method for adding a global declaration to a `SimpLemmas` datastructure. -/
def SimpLemmas.addConst (s : SimpLemmas) (declName : Name) (post : Bool := true) (prio : Nat := eval_prio default) : MetaM SimpLemmas := do
  let simpLemmas ← mkSimpLemmasFromConst declName post prio
  return simpLemmas.foldl addSimpLemmaEntry s

def SimpLemma.getValue (simpLemma : SimpLemma) : MetaM Expr := do
  if simpLemma.proof.isConst && simpLemma.levelParams.isEmpty then
    let info ← getConstInfo simpLemma.proof.constName!
    if info.levelParams.isEmpty then
      return simpLemma.proof
    else
      return simpLemma.proof.updateConst! (← info.levelParams.mapM (fun _ => mkFreshLevelMVar))
  else
    let us ← simpLemma.levelParams.mapM fun _ => mkFreshLevelMVar
    simpLemma.proof.instantiateLevelParamsArray simpLemma.levelParams us

private def preprocessProof (val : Expr) : MetaM (Array Expr) := do
  let type ← inferType val
  checkTypeIsProp type
  let ps ← preprocess val type
  return ps.toArray.map fun (val, _) => val

/- Auxiliary method for creating simp lemmas from a proof term `val`. -/
def mkSimpLemmas (levelParams : Array Name) (proof : Expr) (post : Bool := true) (prio : Nat := eval_prio default) (name? : Option Name := none): MetaM (Array SimpLemma) :=
  withReducible do
    (← preprocessProof proof).mapM fun val => mkSimpLemmaCore val levelParams val post prio name?

/- Auxiliary method for adding a local simp lemma to a `SimpLemmas` datastructure. -/
def SimpLemmas.add (s : SimpLemmas) (levelParams : Array Name) (proof : Expr) (post : Bool := true) (prio : Nat := eval_prio default) (name? : Option Name := none): MetaM SimpLemmas := do
  if proof.isConst then
    s.addConst proof.constName! post prio
  else
    let simpLemmas ← mkSimpLemmas levelParams proof post prio (← getName? proof)
    return simpLemmas.foldl addSimpLemmaEntry s
where
  getName? (e : Expr) : MetaM (Option Name) := do
    match name? with
    | some _ => return name?
    | none   =>
      let f := e.getAppFn
      if f.isConst then
        return f.constName!
      else if f.isFVar then
        let localDecl ← getFVarLocalDecl f
        return localDecl.userName
      else
        return none

end Lean.Meta


example (f : α → α) (a b : α) (h : HEq a b) : f a = f b := by
  subst h
  rfl


/-
Copyright (c) 2021 Gabriel Ebner. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Gabriel Ebner
-/
import Lean
import Mathlib.Logic.Nonempty

/-!
# Once-per-file cache for tactics

This file defines cache data structures for tactics
that are initialized the first time they are accessed.
Since Lean 4 starts one process per file,
these caches are once-per-file
and can for example be used to cache information
about the imported modules.

The `Cache α` data structure is
the most generic version we define.
It is created using `Cache.mk f`
where `f : MetaM α` performs
the initialization of the cache:
```
initialize numberOfImports : Cache Nat ← Cache.mk do
  (← getEnv).imports.size

-- (does not work in the same module where the cache is defined)
#eval show MetaM Nat from numberOfImports.get
```

The `DeclCache α` data structure computes
a fold over the environment's constants:
`DeclCache.mk empty f` constructs such a cache
where `empty : α` and `f : Name → ConstantInfo → α → MetaM α`.
The result of the constants in the imports is cached
between tactic invocations,
while for constants defined in the same file
`f` is evaluated again every time.
This kind of cache can be used e.g.
to populate discrimination trees.
-/

open Lean Meta

namespace Tactic

/-- Once-per-file cache. -/
def Cache (α : Type) :=
  IO.Ref <| Sum (MetaM α) <|
    Task <| Except Exception α

instance : Nonempty (Cache α) :=
  inferInstanceAs <| Nonempty (IO.Ref _)

/-- Creates a cache with an initialization function. -/
def Cache.mk (init : MetaM α) : IO (Cache α) :=
  IO.mkRef <| Sum.inl init

/--
Access the cache.
Calling this function for the first time
will initialize the cache with the function
provided in the constructor.
-/
def Cache.get [Monad m] [MonadEnv m] [MonadOptions m] [MonadLiftT BaseIO m] [MonadExcept Exception m]
    (cache : Cache α) : m α := do
  let t ← match ← show BaseIO _ from ST.Ref.get cache with
    | Sum.inr t => pure t
    | Sum.inl init =>
      let env ← getEnv
      let options ← getOptions -- TODO: sanitize options?
      -- Default heartbeats to a reasonable value.
      -- otherwise librarySearch times out on mathlib
      -- TODO: add customization option
      let options := Core.maxHeartbeats.set options <|
        options.get? Core.maxHeartbeats.name |>.getD 1000000
      let res ← EIO.asTask do
        let metaCtx : Meta.Context := {}
        let metaState : Meta.State := {}
        let coreCtx : Core.Context := {options}
        let coreState : Core.State := {env}
        pure (← ((init ‹_›).run ‹_› ‹_›).run ‹_›).1.1
      show BaseIO _ from cache.set (Sum.inr res)
      pure res
  match t.get with
    | Except.ok res => pure res
    | Except.error err => throw err

/--
Cached fold over the environment's declarations,
where a given function is applied to `α` for every constant.
-/
def DeclCache (α : Type) :=
  Cache α × (Name → ConstantInfo → α → MetaM α)

instance : Nonempty (DeclCache α) :=
  inferInstanceAs <| Nonempty (_ × _)

/--
Creates a `DeclCache`.
The cached structure `α` is initialized with `empty`,
and then `addDecl` is called for every constant in the environment.
Calls to `addDecl` for imported constants are cached.
-/
def DeclCache.mk (profilingName : String) (empty : α) (addDecl : Name → ConstantInfo → α → MetaM α) : IO (DeclCache α) := do
  let cache ← Cache.mk do
    profileitM Exception profilingName (← getOptions) do
    let mut a := empty
    for (n, c) in (← getEnv).constants.map₁.toList do
      a ← addDecl n c a
    return a
  pure (cache, addDecl)

/--
Access the cache.
Calling this function for the first time
will initialize the cache with the function
provided in the constructor.
-/
def DeclCache.get (cache : DeclCache α) : MetaM α := do
  let mut a ← cache.1.get
  for (n, c) in (← getEnv).constants.map₂.toList do
    a ← cache.2 n c a
  return a