/-
Copyright (c) 2018 Sean Leather. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sean Leather, Mario Carneiro
-/
import data.list.alist
import data.finset.sigma
import data.part
/-!
# Finite maps over `multiset`

> THIS FILE IS SYNCHRONIZED WITH MATHLIB4.
> Any changes to this file require a corresponding PR to mathlib4.
-/

universes u v w
open list
variables {α : Type u} {β : α → Type v}

/-! ### multisets of sigma types-/

namespace multiset

/-- Multiset of keys of an association multiset. -/
def keys (s : multiset (sigma β)) : multiset α :=
s.map sigma.fst

@[simp] theorem coe_keys {l : list (sigma β)} :
  keys (l : multiset (sigma β)) = (l.keys : multiset α) :=
rfl

/-- `nodupkeys s` means that `s` has no duplicate keys. -/
def nodupkeys (s : multiset (sigma β)) : Prop :=
quot.lift_on s list.nodupkeys (λ s t p, propext $ perm_nodupkeys p)

@[simp] theorem coe_nodupkeys {l : list (sigma β)} : @nodupkeys α β l ↔ l.nodupkeys := iff.rfl

lemma nodup_keys {m : multiset (Σ a, β a)} : m.keys.nodup ↔ m.nodupkeys :=
by { rcases m with ⟨l⟩, refl }

alias nodup_keys ↔ _ nodupkeys.nodup_keys

lemma nodupkeys.nodup {m : multiset (Σ a, β a)} (h : m.nodupkeys) : m.nodup :=
h.nodup_keys.of_map _

end multiset

/-! ### finmap -/

/-- `finmap β` is the type of finite maps over a multiset. It is effectively
  a quotient of `alist β` by permutation of the underlying list. -/
structure finmap (β : α → Type v) : Type (max u v) :=
(entries : multiset (sigma β))
(nodupkeys : entries.nodupkeys)

/-- The quotient map from `alist` to `finmap`. -/
def alist.to_finmap (s : alist β) : finmap β := ⟨s.entries, s.nodupkeys⟩

local notation (name := to_finmap) `⟦`:max a `⟧`:0 := alist.to_finmap a

theorem alist.to_finmap_eq {s₁ s₂ : alist β} :
  ⟦s₁⟧ = ⟦s₂⟧ ↔ s₁.entries ~ s₂.entries :=
by cases s₁; cases s₂; simp [alist.to_finmap]

@[simp] theorem alist.to_finmap_entries (s : alist β) : ⟦s⟧.entries = s.entries := rfl

/-- Given `l : list (sigma β)`, create a term of type `finmap β` by removing
entries with duplicate keys. -/
def list.to_finmap [decidable_eq α] (s : list (sigma β)) : finmap β := s.to_alist.to_finmap

namespace finmap
open alist

lemma nodup_entries (f : finmap β) : f.entries.nodup := f.nodupkeys.nodup

/-! ### lifting from alist -/

/-- Lift a permutation-respecting function on `alist` to `finmap`. -/
@[elab_as_eliminator] def lift_on
  {γ} (s : finmap β) (f : alist β → γ)
  (H : ∀ a b : alist β, a.entries ~ b.entries → f a = f b) : γ :=
begin
  refine (quotient.lift_on s.1 (λ l, (⟨_, λ nd, f ⟨l, nd⟩⟩ : part γ))
    (λ l₁ l₂ p, part.ext' (perm_nodupkeys p) _) : part γ).get _,
  { exact λ h₁ h₂, H _ _ (by exact p) },
  { have := s.nodupkeys, rcases s.entries with ⟨l⟩, exact id }
end

@[simp] theorem lift_on_to_finmap {γ} (s : alist β) (f : alist β → γ) (H) :
  lift_on ⟦s⟧ f H = f s := by cases s; refl

/-- Lift a permutation-respecting function on 2 `alist`s to 2 `finmap`s. -/
@[elab_as_eliminator] def lift_on₂
  {γ} (s₁ s₂ : finmap β) (f : alist β → alist β → γ)
  (H : ∀ a₁ b₁ a₂ b₂ : alist β, a₁.entries ~ a₂.entries → b₁.entries ~ b₂.entries →
    f a₁ b₁ = f a₂ b₂) : γ :=
lift_on s₁
  (λ l₁, lift_on s₂ (f l₁) (λ b₁ b₂ p, H _ _ _ _ (perm.refl _) p))
  (λ a₁ a₂ p, have H' : f a₁ = f a₂ := funext (λ _, H _ _ _ _ p (perm.refl _)), by simp only [H'])

@[simp] theorem lift_on₂_to_finmap {γ} (s₁ s₂ : alist β) (f : alist β → alist β → γ) (H) :
  lift_on₂ ⟦s₁⟧ ⟦s₂⟧ f H = f s₁ s₂ :=
by cases s₁; cases s₂; refl

/-! ### induction -/

@[elab_as_eliminator] theorem induction_on
  {C : finmap β → Prop} (s : finmap β) (H : ∀ (a : alist β), C ⟦a⟧) : C s :=
by rcases s with ⟨⟨a⟩, h⟩; exact H ⟨a, h⟩

@[elab_as_eliminator] theorem induction_on₂ {C : finmap β → finmap β → Prop}
  (s₁ s₂ : finmap β) (H : ∀ (a₁ a₂ : alist β), C ⟦a₁⟧ ⟦a₂⟧) : C s₁ s₂ :=
induction_on s₁ $ λ l₁, induction_on s₂ $ λ l₂, H l₁ l₂

@[elab_as_eliminator] theorem induction_on₃ {C : finmap β →  finmap β → finmap β → Prop}
  (s₁ s₂ s₃ : finmap β) (H : ∀ (a₁ a₂ a₃ : alist β), C ⟦a₁⟧ ⟦a₂⟧ ⟦a₃⟧) : C s₁ s₂ s₃ :=
induction_on₂ s₁ s₂ $ λ l₁ l₂, induction_on s₃ $ λ l₃, H l₁ l₂ l₃

/-! ### extensionality -/

@[ext] theorem ext : ∀ {s t : finmap β}, s.entries = t.entries → s = t
| ⟨l₁, h₁⟩ ⟨l₂, h₂⟩ H := by congr'

@[simp] theorem ext_iff {s t : finmap β} : s.entries = t.entries ↔ s = t :=
⟨ext, congr_arg _⟩

/-! ### mem -/

/-- The predicate `a ∈ s` means that `s` has a value associated to the key `a`. -/
instance : has_mem α (finmap β) := ⟨λ a s, a ∈ s.entries.keys⟩

theorem mem_def {a : α} {s : finmap β} :
  a ∈ s ↔ a ∈ s.entries.keys := iff.rfl

@[simp] theorem mem_to_finmap {a : α} {s : alist β} :
  a ∈ ⟦s⟧ ↔ a ∈ s := iff.rfl

/-! ### keys -/

/-- The set of keys of a finite map. -/
def keys (s : finmap β) : finset α :=
⟨s.entries.keys, s.nodupkeys.nodup_keys⟩

@[simp] theorem keys_val (s : alist β) : (keys ⟦s⟧).val = s.keys := rfl

@[simp] theorem keys_ext {s₁ s₂ : alist β} :
  keys ⟦s₁⟧ = keys ⟦s₂⟧ ↔ s₁.keys ~ s₂.keys :=
by simp [keys, alist.keys]

theorem mem_keys {a : α} {s : finmap β} : a ∈ s.keys ↔ a ∈ s :=
induction_on s $ λ s, alist.mem_keys

/-! ### empty -/

/-- The empty map. -/
instance : has_emptyc (finmap β) := ⟨⟨0, nodupkeys_nil⟩⟩

instance : inhabited (finmap β) := ⟨∅⟩

@[simp] theorem empty_to_finmap : (⟦∅⟧ : finmap β) = ∅ := rfl

@[simp] theorem to_finmap_nil [decidable_eq α] : ([].to_finmap : finmap β) = ∅ := rfl

theorem not_mem_empty {a : α} : a ∉ (∅ : finmap β) :=
multiset.not_mem_zero a

@[simp] theorem keys_empty : (∅ : finmap β).keys = ∅ := rfl

/-! ### singleton -/

/-- The singleton map. -/
def singleton (a : α) (b : β a) : finmap β := ⟦alist.singleton a b⟧

@[simp] theorem keys_singleton (a : α) (b : β a) :
  (singleton a b).keys = {a} := rfl

@[simp] lemma mem_singleton (x y : α) (b : β y) : x ∈ singleton y b ↔ x = y :=
by simp only [singleton]; erw [mem_cons_eq, mem_nil_iff, or_false]

section

variables [decidable_eq α]

instance has_decidable_eq [∀ a, decidable_eq (β a)] : decidable_eq (finmap β)
| s₁ s₂ := decidable_of_iff _ ext_iff

/-! ### lookup -/

/-- Look up the value associated to a key in a map. -/
def lookup (a : α) (s : finmap β) : option (β a) :=
lift_on s (lookup a) (λ s t, perm_lookup)

@[simp] theorem lookup_to_finmap (a : α) (s : alist β) :
  lookup a ⟦s⟧ = s.lookup a := rfl

@[simp] theorem lookup_list_to_finmap (a : α) (s : list (sigma β)) :
  lookup a s.to_finmap = s.lookup a :=
by rw [list.to_finmap, lookup_to_finmap, lookup_to_alist]

@[simp] theorem lookup_empty (a) : lookup a (∅ : finmap β) = none :=
rfl

theorem lookup_is_some {a : α} {s : finmap β} :
  (s.lookup a).is_some ↔ a ∈ s :=
induction_on s $ λ s, alist.lookup_is_some

theorem lookup_eq_none {a} {s : finmap β} : lookup a s = none ↔ a ∉ s :=
induction_on s $ λ s, alist.lookup_eq_none

lemma mem_lookup_iff {f : finmap β} {a : α} {b : β a} :
  b ∈ f.lookup a ↔ sigma.mk a b ∈ f.entries :=
by { rcases f with ⟨⟨l⟩, hl⟩, exact list.mem_lookup_iff hl }

/-- A version of `finmap.mem_lookup_iff` with LHS in the simp-normal form. -/
lemma lookup_eq_some_iff {f : finmap β} {a : α} {b : β a} :
  f.lookup a = some b ↔ sigma.mk a b ∈ f.entries :=
mem_lookup_iff

@[simp] lemma sigma_keys_lookup (f : finmap β) :
  f.keys.sigma (λ i, (f.lookup i).to_finset) = ⟨f.entries, f.nodup_entries⟩ :=
begin
  ext x,
  have : x ∈ f.entries → x.fst ∈ f.keys, from multiset.mem_map_of_mem _,
  simpa [lookup_eq_some_iff]
end

@[simp] lemma lookup_singleton_eq {a : α} {b : β a} : (singleton a b).lookup a = some b :=
by rw [singleton, lookup_to_finmap, alist.singleton, alist.lookup, lookup_cons_eq]

instance (a : α) (s : finmap β) : decidable (a ∈ s) :=
decidable_of_iff _ lookup_is_some

lemma mem_iff {a : α} {s : finmap β} : a ∈ s ↔ ∃ b, s.lookup a = some b :=
induction_on s $ λ s,
iff.trans list.mem_keys $ exists_congr $ λ b,
(list.mem_lookup_iff s.nodupkeys).symm

lemma mem_of_lookup_eq_some {a : α} {b : β a} {s : finmap β} (h : s.lookup a = some b) : a ∈ s :=
mem_iff.mpr ⟨_, h⟩

theorem ext_lookup {s₁ s₂ : finmap β} : (∀ x, s₁.lookup x = s₂.lookup x) → s₁ = s₂ :=
induction_on₂ s₁ s₂ $ λ s₁ s₂ h,
begin
  simp only [alist.lookup, lookup_to_finmap] at h,
  rw [alist.to_finmap_eq],
  apply lookup_ext s₁.nodupkeys s₂.nodupkeys,
  intros x y,
  rw h,
end

/-- An equivalence between `finmap β` and pairs `(keys : finset α, lookup : Π a, option (β a))` such
that `(lookup a).is_some ↔ a ∈ keys`. -/
@[simps apply_coe_fst apply_coe_snd]
def keys_lookup_equiv :
  finmap β ≃ {f : finset α × (Π a, option (β a)) // ∀ i, (f.2 i).is_some ↔ i ∈ f.1} :=
{ to_fun := λ f, ⟨(f.keys, λ i, f.lookup i), λ i, lookup_is_some⟩,
  inv_fun := λ f, ⟨(f.1.1.sigma $ λ i, (f.1.2 i).to_finset).val,
    begin
      refine multiset.nodup_keys.1 ((finset.nodup _).map_on _),
      simp only [finset.mem_val, finset.mem_sigma, option.mem_to_finset, option.mem_def],
      rintro ⟨i, x⟩ ⟨hi, hx⟩ ⟨j, y⟩ ⟨hj, hy⟩ (rfl : i = j),
      obtain rfl : x = y, from option.some.inj (hx.symm.trans hy),
      refl
    end⟩,
  left_inv := λ f, ext $ by simp,
  right_inv := λ ⟨⟨s, f⟩, hf⟩,
    begin
      ext : 2; dsimp [keys],
      { ext1 i,
        have : i ∈ s → (∃ x, f i = some x),
          from λ hi, ⟨option.get _, option.get_mem $ (hf i).2 hi⟩,
        simpa [multiset.keys] },
      { ext i x : 2,
        simp only [option.mem_def, lookup_eq_some_iff, finset.mem_val, finset.mem_sigma,
          option.mem_to_finset, and_iff_right_iff_imp, ← hf],
        exact λ h, option.is_some_iff_exists.2 ⟨_, h⟩ }
    end }

@[simp] lemma keys_lookup_equiv_symm_apply_keys :
  ∀ f : {f : finset α × (Π a, option (β a)) // ∀ i, (f.2 i).is_some ↔ i ∈ f.1},
    (keys_lookup_equiv.symm f).keys = (f : finset α × Π a, option (β a)).1 :=
keys_lookup_equiv.surjective.forall.2 $ λ f,
  by simp only [equiv.symm_apply_apply, keys_lookup_equiv_apply_coe_fst]

@[simp] lemma keys_lookup_equiv_symm_apply_lookup :
  ∀ (f : {f : finset α × (Π a, option (β a)) // ∀ i, (f.2 i).is_some ↔ i ∈ f.1}) a,
    (keys_lookup_equiv.symm f).lookup a = (f : finset α × Π a, option (β a)).2 a :=
keys_lookup_equiv.surjective.forall.2 $ λ f a,
  by simp only [equiv.symm_apply_apply, keys_lookup_equiv_apply_coe_snd]

/-! ### replace -/

/-- Replace a key with a given value in a finite map.
  If the key is not present it does nothing. -/
def replace (a : α) (b : β a) (s : finmap β) : finmap β :=
lift_on s (λ t, ⟦replace a b t⟧) $
λ s₁ s₂ p, to_finmap_eq.2 $ perm_replace p

@[simp] theorem replace_to_finmap (a : α) (b : β a) (s : alist β) :
  replace a b ⟦s⟧ = ⟦s.replace a b⟧ := by simp [replace]

@[simp] theorem keys_replace (a : α) (b : β a) (s : finmap β) :
  (replace a b s).keys = s.keys :=
induction_on s $ λ s, by simp

@[simp] theorem mem_replace {a a' : α} {b : β a} {s : finmap β} :
  a' ∈ replace a b s ↔ a' ∈ s :=
induction_on s $ λ s, by simp

end

/-! ### foldl -/

/-- Fold a commutative function over the key-value pairs in the map -/
def foldl {δ : Type w} (f : δ → Π a, β a → δ)
  (H : ∀ d a₁ b₁ a₂ b₂, f (f d a₁ b₁) a₂ b₂ = f (f d a₂ b₂) a₁ b₁)
  (d : δ) (m : finmap β) : δ :=
m.entries.foldl (λ d s, f d s.1 s.2) (λ d s t, H _ _ _ _ _) d

/-- `any f s` returns `tt` iff there exists a value `v` in `s` such that `f v = tt`. -/
def any (f : Π x, β x → bool) (s : finmap β) : bool :=
s.foldl (λ x y z, x || f y z) (by { intros, simp_rw [bool.bor_assoc, bool.bor_comm] }) ff

/-- `all f s` returns `tt` iff `f v = tt` for all values `v` in `s`. -/
def all (f : Π x, β x → bool) (s : finmap β) : bool :=
s.foldl (λ x y z, x && f y z) (by { intros, simp_rw [bool.band_assoc, bool.band_comm] }) tt

/-! ### erase -/

section

variables [decidable_eq α]

/-- Erase a key from the map. If the key is not present it does nothing. -/
def erase (a : α) (s : finmap β) : finmap β :=
lift_on s (λ t, ⟦erase a t⟧) $
λ s₁ s₂ p, to_finmap_eq.2 $ perm_erase p

@[simp] theorem erase_to_finmap (a : α) (s : alist β) :
  erase a ⟦s⟧ = ⟦s.erase a⟧ := by simp [erase]

@[simp]

@[simp] theorem keys_erase (a : α) (s : finmap β) :
  (erase a s).keys = s.keys.erase a :=
induction_on s $ λ s, by simp

@[simp] theorem mem_erase {a a' : α} {s : finmap β} : a' ∈ erase a s ↔ a' ≠ a ∧ a' ∈ s :=
induction_on s $ λ s, by simp

theorem not_mem_erase_self {a : α} {s : finmap β} : ¬ a ∈ erase a s :=
by rw [mem_erase, not_and_distrib, not_not]; left; refl

@[simp] theorem lookup_erase (a) (s : finmap β) : lookup a (erase a s) = none :=
induction_on s $ lookup_erase a

@[simp] theorem lookup_erase_ne {a a'} {s : finmap β} (h : a ≠ a') :
  lookup a (erase a' s) = lookup a s :=
induction_on s $ λ s, lookup_erase_ne h

theorem erase_erase {a a' : α} {s : finmap β} : erase a (erase a' s) = erase a' (erase a s) :=
induction_on s $ λ s, ext (by simp only [erase_erase, erase_to_finmap])

/-! ### sdiff -/

/-- `sdiff s s'` consists of all key-value pairs from `s` and `s'` where the keys are in `s` or
`s'` but not both. -/
def sdiff (s s' : finmap β) : finmap β :=
s'.foldl (λ s x _, s.erase x) (λ a₀ a₁ _ a₂ _, erase_erase) s

instance : has_sdiff (finmap β) := ⟨sdiff⟩

/-! ### insert -/

/-- Insert a key-value pair into a finite map, replacing any existing pair with
  the same key. -/
def insert (a : α) (b : β a) (s : finmap β) : finmap β :=
lift_on s (λ t, ⟦insert a b t⟧) $
λ s₁ s₂ p, to_finmap_eq.2 $ perm_insert p

@[simp] theorem insert_to_finmap (a : α) (b : β a) (s : alist β) :
  insert a b ⟦s⟧ = ⟦s.insert a b⟧ := by simp [insert]

theorem insert_entries_of_neg {a : α} {b : β a} {s : finmap β} : a ∉ s →
  (insert a b s).entries = ⟨a, b⟩ ::ₘ s.entries :=
induction_on s $ λ s h,
by simp [insert_entries_of_neg (mt mem_to_finmap.1 h)]

@[simp] theorem mem_insert {a a' : α} {b' : β a'} {s : finmap β} :
  a ∈ insert a' b' s ↔ a = a' ∨ a ∈ s :=
induction_on s mem_insert

@[simp] theorem lookup_insert {a} {b : β a} (s : finmap β) :
  lookup a (insert a b s) = some b :=
induction_on s $ λ s,
by simp only [insert_to_finmap, lookup_to_finmap, lookup_insert]

@[simp] theorem lookup_insert_of_ne {a a'} {b : β a} (s : finmap β) (h : a' ≠ a) :
  lookup a' (insert a b s) = lookup a' s :=
induction_on s $ λ s,
by simp only [insert_to_finmap, lookup_to_finmap, lookup_insert_ne h]

@[simp] theorem insert_insert {a} {b b' : β a} (s : finmap β) :
  (s.insert a b).insert a b' = s.insert a b' :=
induction_on s $ λ s,
by simp only [insert_to_finmap, insert_insert]

theorem insert_insert_of_ne {a a'} {b : β a} {b' : β a'} (s : finmap β) (h : a ≠ a') :
  (s.insert a b).insert a' b' = (s.insert a' b').insert a b :=
induction_on s $ λ s,
by simp only [insert_to_finmap, alist.to_finmap_eq, insert_insert_of_ne _ h]

theorem to_finmap_cons (a : α) (b : β a) (xs : list (sigma β)) :
  list.to_finmap (⟨a,b⟩ :: xs) = insert a b xs.to_finmap := rfl

theorem mem_list_to_finmap (a : α) (xs : list (sigma β)) :
  a ∈ xs.to_finmap ↔ (∃ b : β a, sigma.mk a b ∈ xs) :=
by { induction xs with x xs; [skip, cases x];
     simp only [to_finmap_cons, *, not_mem_empty, exists_or_distrib, not_mem_nil, to_finmap_nil,
                exists_false, mem_cons_iff, mem_insert, exists_and_distrib_left];
     apply or_congr _ iff.rfl,
     conv { to_lhs, rw ← and_true (a = x_fst) },
     apply and_congr_right, rintro ⟨⟩, simp only [exists_eq, heq_iff_eq] }

@[simp] theorem insert_singleton_eq {a : α} {b b' : β a} :
  insert a b (singleton a b') = singleton a b :=
by simp only [singleton, finmap.insert_to_finmap, alist.insert_singleton_eq]

/-! ### extract -/

/-- Erase a key from the map, and return the corresponding value, if found. -/
def extract (a : α) (s : finmap β) : option (β a) × finmap β :=
lift_on s (λ t, prod.map id to_finmap (extract a t)) $
λ s₁ s₂ p, by simp [perm_lookup p, to_finmap_eq, perm_erase p]

@[simp] theorem extract_eq_lookup_erase (a : α) (s : finmap β) :
  extract a s = (lookup a s, erase a s) :=
induction_on s $ λ s, by simp [extract]

/-! ### union -/

/-- `s₁ ∪ s₂` is the key-based union of two finite maps. It is left-biased: if
there exists an `a ∈ s₁`, `lookup a (s₁ ∪ s₂) = lookup a s₁`. -/
def union (s₁ s₂ : finmap β) : finmap β :=
lift_on₂ s₁ s₂ (λ s₁ s₂, ⟦s₁ ∪ s₂⟧) $
λ s₁ s₂ s₃ s₄ p₁₃ p₂₄, to_finmap_eq.mpr $ perm_union p₁₃ p₂₄

instance : has_union (finmap β) := ⟨union⟩

@[simp] theorem mem_union {a} {s₁ s₂ : finmap β} :
  a ∈ s₁ ∪ s₂ ↔ a ∈ s₁ ∨ a ∈ s₂ :=
induction_on₂ s₁ s₂ $ λ _ _, mem_union

@[simp] theorem union_to_finmap (s₁ s₂ : alist β) : ⟦s₁⟧ ∪ ⟦s₂⟧ = ⟦s₁ ∪ s₂⟧ :=
by simp [(∪), union]

theorem keys_union {s₁ s₂ : finmap β} : (s₁ ∪ s₂).keys = s₁.keys ∪ s₂.keys :=
induction_on₂ s₁ s₂ $ λ s₁ s₂, finset.ext $ by simp [keys]

@[simp] theorem lookup_union_left {a} {s₁ s₂ : finmap β} :
  a ∈ s₁ → lookup a (s₁ ∪ s₂) = lookup a s₁ :=
induction_on₂ s₁ s₂ $ λ s₁ s₂, lookup_union_left

@[simp] theorem lookup_union_right {a} {s₁ s₂ : finmap β} :
  a ∉ s₁ → lookup a (s₁ ∪ s₂) = lookup a s₂ :=
induction_on₂ s₁ s₂ $ λ s₁ s₂, lookup_union_right

theorem lookup_union_left_of_not_in {a} {s₁ s₂ : finmap β} (h : a ∉ s₂) :
  lookup a (s₁ ∪ s₂) = lookup a s₁ :=
begin
  by_cases h' : a ∈ s₁,
  { rw lookup_union_left h' },
  { rw [lookup_union_right h', lookup_eq_none.mpr h, lookup_eq_none.mpr h'] }
end

@[simp] theorem mem_lookup_union {a} {b : β a} {s₁ s₂ : finmap β} :
  b ∈ lookup a (s₁ ∪ s₂) ↔ b ∈ lookup a s₁ ∨ a ∉ s₁ ∧ b ∈ lookup a s₂ :=
induction_on₂ s₁ s₂ $ λ s₁ s₂, mem_lookup_union

theorem mem_lookup_union_middle {a} {b : β a} {s₁ s₂ s₃ : finmap β} :
  b ∈ lookup a (s₁ ∪ s₃) → a ∉ s₂ → b ∈ lookup a (s₁ ∪ s₂ ∪ s₃) :=
induction_on₃ s₁ s₂ s₃ $ λ s₁ s₂ s₃, mem_lookup_union_middle

theorem insert_union {a} {b : β a} {s₁ s₂ : finmap β} :
  insert a b (s₁ ∪ s₂) = insert a b s₁ ∪ s₂ :=
induction_on₂ s₁ s₂ $ λ a₁ a₂, by simp [insert_union]

theorem union_assoc {s₁ s₂ s₃ : finmap β} : (s₁ ∪ s₂) ∪ s₃ = s₁ ∪ (s₂ ∪ s₃) :=
induction_on₃ s₁ s₂ s₃ $ λ s₁ s₂ s₃,
by simp only [alist.to_finmap_eq, union_to_finmap, alist.union_assoc]

@[simp] theorem empty_union {s₁ : finmap β} : ∅ ∪ s₁ = s₁ :=
induction_on s₁ $ λ s₁, by rw ← empty_to_finmap;
  simp [- empty_to_finmap, alist.to_finmap_eq, union_to_finmap, alist.union_assoc]

@[simp] theorem union_empty {s₁ : finmap β} : s₁ ∪ ∅ = s₁ :=
induction_on s₁ $ λ s₁, by rw ← empty_to_finmap;
  simp [- empty_to_finmap, alist.to_finmap_eq, union_to_finmap, alist.union_assoc]

theorem erase_union_singleton (a : α) (b : β a) (s : finmap β) (h : s.lookup a = some b) :
  s.erase a ∪ singleton a b = s :=
ext_lookup
(λ x, by { by_cases h' : x = a,
      { subst a, rw [lookup_union_right not_mem_erase_self, lookup_singleton_eq, h], },
      { have : x ∉ singleton a b, { rwa mem_singleton },
        rw [lookup_union_left_of_not_in this, lookup_erase_ne h'] } } )

end

/-! ### disjoint -/

/-- `disjoint s₁ s₂` holds if `s₁` and `s₂` have no keys in common. -/
def disjoint (s₁ s₂ : finmap β) : Prop :=
∀ x ∈ s₁, ¬ x ∈ s₂

lemma disjoint_empty (x : finmap β) : disjoint ∅ x .

@[symm]
lemma disjoint.symm (x y : finmap β) (h : disjoint x y) : disjoint y x :=
λ p hy hx, h p hx hy

lemma disjoint.symm_iff (x y : finmap β) : disjoint x y ↔ disjoint y x :=
⟨disjoint.symm x y, disjoint.symm y x⟩

section

variables [decidable_eq α]

instance : decidable_rel (@disjoint α β) :=
λ x y, by dsimp only [disjoint]; apply_instance

lemma disjoint_union_left (x y z : finmap β) : disjoint (x ∪ y) z ↔ disjoint x z ∧ disjoint y z :=
by simp [disjoint, finmap.mem_union, or_imp_distrib, forall_and_distrib]

lemma disjoint_union_right (x y z : finmap β) : disjoint x (y ∪ z) ↔ disjoint x y ∧ disjoint x z :=
by rw [disjoint.symm_iff, disjoint_union_left, disjoint.symm_iff _ x, disjoint.symm_iff _ x]

theorem union_comm_of_disjoint {s₁ s₂ : finmap β} : disjoint s₁ s₂ → s₁ ∪ s₂ = s₂ ∪ s₁ :=
induction_on₂ s₁ s₂ $ λ s₁ s₂,
by { intros h, simp only [alist.to_finmap_eq, union_to_finmap, alist.union_comm_of_disjoint h] }

theorem union_cancel {s₁ s₂ s₃ : finmap β} (h : disjoint s₁ s₃) (h' : disjoint s₂ s₃) :
  s₁ ∪ s₃ = s₂ ∪ s₃ ↔ s₁ = s₂ :=
⟨λ h'', begin
          apply ext_lookup, intro x,
          have : (s₁ ∪ s₃).lookup x = (s₂ ∪ s₃).lookup x, from h'' ▸ rfl,
          by_cases hs₁ : x ∈ s₁,
          { rwa [lookup_union_left hs₁, lookup_union_left_of_not_in (h _ hs₁)] at this, },
          { by_cases hs₂ : x ∈ s₂,
            { rwa [lookup_union_left_of_not_in (h' _ hs₂), lookup_union_left hs₂] at this, },
            { rw [lookup_eq_none.mpr hs₁, lookup_eq_none.mpr hs₂] } }
        end,
 λ h, h ▸ rfl⟩

end

end finmap


/-
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 data.array.lemmas
import data.list.join
import data.list.range
import data.pnat.basic

/-!
# Hash maps

Defines a hash map data structure, representing a finite key-value map
with a value type that may depend on the key type.  The structure
requires a `nat`-valued hash function to associate keys to buckets.

## Main definitions

* `hash_map`: constructed with `mk_hash_map`.

## Implementation details

A hash map with key type `α` and (dependent) value type `β : α → Type*`
consists of an array of *buckets*, which are lists containing
key/value pairs for that bucket.  The hash function is taken modulo `n`
to assign keys to their respective bucket.  Because of this, some care
should be put into the hash function to ensure it evenly distributes
keys.

The bucket array is an `array`.  These have special VM support for
in-place modification if there is only ever one reference to them.  If
one takes special care to never keep references to old versions of a
hash map alive after updating it, then the hash map will be modified
in-place.  In this documentation, when we say a hash map is modified
in-place, we are assuming the API is being used in this manner.

When inserting (`hash_map.insert`), if the number of stored pairs (the
*size*) is going to exceed the number of buckets, then a new hash map
is first created with double the number of buckets and everything in
the old hash map is reinserted along with the new key/value pair.
Otherwise, the bucket array is modified in-place.  The amortized
running time of inserting $$n$$ elements into a hash map is $$O(n)$$.

When removing (`hash_map.erase`), the hash map is modified in-place.
The implementation does not reduce the number of buckets in the hash
map if the size gets too low.

## Tags

hash map

-/

universes u v w

/-- `bucket_array α β` is the underlying data type for `hash_map α β`,
  an array of linked lists of key-value pairs. -/
def bucket_array (α : Type u) (β : α → Type v) (n : ℕ+) :=
array n (list Σ a, β a)

/-- Make a hash_map index from a `nat` hash value and a (positive) buffer size -/
def hash_map.mk_idx (n : ℕ+) (i : nat) : fin n :=
⟨i % n, nat.mod_lt _ n.2⟩

namespace bucket_array
section
parameters {α : Type u} {β : α → Type v} (hash_fn : α → nat)
variables {n : ℕ+} (data : bucket_array α β n)

instance : inhabited (bucket_array α β n) :=
⟨mk_array _ []⟩

/-- Read the bucket corresponding to an element -/
def read (a : α) : list Σ a, β a :=
let bidx := hash_map.mk_idx n (hash_fn a) in
data.read bidx

/-- Write the bucket corresponding to an element -/
def write (a : α) (l : list Σ a, β a) : bucket_array α β n :=
let bidx := hash_map.mk_idx n (hash_fn a) in
data.write bidx l

/-- Modify (read, apply `f`, and write) the bucket corresponding to an element -/
def modify (a : α) (f : list (Σ a, β a) → list (Σ a, β a)) : bucket_array α β n :=
let bidx := hash_map.mk_idx n (hash_fn a) in
array.write data bidx (f (array.read data bidx))

/-- The list of all key-value pairs in the bucket list -/
def as_list : list Σ a, β a := data.to_list.join

theorem mem_as_list {a : Σ a, β a} : a ∈ data.as_list ↔ ∃i, a ∈ array.read data i :=
have (∃ (l : list (Σ (a : α), β a)) (i : fin (n.val)), a ∈ l ∧ array.read data i = l) ↔
  ∃ (i : fin (n.val)), a ∈ array.read data i,
by rw exists_swap; exact exists_congr (λ i, by simp),
by simp [as_list]; simpa [array.mem.def, and_comm]

/-- Fold a function `f` over the key-value pairs in the bucket list -/
def foldl {δ : Type w} (d : δ) (f : δ → Π a, β a → δ) : δ :=
data.foldl d (λ b d, b.foldl (λ r a, f r a.1 a.2) d)

theorem foldl_eq {δ : Type w} (d : δ) (f : δ → Π a, β a → δ) :
  data.foldl d f = data.as_list.foldl (λ r a, f r a.1 a.2) d :=
by rw [foldl, as_list, list.foldl_join, ← array.to_list_foldl]

end
end bucket_array

namespace hash_map
section
parameters {α : Type u} {β : α → Type v} (hash_fn : α → nat)

/-- Insert the pair `⟨a, b⟩` into the correct location in the bucket array
  (without checking for duplication) -/
def reinsert_aux {n} (data : bucket_array α β n) (a : α) (b : β a) : bucket_array α β n :=
data.modify hash_fn a (λl, ⟨a, b⟩ :: l)

theorem mk_as_list (n : ℕ+) : bucket_array.as_list (mk_array n [] : bucket_array α β n) = [] :=
list.eq_nil_iff_forall_not_mem.mpr $ λ x m,
let ⟨i, h⟩ := (bucket_array.mem_as_list _).1 m in h

parameter [decidable_eq α]

/-- Search a bucket for a key `a` and return the value -/
def find_aux (a : α) : list (Σ a, β a) → option (β a)
| []          := none
| (⟨a',b⟩::t) := if h : a' = a then some (eq.rec_on h b) else find_aux t

theorem find_aux_iff {a : α} {b : β a} :
  Π {l : list Σ a, β a}, (l.map sigma.fst).nodup → (find_aux a l = some b ↔ sigma.mk a b ∈ l)
| []          nd := ⟨λn, by injection n, false.elim⟩
| (⟨a',b'⟩::t) nd := begin
  by_cases a' = a,
  { clear find_aux_iff, subst h,
    suffices : b' = b ↔ b' = b ∨ sigma.mk a' b ∈ t, {simpa [find_aux, eq_comm]},
    refine (or_iff_left_of_imp (λ m, _)).symm,
    have : a' ∉ t.map sigma.fst, from list.not_mem_of_nodup_cons nd,
    exact this.elim (list.mem_map_of_mem sigma.fst m) },
  { have : sigma.mk a b ≠ ⟨a', b'⟩,
    { intro e, injection e with e, exact h e.symm },
    simp at nd, simp [find_aux, h, ne.symm h, find_aux_iff, nd] }
end

/-- Returns `tt` if the bucket `l` contains the key `a` -/
def contains_aux (a : α) (l : list Σ a, β a) : bool :=
(find_aux a l).is_some

theorem contains_aux_iff {a : α} {l : list Σ a, β a} (nd : (l.map sigma.fst).nodup) :
  contains_aux a l ↔ a ∈ l.map sigma.fst :=
begin
  unfold contains_aux,
  cases h : find_aux a l with b; simp,
  { assume (b : β a) (m : sigma.mk a b ∈ l),
    rw (find_aux_iff nd).2 m at h,
    contradiction },
  { show ∃ (b : β a), sigma.mk a b ∈ l,
    exact ⟨_, (find_aux_iff nd).1 h⟩ },
end

/-- Modify a bucket to replace a value in the list. Leaves the list
 unchanged if the key is not found. -/
def replace_aux (a : α) (b : β a) : list (Σ a, β a) → list (Σ a, β a)
| []            := []
| (⟨a', b'⟩::t) := if a' = a then ⟨a, b⟩::t else ⟨a', b'⟩ :: replace_aux t

/-- Modify a bucket to remove a key, if it exists. -/
def erase_aux (a : α) : list (Σ a, β a) → list (Σ a, β a)
| []            := []
| (⟨a', b'⟩::t) := if a' = a then t else ⟨a', b'⟩ :: erase_aux t

/-- The predicate `valid bkts sz` means that `bkts` satisfies the `hash_map`
  invariants: There are exactly `sz` elements in it, every pair is in the
  bucket determined by its key and the hash function, and no key appears
  multiple times in the list. -/
structure valid {n} (bkts : bucket_array α β n) (sz : nat) : Prop :=
(len : bkts.as_list.length = sz)
(idx : ∀ {i} {a : Σ a, β a}, a ∈ array.read bkts i →
  mk_idx n (hash_fn a.1) = i)
(nodup : ∀i, ((array.read bkts i).map sigma.fst).nodup)

theorem valid.idx_enum {n} {bkts : bucket_array α β n} {sz : nat} (v : valid bkts sz)
  {i l} (he : (i, l) ∈ bkts.to_list.enum) {a} {b : β a} (hl : sigma.mk a b ∈ l) :
  ∃ h, mk_idx n (hash_fn a) = ⟨i, h⟩ :=
(array.mem_to_list_enum.mp he).imp (λ h e, by subst e; exact v.idx hl)

theorem valid.idx_enum_1 {n} {bkts : bucket_array α β n} {sz : nat} (v : valid bkts sz)
  {i l} (he : (i, l) ∈ bkts.to_list.enum) {a} {b : β a} (hl : sigma.mk a b ∈ l) :
  (mk_idx n (hash_fn a)).1 = i :=
let ⟨h, e⟩ := v.idx_enum _ he hl in by rw e; refl

theorem valid.as_list_nodup {n} {bkts : bucket_array α β n} {sz : nat} (v : valid bkts sz) :
  (bkts.as_list.map sigma.fst).nodup :=
begin
  suffices : (bkts.to_list.map (list.map sigma.fst)).pairwise list.disjoint,
  { suffices : ∀ l, array.mem l bkts → (l.map sigma.fst).nodup,
      by simpa [bucket_array.as_list, list.nodup_join, *],
    rintros l ⟨i, rfl⟩,
    apply v.nodup },
  rw [← list.enum_map_snd bkts.to_list, list.pairwise_map, list.pairwise_map],
  have : (bkts.to_list.enum.map prod.fst).nodup := by simp [list.nodup_range],
  refine list.pairwise.imp_of_mem _ ((list.pairwise_map _).1 this),
  rw prod.forall, intros i l₁,
  rw prod.forall, intros j l₂ me₁ me₂ ij,
  simp [list.disjoint], intros a b ml₁ b' ml₂,
  apply ij, rwa [← v.idx_enum_1 _ me₁ ml₁, ← v.idx_enum_1 _ me₂ ml₂]
end

theorem mk_valid (n : ℕ+) : @valid n (mk_array n []) 0 :=
⟨by simp [mk_as_list], λ i a h, by cases h, λ i, list.nodup_nil⟩

theorem valid.find_aux_iff {n} {bkts : bucket_array α β n} {sz : nat} (v : valid bkts sz) {a : α}
  {b : β a} :
  find_aux a (bkts.read hash_fn a) = some b ↔ sigma.mk a b ∈ bkts.as_list :=
(find_aux_iff (v.nodup _)).trans $
by rw bkts.mem_as_list; exact ⟨λ h, ⟨_, h⟩, λ ⟨i, h⟩, (v.idx h).symm ▸ h⟩

theorem valid.contains_aux_iff {n} {bkts : bucket_array α β n} {sz : nat} (v : valid bkts sz)
  (a : α) :
  contains_aux a (bkts.read hash_fn a) ↔ a ∈ bkts.as_list.map sigma.fst :=
by simp [contains_aux, option.is_some_iff_exists, v.find_aux_iff hash_fn]

section
  parameters {n : ℕ+} {bkts : bucket_array α β n}
             {bidx : fin n} {f : list (Σ a, β a) → list (Σ a, β a)}
             (u v1 v2 w : list Σ a, β a)

  local notation `L` := array.read bkts bidx
  private def bkts' : bucket_array α β n := array.write bkts bidx (f L)

  variables (hl : L = u ++ v1 ++ w)
            (hfl : f L = u ++ v2 ++ w)
  include hl hfl

  theorem append_of_modify :
  ∃ u' w', bkts.as_list = u' ++ v1 ++ w' ∧ bkts'.as_list = u' ++ v2 ++ w' :=
  begin
    unfold bucket_array.as_list,
    have h : (bidx : ℕ) < bkts.to_list.length, { simp only [bidx.is_lt, array.to_list_length] },
    refine ⟨(bkts.to_list.take bidx).join ++ u, w ++ (bkts.to_list.drop (bidx+1)).join, _, _⟩,
    { conv { to_lhs,
        rw [← list.take_append_drop bidx bkts.to_list, list.drop_eq_nth_le_cons h],
        simp [hl] }, simp },
    { conv { to_lhs,
        rw [bkts', array.write_to_list, list.update_nth_eq_take_cons_drop _ h],
        simp [hfl] }, simp }
  end

  variables (hvnd : (v2.map sigma.fst).nodup)
            (hal : ∀ (a : Σ a, β a), a ∈ v2 → mk_idx n (hash_fn a.1) = bidx)
            (djuv : (u.map sigma.fst).disjoint (v2.map sigma.fst))
            (djwv : (w.map sigma.fst).disjoint (v2.map sigma.fst))
  include hvnd hal djuv djwv

  theorem valid.modify {sz : ℕ} (v : valid bkts sz) :
    v1.length ≤ sz + v2.length ∧ valid bkts' (sz + v2.length - v1.length) :=
  begin
    rcases append_of_modify u v1 v2 w hl hfl with ⟨u', w', e₁, e₂⟩,
    rw [← v.len, e₁],
    suffices : valid bkts' (u' ++ v2 ++ w').length,
    { simpa [ge, add_comm, add_left_comm, nat.le_add_right, add_tsub_cancel_left] },
    refine ⟨congr_arg _ e₂, λ i a, _, λ i, _⟩,
    { by_cases bidx = i,
      { subst i, rw [bkts', array.read_write, hfl],
        have := @valid.idx _ _ _ v bidx a,
        simp only [hl, list.mem_append, or_imp_distrib, forall_and_distrib] at this ⊢,
        exact ⟨⟨this.1.1, hal _⟩, this.2⟩ },
      { rw [bkts', array.read_write_of_ne _ _ h], apply v.idx } },
    { by_cases bidx = i,
      { subst i, rw [bkts', array.read_write, hfl],
        have := @valid.nodup _ _ _ v bidx,
        simp [hl, list.nodup_append] at this,
        simp [list.nodup_append, this, hvnd, djuv, djwv.symm] },
      { rw [bkts', array.read_write_of_ne _ _ h], apply v.nodup } }
  end
end

theorem valid.replace_aux (a : α) (b : β a) : Π (l : list (Σ a, β a)), a ∈ l.map sigma.fst →
  ∃ (u w : list Σ a, β a) b', l = u ++ [⟨a, b'⟩] ++ w ∧ replace_aux a b l = u ++ [⟨a, b⟩] ++ w
| []            := false.elim
| (⟨a', b'⟩::t) := begin
  by_cases e : a' = a,
  { subst a',
    suffices : ∃ (u w : list Σ a, β a) (b'' : β a),
      (sigma.mk a b') :: t = u ++ ⟨a, b''⟩ :: w ∧
      replace_aux a b (⟨a, b'⟩ :: t) = u ++ ⟨a, b⟩ :: w, {simpa},
    refine ⟨[], t, b', _⟩, simp [replace_aux] },
  { suffices : ∀ (x : β a) (_ : sigma.mk a x ∈ t), ∃ u w (b'' : β a),
      (sigma.mk a' b') :: t = u ++ ⟨a, b''⟩ :: w ∧
      (sigma.mk a' b') :: (replace_aux a b t) = u ++ ⟨a, b⟩ :: w,
    { simpa [replace_aux, ne.symm e, e] },
    intros x m,
    have IH : ∀ (x : β a) (_ : sigma.mk a x ∈ t), ∃ u w (b'' : β a),
      t = u ++ ⟨a, b''⟩ :: w ∧ replace_aux a b t = u ++ ⟨a, b⟩ :: w,
    { simpa using valid.replace_aux t },
    rcases IH x m with ⟨u, w, b'', hl, hfl⟩,
    exact ⟨⟨a', b'⟩ :: u, w, b'', by simp [hl, hfl.symm, ne.symm e]⟩ }
end

theorem valid.replace {n : ℕ+}
  {bkts : bucket_array α β n} {sz : ℕ} (a : α) (b : β a)
  (Hc : contains_aux a (bkts.read hash_fn a))
  (v : valid bkts sz) : valid (bkts.modify hash_fn a (replace_aux a b)) sz :=
begin
  have nd := v.nodup (mk_idx n (hash_fn a)),
  rcases hash_map.valid.replace_aux a b (array.read bkts (mk_idx n (hash_fn a)))
    ((contains_aux_iff nd).1 Hc) with ⟨u, w, b', hl, hfl⟩,
  simp [hl, list.nodup_append] at nd,
  refine (v.modify hash_fn
    u [⟨a, b'⟩] [⟨a, b⟩] w hl hfl (list.nodup_singleton _)
    (λa' e, by simp at e; rw e)
    (λa' e1 e2, _)
    (λa' e1 e2, _)).2;
  { revert e1, simp [-sigma.exists] at e2, subst a', simp [nd] }
end

theorem valid.insert {n : ℕ+}
  {bkts : bucket_array α β n} {sz : ℕ} (a : α) (b : β a)
  (Hnc : ¬ contains_aux a (bkts.read hash_fn a))
  (v : valid bkts sz) : valid (reinsert_aux bkts a b) (sz+1) :=
begin
  have nd := v.nodup (mk_idx n (hash_fn a)),
  refine (v.modify hash_fn
    [] [] [⟨a, b⟩] (bkts.read hash_fn a) rfl rfl (list.nodup_singleton _)
    (λa' e, by simp at e; rw e)
    (λa', false.elim)
    (λa' e1 e2, _)).2,
  simp [-sigma.exists] at e2, subst a',
  exact Hnc ((contains_aux_iff nd).2 e1)
end

theorem valid.erase_aux (a : α) : Π (l : list (Σ a, β a)), a ∈ l.map sigma.fst →
  ∃ (u w : list Σ a, β a) b, l = u ++ [⟨a, b⟩] ++ w ∧ erase_aux a l = u ++ [] ++ w
| []            := false.elim
| (⟨a', b'⟩::t) := begin
  by_cases e : a' = a,
  { subst a',
    simpa [erase_aux, and_comm] using show ∃ u w (x : β a),
      t = u ++ w ∧ (sigma.mk a b') :: t = u ++ ⟨a, x⟩ :: w,
      from ⟨[], t, b', by simp⟩ },
  { simp [erase_aux, e, ne.symm e],
    suffices : ∀ (b : β a) (_ : sigma.mk a b ∈ t), ∃ u w (x : β a),
      (sigma.mk a' b') :: t = u ++ ⟨a, x⟩ :: w ∧
      (sigma.mk a' b') :: (erase_aux a t) = u ++ w,
    { simpa [replace_aux, ne.symm e, e] },
    intros b m,
    have IH : ∀ (x : β a) (_ : sigma.mk a x ∈ t), ∃ u w (x : β a),
      t = u ++ ⟨a, x⟩ :: w ∧ erase_aux a t = u ++ w,
    { simpa using valid.erase_aux t },
    rcases IH b m with ⟨u, w, b'', hl, hfl⟩,
    exact ⟨⟨a', b'⟩ :: u, w, b'', by simp [hl, hfl.symm]⟩ }
end

theorem valid.erase {n} {bkts : bucket_array α β n} {sz}
  (a : α) (Hc : contains_aux a (bkts.read hash_fn a))
  (v : valid bkts sz) : valid (bkts.modify hash_fn a (erase_aux a)) (sz-1) :=
begin
  have nd := v.nodup (mk_idx n (hash_fn a)),
  rcases hash_map.valid.erase_aux a (array.read bkts (mk_idx n (hash_fn a)))
    ((contains_aux_iff nd).1 Hc) with ⟨u, w, b, hl, hfl⟩,
  refine (v.modify hash_fn u [⟨a, b⟩] [] w hl hfl list.nodup_nil _ _ _).2;
  simp
end

end
end hash_map

/-- A hash map data structure, representing a finite key-value map
  with key type `α` and value type `β` (which may depend on `α`). -/
structure hash_map (α : Type u) [decidable_eq α] (β : α → Type v) :=
(hash_fn : α → nat)
(size : ℕ)
(nbuckets : ℕ+)
(buckets : bucket_array α β nbuckets)
(is_valid : hash_map.valid hash_fn buckets size)

/-- Construct an empty hash map with buffer size `nbuckets` (default 8). -/
def mk_hash_map {α : Type u} [decidable_eq α] {β : α → Type v} (hash_fn : α → nat) (nbuckets := 8) :
  hash_map α β :=
let n := if nbuckets = 0 then 8 else nbuckets in
let nz : n > 0 := by abstract { cases nbuckets; simp [if_pos, nat.succ_ne_zero] } in
{ hash_fn  := hash_fn,
  size     := 0,
  nbuckets := ⟨n, nz⟩,
  buckets  := mk_array n [],
  is_valid := hash_map.mk_valid _ _ }

namespace hash_map
variables {α : Type u} {β : α → Type v} [decidable_eq α]

/-- Return the value corresponding to a key, or `none` if not found -/
def find (m : hash_map α β) (a : α) : option (β a) :=
find_aux a (m.buckets.read m.hash_fn a)

/-- Return `tt` if the key exists in the map -/
def contains (m : hash_map α β) (a : α) : bool :=
(m.find a).is_some

instance : has_mem α (hash_map α β) := ⟨λa m, m.contains a⟩

/-- Fold a function over the key-value pairs in the map -/
def fold {δ : Type w} (m : hash_map α β) (d : δ) (f : δ → Π a, β a → δ) : δ :=
m.buckets.foldl d f

/-- The list of key-value pairs in the map -/
def entries (m : hash_map α β) : list Σ a, β a :=
m.buckets.as_list

/-- The list of keys in the map -/
def keys (m : hash_map α β) : list α :=
m.entries.map sigma.fst

theorem find_iff (m : hash_map α β) (a : α) (b : β a) :
  m.find a = some b ↔ sigma.mk a b ∈ m.entries :=
m.is_valid.find_aux_iff _

theorem contains_iff (m : hash_map α β) (a : α) :
  m.contains a ↔ a ∈ m.keys :=
m.is_valid.contains_aux_iff _ _

theorem entries_empty (hash_fn : α → nat) (n) :
  (@mk_hash_map α _ β hash_fn n).entries = [] :=
mk_as_list _

theorem find_empty (hash_fn : α → nat) (n a) :
  (@mk_hash_map α _ β hash_fn n).find a = none :=
by induction h : (@mk_hash_map α _ β hash_fn n).find a; [refl,
   { have := (find_iff _ _ _).1 h, rw entries_empty at this, contradiction }]

theorem not_contains_empty (hash_fn : α → nat) (n a) :
  ¬ (@mk_hash_map α _ β hash_fn n).contains a :=
by apply bool_iff_false.2; dsimp [contains]; rw [find_empty]; refl

theorem insert_lemma (hash_fn : α → nat) {n n'}
  {bkts : bucket_array α β n} {sz} (v : valid hash_fn bkts sz) :
  valid hash_fn (bkts.foldl (mk_array _ [] : bucket_array α β n') (reinsert_aux hash_fn)) sz :=
begin
  suffices : ∀ (l : list Σ a, β a) (t : bucket_array α β n') sz,
    valid hash_fn t sz → ((l ++ t.as_list).map sigma.fst).nodup →
    valid hash_fn (l.foldl (λr (a : Σ a, β a), reinsert_aux hash_fn r a.1 a.2) t) (sz + l.length),
  { have p := this bkts.as_list _ _ (mk_valid _ _),
    rw [mk_as_list, list.append_nil, zero_add, v.len] at p,
    rw bucket_array.foldl_eq,
    exact p (v.as_list_nodup _) },
  intro l, induction l with c l IH; intros t sz v nd, {exact v},
  rw show sz + (c :: l).length = sz + 1 + l.length, by simp [add_comm, add_assoc],
  rcases (show (l.map sigma.fst).nodup ∧
      ((bucket_array.as_list t).map sigma.fst).nodup ∧
      c.fst ∉ l.map sigma.fst ∧
      c.fst ∉ (bucket_array.as_list t).map sigma.fst ∧
      (l.map sigma.fst).disjoint ((bucket_array.as_list t).map sigma.fst),
    by simpa [list.nodup_append, not_or_distrib, and_comm, and.left_comm] using nd)
    with ⟨nd1, nd2, nm1, nm2, dj⟩,
  have v' := v.insert _ _ c.2 (λHc, nm2 $ (v.contains_aux_iff _ c.1).1 Hc),
  apply IH _ _ v',
  suffices : ∀ ⦃a : α⦄ (b : β a), sigma.mk a b ∈ l →
    ∀ (b' : β a), sigma.mk a b' ∈ (reinsert_aux hash_fn t c.1 c.2).as_list → false,
  { simpa [list.nodup_append, nd1, v'.as_list_nodup _, list.disjoint] },
  intros a b m1 b' m2,
  rcases (reinsert_aux hash_fn t c.1 c.2).mem_as_list.1 m2 with ⟨i, im⟩,
  have : sigma.mk a b' ∉ array.read t i,
  { intro m3,
    have : a ∈ list.map sigma.fst t.as_list :=
      list.mem_map_of_mem sigma.fst (t.mem_as_list.2 ⟨_, m3⟩),
    exact dj (list.mem_map_of_mem sigma.fst m1) this },
  by_cases h : mk_idx n' (hash_fn c.1) = i,
  { subst h,
    have e : sigma.mk a b' = ⟨c.1, c.2⟩,
    { simpa [reinsert_aux, bucket_array.modify, array.read_write, this] using im },
    injection e with e, subst a,
    exact nm1.elim (@list.mem_map_of_mem _ _ sigma.fst _ _ m1) },
  { apply this,
    simpa [reinsert_aux, bucket_array.modify, array.read_write_of_ne _ _ h] using im }
end

/-- Insert a key-value pair into the map. (Modifies `m` in-place when applicable) -/
def insert : Π (m : hash_map α β) (a : α) (b : β a), hash_map α β
| ⟨hash_fn, size, n, buckets, v⟩ a b :=
let bkt := buckets.read hash_fn a in
if hc : contains_aux a bkt then
{ hash_fn  := hash_fn,
  size     := size,
  nbuckets := n,
  buckets  := buckets.modify hash_fn a (replace_aux a b),
  is_valid := v.replace _ a b hc }
else
let size'    := size + 1,
    buckets' := buckets.modify hash_fn a (λl, ⟨a, b⟩::l),
    valid'   := v.insert _ a b hc in
if size' ≤ n then
{ hash_fn  := hash_fn,
  size     := size',
  nbuckets := n,
  buckets  := buckets',
  is_valid := valid' }
else
let n'        : ℕ+ := ⟨n * 2, mul_pos n.2 dec_trivial⟩,
    buckets'' : bucket_array α β n' :=
                buckets'.foldl (mk_array _ []) (reinsert_aux hash_fn) in
{ hash_fn  := hash_fn,
  size     := size',
  nbuckets := n',
  buckets  := buckets'',
  is_valid := insert_lemma _ valid' }

theorem mem_insert : Π (m : hash_map α β) (a b a' b'),
  (sigma.mk a' b' : sigma β) ∈ (m.insert a b).entries ↔
  if a = a' then b == b' else sigma.mk a' b' ∈ m.entries
| ⟨hash_fn, size, n, bkts, v⟩ a b a' b' := begin
  let bkt := bkts.read hash_fn a,
  have nd : (bkt.map sigma.fst).nodup := v.nodup (mk_idx n (hash_fn a)),
  have lem : Π (bkts' : bucket_array α β n) (v1 u w)
    (hl : bucket_array.as_list bkts = u ++ v1 ++ w)
    (hfl : bucket_array.as_list bkts' = u ++ [⟨a, b⟩] ++ w)
    (veq : (v1 = [] ∧ ¬ contains_aux a bkt) ∨ ∃b'', v1 = [⟨a, b''⟩]),
    sigma.mk a' b' ∈ bkts'.as_list ↔
    if a = a' then b == b' else sigma.mk a' b' ∈ bkts.as_list,
  { intros bkts' v1 u w hl hfl veq,
    rw [hl, hfl],
    by_cases h : a = a',
    { subst a',
      suffices : b = b' ∨ sigma.mk a b' ∈ u ∨ sigma.mk a b' ∈ w ↔ b = b',
      { simpa [eq_comm, or.left_comm] },
      refine or_iff_left_of_imp (not.elim $ not_or_distrib.2 _),
      rcases veq with ⟨rfl, Hnc⟩ | ⟨b'', rfl⟩,
      { have na := (not_iff_not_of_iff $ v.contains_aux_iff _ _).1 Hnc,
        simp [hl, not_or_distrib] at na, simp [na] },
      { have nd' := v.as_list_nodup _,
        simp [hl, list.nodup_append] at nd', simp [nd'] } },
    { suffices : sigma.mk a' b' ∉ v1, {simp [h, ne.symm h, this]},
      rcases veq with ⟨rfl, Hnc⟩ | ⟨b'', rfl⟩; simp [ne.symm h] } },
  by_cases Hc : (contains_aux a bkt : Prop),
  { rcases hash_map.valid.replace_aux a b (array.read bkts (mk_idx n (hash_fn a)))
      ((contains_aux_iff nd).1 Hc) with ⟨u', w', b'', hl', hfl'⟩,
    rcases (append_of_modify u' [⟨a, b''⟩] [⟨a, b⟩] w' hl' hfl') with ⟨u, w, hl, hfl⟩,
    simpa [insert, @dif_pos (contains_aux a bkt) _ Hc]
      using lem _ _ u w hl hfl (or.inr ⟨b'', rfl⟩) },
  { let size' := size + 1,
    let bkts' := bkts.modify hash_fn a (λl, ⟨a, b⟩::l),
    have mi : sigma.mk a' b' ∈ bkts'.as_list ↔
        if a = a' then b == b' else sigma.mk a' b' ∈ bkts.as_list :=
      let ⟨u, w, hl, hfl⟩ := append_of_modify [] [] [⟨a, b⟩] _ rfl rfl in
      lem bkts' _ u w hl hfl $ or.inl ⟨rfl, Hc⟩,
    simp [insert, @dif_neg (contains_aux a bkt) _ Hc],
    by_cases h : size' ≤ n,
    { simpa [show size' ≤ n, from h] using mi },
    { let n' : ℕ+ := ⟨n * 2, mul_pos n.2 dec_trivial⟩,
      let bkts'' : bucket_array α β n' := bkts'.foldl (mk_array _ []) (reinsert_aux hash_fn),
      suffices : sigma.mk a' b' ∈ bkts''.as_list ↔ sigma.mk a' b' ∈ bkts'.as_list.reverse,
      { simpa [show ¬ size' ≤ n, from h, mi] },
      rw [show bkts'' = bkts'.as_list.foldl _ _, from bkts'.foldl_eq _ _,
          ← list.foldr_reverse],
      induction bkts'.as_list.reverse with a l IH,
      { simp [mk_as_list] },
      { cases a with a'' b'',
        let B := l.foldr (λ (y : sigma β) (x : bucket_array α β n'),
          reinsert_aux hash_fn x y.1 y.2) (mk_array n' []),
        rcases append_of_modify [] [] [⟨a'', b''⟩] _ rfl rfl with ⟨u, w, hl, hfl⟩,
        simp [IH.symm, or.left_comm, show B.as_list = _, from hl,
              show (reinsert_aux hash_fn B a'' b'').as_list = _, from hfl] } } }
end

theorem find_insert_eq (m : hash_map α β) (a : α) (b : β a) : (m.insert a b).find a = some b :=
(find_iff (m.insert a b) a b).2 $ (mem_insert m a b a b).2 $ by rw if_pos rfl

theorem find_insert_ne (m : hash_map α β) (a a' : α) (b : β a) (h : a ≠ a') :
  (m.insert a b).find a' = m.find a' :=
option.eq_of_eq_some $ λb',
let t := mem_insert m a b a' b' in
(find_iff _ _ _).trans $ iff.trans (by rwa if_neg h at t) (find_iff _ _ _).symm

theorem find_insert (m : hash_map α β) (a' a : α) (b : β a) :
  (m.insert a b).find a' = if h : a = a' then some (eq.rec_on h b) else m.find a' :=
if h : a = a' then by rw dif_pos h; exact
  match a', h with ._, rfl := find_insert_eq m a b end
else by rw dif_neg h; exact find_insert_ne m a a' b h

/-- Insert a list of key-value pairs into the map. (Modifies `m` in-place when applicable) -/
def insert_all (l : list (Σ a, β a)) (m : hash_map α β) : hash_map α β :=
l.foldl (λ m ⟨a, b⟩, insert m a b) m

/-- Construct a hash map from a list of key-value pairs. -/
def of_list (l : list (Σ a, β a)) (hash_fn) : hash_map α β :=
insert_all l (mk_hash_map hash_fn (2 * l.length))

/-- Remove a key from the map. (Modifies `m` in-place when applicable) -/
def erase (m : hash_map α β) (a : α) : hash_map α β :=
match m with ⟨hash_fn, size, n, buckets, v⟩ :=
  if hc : contains_aux a (buckets.read hash_fn a) then
  { hash_fn  := hash_fn,
    size     := size - 1,
    nbuckets := n,
    buckets  := buckets.modify hash_fn a (erase_aux a),
    is_valid := v.erase _ a hc }
  else m
end

theorem mem_erase : Π (m : hash_map α β) (a a' b'),
  (sigma.mk a' b' : sigma β) ∈ (m.erase a).entries ↔
  a ≠ a' ∧ sigma.mk a' b' ∈ m.entries
| ⟨hash_fn, size, n, bkts, v⟩ a a' b' := begin
  let bkt := bkts.read hash_fn a,
  by_cases Hc : (contains_aux a bkt : Prop),
  { let bkts' := bkts.modify hash_fn a (erase_aux a),
    suffices : sigma.mk a' b' ∈ bkts'.as_list ↔ a ≠ a' ∧ sigma.mk a' b' ∈ bkts.as_list,
    { simpa [erase, @dif_pos (contains_aux a bkt) _ Hc] },
    have nd := v.nodup (mk_idx n (hash_fn a)),
    rcases valid.erase_aux a bkt ((contains_aux_iff nd).1 Hc) with ⟨u', w', b, hl', hfl'⟩,
    rcases append_of_modify u' [⟨a, b⟩] [] _ hl' hfl' with ⟨u, w, hl, hfl⟩,
    suffices : ∀_:sigma.mk a' b' ∈ u ∨ sigma.mk a' b' ∈ w, a ≠ a',
    { have : sigma.mk a' b' ∈ u ∨ sigma.mk a' b' ∈ w ↔ (¬a = a' ∧ a' = a) ∧ b' == b ∨
        ¬a = a' ∧ (sigma.mk a' b' ∈ u ∨ sigma.mk a' b' ∈ w),
      { simp [eq_comm, not_and_self_iff, and_iff_right_of_imp this] },
      simpa [hl, show bkts'.as_list = _, from hfl, and_or_distrib_left,
             and_comm, and.left_comm, or.left_comm] },
    intros m e, subst a', revert m, apply not_or_distrib.2,
    have nd' := v.as_list_nodup _,
    simp [hl, list.nodup_append] at nd', simp [nd'] },
  { suffices : ∀_:sigma.mk a' b' ∈ bucket_array.as_list bkts, a ≠ a',
    { simp [erase, @dif_neg (contains_aux a bkt) _ Hc, entries, and_iff_right_of_imp this] },
    intros m e, subst a',
    exact Hc ((v.contains_aux_iff _ _).2 (list.mem_map_of_mem sigma.fst m)) }
end

theorem find_erase_eq (m : hash_map α β) (a : α) : (m.erase a).find a = none :=
begin
  cases h : (m.erase a).find a with b, {refl},
  exact absurd rfl ((mem_erase m a a b).1 ((find_iff (m.erase a) a b).1 h)).left
end

theorem find_erase_ne (m : hash_map α β) (a a' : α) (h : a ≠ a') :
  (m.erase a).find a' = m.find a' :=
option.eq_of_eq_some $ λb',
(find_iff _ _ _).trans $ (mem_erase m a a' b').trans $
  (and_iff_right h).trans (find_iff _ _ _).symm

theorem find_erase (m : hash_map α β) (a' a : α) :
  (m.erase a).find a' = if a = a' then none else m.find a' :=
if h : a = a' then by subst a'; simp [find_erase_eq m a]
else by rw if_neg h; exact find_erase_ne m a a' h

section string
variables [has_to_string α] [∀ a, has_to_string (β a)]
open prod
private def key_data_to_string (a : α) (b : β a) (first : bool) : string :=
(if first then "" else ", ") ++ sformat!"{a} ← {b}"

private def to_string (m : hash_map α β) : string :=
"⟨" ++ (fst (fold m ("", tt) (λ p a b, (fst p ++ key_data_to_string a b (snd p), ff)))) ++ "⟩"

instance : has_to_string (hash_map α β) :=
⟨to_string⟩

end string

section format
open format prod
variables [has_to_format α] [∀ a, has_to_format (β a)]

private meta def format_key_data (a : α) (b : β a) (first : bool) : format :=
(if first then to_fmt "" else to_fmt "," ++ line) ++
  to_fmt a ++ space ++ to_fmt "←" ++ space ++ to_fmt b

private meta def to_format (m : hash_map α β) : format :=
group $ to_fmt "⟨" ++
  nest 1 (fst (fold m (to_fmt "", tt) (λ p a b, (fst p ++ format_key_data a b (snd p), ff)))) ++
  to_fmt "⟩"

meta instance : has_to_format (hash_map α β) :=
⟨to_format⟩
end format

/-- `hash_map` with key type `nat` and value type that may vary. -/
instance {β : ℕ → Type*} : inhabited (hash_map ℕ β) := ⟨mk_hash_map id⟩

end hash_map


/- Author: E.W.Ayers.
   This section roughly follows Chapter 3, §1, §2 of Sheaves in Geology and Logic by Saunders Maclane and Ieke M.
 -/

import sieve
import category.pullbacks

universes u v w
namespace category_theory

open category_theory limits order lattice

/-- A set of sieves for every object in the category: a candidate to be a Grothendieck topology. -/
def sieve_set (C : Type u) [category.{v} C] := Π (X : C), set (sieve X)

def sieve_set.trivial (C : Type u) [category.{v} C] : sieve_set C := λ X, {⊤}

lemma mem_trivial (C : Type u) [category.{v} C] {X : C} (S : sieve X) :
  S ∈ sieve_set.trivial C X ↔ S = ⊤ :=
set.mem_singleton_iff

/-- A sieve on `X` is dense if for any arrow `f : Y ⟶ X`, there is a `g : Z ⟶ Y` with `g ≫ f ∈ S`. -/
def sieve_set.dense (C : Type u) [category.{v} C] : sieve_set C :=
λ X, {S | ∀ {Y : C} (f : Y ⟶ X), ∃ Z (g : Z ⟶ Y), over.mk (g ≫ f) ∈ S.arrows }

/-- The atomic sieve_set just contains all of the non-empty sieves. -/
def sieve_set.atomic (C : Type u) [category.{v} C] : sieve_set C :=
λ X, {S | ∃ {Y} (f : Y ⟶ X), over.mk f ∈ S.arrows}

open sieve category

/--
Definition of a Grothendieck Topology: a set of sieves `J X` on each object `X` satisfying three axioms:
1. For every object `X`, the maximal sieve is in `J X`.
2. If `S ∈ J X` then its pullback along any `h : Y ⟶ X` is in `J Y`.
3. If `S ∈ J X` and `R` is a sieve on `X`, then provided that the pullback of `R` along any arrow
   `f : Y ⟶ X` in `S` is in `J Y`, we have that `R` itself is in `J X`.
-/
class grothendieck {C : Type u} [category.{v} C] (J : sieve_set C) : Prop :=
(max : ∀ X, ⊤ ∈ J X)
(stab : ∀ {X Y} (S ∈ J X) (h : Y ⟶ X), sieve.pullback S h ∈ J Y)
(trans : ∀ ⦃X⦄ (S : sieve X) (hS : S ∈ J X) (R : sieve X), (∀ {Y} (f : Y ⟶ X), over.mk f ∈ S.arrows → R.pullback f ∈ J Y) → R ∈ J X)

/-- A site is a category equipped with a grothendieck topology. -/
structure Site :=
(C : Type u)
[𝒞 : category.{v} C]
(J : sieve_set C)
[g : grothendieck J]

namespace grothendieck
variables {C : Type u} [category.{v} C]
variables {X Y : C} {S R : sieve X}
variables {J : sieve_set C} [grothendieck J]

def superset_covering (Hss : S ≤ R) (sjx : S ∈ J X) : R ∈ J X :=
begin
  apply grothendieck.trans _ sjx,
  intros Y h hh,
  dsimp,
  have : S.pullback h ≤ R.pullback h,
    apply pullback_le_map Hss,
  have : S.pullback h = ⊤,
    rw ← id_mem_iff_eq_top,
    simpa,
  have : R.pullback h = ⊤,
    apply top_unique,
    rwa ← this,
  rw this,
  apply grothendieck.max,
end

def covers (J : sieve_set C) (S : sieve X) (f : Y ⟶ X) : Prop := S.pullback f ∈ J Y

lemma arrow_max (f : Y ⟶ X) (S : sieve X) [grothendieck J] (hf : over.mk f ∈ S.arrows) : covers J S f :=
begin
  rw [covers, (pullback_eq_top_iff_mem f).1 hf],
  apply grothendieck.max,
end
lemma arrow_stab (f : Y ⟶ X) (S : sieve X) (h : covers J S f) {Z : C} (g : Z ⟶ Y) : covers J S (g ≫ f) :=
begin
  rw [covers, pullback_comp],
  apply grothendieck.stab,
  apply h,
end
lemma arrow_trans (f : Y ⟶ X) (S R : sieve X) (h : covers J S f) : (∀ {Z : C} (g : Z ⟶ X), over.mk g ∈ S.arrows → covers J R g) → covers J R f :=
begin
  intro k,
  apply grothendieck.trans (S.pullback f) h,
  intros Z g hg,
  rw ← pullback_comp,
  apply k (g ≫ f) hg,
end

lemma intersection_covering (rj : R ∈ J X) (sj : S ∈ J X) : R ⊓ S ∈ J X :=
begin
  apply grothendieck.trans R rj,
  intros Y f Hf,
  have : S.pullback f ≤ (R ⊓ S).pullback f,
    intros Z g hg,
    refine ⟨downward_closed _ Hf _, hg⟩,
  apply superset_covering this,
  apply grothendieck.stab _ sj,
  apply_instance,
end

lemma arrow_intersect (f : Y ⟶ X) (S R : sieve X) (hS : covers J S f) (hR : covers J R f) : covers J (S ⊓ R) f :=
begin
  rw [covers, pullback_inter],
  apply intersection_covering;
  assumption
end

open sieve_set

instance trivial.grothendieck : grothendieck (sieve_set.trivial C) :=
{ max := λ X, set.mem_singleton _,
  stab := λ X Y S HS h,
  begin
    rw mem_trivial at *,
    rw [HS, pullback_top],
  end,
  trans := λ X S HS R HR,
  begin
    rw [mem_trivial, ← id_mem_iff_eq_top, pullback_eq_top_iff_mem],
    simp only [mem_trivial] at HR,
    apply HR,
    rwa [id_mem_iff_eq_top, ← mem_trivial],
  end }

instance dense.grothendieck : grothendieck (dense C) :=
{ max := λ X Y f, ⟨Y, 𝟙 Y, ⟨⟩⟩,
  stab :=
    begin
      intros X Y S H h Z f,
      rcases H (f ≫ h) with ⟨W, g, H⟩,
      refine ⟨W, g, _⟩,
      simpa,
    end,
  trans :=
    begin
      intros X S H₁ R H₂ Y f,
      rcases H₁ f with ⟨Z,g,H₃⟩,
      rcases H₂ _ H₃ (𝟙 Z) with ⟨W,h,H₄⟩,
      refine ⟨W, (h ≫ 𝟙 Z ≫ g), _⟩,
      simpa using H₄,
    end }

/--
A category satisfies the right Ore condition if any span can be completed to a
commutative square.
NB. Any category with pullbacks obviously satisfies the right Ore condition.
-/
def right_ore_condition (C : Type u) [category.{v} C] : Prop :=
∀ {X Y Z : C} (yx : Y ⟶ X) (zx : Z ⟶ X), ∃ W (wy : W ⟶ Y) (wz : W ⟶ Z), wy ≫ yx = wz ≫ zx

/--
The atomic sieveset is a grothendieck topology when it
satisfies the 'square' property. Which says that every span `Y ⟶ X ⟵ Z` forms a commuting
diagram.
-/
instance atomic.grothendieck
  (hro : right_ore_condition C)
  : grothendieck (atomic C) :=
{ max := λ X, ⟨_, 𝟙 _, ⟨⟩⟩,
  stab :=
  begin
    rintros X Y S ⟨Z, f, hf⟩ h,
    rcases hro h f with ⟨W, g, k, comm⟩,
    refine ⟨_, g, _⟩,
    simp [mem_pullback, comm, hf],
  end,
  trans :=
  begin
    rintros X S ⟨Y, f, hf⟩ R h,
    rcases h f hf with ⟨Z, g, hg⟩,
    exact ⟨_, _, hg⟩,
  end }

open opposite

def matching_family (P : Cᵒᵖ ⥤ Type v) (S : sieve X) :=
S.as_functor ⟶ P

def amalgamation {P : Cᵒᵖ ⥤ Type v} {S : sieve X} (γ : matching_family P S) :=
{α : yoneda.obj X ⟶ P // sieve.functor_inclusion S ≫ α = γ}

@[derive subsingleton]
def sheaf_condition (J : sieve_set C) [grothendieck J] (P : Cᵒᵖ ⥤ Type v) : Type (max u v) :=
Π (X : C) (S : sieve X) (γ : matching_family P S), S ∈ J X → unique (amalgamation γ)

def matching_family' (P : Cᵒᵖ ⥤ Type v) {c : C} (S : sieve c) :=
{x : Π {d : C} (f : d ⟶ c), over.mk f ∈ S.arrows → P.obj (opposite.op d) // ∀ {d e : C} (f : d ⟶ c) (g : e ⟶ d) (h : over.mk f ∈ S.arrows), x (g ≫ f) (sieve.downward_closed _ h _) = P.map g.op (x f h)}

def amalgamation' {P : Cᵒᵖ ⥤ Type v} {c : C} {S : sieve c} (γ : matching_family' P S) :=
{y : P.obj (opposite.op c) // ∀ {d : C} (f : d ⟶ c) (hf : over.mk f ∈ S.arrows), P.map f.op y = γ.1 f hf}

@[derive subsingleton]
def sheaf_condition' (J : sieve_set C) [grothendieck J] (P : Cᵒᵖ ⥤ Type v) : Type (max u v) :=
Π (c : C) (S : sieve c) (γ : matching_family' P S), S ∈ J c → unique (amalgamation' γ)

def matching_family'_equiv_matching_family (P : Cᵒᵖ ⥤ Type v) : matching_family' P S ≃ matching_family P S :=
{ to_fun := λ x, ⟨λ _ t, x.1 _ t.2, λ c c' f, funext $ λ t, x.2 _ _ t.2⟩,
  inv_fun := λ x, ⟨λ d f hf, x.app _ ⟨f, hf⟩, λ d d' f g h, congr_fun (x.2 g.op) ⟨f, h⟩⟩,
  left_inv := λ _, subtype.ext $ funext $ λ _, funext $ λ _, funext $ λ _, rfl,
  right_inv := λ _, by { ext _ ⟨_, _⟩, refl } }

def amalgamation'_equiv_amalgamation (P : Cᵒᵖ ⥤ Type v) (x : matching_family' P S) :
  amalgamation (matching_family'_equiv_matching_family P x) ≃ (amalgamation' x) :=
{ to_fun := λ γ,
  { val := γ.1.app _ (𝟙 X),
    property := λ d f hf,
    begin
      have := congr_fun (γ.1.naturality f.op) (𝟙 _),
      dsimp at this,
      erw ← this,
      rw comp_id,
      have q := congr_arg (λ t, nat_trans.app t (opposite.op d)) γ.2,
      dsimp at q,
      have := congr_fun q ⟨f, hf⟩,
      exact this,
    end },
  inv_fun := λ γ,
  { val :=
    { app := λ c f, P.map f.op γ.1,
      naturality' := λ c c' f, funext $ λ g,
      begin
        dsimp at g,
        dsimp,
        rw P.map_comp,
        refl,
      end },
    property :=
    begin
      ext c ⟨f, hf⟩,
      apply γ.2,
    end },
  left_inv :=
  begin
    rintro ⟨γ₁, γ₂⟩,
    ext d f,
    dsimp,
    dsimp at f,
    have := congr_fun (γ₁.naturality f.op) (𝟙 X),
    dsimp at this,
    rw [← this, comp_id],
  end,
  right_inv :=
  begin
    rintro ⟨γ₁, γ₂⟩,
    ext1,
    dsimp,
    rw P.map_id,
    refl,
  end }

def sheaf'_equiv_sheaf (J : sieve_set C) [grothendieck J] (P : Cᵒᵖ ⥤ Type v) :
  sheaf_condition J P ≅ sheaf_condition' J P :=
{ hom :=
  begin
    intros h c S γ hS,
    apply equiv.unique (amalgamation'_equiv_amalgamation _ _).symm,
    apply h _ _ _ hS,
  end,
  inv :=
  begin
    intros h c S γ hS,
    haveI := h _ _ ((matching_family'_equiv_matching_family P).symm γ) hS,
    have := equiv.unique (amalgamation'_equiv_amalgamation P ((matching_family'_equiv_matching_family P).symm γ)),
    simpa using this,
  end }

end grothendieck

end category_theory


import data.list
import tactic

universes u v w

theorem list.tfae_prf {a b : Prop} {l : list Prop} (h : list.tfae l) (ha : a ∈ l) (hb : b ∈ l) (ha_prf : a) : b :=
(h a ha b hb).mp ha_prf

lemma and_iff_right_of_left_if_right {p q : Prop} (h : p → q) : q ∧ p ↔ p :=
⟨λ h₂, h₂.right, λ h₂, ⟨h h₂, h₂⟩⟩

lemma choice_2_arg {α : Sort u} {β : Sort v} {γ : α → β → Sort w}
{r : Π (x : α) (y : β), γ x y → Prop} (h : ∀ (x : α) (y : β), ∃ (z : γ x y), r x y z) :
∃ (f : Π (x : α) (y : β), γ x y), ∀ (x : α) (y : β), r x y (f x y) :=
begin
  let γ' : (pprod α β) → Sort w := (λ z, γ z.fst z.snd),
  let r' := λ (x : pprod α β) (z : γ' x), r x.fst x.snd z,
  have h' : ∀ (x : pprod α β), ∃ z : γ' x, r' x z := (λ x, h x.fst x.snd),
  rcases classical.axiom_of_choice h' with ⟨f, hf⟩,
  let f := (λ x y, f ⟨x, y⟩),
  existsi f, intros x y, exact hf ⟨x, y⟩,
end

lemma subst_right_of_and {α : Sort u} {p : Prop} {a b c : α} (h : p → b = c) : p ∧ a = b ↔ p ∧ a = c :=
⟨ assume h₂, ⟨h₂.left, (h h₂.left) ▸ h₂.right⟩,
  assume h₂, ⟨h₂.left, (h h₂.left).symm ▸ h₂.right⟩ ⟩

lemma exists_unique_eq {α : Sort u} {β : Sort v} (f : α → β) : ∀ x : α, ∃! y : β, (λ (a : α) (b : β), b = f a) x y :=
λ x, ⟨f x, rfl, λ y, assume hy, hy⟩

lemma rel_congr_left {α : Sort u} {r : α → α → Prop} (symm : symmetric r) (trans : transitive r)
  {x y : α} (xy : r x y) {z : α} : r x z ↔ r y z :=
⟨λ xz, trans (symm xy) xz, λ yz, trans xy yz⟩

lemma iff_of_not_of_not {p : Prop} (np : ¬ p) {q : Prop} (nq : ¬ q) : p ↔ q :=
⟨λ hp, false.elim (np hp), λ hq, false.elim (nq hq)⟩

lemma ite_prop {α : Sort u} {p : α → Prop} {x y : α} (px : p x) (py : p y) {q : Prop} [h : decidable q] :
p (if q then x else y) :=
begin
  by_cases q,
    rw if_pos h, exact px,
  rw if_neg h, exact py,
end


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

! This file was ported from Lean 3 source module category_theory.extensive
! leanprover-community/mathlib commit ac3ae212f394f508df43e37aa093722fa9b65d31
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathbin.CategoryTheory.Limits.Shapes.CommSq
import Mathbin.CategoryTheory.Limits.Shapes.StrictInitial
import Mathbin.CategoryTheory.Limits.Shapes.Types
import Mathbin.Topology.Category.Top.Limits
import Mathbin.CategoryTheory.Limits.FunctorCategory

/-!

# Extensive categories

## Main definitions
- `category_theory.is_van_kampen_colimit`: A (colimit) cocone over a diagram `F : J ⥤ C` is van
  Kampen if for every cocone `c'` over the pullback of the diagram `F' : J ⥤ C'`,
  `c'` is colimiting iff `c'` is the pullback of `c`.
- `category_theory.finitary_extensive`: A category is (finitary) extensive if it has finite
  coproducts, and binary coproducts are van Kampen.

## Main Results
- `category_theory.has_strict_initial_objects_of_finitary_extensive`: The initial object
  in extensive categories is strict.
- `category_theory.finitary_extensive.mono_inr_of_is_colimit`: Coproduct injections are monic in
  extensive categories.
- `category_theory.binary_cofan.is_pullback_initial_to_of_is_van_kampen`: In extensive categories,
  sums are disjoint, i.e. the pullback of `X ⟶ X ⨿ Y` and `Y ⟶ X ⨿ Y` is the initial object.
- `category_theory.types.finitary_extensive`: The category of types is extensive.

## TODO

Show that the following are finitary extensive:
- the categories of sheaves over a site
- `Scheme`
- `AffineScheme` (`CommRingᵒᵖ`)

## References
- https://ncatlab.org/nlab/show/extensive+category
- [Carboni et al, Introduction to extensive and distributive categories][CARBONI1993145]

-/

open CategoryTheory.Limits

namespace CategoryTheory

universe v' u' v u

variable {J : Type v'} [Category.{u'} J] {C : Type u} [Category.{v} C]

/-- A natural transformation is equifibered if every commutative square of the following form is
a pullback.
```
F(X) → F(Y)
 ↓      ↓
G(X) → G(Y)
```
-/
def NatTrans.Equifibered {F G : J ⥤ C} (α : F ⟶ G) : Prop :=
  ∀ ⦃i j : J⦄ (f : i ⟶ j), IsPullback (F.map f) (α.app i) (α.app j) (G.map f)
#align category_theory.nat_trans.equifibered CategoryTheory.NatTrans.Equifibered

theorem NatTrans.equifibered_of_isIso {F G : J ⥤ C} (α : F ⟶ G) [IsIso α] : α.Equifibered :=
  fun _ _ f => IsPullback.of_vert_isIso ⟨NatTrans.naturality _ f⟩
#align category_theory.nat_trans.equifibered_of_is_iso CategoryTheory.NatTrans.equifibered_of_isIso

theorem NatTrans.Equifibered.comp {F G H : J ⥤ C} {α : F ⟶ G} {β : G ⟶ H} (hα : α.Equifibered)
    (hβ : β.Equifibered) : (α ≫ β).Equifibered := fun i j f => (hα f).paste_vert (hβ f)
#align category_theory.nat_trans.equifibered.comp CategoryTheory.NatTrans.Equifibered.comp

/-- A (colimit) cocone over a diagram `F : J ⥤ C` is universal if it is stable under pullbacks. -/
def IsUniversalColimit {F : J ⥤ C} (c : Cocone F) : Prop :=
  ∀ ⦃F' : J ⥤ C⦄ (c' : Cocone F') (α : F' ⟶ F) (f : c'.pt ⟶ c.pt)
    (h : α ≫ c.ι = c'.ι ≫ (Functor.const J).map f) (hα : α.Equifibered),
    (∀ j : J, IsPullback (c'.ι.app j) (α.app j) f (c.ι.app j)) → Nonempty (IsColimit c')
#align category_theory.is_universal_colimit CategoryTheory.IsUniversalColimit

/-- A (colimit) cocone over a diagram `F : J ⥤ C` is van Kampen if for every cocone `c'` over the
pullback of the diagram `F' : J ⥤ C'`, `c'` is colimiting iff `c'` is the pullback of `c`.

TODO: Show that this is iff the functor `C ⥤ Catᵒᵖ` sending `x` to `C/x` preserves it.
TODO: Show that this is iff the inclusion functor `C ⥤ Span(C)` preserves it.
-/
def IsVanKampenColimit {F : J ⥤ C} (c : Cocone F) : Prop :=
  ∀ ⦃F' : J ⥤ C⦄ (c' : Cocone F') (α : F' ⟶ F) (f : c'.pt ⟶ c.pt)
    (h : α ≫ c.ι = c'.ι ≫ (Functor.const J).map f) (hα : α.Equifibered),
    Nonempty (IsColimit c') ↔ ∀ j : J, IsPullback (c'.ι.app j) (α.app j) f (c.ι.app j)
#align category_theory.is_van_kampen_colimit CategoryTheory.IsVanKampenColimit

theorem IsVanKampenColimit.is_universal {F : J ⥤ C} {c : Cocone F} (H : IsVanKampenColimit c) :
    IsUniversalColimit c := fun _ c' α f h hα => (H c' α f h hα).mpr
#align category_theory.is_van_kampen_colimit.is_universal CategoryTheory.IsVanKampenColimit.is_universal

/-- A van Kampen colimit is a colimit. -/
noncomputable def IsVanKampenColimit.isColimit {F : J ⥤ C} {c : Cocone F}
    (h : IsVanKampenColimit c) : IsColimit c :=
  by
  refine'
    ((h c (𝟙 F) (𝟙 c.X : _) (by rw [Functor.map_id, category.comp_id, category.id_comp])
            (nat_trans.equifibered_of_is_iso _)).mpr
        fun j => _).some
  haveI : is_iso (𝟙 c.X) := inferInstance
  exact is_pullback.of_vert_is_iso ⟨by erw [nat_trans.id_app, category.comp_id, category.id_comp]⟩
#align category_theory.is_van_kampen_colimit.is_colimit CategoryTheory.IsVanKampenColimit.isColimit

theorem IsInitial.isVanKampenColimit [HasStrictInitialObjects C] {X : C} (h : IsInitial X) :
    IsVanKampenColimit (asEmptyCocone X) :=
  by
  intro F' c' α f hf hα
  have : F' = functor.empty C := by apply functor.hext <;> rintro ⟨⟨⟩⟩
  subst this
  haveI := h.is_iso_to f
  refine'
    ⟨by rintro _ ⟨⟨⟩⟩, fun _ =>
      ⟨is_colimit.of_iso_colimit h (cocones.ext (as_iso f).symm <| by rintro ⟨⟨⟩⟩)⟩⟩
#align category_theory.is_initial.is_van_kampen_colimit CategoryTheory.IsInitial.isVanKampenColimit

section Extensive

variable {X Y : C}

/-- A category is (finitary) extensive if it has finite coproducts,
and binary coproducts are van Kampen.

TODO: Show that this is iff all finite coproducts are van Kampen. -/
class FinitaryExtensive (C : Type u) [Category.{v} C] : Prop where
  [HasFiniteCoproducts : HasFiniteCoproducts C]
  van_kampen' : ∀ {X Y : C} (c : BinaryCofan X Y), IsColimit c → IsVanKampenColimit c
#align category_theory.finitary_extensive CategoryTheory.FinitaryExtensive

attribute [instance] finitary_extensive.has_finite_coproducts

theorem FinitaryExtensive.van_kampen [FinitaryExtensive C] {F : Discrete WalkingPair ⥤ C}
    (c : Cocone F) (hc : IsColimit c) : IsVanKampenColimit c :=
  by
  let X := F.obj ⟨walking_pair.left⟩
  let Y := F.obj ⟨walking_pair.right⟩
  have : F = pair X Y := by
    apply functor.hext
    · rintro ⟨⟨⟩⟩ <;> rfl
    · rintro ⟨⟨⟩⟩ ⟨j⟩ ⟨⟨rfl : _ = j⟩⟩ <;> simpa
  clear_value X Y
  subst this
  exact finitary_extensive.van_kampen' c hc
#align category_theory.finitary_extensive.van_kampen CategoryTheory.FinitaryExtensive.van_kampen

theorem map_pair_equifibered {F F' : Discrete WalkingPair ⥤ C} (α : F ⟶ F') : α.Equifibered :=
  by
  rintro ⟨⟨⟩⟩ ⟨j⟩ ⟨⟨rfl : _ = j⟩⟩
  all_goals
    dsimp; simp only [discrete.functor_map_id]
    exact is_pullback.of_horiz_is_iso ⟨by simp only [category.comp_id, category.id_comp]⟩
#align category_theory.map_pair_equifibered CategoryTheory.map_pair_equifibered

theorem BinaryCofan.is_van_kampen_iff (c : BinaryCofan X Y) :
    IsVanKampenColimit c ↔
      ∀ {X' Y' : C} (c' : BinaryCofan X' Y') (αX : X' ⟶ X) (αY : Y' ⟶ Y) (f : c'.pt ⟶ c.pt)
        (hαX : αX ≫ c.inl = c'.inl ≫ f) (hαY : αY ≫ c.inr = c'.inr ≫ f),
        Nonempty (IsColimit c') ↔ IsPullback c'.inl αX f c.inl ∧ IsPullback c'.inr αY f c.inr :=
  by
  constructor
  · introv H hαX hαY
    rw [H c' (map_pair αX αY) f (by ext ⟨⟨⟩⟩ <;> dsimp <;> assumption) (map_pair_equifibered _)]
    constructor
    · intro H
      exact ⟨H _, H _⟩
    · rintro H ⟨⟨⟩⟩
      exacts[H.1, H.2]
  · introv H F' hα h
    let X' := F'.obj ⟨walking_pair.left⟩
    let Y' := F'.obj ⟨walking_pair.right⟩
    have : F' = pair X' Y' := by
      apply functor.hext
      · rintro ⟨⟨⟩⟩ <;> rfl
      · rintro ⟨⟨⟩⟩ ⟨j⟩ ⟨⟨rfl : _ = j⟩⟩ <;> simpa
    clear_value X' Y'
    subst this
    change binary_cofan X' Y' at c'
    rw [H c' _ _ _ (nat_trans.congr_app hα ⟨walking_pair.left⟩)
        (nat_trans.congr_app hα ⟨walking_pair.right⟩)]
    constructor
    · rintro H ⟨⟨⟩⟩
      exacts[H.1, H.2]
    · intro H
      exact ⟨H _, H _⟩
#align category_theory.binary_cofan.is_van_kampen_iff CategoryTheory.BinaryCofan.is_van_kampen_iff

theorem BinaryCofan.is_van_kampen_mk {X Y : C} (c : BinaryCofan X Y)
    (cofans : ∀ X Y : C, BinaryCofan X Y) (colimits : ∀ X Y, IsColimit (cofans X Y))
    (cones : ∀ {X Y Z : C} (f : X ⟶ Z) (g : Y ⟶ Z), PullbackCone f g)
    (limits : ∀ {X Y Z : C} (f : X ⟶ Z) (g : Y ⟶ Z), IsLimit (cones f g))
    (h₁ :
      ∀ {X' Y' : C} (αX : X' ⟶ X) (αY : Y' ⟶ Y) (f : (cofans X' Y').pt ⟶ c.pt)
        (hαX : αX ≫ c.inl = (cofans X' Y').inl ≫ f) (hαY : αY ≫ c.inr = (cofans X' Y').inr ≫ f),
        IsPullback (cofans X' Y').inl αX f c.inl ∧ IsPullback (cofans X' Y').inr αY f c.inr)
    (h₂ :
      ∀ {Z : C} (f : Z ⟶ c.pt),
        IsColimit (BinaryCofan.mk (cones f c.inl).fst (cones f c.inr).fst)) :
    IsVanKampenColimit c := by
  rw [binary_cofan.is_van_kampen_iff]
  introv hX hY
  constructor
  · rintro ⟨h⟩
    let e := h.cocone_point_unique_up_to_iso (colimits _ _)
    obtain ⟨hl, hr⟩ := h₁ αX αY (e.inv ≫ f) (by simp [hX]) (by simp [hY])
    constructor
    · rw [← category.id_comp αX, ← iso.hom_inv_id_assoc e f]
      have : c'.inl ≫ e.hom = 𝟙 X' ≫ (cofans X' Y').inl :=
        by
        dsimp
        simp
      haveI : is_iso (𝟙 X') := inferInstance
      exact (is_pullback.of_vert_is_iso ⟨this⟩).paste_vert hl
    · rw [← category.id_comp αY, ← iso.hom_inv_id_assoc e f]
      have : c'.inr ≫ e.hom = 𝟙 Y' ≫ (cofans X' Y').inr :=
        by
        dsimp
        simp
      haveI : is_iso (𝟙 Y') := inferInstance
      exact (is_pullback.of_vert_is_iso ⟨this⟩).paste_vert hr
  · rintro ⟨H₁, H₂⟩
    refine' ⟨is_colimit.of_iso_colimit _ <| (iso_binary_cofan_mk _).symm⟩
    let e₁ : X' ≅ _ := H₁.is_limit.cone_point_unique_up_to_iso (limits _ _)
    let e₂ : Y' ≅ _ := H₂.is_limit.cone_point_unique_up_to_iso (limits _ _)
    have he₁ : c'.inl = e₁.hom ≫ (cones f c.inl).fst := by simp
    have he₂ : c'.inr = e₂.hom ≫ (cones f c.inr).fst := by simp
    rw [he₁, he₂]
    apply binary_cofan.is_colimit_comp_right_iso (binary_cofan.mk _ _)
    apply binary_cofan.is_colimit_comp_left_iso (binary_cofan.mk _ _)
    exact h₂ f
#align category_theory.binary_cofan.is_van_kampen_mk CategoryTheory.BinaryCofan.is_van_kampen_mk

theorem BinaryCofan.mono_inr_of_is_van_kampen [HasInitial C] {X Y : C} {c : BinaryCofan X Y}
    (h : IsVanKampenColimit c) : Mono c.inr :=
  by
  refine' pullback_cone.mono_of_is_limit_mk_id_id _ (is_pullback.is_limit _)
  refine'
    (h (binary_cofan.mk (initial.to Y) (𝟙 Y)) (map_pair (initial.to X) (𝟙 Y)) c.inr _
          (map_pair_equifibered _)).mp
      ⟨_⟩ ⟨walking_pair.right⟩
  · ext ⟨⟨⟩⟩ <;> dsimp <;> simp
  ·
    exact
      ((binary_cofan.is_colimit_iff_is_iso_inr initial_is_initial _).mpr
          (by
            dsimp
            infer_instance)).some
#align category_theory.binary_cofan.mono_inr_of_is_van_kampen CategoryTheory.BinaryCofan.mono_inr_of_is_van_kampen

theorem FinitaryExtensive.mono_inr_of_isColimit [FinitaryExtensive C] {c : BinaryCofan X Y}
    (hc : IsColimit c) : Mono c.inr :=
  BinaryCofan.mono_inr_of_is_van_kampen (FinitaryExtensive.van_kampen c hc)
#align category_theory.finitary_extensive.mono_inr_of_is_colimit CategoryTheory.FinitaryExtensive.mono_inr_of_isColimit

theorem FinitaryExtensive.mono_inl_of_isColimit [FinitaryExtensive C] {c : BinaryCofan X Y}
    (hc : IsColimit c) : Mono c.inl :=
  FinitaryExtensive.mono_inr_of_isColimit (BinaryCofan.isColimitFlip hc)
#align category_theory.finitary_extensive.mono_inl_of_is_colimit CategoryTheory.FinitaryExtensive.mono_inl_of_isColimit

instance [FinitaryExtensive C] (X Y : C) : Mono (coprod.inl : X ⟶ X ⨿ Y) :=
  (FinitaryExtensive.mono_inl_of_isColimit (coprodIsCoprod X Y) : _)

instance [FinitaryExtensive C] (X Y : C) : Mono (coprod.inr : Y ⟶ X ⨿ Y) :=
  (FinitaryExtensive.mono_inr_of_isColimit (coprodIsCoprod X Y) : _)

theorem BinaryCofan.isPullback_initial_to_of_is_van_kampen [HasInitial C] {c : BinaryCofan X Y}
    (h : IsVanKampenColimit c) : IsPullback (initial.to _) (initial.to _) c.inl c.inr :=
  by
  refine'
    ((h (binary_cofan.mk (initial.to Y) (𝟙 Y)) (map_pair (initial.to X) (𝟙 Y)) c.inr _
            (map_pair_equifibered _)).mp
        ⟨_⟩ ⟨walking_pair.left⟩).flip
  · ext ⟨⟨⟩⟩ <;> dsimp <;> simp
  ·
    exact
      ((binary_cofan.is_colimit_iff_is_iso_inr initial_is_initial _).mpr
          (by
            dsimp
            infer_instance)).some
#align category_theory.binary_cofan.is_pullback_initial_to_of_is_van_kampen CategoryTheory.BinaryCofan.isPullback_initial_to_of_is_van_kampen

theorem FinitaryExtensive.isPullback_initial_to_binaryCofan [FinitaryExtensive C]
    {c : BinaryCofan X Y} (hc : IsColimit c) :
    IsPullback (initial.to _) (initial.to _) c.inl c.inr :=
  BinaryCofan.isPullback_initial_to_of_is_van_kampen (FinitaryExtensive.van_kampen c hc)
#align category_theory.finitary_extensive.is_pullback_initial_to_binary_cofan CategoryTheory.FinitaryExtensive.isPullback_initial_to_binaryCofan

theorem has_strict_initial_of_is_universal [HasInitial C]
    (H : IsUniversalColimit (BinaryCofan.mk (𝟙 (⊥_ C)) (𝟙 (⊥_ C)))) : HasStrictInitialObjects C :=
  hasStrictInitialObjects_of_initial_is_strict
    (by
      intro A f
      suffices is_colimit (binary_cofan.mk (𝟙 A) (𝟙 A))
        by
        obtain ⟨l, h₁, h₂⟩ := limits.binary_cofan.is_colimit.desc' this (f ≫ initial.to A) (𝟙 A)
        rcases(category.id_comp _).symm.trans h₂ with rfl
        exact ⟨⟨_, ((category.id_comp _).symm.trans h₁).symm, initial_is_initial.hom_ext _ _⟩⟩
      refine'
        (H (binary_cofan.mk (𝟙 _) (𝟙 _)) (map_pair f f) f (by ext ⟨⟨⟩⟩ <;> dsimp <;> simp)
            (map_pair_equifibered _) _).some
      rintro ⟨⟨⟩⟩ <;> dsimp <;>
        exact is_pullback.of_horiz_is_iso ⟨(category.id_comp _).trans (category.comp_id _).symm⟩)
#align category_theory.has_strict_initial_of_is_universal CategoryTheory.has_strict_initial_of_is_universal

instance (priority := 100) hasStrictInitialObjects_of_finitaryExtensive [FinitaryExtensive C] :
    HasStrictInitialObjects C :=
  has_strict_initial_of_is_universal
    (FinitaryExtensive.van_kampen _
        ((BinaryCofan.isColimit_iff_isIso_inr initialIsInitial _).mpr
            (by
              dsimp
              infer_instance)).some).is_universal
#align category_theory.has_strict_initial_objects_of_finitary_extensive CategoryTheory.hasStrictInitialObjects_of_finitaryExtensive

theorem finitaryExtensive_iff_of_isTerminal (C : Type u) [Category.{v} C] [HasFiniteCoproducts C]
    (T : C) (HT : IsTerminal T) (c₀ : BinaryCofan T T) (hc₀ : IsColimit c₀) :
    FinitaryExtensive C ↔ IsVanKampenColimit c₀ :=
  by
  refine' ⟨fun H => H.2 c₀ hc₀, fun H => _⟩
  constructor
  simp_rw [binary_cofan.is_van_kampen_iff] at H⊢
  intro X Y c hc X' Y' c' αX αY f hX hY
  obtain ⟨d, hd, hd'⟩ :=
    limits.binary_cofan.is_colimit.desc' hc (HT.from _ ≫ c₀.inl) (HT.from _ ≫ c₀.inr)
  rw [H c' (αX ≫ HT.from _) (αY ≫ HT.from _) (f ≫ d) (by rw [← reassoc_of hX, hd, category.assoc])
      (by rw [← reassoc_of hY, hd', category.assoc])]
  obtain ⟨hl, hr⟩ := (H c (HT.from _) (HT.from _) d hd.symm hd'.symm).mp ⟨hc⟩
  rw [hl.paste_vert_iff hX.symm, hr.paste_vert_iff hY.symm]
#align category_theory.finitary_extensive_iff_of_is_terminal CategoryTheory.finitaryExtensive_iff_of_isTerminal

instance types.finitaryExtensive : FinitaryExtensive (Type u) :=
  by
  rw [finitary_extensive_iff_of_is_terminal (Type u) PUnit types.is_terminal_punit _
      (types.binary_coproduct_colimit _ _)]
  apply
    binary_cofan.is_van_kampen_mk _ _ (fun X Y => types.binary_coproduct_colimit X Y) _
      fun X Y Z f g => (limits.types.pullback_limit_cone f g).2
  · intros
    constructor
    · refine' ⟨⟨hαX.symm⟩, ⟨pullback_cone.is_limit_aux' _ _⟩⟩
      intro s
      have : ∀ x, ∃! y, s.fst x = Sum.inl y := by
        intro x
        cases h : s.fst x
        · simp_rw [sum.inl_injective.eq_iff]
          exact exists_unique_eq'
        · apply_fun f  at h
          cases ((congr_fun s.condition x).symm.trans h).trans (congr_fun hαY val : _).symm
      delta ExistsUnique at this
      choose l hl hl'
      exact
        ⟨l, (funext hl).symm, types.is_terminal_punit.hom_ext _ _, fun l' h₁ h₂ =>
          funext fun x => hl' x (l' x) (congr_fun h₁ x).symm⟩
    · refine' ⟨⟨hαY.symm⟩, ⟨pullback_cone.is_limit_aux' _ _⟩⟩
      intro s
      dsimp
      have : ∀ x, ∃! y, s.fst x = Sum.inr y := by
        intro x
        cases h : s.fst x
        · apply_fun f  at h
          cases ((congr_fun s.condition x).symm.trans h).trans (congr_fun hαX val : _).symm
        · simp_rw [sum.inr_injective.eq_iff]
          exact exists_unique_eq'
      delta ExistsUnique at this
      choose l hl hl'
      exact
        ⟨l, (funext hl).symm, types.is_terminal_punit.hom_ext _ _, fun l' h₁ h₂ =>
          funext fun x => hl' x (l' x) (congr_fun h₁ x).symm⟩
  · intro Z f
    dsimp [limits.types.binary_coproduct_cocone]
    delta types.pullback_obj
    have : ∀ x, f x = Sum.inl PUnit.unit ∨ f x = Sum.inr PUnit.unit :=
      by
      intro x
      rcases f x with (⟨⟨⟩⟩ | ⟨⟨⟩⟩)
      exacts[Or.inl rfl, Or.inr rfl]
    let eX : { p : Z × PUnit // f p.fst = Sum.inl p.snd } ≃ { x : Z // f x = Sum.inl PUnit.unit } :=
      ⟨fun p => ⟨p.1.1, by convert p.2⟩, fun x => ⟨⟨_, _⟩, x.2⟩, fun _ => by ext <;> rfl, fun _ =>
        by ext <;> rfl⟩
    let eY : { p : Z × PUnit // f p.fst = Sum.inr p.snd } ≃ { x : Z // f x = Sum.inr PUnit.unit } :=
      ⟨fun p => ⟨p.1.1, p.2.trans (congr_arg Sum.inr <| Subsingleton.elim _ _)⟩, fun x =>
        ⟨⟨_, _⟩, x.2⟩, fun _ => by ext <;> rfl, fun _ => by ext <;> rfl⟩
    fapply binary_cofan.is_colimit_mk
    ·
      exact fun s x =>
        dite _ (fun h => s.inl <| eX.symm ⟨x, h⟩) fun h =>
          s.inr <| eY.symm ⟨x, (this x).resolve_left h⟩
    · intro s
      ext ⟨⟨x, ⟨⟩⟩, _⟩
      dsimp
      split_ifs <;> rfl
    · intro s
      ext ⟨⟨x, ⟨⟩⟩, hx⟩
      dsimp
      split_ifs
      · cases h.symm.trans hx
      · rfl
    · intro s m e₁ e₂
      ext x
      split_ifs
      · rw [← e₁]
        rfl
      · rw [← e₂]
        rfl
#align category_theory.types.finitary_extensive CategoryTheory.types.finitaryExtensive

section TopCat

/-- (Implementation) An auxiliary lemma for the proof that `Top` is finitary extensive. -/
def finitaryExtensiveTopAux (Z : TopCat.{u}) (f : Z ⟶ TopCat.of (Sum PUnit.{u + 1} PUnit.{u + 1})) :
    IsColimit
      (BinaryCofan.mk
        (TopCat.pullbackFst f (TopCat.binaryCofan (TopCat.of PUnit) (TopCat.of PUnit)).inl)
        (TopCat.pullbackFst f (TopCat.binaryCofan (TopCat.of PUnit) (TopCat.of PUnit)).inr)) :=
  by
  have : ∀ x, f x = Sum.inl PUnit.unit ∨ f x = Sum.inr PUnit.unit :=
    by
    intro x
    rcases f x with (⟨⟨⟩⟩ | ⟨⟨⟩⟩)
    exacts[Or.inl rfl, Or.inr rfl]
  let eX : { p : Z × PUnit // f p.fst = Sum.inl p.snd } ≃ { x : Z // f x = Sum.inl PUnit.unit } :=
    ⟨fun p => ⟨p.1.1, p.2.trans (congr_arg Sum.inl <| Subsingleton.elim _ _)⟩, fun x =>
      ⟨⟨_, _⟩, x.2⟩, fun _ => by ext <;> rfl, fun _ => by ext <;> rfl⟩
  let eY : { p : Z × PUnit // f p.fst = Sum.inr p.snd } ≃ { x : Z // f x = Sum.inr PUnit.unit } :=
    ⟨fun p => ⟨p.1.1, p.2.trans (congr_arg Sum.inr <| Subsingleton.elim _ _)⟩, fun x =>
      ⟨⟨_, _⟩, x.2⟩, fun _ => by ext <;> rfl, fun _ => by ext <;> rfl⟩
  fapply binary_cofan.is_colimit_mk
  · refine' fun s =>
      ⟨fun x =>
        dite _ (fun h => s.inl <| eX.symm ⟨x, h⟩) fun h =>
          s.inr <| eY.symm ⟨x, (this x).resolve_left h⟩,
        _⟩
    rw [continuous_iff_continuousAt]
    intro x
    by_cases f x = Sum.inl PUnit.unit
    · revert h x
      apply (IsOpen.continuousOn_iff _).mp
      · rw [continuousOn_iff_continuous_restrict]
        convert_to Continuous fun x : { x | f x = Sum.inl PUnit.unit } =>
            s.inl ⟨(x, PUnit.unit), x.2⟩
        · ext ⟨x, hx⟩
          exact dif_pos hx
        continuity
      · convert f.2.1 _ openEmbedding_inl.open_range
        ext x
        exact
          ⟨fun h => ⟨_, h.symm⟩, fun ⟨e, h⟩ =>
            h.symm.trans (congr_arg Sum.inl <| Subsingleton.elim _ _)⟩
    · revert h x
      apply (IsOpen.continuousOn_iff _).mp
      · rw [continuousOn_iff_continuous_restrict]
        convert_to Continuous fun x : { x | f x ≠ Sum.inl PUnit.unit } =>
            s.inr ⟨(x, PUnit.unit), (this _).resolve_left x.2⟩
        · ext ⟨x, hx⟩
          exact dif_neg hx
        continuity
      · convert f.2.1 _ openEmbedding_inr.open_range
        ext x
        change f x ≠ Sum.inl PUnit.unit ↔ f x ∈ Set.range Sum.inr
        trans f x = Sum.inr PUnit.unit
        ·
          rcases f x with (⟨⟨⟩⟩ | ⟨⟨⟩⟩) <;>
            simp only [iff_self_iff, eq_self_iff_true, not_true, Ne.def, not_false_iff]
        ·
          exact
            ⟨fun h => ⟨_, h.symm⟩, fun ⟨e, h⟩ =>
              h.symm.trans (congr_arg Sum.inr <| Subsingleton.elim _ _)⟩
  · intro s
    ext ⟨⟨x, ⟨⟩⟩, _⟩
    change dite _ _ _ = _
    split_ifs <;> rfl
  · intro s
    ext ⟨⟨x, ⟨⟩⟩, hx⟩
    change dite _ _ _ = _
    split_ifs
    · cases h.symm.trans hx
    · rfl
  · intro s m e₁ e₂
    ext x
    change m x = dite _ _ _
    split_ifs
    · rw [← e₁]
      rfl
    · rw [← e₂]
      rfl
#align category_theory.finitary_extensive_Top_aux CategoryTheory.finitaryExtensiveTopAux

instance : FinitaryExtensive TopCat.{u} :=
  by
  rw [finitary_extensive_iff_of_is_terminal TopCat.{u} _ TopCat.isTerminalPunit _
      (TopCat.binaryCofanIsColimit _ _)]
  apply
    binary_cofan.is_van_kampen_mk _ _ (fun X Y => TopCat.binaryCofanIsColimit X Y) _
      fun X Y Z f g => TopCat.pullbackConeIsLimit f g
  · intros
    constructor
    · refine' ⟨⟨hαX.symm⟩, ⟨pullback_cone.is_limit_aux' _ _⟩⟩
      intro s
      have : ∀ x, ∃! y, s.fst x = Sum.inl y := by
        intro x
        cases h : s.fst x
        · simp_rw [sum.inl_injective.eq_iff]
          exact exists_unique_eq'
        · apply_fun f  at h
          cases
            ((concrete_category.congr_hom s.condition x).symm.trans h).trans
              (concrete_category.congr_hom hαY val : _).symm
      delta ExistsUnique at this
      choose l hl hl'
      refine'
        ⟨⟨l, _⟩, ContinuousMap.ext fun a => (hl a).symm, Top.is_terminal_punit.hom_ext _ _,
          fun l' h₁ h₂ =>
          ContinuousMap.ext fun x => hl' x (l' x) (concrete_category.congr_hom h₁ x).symm⟩
      apply embedding_inl.to_inducing.continuous_iff.mpr
      convert s.fst.2 using 1
      exact (funext hl).symm
    · refine' ⟨⟨hαY.symm⟩, ⟨pullback_cone.is_limit_aux' _ _⟩⟩
      intro s
      dsimp
      have : ∀ x, ∃! y, s.fst x = Sum.inr y := by
        intro x
        cases h : s.fst x
        · apply_fun f  at h
          cases
            ((concrete_category.congr_hom s.condition x).symm.trans h).trans
              (concrete_category.congr_hom hαX val : _).symm
        · simp_rw [sum.inr_injective.eq_iff]
          exact exists_unique_eq'
      delta ExistsUnique at this
      choose l hl hl'
      refine'
        ⟨⟨l, _⟩, ContinuousMap.ext fun a => (hl a).symm, Top.is_terminal_punit.hom_ext _ _,
          fun l' h₁ h₂ =>
          ContinuousMap.ext fun x => hl' x (l' x) (concrete_category.congr_hom h₁ x).symm⟩
      apply embedding_inr.to_inducing.continuous_iff.mpr
      convert s.fst.2 using 1
      exact (funext hl).symm
  · intro Z f
    exact finitary_extensive_Top_aux Z f

end TopCat

section Functor

universe v'' u''

variable {D : Type u''} [Category.{v''} D]

theorem NatTrans.Equifibered.whiskerRight {F G : J ⥤ C} {α : F ⟶ G} (hα : α.Equifibered) (H : C ⥤ D)
    [PreservesLimitsOfShape WalkingCospan H] : (whiskerRight α H).Equifibered := fun i j f =>
  (hα f).map H
#align category_theory.nat_trans.equifibered.whisker_right CategoryTheory.NatTrans.Equifibered.whiskerRight

theorem IsVanKampenColimit.of_iso {F : J ⥤ C} {c c' : Cocone F} (H : IsVanKampenColimit c)
    (e : c ≅ c') : IsVanKampenColimit c' :=
  by
  intro F' c'' α f h hα
  have : c'.ι ≫ (Functor.Const J).map e.inv.hom = c.ι :=
    by
    ext j
    exact e.inv.2 j
  rw [H c'' α (f ≫ e.inv.1) (by rw [functor.map_comp, ← reassoc_of h, this]) hα]
  apply forall_congr'
  intro j
  conv_lhs => rw [← category.comp_id (α.app j)]
  haveI : is_iso e.inv.hom := functor.map_is_iso (cocones.forget _) e.inv
  exact (is_pullback.of_vert_is_iso ⟨by simp⟩).paste_vert_iff (nat_trans.congr_app h j).symm
#align category_theory.is_van_kampen_colimit.of_iso CategoryTheory.IsVanKampenColimit.of_iso

theorem IsVanKampenColimit.of_map {D : Type _} [Category D] (G : C ⥤ D) {F : J ⥤ C} {c : Cocone F}
    [PreservesLimitsOfShape WalkingCospan G] [ReflectsLimitsOfShape WalkingCospan G]
    [PreservesColimitsOfShape J G] [ReflectsColimitsOfShape J G]
    (H : IsVanKampenColimit (G.mapCocone c)) : IsVanKampenColimit c :=
  by
  intro F' c' α f h hα
  refine'
    (Iff.trans _
          (H (G.map_cocone c') (whisker_right α G) (G.map f)
            (by
              ext j
              simpa using G.congr_map (nat_trans.congr_app h j))
            (hα.whisker_right G))).trans
      (forall_congr' fun j => _)
  · exact ⟨fun h => ⟨is_colimit_of_preserves G h.some⟩, fun h => ⟨is_colimit_of_reflects G h.some⟩⟩
  · exact is_pullback.map_iff G (nat_trans.congr_app h.symm j)
#align category_theory.is_van_kampen_colimit.of_map CategoryTheory.IsVanKampenColimit.of_map

theorem isVanKampenColimit_of_evaluation [HasPullbacks D] [HasColimitsOfShape J D] (F : J ⥤ C ⥤ D)
    (c : Cocone F) (hc : ∀ x : C, IsVanKampenColimit (((evaluation C D).obj x).mapCocone c)) :
    IsVanKampenColimit c := by
  intro F' c' α f e hα
  have := fun x =>
    hc x (((evaluation C D).obj x).mapCocone c') (whisker_right α _)
      (((evaluation C D).obj x).map f)
      (by
        ext y
        dsimp
        exact nat_trans.congr_app (nat_trans.congr_app e y) x)
      (hα.whisker_right _)
  constructor
  · rintro ⟨hc'⟩ j
    refine' ⟨⟨(nat_trans.congr_app e j).symm⟩, ⟨evaluation_jointly_reflects_limits _ _⟩⟩
    refine' fun x => (is_limit_map_cone_pullback_cone_equiv _ _).symm _
    exact ((this x).mp ⟨preserves_colimit.preserves hc'⟩ _).IsLimit
  ·
    exact fun H =>
      ⟨evaluation_jointly_reflects_colimits _ fun x =>
          ((this x).mpr fun j => (H j).map ((evaluation C D).obj x)).some⟩
#align category_theory.is_van_kampen_colimit_of_evaluation CategoryTheory.isVanKampenColimit_of_evaluation

instance [HasPullbacks C] [FinitaryExtensive C] : FinitaryExtensive (D ⥤ C) :=
  haveI : has_finite_coproducts (D ⥤ C) := ⟨fun n => limits.functor_category_has_colimits_of_shape⟩
  ⟨fun X Y c hc =>
    is_van_kampen_colimit_of_evaluation _ c fun x =>
      finitary_extensive.van_kampen _ <| preserves_colimit.preserves hc⟩

theorem finitaryExtensive_of_preserves_and_reflects (F : C ⥤ D) [FinitaryExtensive D]
    [HasFiniteCoproducts C] [PreservesLimitsOfShape WalkingCospan F]
    [ReflectsLimitsOfShape WalkingCospan F] [PreservesColimitsOfShape (Discrete WalkingPair) F]
    [ReflectsColimitsOfShape (Discrete WalkingPair) F] : FinitaryExtensive C :=
  ⟨fun X Y c hc => (FinitaryExtensive.van_kampen _ (isColimitOfPreserves F hc)).of_map F⟩
#align category_theory.finitary_extensive_of_preserves_and_reflects CategoryTheory.finitaryExtensive_of_preserves_and_reflects

theorem finitaryExtensive_of_preserves_and_reflects_isomorphism (F : C ⥤ D) [FinitaryExtensive D]
    [HasFiniteCoproducts C] [HasPullbacks C] [PreservesLimitsOfShape WalkingCospan F]
    [PreservesColimitsOfShape (Discrete WalkingPair) F] [ReflectsIsomorphisms F] :
    FinitaryExtensive C :=
  by
  haveI : reflects_limits_of_shape walking_cospan F :=
    reflects_limits_of_shape_of_reflects_isomorphisms
  haveI : reflects_colimits_of_shape (discrete walking_pair) F :=
    reflects_colimits_of_shape_of_reflects_isomorphisms
  exact finitary_extensive_of_preserves_and_reflects F
#align category_theory.finitary_extensive_of_preserves_and_reflects_isomorphism CategoryTheory.finitaryExtensive_of_preserves_and_reflects_isomorphism

end Functor

end Extensive

end CategoryTheory