/-
Copyright (c) 2022 Devon Tuma. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Devon Tuma
-/
import computational_monads.simulation_semantics.constructions.stateless_oracle
import computational_monads.constructions.uniform_select

/-!
# Coercions Between Computations With Additional Oracles

This file defines a `is_sub_spec` relation for pairs of `oracle_spec` where one can be
thought of as an extension of the other with additional oracles.
The definition consists of a function from query inputs in the original oracle to a
computation using the new set of oracles, such that the result of the mapping
doesn't affect the underlying probability distribution on the oracle call.

We use the notation `spec ⊂ₒ spec'` to represent that one set of oracles is a subset of another,
where the non-exclusive subset symbol reflects that we avoid defining this instance reflexively.
This decision is based on the `is_coe` construction, where we don't want to coerce a computation
to itself by calling a reflexive `is_sub_spec` construction.

We define the map to output a computation rather than a new set of oracle inputs in the new spec
to avoid type checking issues, as the `query` output type will not be definitionally equal
to the `query` output type in the original `oracle_spec`, causing issues in defining `has_coe`.
In practice the mapping will still usually output a `query` call,
and the equality between the underlying distributions is generally sufficient.

From this definition we construct a `is_coe` instance to coerce a computation with one set of
oracles to one with a larger set of oracles, using the `is_sub_spec` to call `simulate'`.
We show that this coercion has no effect on `support`, `eval_dist`, or `prob_event`.
-/

variables {α β γ : Type}

namespace oracle_spec

open oracle_comp

/-- Example of a computation that naturally should work, but doesn't without `sub_spec` coercions.
The fundamental issue being that the type system doesn't have a sense of "additional" oracles.
In this case, performing a validity check on the adversaries results isn't easily possible.
Note the actual version is commented out, only the un-checked version will compile. -/
example {regular_spec adversary_spec : oracle_spec}
  (adversary : oracle_comp (regular_spec ++ adversary_spec) α)
  (validity_check : α → oracle_comp regular_spec bool) :
  oracle_comp (regular_spec ++ adversary_spec) (option α) :=
-- do { x ← adversary, b ← validity_check x, return (if b = tt then some x else none) }
do { x ← adversary, return x }

/-- Relation defining an inclusion of one set of oracles into another, where the mapping
doesn't affect the underlying probability distribution of the computation.
Informally, `sub_spec ⊂ₒ super_spec` means that for any query to an oracle of `sub_spec`,
it can be perfectly simulated by a computation using the oracles of `super_spec`. -/
class is_sub_spec (sub_spec super_spec : oracle_spec) :=
(to_fun (i : sub_spec.ι) (t : sub_spec.domain i) : oracle_comp super_spec (sub_spec.range i))
(eval_dist_to_fun' : ∀ i t, ⁅to_fun i t⁆ = ⁅query i t⁆)

infixl ` ⊂ₒ `:65 := is_sub_spec

namespace is_sub_spec

variables (sub_spec super_spec : oracle_spec) [h : sub_spec ⊂ₒ super_spec]
  (i : sub_spec.ι) (t : sub_spec.domain i)

@[simp] lemma support_to_fun : (h.to_fun i t).support = ⊤ :=
by rw [← support_eval_dist, h.eval_dist_to_fun', support_eval_dist, support_query]

@[simp] lemma fin_support_to_fun [∀ i t, (h.to_fun i t).decidable] :
  (h.to_fun i t).fin_support = ⊤ :=
by simp only [fin_support_eq_iff_support_eq_coe, finset.top_eq_univ,
  support_to_fun, set.top_eq_univ, finset.coe_univ]

@[simp] lemma eval_dist_to_fun : ⁅h.to_fun i t⁆ = pmf.uniform_of_fintype (sub_spec.range i) :=
by rw [h.eval_dist_to_fun', eval_dist_query]

@[simp] lemma prob_event_to_fun (e : set (sub_spec.range i)) :
  ⁅e | h.to_fun i t⁆ = ⁅e | query i t⁆ :=
prob_event_eq_of_eval_dist_eq (h.eval_dist_to_fun' i t) e

end is_sub_spec

end oracle_spec

namespace oracle_comp

open oracle_spec

/-- Given a `is_sub_spec` instance between `sub_spec` and `super_spec`, we can coerce a computation
with oracles `sub_spec` to one with `super_spec` by simulating with `is_sub_spec.to_fun`. -/
instance coe_sub_spec (sub_spec super_spec : oracle_spec) [h : sub_spec ⊂ₒ super_spec] (α : Type) :
  has_coe (oracle_comp sub_spec α) (oracle_comp super_spec α) :=
{coe := λ oa, default_simulate' ⟪λ i t, h.to_fun i t⟫ oa}

lemma coe_sub_spec_def {sub_spec super_spec : oracle_spec} [h : sub_spec ⊂ₒ super_spec]
  (oa : oracle_comp sub_spec α) : (↑oa : oracle_comp super_spec α) =
    default_simulate' ⟪λ i t, h.to_fun i t⟫ oa := rfl

section coe_sub_spec

variables (sub_spec super_spec : oracle_spec) [h : sub_spec ⊂ₒ super_spec]
  (a : α) (oa : oracle_comp sub_spec α) (ob : α → oracle_comp sub_spec β)
  (i : sub_spec.ι) (t : sub_spec.domain i) (e : set α)
include h

instance coe_sub_spec.decidable [∀ i t, (@is_sub_spec.to_fun sub_spec super_spec h i t).decidable]
  (oa : oracle_comp sub_spec α) [oa.decidable] : (↑oa : oracle_comp super_spec α).decidable :=
simulate'.decidable _ oa ()

lemma coe_sub_spec_return : (↑(return a : oracle_comp sub_spec α) : oracle_comp super_spec α) =
  prod.fst <$> return (a, ()) := rfl

lemma coe_sub_spec_bind : (↑(oa >>= ob) : oracle_comp super_spec β) =
  prod.fst <$> (default_simulate ⟪λ i t, h.to_fun i t⟫ oa >>=
    λ x, simulate ⟪λ i t, h.to_fun i t⟫ (ob x.1) x.2) :=
by rw [coe_sub_spec_def, default_simulate', simulate'_bind]

lemma coe_sub_spec_query : (↑(query i t) : oracle_comp super_spec (sub_spec.range i)) =
  prod.fst <$> (h.to_fun i t >>= λ u, return (u, ())) :=
by rw [coe_sub_spec_def, default_simulate', simulate'_query, stateless_oracle.apply_eq]

/-- `support` is unchanged after coercing a computation via a sub-spec instance. -/
@[simp] lemma support_coe_sub_spec : (↑oa : oracle_comp super_spec α).support = oa.support :=
stateless_oracle.support_simulate'_eq_support _ _ ()
  (λ i t, is_sub_spec.support_to_fun sub_spec super_spec i t)

/-- `fin_support` is unchanged after coercing a computation via a sub-spec instance. -/
@[simp] lemma fin_support_coe_sub_spec [∀ i t, (@is_sub_spec.to_fun sub_spec super_spec _ i t).decidable]
  [oa.decidable] : (↑oa : oracle_comp super_spec α).fin_support = oa.fin_support :=
by rw [fin_support_eq_fin_support_iff_support_eq_support, support_coe_sub_spec]

/-- `eval_dist` is unchanged after coercing a computation via a sub-spec instance. -/
@[simp] lemma eval_dist_coe_sub_spec : ⁅(↑oa : oracle_comp super_spec α)⁆ = ⁅oa⁆ :=
stateless_oracle.eval_dist_simulate'_eq_eval_dist _ _ ()
  (λ i t, is_sub_spec.eval_dist_to_fun sub_spec super_spec i t)

/-- `prob_event` is unchanged after coercing a computation via a sub-spec instance. -/
@[simp] lemma prob_event_coe_sub_spec : ⁅e | (↑oa : oracle_comp super_spec α)⁆ = ⁅e | oa⁆ :=
stateless_oracle.prob_event_simulate'_eq_prob_event _ _ ()
  (λ i t, is_sub_spec.eval_dist_to_fun sub_spec super_spec i t) e

end coe_sub_spec

section simulate_coe_sub_spec

variables {sub_spec super_spec spec : oracle_spec} [h : sub_spec ⊂ₒ super_spec] {S S' : Type}
  (so : sim_oracle sub_spec spec S) (so' : sim_oracle super_spec spec S')
  (s : S) (s' : S') (a : α) (oa : oracle_comp sub_spec α) (ob : α → oracle_comp sub_spec β)
  (i : sub_spec.ι) (t : sub_spec.domain i)
include h

section support

@[simp] lemma support_simulate_coe_sub_spec_return :
  (simulate so' (return a : oracle_comp sub_spec α) s').support = {(a, s')} :=
by rw [coe_sub_spec_return, simulate_map, simulate_return, support_map, support_return,
  set.image_singleton, prod.map, id.def]

@[simp] lemma support_simulate'_coe_sub_spec_return :
  (simulate' so' (return a : oracle_comp sub_spec α) s').support = {a} :=
by simp only [support_simulate', support_simulate_coe_sub_spec_return, set.image_singleton]

@[simp] lemma support_simulate_coe_sub_spec_bind :
  (simulate so' (↑(oa >>= ob) : oracle_comp super_spec β) s').support =
    ⋃ x ∈ (simulate so' (↑oa : oracle_comp super_spec α) s').support,
      (simulate so' ↑(ob $ prod.fst x) x.2).support :=
calc (simulate so' (↑(oa >>= ob) : oracle_comp super_spec β) s').support =
  (simulate so' ↑oa s' >>= λ (x : α × S'), simulate so' ↑(ob x.1) x.2).support :
    by simp_rw [coe_sub_spec_def, default_simulate', simulate', simulate_bind,
      support_simulate_map_bind, simulate_map, support_bind_map, support_map,
      simulate_eq_default_simulate, prod.map_snd, prod.map_fst, id.def]
  ... = ⋃ x ∈ (simulate so' (↑oa : oracle_comp super_spec α) s').support,
    (simulate so' ↑(ob $ prod.fst x) x.2).support : by rw [support_bind]

@[simp] lemma support_simulate'_coe_sub_spec_bind :
  (simulate' so' (↑(oa >>= ob) : oracle_comp super_spec β) s').support =
    ⋃ x ∈ (simulate so' (↑oa : oracle_comp super_spec α) s').support,
      (simulate' so' ↑(ob $ prod.fst x) x.2).support :=
by simp only [support_simulate', support_simulate_coe_sub_spec_bind, set.image_Union]

@[simp] lemma support_simulate_coe_sub_spec_query :
  (simulate so' (↑(query i t) : oracle_comp super_spec (sub_spec.range i)) s').support =
    (simulate so' (h.to_fun i t) s').support :=
by simp_rw [coe_sub_spec_def, default_simulate', simulate'_query, stateless_oracle.apply_eq,
  support_simulate_map, support_simulate_bind, support_simulate_return, set.image_Union,
  set.image_singleton, prod.map_mk, id.def, prod.mk.eta, set.bUnion_of_singleton]

/-- Given two simulation oracles `so : sim_oracle spec spec'' S` and
`so' : sim_oracle spec' spec'' : S'` with the starting specs satisfying `spec ⊂ₒ spec'`,
and a function `f : S → S'` between their states, if the support after simulating the
sub-spec coersion function with the second oracle looks like the support after simulating with the
first oracle then applying `f`, then simulating the coercion of any computation with the second
oracle has the same support as simulating the uncoerced computation and mapping by `f`. -/
lemma support_simulate_coe_sub_spec (f : S → S') (hf : ∀ i t s,
  (simulate so' (h.to_fun i t) (f s)).support = prod.map id f '' (so i (t, s)).support) :
  (simulate so' (↑oa : oracle_comp super_spec α) (f s)).support =
    (prod.map id f) '' (simulate so oa s).support :=
begin
  induction oa using oracle_comp.induction_on with α a α β oa ob hoa hob i t generalizing s,
  { simpa only [support_simulate_coe_sub_spec_return,
      support_simulate_return, set.image_singleton] },
  { ext y,
    simp_rw [support_simulate_coe_sub_spec_bind, hoa, support_simulate_bind,
      set.image_Union, ← hob, set.mem_Union],
    refine ⟨λ h, let ⟨x, ⟨y', hy', hxy⟩, hx⟩ := h in ⟨y', hy', by simpa only [← hxy] using hx⟩,
      λ h, let ⟨x, hy, hx⟩ := h in ⟨(x.1, f x.2), ⟨x, hy, rfl⟩, hx⟩⟩ },
  { rw [support_simulate_coe_sub_spec_query, hf, support_simulate_query] }
end

/-- Version of `support_simulate_coe_sub_spec` for `simulate'`. In this case we get exact equality
between the support of the simulations, since the oracle states are irrelevent. -/
lemma support_simulate'_coe_sub_spec (f : S → S') (hf : ∀ i t s,
  (simulate so' (h.to_fun i t) (f s)).support = prod.map id f '' (so i (t, s)).support) :
  (simulate' so' (↑oa : oracle_comp super_spec α) (f s)).support = (simulate' so oa s).support :=
by simp only [support_simulate_coe_sub_spec so so' s oa f hf,
  set.image_image, support_simulate', prod_map, id.def]

end support

section fin_support

end fin_support

section eval_dist

@[simp] lemma eval_dist_simulate_coe_sub_spec_return :
  ⁅simulate so' ↑(return a : oracle_comp sub_spec α) s'⁆ = pmf.pure (a, s') :=
by simp only [coe_sub_spec_return, simulate_map, simulate_return, eval_dist_map, eval_dist_return,
  pmf.map_pure, prod.map_mk, id.def]

@[simp] lemma eval_dist_simulate_coe_sub_spec_bind :
  ⁅simulate so' (↑(oa >>= ob) : oracle_comp super_spec β) s'⁆ =
    ⁅simulate so' ↑oa s'⁆.bind (λ x, ⁅simulate so' ↑(ob $ prod.fst x) x.2⁆) :=
calc ⁅simulate so' (↑(oa >>= ob) : oracle_comp super_spec β) s'⁆
  = ⁅simulate so' (default_simulate ⟪h.to_fun⟫ oa) s'⁆.bind (λ (x : (α × unit) × S'),
      ⁅simulate so' (simulate ⟪h.to_fun⟫ (ob x.1.1) x.1.2) x.2⁆.map (prod.map prod.fst id)) :
    by simp only [coe_sub_spec_bind, simulate_map, simulate_bind,
      eval_dist_map, eval_dist_bind, pmf.map_bind]
  ... = (⁅simulate so' (default_simulate ⟪h.to_fun⟫ oa) s'⁆.map (prod.map prod.fst id)).bind
      (λ x, ⁅simulate so' (default_simulate ⟪h.to_fun⟫ (ob x.1)) x.2⁆.map (prod.map prod.fst id)) :
    symm (trans (pmf.bind_map _ _ _) (congr_arg (λ _, pmf.bind _ _) (funext $ λ x, by simp only
      [function.comp_app, prod_map, id.def, stateless_oracle.simulate_eq_default_simulate])))
  ... = ⁅simulate so' (default_simulate' ⟪h.to_fun⟫ oa) s'⁆.bind (λ (x : α × S'),
      ⁅simulate so' (default_simulate ⟪h.to_fun⟫ (ob x.1)) x.2⁆.map (prod.map prod.fst id)) :
    by rw [default_simulate', simulate', eval_dist_simulate_map]
  ... = ⁅simulate so' ↑oa s'⁆.bind (λ (x : α × S'), ⁅simulate so' ↑(ob $ prod.fst x) x.2⁆) :
    by simp only [coe_sub_spec_def, simulate', eval_dist_simulate_map, default_simulate']

@[simp] lemma eval_dist_simulate_coe_sub_spec_query :
  ⁅simulate so' (↑(query i t) : oracle_comp super_spec _) s'⁆ =
    ⁅simulate so' (h.to_fun i t) s'⁆ :=
by simp only [coe_sub_spec_query, eval_dist_simulate_map_bind, eval_dist_simulate_return,
  pmf.map_pure, prod.map_mk, id.def, prod.mk.eta, pmf.bind_pure]

/-- Given two simulation oracles `so : sim_oracle spec spec'' S` and
`so' : sim_oracle spec' spec'' : S'` with the starting specs satisfying `spec ⊂ₒ spec'`,
and a function `f : S → S'` between their states, if the distribution after simulating the
sub-spec coersion function with the second oracle looks like the distribution after simulating with
the first oracle then applying `f`, then simulating the coercion of any computation with the second
oracle has the same distribution as simulating the uncoerced computation and mapping by `f`. -/
lemma eval_dist_simulate_coe_sub_spec (f : S → S') (hf : ∀ i t s,
  ⁅simulate so' (h.to_fun i t) (f s)⁆ = pmf.map (prod.map id f) ⁅so i (t, s)⁆) :
  ⁅simulate so' (↑oa : oracle_comp super_spec α) (f s)⁆ =
    ⁅simulate so oa s⁆.map (prod.map id f) :=
begin
  induction oa using oracle_comp.induction_on with α a α β oa ob hoa hob i t generalizing s,
  { simp only [eval_dist_simulate_coe_sub_spec_return, simulate_return, eval_dist_return,
      pmf.map_pure, prod.map_mk, id.def] },
  { simp_rw [eval_dist_simulate_coe_sub_spec_bind, hoa,
      eval_dist_simulate_bind, pmf.map_bind, pmf.bind_map],
    refine congr_arg (λ _, pmf.bind _ _) (funext $ λ x, (hob _ _)) },
  { rw [eval_dist_simulate_query, eval_dist_simulate_coe_sub_spec_query, hf] }
end

/-- Version of `eval_dist_simulate_coe_sub_spec` for `simulate'`. In this case we get exact
equality between the distributions of the simulations, since the oracle states are irrelevent. -/
lemma eval_dist_simulate'_coe_sub_spec (f : S → S') (hf : ∀ i t s,
  ⁅simulate so' (h.to_fun i t) (f s)⁆ = pmf.map (prod.map id f) ⁅so i (t, s)⁆) :
  ⁅simulate' so' (↑oa : oracle_comp super_spec α) (f s)⁆ = ⁅simulate' so oa s⁆ :=
by simp only [eval_dist_simulate', eval_dist_simulate_coe_sub_spec so so' s oa f hf, pmf.map_comp,
  prod.map_fst', function.comp.left_id]

end eval_dist

section prob_event

/-- Extension of `eval_dist_simulate_coe_sub_spec` to `prob_event`. We keep the same hypothesis
about `eval_dist` rather than a one in terms of `prob_event` for simplicity. -/
lemma prob_event_simulate_coe_sub_spec (e : set (α × S')) (f : S → S') (hf : ∀ i t s,
  ⁅simulate so' (h.to_fun i t) (f s)⁆ = pmf.map (prod.map id f) ⁅so i (t, s)⁆) :
  ⁅e | simulate so' (↑oa : oracle_comp super_spec α) (f s)⁆ =
    ⁅e | prod.map id f <$> simulate so oa s⁆ :=
by simp_rw [prob_event_eq_tsum_indicator, eval_dist_map,
  eval_dist_simulate_coe_sub_spec so so' s oa f hf]

/-- Extension of `eval_dist_simulate'_coe_sub_spec` to `prob_event`. We keep the same hypothesis
about `eval_dist` rather than a one in terms of `prob_event` for simplicity. -/
lemma prob_event_simulate'_coe_sub_spec (e : set α) (f : S → S') (hf : ∀ i t s,
  ⁅simulate so' (h.to_fun i t) (f s)⁆ = pmf.map (prod.map id f) ⁅so i (t, s)⁆) :
  ⁅e | simulate' so' (↑oa : oracle_comp super_spec α) (f s)⁆ = ⁅e | simulate' so oa s⁆ :=
by simpa only [prob_event_simulate', prob_event_simulate_coe_sub_spec so so' s oa _ f hf,
  prob_event_map, set.preimage_preimage, prod.map_fst]

end prob_event

end simulate_coe_sub_spec

end oracle_comp


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Mario Carneiro
-/
prelude
import Fixtures.Termination.Init.Tactics
set_option linter.all false -- prevent error messages from runFrontend

/-! # SizeOf -/

/--
`SizeOf` is a typeclass automatically derived for every inductive type,
which equips the type with a "size" function to `Nat`.
The default instance defines each constructor to be `1` plus the sum of the
sizes of all the constructor fields.

This is used for proofs by well-founded induction, since every field of the
constructor has a smaller size than the constructor itself,
and in many cases this will suffice to do the proof that a recursive function
is only called on smaller values.
If the default proof strategy fails, it is recommended to supply a custom
size measure using the `termination_by` argument on the function definition.
-/
class SizeOf (α : Sort u) where
  /-- The "size" of an element, a natural number which decreases on fields of
  each inductive type. -/
  sizeOf : α → Nat

export SizeOf (sizeOf)

/-!
Declare `SizeOf` instances and theorems for types declared before `SizeOf`.
From now on, the inductive compiler will automatically generate `SizeOf` instances and theorems.
-/

/--
Every type `α` has a default `SizeOf` instance that just returns `0`
for every element of `α`.
-/
protected def default.sizeOf (α : Sort u) : α → Nat
  | _ => 0

instance (priority := low) (α : Sort u) : SizeOf α where
  sizeOf := default.sizeOf α

@[simp] theorem sizeOf_default (n : α) : sizeOf n = 0 := rfl

instance : SizeOf Nat where
  sizeOf n := n

@[simp] theorem sizeOf_nat (n : Nat) : sizeOf n = n := rfl

instance [SizeOf α] : SizeOf (Unit → α) where
  sizeOf f := sizeOf (f ())

@[simp] theorem sizeOf_thunk [SizeOf α] (f : Unit → α) : sizeOf f = sizeOf (f ()) :=
  rfl

deriving instance SizeOf for PUnit
deriving instance SizeOf for Prod
deriving instance SizeOf for PProd
deriving instance SizeOf for MProd
deriving instance SizeOf for Bool
deriving instance SizeOf for Subtype
deriving instance SizeOf for PLift
deriving instance SizeOf for ULift
deriving instance SizeOf for Decidable
deriving instance SizeOf for Fin
deriving instance SizeOf for UInt8
deriving instance SizeOf for UInt16
deriving instance SizeOf for UInt32
deriving instance SizeOf for UInt64
deriving instance SizeOf for USize
deriving instance SizeOf for Char
deriving instance SizeOf for Option
deriving instance SizeOf for List
deriving instance SizeOf for String
deriving instance SizeOf for String.Pos
deriving instance SizeOf for Substring
deriving instance SizeOf for Array
deriving instance SizeOf for Except
deriving instance SizeOf for EStateM.Result

@[simp] theorem Unit.sizeOf (u : Unit) : sizeOf u = 1 := rfl
@[simp] theorem Unit.sizeOf' (u : Unit) : SizeOf.sizeOf u = 1 := by cases u <;> rfl
@[simp] theorem Bool.sizeOf_eq_one (b : Bool) : sizeOf b = 1 := by cases b <;> rfl

namespace Lean

/--
We manually define the `Lean.Name` instance because we use
an opaque function for computing the hashcode field.
-/
protected noncomputable def Name.sizeOf : Name → Nat
  | anonymous => 1
  | str p s   => 1 + Name.sizeOf p + sizeOf s
  | num p n   => 1 + Name.sizeOf p + sizeOf n

noncomputable instance : SizeOf Name where
  sizeOf n := n.sizeOf

@[simp] theorem Name.anonymous.sizeOf_spec : sizeOf anonymous = 1 :=
  rfl
@[simp] theorem Name.str.sizeOf_spec (p : Name) (s : String) : sizeOf (str p s) = 1 + sizeOf p + sizeOf s :=
  rfl
@[simp] theorem Name.num.sizeOf_spec (p : Name) (n : Nat) : sizeOf (num p n) = 1 + sizeOf p + sizeOf n :=
  rfl

deriving instance SizeOf for SourceInfo
deriving instance SizeOf for Syntax
deriving instance SizeOf for TSyntax
deriving instance SizeOf for Syntax.SepArray
deriving instance SizeOf for Syntax.TSepArray
deriving instance SizeOf for ParserDescr
deriving instance SizeOf for MacroScopesView
deriving instance SizeOf for Macro.Context
deriving instance SizeOf for Macro.Exception
deriving instance SizeOf for Macro.State
deriving instance SizeOf for Macro.Methods

end Lean


/-
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

! This file was ported from Lean 3 source module data.array.lemmas
! leanprover-community/mathlib commit 78314d08d707a6338079f00094bbdb90bf11fc41
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathbin.Control.Traversable.Equiv
import Mathbin.Data.Vector.Basic

universe u v w

namespace DArray

variable {n : ℕ} {α : Fin n → Type u}

instance [∀ i, Inhabited (α i)] : Inhabited (DArray n α) :=
  ⟨⟨default⟩⟩

end DArray

namespace Array'

instance {n α} [Inhabited α] : Inhabited (Array' n α) :=
  DArray.inhabited

theorem toList_of_hEq {n₁ n₂ α} {a₁ : Array' n₁ α} {a₂ : Array' n₂ α} (hn : n₁ = n₂)
    (ha : HEq a₁ a₂) : a₁.toList = a₂.toList := by congr <;> assumption
#align array.to_list_of_heq Array'.toList_of_hEq

-- rev_list
section RevList

variable {n : ℕ} {α : Type u} {a : Array' n α}

theorem rev_list_reverse_aux :
    ∀ (i) (h : i ≤ n) (t : List α),
      (a.iterateAux (fun _ => (· :: ·)) i h []).reverseAux t =
        a.revIterateAux (fun _ => (· :: ·)) i h t
  | 0, h, t => rfl
  | i + 1, h, t => rev_list_reverse_aux i _ _
#align array.rev_list_reverse_aux Array'.rev_list_reverse_aux

@[simp]
theorem revList_reverse : a.revList.reverse = a.toList :=
  rev_list_reverse_aux _ _ _
#align array.rev_list_reverse Array'.revList_reverse

@[simp]
theorem toList_reverse : a.toList.reverse = a.revList := by
  rw [← rev_list_reverse, List.reverse_reverse]
#align array.to_list_reverse Array'.toList_reverse

end RevList

-- mem
section Mem

variable {n : ℕ} {α : Type u} {v : α} {a : Array' n α}

theorem Mem.def : v ∈ a ↔ ∃ i, a.read i = v :=
  Iff.rfl
#align array.mem.def Array'.Mem.def

theorem mem_rev_list_aux :
    ∀ {i} (h : i ≤ n),
      (∃ j : Fin n, (j : ℕ) < i ∧ read a j = v) ↔ v ∈ a.iterateAux (fun _ => (· :: ·)) i h []
  | 0, _ => ⟨fun ⟨i, n, _⟩ => absurd n i.val.not_lt_zero, False.elim⟩
  | i + 1, h =>
    let IH := mem_rev_list_aux (le_of_lt h)
    ⟨fun ⟨j, ji1, e⟩ =>
      Or.elim (lt_or_eq_of_le <| Nat.le_of_succ_le_succ ji1)
        (fun ji => List.mem_cons_of_mem _ <| IH.1 ⟨j, ji, e⟩) fun je => by
        simp [DArray.iterateAux] <;> apply Or.inl <;> unfold read at e <;>
            have H : j = ⟨i, h⟩ := Fin.eq_of_veq je <;>
          rwa [← H, e],
      fun m => by
      simp [DArray.iterateAux, List.Mem] at m
      cases' m with e m'
      exact ⟨⟨i, h⟩, Nat.lt_succ_self _, Eq.symm e⟩
      exact
        let ⟨j, ji, e⟩ := IH.2 m'
        ⟨j, Nat.le_succ_of_le ji, e⟩⟩
#align array.mem_rev_list_aux Array'.mem_rev_list_aux

@[simp]
theorem mem_revList : v ∈ a.revList ↔ v ∈ a :=
  Iff.symm <|
    Iff.trans
      (exists_congr fun j =>
        Iff.symm <| show j.1 < n ∧ read a j = v ↔ read a j = v from and_iff_right j.2)
      (mem_rev_list_aux _)
#align array.mem_rev_list Array'.mem_revList

@[simp]
theorem mem_toList : v ∈ a.toList ↔ v ∈ a := by
  rw [← rev_list_reverse] <;> exact list.mem_reverse.trans mem_rev_list
#align array.mem_to_list Array'.mem_toList

end Mem

-- foldr
section Foldr

variable {n : ℕ} {α : Type u} {β : Type w} {b : β} {f : α → β → β} {a : Array' n α}

theorem rev_list_foldr_aux :
    ∀ {i} (h : i ≤ n),
      (DArray.iterateAux a (fun _ => (· :: ·)) i h []).foldr f b =
        DArray.iterateAux a (fun _ => f) i h b
  | 0, h => rfl
  | j + 1, h => congr_arg (f (read a ⟨j, h⟩)) (rev_list_foldr_aux _)
#align array.rev_list_foldr_aux Array'.rev_list_foldr_aux

theorem revList_foldr : a.revList.foldr f b = a.foldl b f :=
  rev_list_foldr_aux _
#align array.rev_list_foldr Array'.revList_foldr

end Foldr

-- foldl
section Foldl

variable {n : ℕ} {α : Type u} {β : Type w} {b : β} {f : β → α → β} {a : Array' n α}

theorem toList_foldl : a.toList.foldl f b = a.foldl b (Function.swap f) := by
  rw [← rev_list_reverse, List.foldl_reverse, rev_list_foldr]
#align array.to_list_foldl Array'.toList_foldl

end Foldl

-- length
section Length

variable {n : ℕ} {α : Type u}

theorem rev_list_length_aux (a : Array' n α) (i h) :
    (a.iterateAux (fun _ => (· :: ·)) i h []).length = i := by
  induction i <;> simp [*, DArray.iterateAux]
#align array.rev_list_length_aux Array'.rev_list_length_aux

@[simp]
theorem revList_length (a : Array' n α) : a.revList.length = n :=
  rev_list_length_aux a _ _
#align array.rev_list_length Array'.revList_length

@[simp]
theorem toList_length (a : Array' n α) : a.toList.length = n := by
  rw [← rev_list_reverse, List.length_reverse, rev_list_length]
#align array.to_list_length Array'.toList_length

end Length

-- nth
section Nth

variable {n : ℕ} {α : Type u} {a : Array' n α}

theorem to_list_nthLe_aux (i : ℕ) (ih : i < n) :
    ∀ (j) {jh t h'},
      (∀ k tl, j + k = i → List.nthLe t k tl = a.read ⟨i, ih⟩) →
        (a.revIterateAux (fun _ => (· :: ·)) j jh t).nthLe i h' = a.read ⟨i, ih⟩
  | 0, _, _, _, al => al i _ <| zero_add _
  | j + 1, jh, t, h', al =>
    to_list_nth_le_aux j fun k tl hjk =>
      show List.nthLe (a.read ⟨j, jh⟩ :: t) k tl = a.read ⟨i, ih⟩ from
        match k, hjk, tl with
        | 0, e, tl =>
          match i, e, ih with
          | _, rfl, _ => rfl
        | k' + 1, _, tl => by
          simp [List.nthLe] <;> exact al _ _ (by simp [add_comm, add_assoc, *] <;> cc)
#align array.to_list_nth_le_aux Array'.to_list_nthLe_aux

theorem toList_nthLe (i : ℕ) (h h') : List.nthLe a.toList i h' = a.read ⟨i, h⟩ :=
  to_list_nthLe_aux _ _ _ fun k tl => absurd tl k.not_lt_zero
#align array.to_list_nth_le Array'.toList_nthLe

@[simp]
theorem toList_nth_le' (a : Array' n α) (i : Fin n) (h') : List.nthLe a.toList i h' = a.read i := by
  cases i <;> apply to_list_nth_le
#align array.to_list_nth_le' Array'.toList_nth_le'

theorem toList_get? {i v} : List.get? a.toList i = some v ↔ ∃ h, a.read ⟨i, h⟩ = v :=
  by
  rw [List.get?_eq_some']
  have ll := to_list_length a
  constructor <;> intro h <;> cases' h with h e <;> subst v
  · exact ⟨ll ▸ h, (to_list_nth_le _ _ _).symm⟩
  · exact ⟨ll.symm ▸ h, to_list_nth_le _ _ _⟩
#align array.to_list_nth Array'.toList_get?

theorem write_toList {i v} : (a.write i v).toList = a.toList.set i v :=
  List.ext_nthLe (by simp) fun j h₁ h₂ =>
    by
    have h₃ : j < n := by simpa using h₁
    rw [to_list_nth_le _ h₃]
    refine'
      let ⟨_, e⟩ := List.get?_eq_some'.1 _
      e.symm
    by_cases ij : (i : ℕ) = j
    · subst j
      rw [show (⟨(i : ℕ), h₃⟩ : Fin _) = i from Fin.eq_of_veq rfl, Array'.read_write,
        List.get?_set_eq_of_lt]
      simp [h₃]
    · rw [List.get?_set_ne _ _ ij, a.read_write_of_ne, to_list_nth.2 ⟨h₃, rfl⟩]
      exact Fin.ne_of_vne ij
#align array.write_to_list Array'.write_toList

end Nth

-- enum
section Enum

variable {n : ℕ} {α : Type u} {a : Array' n α}

theorem mem_toList_enum {i v} : (i, v) ∈ a.toList.enum ↔ ∃ h, a.read ⟨i, h⟩ = v := by
  simp [List.mem_iff_get?, to_list_nth, and_comm, and_assoc, and_left_comm]
#align array.mem_to_list_enum Array'.mem_toList_enum

end Enum

-- to_array
section ToArray

variable {n : ℕ} {α : Type u}

@[simp]
theorem toList_toArray (a : Array' n α) : HEq a.toList.toArray a :=
  hEq_of_hEq_of_eq
      (@Eq.drecOn
        (fun m (e : a.toList.length = m) =>
          HEq (DArray.mk fun v => a.toList.nthLe v.1 v.2)
            (@DArray.mk m (fun _ => α) fun v => a.toList.nthLe v.1 <| e.symm ▸ v.2))
        a.toList_length HEq.rfl) <|
    DArray.ext fun ⟨i, h⟩ => toList_nthLe i h _
#align array.to_list_to_array Array'.toList_toArray

@[simp]
theorem toArray_toList (l : List α) : l.toArray.toList = l :=
  List.ext_nthLe (toList_length _) fun n h1 h2 => toList_nthLe _ h2 _
#align array.to_array_to_list Array'.toArray_toList

end ToArray

-- push_back
section PushBack

variable {n : ℕ} {α : Type u} {v : α} {a : Array' n α}

theorem pushBack_rev_list_aux :
    ∀ i h h',
      DArray.iterateAux (a.pushBack v) (fun _ => (· :: ·)) i h [] =
        DArray.iterateAux a (fun _ => (· :: ·)) i h' []
  | 0, h, h' => rfl
  | i + 1, h, h' => by
    simp [DArray.iterateAux]
    refine' ⟨_, push_back_rev_list_aux _ _ _⟩
    dsimp [read, DArray.read, push_back]
    rw [dif_neg]; rfl
    exact ne_of_lt h'
#align array.push_back_rev_list_aux Array'.pushBack_rev_list_aux

@[simp]
theorem pushBack_revList : (a.pushBack v).revList = v :: a.revList :=
  by
  unfold push_back rev_list foldl iterate DArray.iterate
  dsimp [DArray.iterateAux, read, DArray.read, push_back]
  rw [dif_pos (Eq.refl n)]
  apply congr_arg
  apply push_back_rev_list_aux
#align array.push_back_rev_list Array'.pushBack_revList

@[simp]
theorem pushBack_toList : (a.pushBack v).toList = a.toList ++ [v] := by
  rw [← rev_list_reverse, ← rev_list_reverse, push_back_rev_list, List.reverse_cons]
#align array.push_back_to_list Array'.pushBack_toList

@[simp]
theorem read_pushBack_left (i : Fin n) : (a.pushBack v).read i.cast_succ = a.read i :=
  by
  cases' i with i hi
  have : ¬i = n := ne_of_lt hi
  simp [push_back, this, Fin.castSucc, Fin.castAdd, Fin.castLe, Fin.castLt, read, DArray.read]
#align array.read_push_back_left Array'.read_pushBack_left

@[simp]
theorem read_pushBack_right : (a.pushBack v).read (Fin.last _) = v :=
  by
  cases' hn : Fin.last n with k hk
  have : k = n := by simpa [Fin.eq_iff_veq] using hn.symm
  simp [push_back, this, Fin.castSucc, Fin.castAdd, Fin.castLe, Fin.castLt, read, DArray.read]
#align array.read_push_back_right Array'.read_pushBack_right

end PushBack

-- foreach
section Foreach

variable {n : ℕ} {α : Type u} {β : Type v} {i : Fin n} {f : Fin n → α → β} {a : Array' n α}

@[simp]
theorem read_foreach : (foreach a f).read i = f i (a.read i) :=
  rfl
#align array.read_foreach Array'.read_foreach

end Foreach

-- map
section Map

variable {n : ℕ} {α : Type u} {β : Type v} {i : Fin n} {f : α → β} {a : Array' n α}

theorem read_map : (a.map f).read i = f (a.read i) :=
  read_foreach
#align array.read_map Array'.read_map

end Map

-- map₂
section Map₂

variable {n : ℕ} {α : Type u} {i : Fin n} {f : α → α → α} {a₁ a₂ : Array' n α}

@[simp]
theorem read_map₂ : (map₂ f a₁ a₂).read i = f (a₁.read i) (a₂.read i) :=
  read_foreach
#align array.read_map₂ Array'.read_map₂

end Map₂

end Array'


/-
Copyright (c) 2018 Simon Hudon. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Simon Hudon

General utility functions for buffers.
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.Lean3Lib.data.buffer
import Mathlib.data.array.lemmas
import Mathlib.control.traversable.instances
import Mathlib.PostPort

universes u_1

namespace Mathlib

namespace buffer

protected instance inhabited {α : Type u_1} : Inhabited (buffer α) :=
  { default := nil }

theorem ext {α : Type u_1} {b₁ : buffer α} {b₂ : buffer α} : to_list b₁ = to_list b₂ → b₁ = b₂ := sorry

protected instance decidable_eq (α : Type u_1) [DecidableEq α] : DecidableEq (buffer α) :=
  id
    fun (_v : buffer α) =>
      sigma.cases_on _v
        fun (fst : ℕ) (snd : array fst α) (w : buffer α) =>
          sigma.cases_on w
            fun (w_fst : ℕ) (w_snd : array w_fst α) =>
              decidable.by_cases
                (fun (ᾰ : fst = w_fst) =>
                  Eq._oldrec
                    (fun (w_snd : array fst α) =>
                      decidable.by_cases (fun (ᾰ : snd = w_snd) => Eq._oldrec (is_true sorry) ᾰ)
                        fun (ᾰ : ¬snd = w_snd) => isFalse sorry)
                    ᾰ w_snd)
                fun (ᾰ : ¬fst = w_fst) => isFalse sorry

@[simp] theorem to_list_append_list {α : Type u_1} {xs : List α} {b : buffer α} : to_list (append_list b xs) = to_list b ++ xs := sorry

@[simp] theorem append_list_mk_buffer {α : Type u_1} {xs : List α} : append_list mk_buffer xs = array.to_buffer (list.to_array xs) := sorry

/-- The natural equivalence between lists and buffers, using
`list.to_buffer` and `buffer.to_list`. -/
def list_equiv_buffer (α : Type u_1) : List α ≃ buffer α :=
  equiv.mk list.to_buffer to_list sorry sorry

protected instance traversable : traversable buffer :=
  equiv.traversable list_equiv_buffer

protected instance is_lawful_traversable : is_lawful_traversable buffer :=
  equiv.is_lawful_traversable list_equiv_buffer


/-
Copyright (c) 2020 Scott Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Bhavik Mehta, Scott Morrison
-/
import category_theory.subobject.basic
import category_theory.preadditive.basic

/-!
# Factoring through subobjects

The predicate `h : P.factors f`, for `P : subobject Y` and `f : X ⟶ Y`
asserts the existence of some `P.factor_thru f : X ⟶ (P : C)` making the obvious diagram commute.

-/

universes v₁ v₂ u₁ u₂

noncomputable theory

open category_theory category_theory.category category_theory.limits

variables {C : Type u₁} [category.{v₁} C] {X Y Z : C}
variables {D : Type u₂} [category.{v₂} D]

namespace category_theory

namespace mono_over

/-- When `f : X ⟶ Y` and `P : mono_over Y`,
`P.factors f` expresses that there exists a factorisation of `f` through `P`.
Given `h : P.factors f`, you can recover the morphism as `P.factor_thru f h`.
-/
def factors {X Y : C} (P : mono_over Y) (f : X ⟶ Y) : Prop := ∃ g : X ⟶ (P : C), g ≫ P.arrow = f

lemma factors_congr {X : C} {f g : mono_over X} {Y : C} (h : Y ⟶ X) (e : f ≅ g) :
  f.factors h ↔ g.factors h :=
⟨λ ⟨u, hu⟩, ⟨u ≫ (((mono_over.forget _).map e.hom)).left, by simp [hu]⟩,
 λ ⟨u, hu⟩, ⟨u ≫ (((mono_over.forget _).map e.inv)).left, by simp [hu]⟩⟩

/-- `P.factor_thru f h` provides a factorisation of `f : X ⟶ Y` through some `P : mono_over Y`,
given the evidence `h : P.factors f` that such a factorisation exists. -/
def factor_thru {X Y : C} (P : mono_over Y) (f : X ⟶ Y) (h : factors P f) : X ⟶ (P : C) :=
classical.some h

end mono_over

namespace subobject

/-- When `f : X ⟶ Y` and `P : subobject Y`,
`P.factors f` expresses that there exists a factorisation of `f` through `P`.
Given `h : P.factors f`, you can recover the morphism as `P.factor_thru f h`.
-/
def factors {X Y : C} (P : subobject Y) (f : X ⟶ Y) : Prop :=
quotient.lift_on' P (λ P, P.factors f)
begin
  rintros P Q ⟨h⟩,
  apply propext,
  split,
  { rintro ⟨i, w⟩,
    exact ⟨i ≫ h.hom.left, by erw [category.assoc, over.w h.hom, w]⟩, },
  { rintro ⟨i, w⟩,
    exact ⟨i ≫ h.inv.left, by erw [category.assoc, over.w h.inv, w]⟩, },
end

@[simp]

lemma mk_factors_self (f : X ⟶ Y) [mono f] : (mk f).factors f := ⟨𝟙 _, by simp⟩

lemma factors_iff {X Y : C} (P : subobject Y) (f : X ⟶ Y) :
  P.factors f ↔ (representative.obj P).factors f :=
quot.induction_on P $ λ a, mono_over.factors_congr _ (representative_iso _).symm

lemma factors_self {X : C} (P : subobject X) : P.factors P.arrow :=
(factors_iff _ _).mpr ⟨𝟙 P, (by simp)⟩

lemma factors_comp_arrow {X Y : C} {P : subobject Y} (f : X ⟶ P) : P.factors (f ≫ P.arrow) :=
(factors_iff _ _).mpr ⟨f, rfl⟩

lemma factors_of_factors_right {X Y Z : C} {P : subobject Z} (f : X ⟶ Y) {g : Y ⟶ Z}
  (h : P.factors g) : P.factors (f ≫ g) :=
begin
  revert P,
  refine quotient.ind' _,
  intro P,
  rintro ⟨g, rfl⟩,
  exact ⟨f ≫ g, by simp⟩,
end

lemma factors_zero [has_zero_morphisms C] {X Y : C} {P : subobject Y} :
  P.factors (0 : X ⟶ Y) :=
(factors_iff _ _).mpr ⟨0, by simp⟩

lemma factors_of_le {Y Z : C} {P Q : subobject Y} (f : Z ⟶ Y) (h : P ≤ Q) :
  P.factors f → Q.factors f :=
by { simp only [factors_iff], exact λ ⟨u, hu⟩, ⟨u ≫ of_le _ _ h, by simp [←hu]⟩ }

/-- `P.factor_thru f h` provides a factorisation of `f : X ⟶ Y` through some `P : subobject Y`,
given the evidence `h : P.factors f` that such a factorisation exists. -/
def factor_thru {X Y : C} (P : subobject Y) (f : X ⟶ Y) (h : factors P f) : X ⟶ P :=
classical.some ((factors_iff _ _).mp h)

@[simp, reassoc] lemma factor_thru_arrow {X Y : C} (P : subobject Y) (f : X ⟶ Y) (h : factors P f) :
  P.factor_thru f h ≫ P.arrow = f :=
classical.some_spec ((factors_iff _ _).mp h)

@[simp] lemma factor_thru_self {X : C} (P : subobject X) (h) :
  P.factor_thru P.arrow h = 𝟙 P :=
by { ext, simp, }

@[simp] lemma factor_thru_mk_self (f : X ⟶ Y) [mono f] :
  (mk f).factor_thru f (mk_factors_self f) = (underlying_iso f).inv :=
by { ext, simp, }

@[simp] lemma factor_thru_comp_arrow {X Y : C} {P : subobject Y} (f : X ⟶ P) (h) :
  P.factor_thru (f ≫ P.arrow) h = f :=
by { ext, simp, }

@[simp] lemma factor_thru_eq_zero [has_zero_morphisms C]
  {X Y : C} {P : subobject Y} {f : X ⟶ Y} {h : factors P f} :
  P.factor_thru f h = 0 ↔ f = 0 :=
begin
  fsplit,
  { intro w,
    replace w := w =≫ P.arrow,
    simpa using w, },
  { rintro rfl,
    ext, simp, },
end

lemma factor_thru_right {X Y Z : C} {P : subobject Z} (f : X ⟶ Y) (g : Y ⟶ Z) (h : P.factors g) :
  f ≫ P.factor_thru g h = P.factor_thru (f ≫ g) (factors_of_factors_right f h) :=
begin
  apply (cancel_mono P.arrow).mp,
  simp,
end

@[simp]
lemma factor_thru_zero
  [has_zero_morphisms C] {X Y : C} {P : subobject Y} (h : P.factors (0 : X ⟶ Y)) :
  P.factor_thru 0 h = 0 :=
by simp

-- `h` is an explicit argument here so we can use
-- `rw factor_thru_le h`, obtaining a subgoal `P.factors f`.
-- (While the reverse direction looks plausible as a simp lemma, it seems to be unproductive.)
lemma factor_thru_of_le
  {Y Z : C} {P Q : subobject Y} {f : Z ⟶ Y} (h : P ≤ Q) (w : P.factors f) :
  Q.factor_thru f (factors_of_le f h w) = P.factor_thru f w ≫ of_le P Q h :=
by { ext, simp, }

section preadditive

variables [preadditive C]

lemma factors_add {X Y : C} {P : subobject Y} (f g : X ⟶ Y) (wf : P.factors f) (wg : P.factors g) :
  P.factors (f + g) :=
(factors_iff _ _).mpr ⟨P.factor_thru f wf + P.factor_thru g wg, by simp⟩

-- This can't be a `simp` lemma as `wf` and `wg` may not exist.
-- However you can `rw` by it to assert that `f` and `g` factor through `P` separately.
lemma factor_thru_add {X Y : C} {P : subobject Y} (f g : X ⟶ Y)
   (w : P.factors (f + g)) (wf : P.factors f) (wg : P.factors g) :
  P.factor_thru (f + g) w = P.factor_thru f wf + P.factor_thru g wg :=
by { ext, simp, }

lemma factors_left_of_factors_add {X Y : C} {P : subobject Y} (f g : X ⟶ Y)
  (w : P.factors (f + g)) (wg : P.factors g) : P.factors f :=
(factors_iff _ _).mpr ⟨P.factor_thru (f + g) w - P.factor_thru g wg, by simp⟩

@[simp]
lemma factor_thru_add_sub_factor_thru_right {X Y : C} {P : subobject Y} (f g : X ⟶ Y)
  (w : P.factors (f + g)) (wg : P.factors g) :
  P.factor_thru (f + g) w - P.factor_thru g wg =
    P.factor_thru f (factors_left_of_factors_add f g w wg) :=
by { ext, simp, }

lemma factors_right_of_factors_add {X Y : C} {P : subobject Y} (f g : X ⟶ Y)
  (w : P.factors (f + g)) (wf : P.factors f) : P.factors g :=
(factors_iff _ _).mpr ⟨P.factor_thru (f + g) w - P.factor_thru f wf, by simp⟩

@[simp]
lemma factor_thru_add_sub_factor_thru_left {X Y : C} {P : subobject Y} (f g : X ⟶ Y)
  (w : P.factors (f + g)) (wf : P.factors f) :
  P.factor_thru (f + g) w - P.factor_thru f wf =
    P.factor_thru g (factors_right_of_factors_add f g w wf) :=
by { ext, simp, }

end preadditive

end subobject

end category_theory


/-
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

/-!
# Irreducible definitions

This file defines an `irreducible_def` command,
which works almost like the `def` command
except that the introduced definition
does not reduce to the value.
Instead, the command
adds a `_def` lemma
which can be used for rewriting.

```
irreducible_def frobnicate (a b : Nat) :=
  a + b

example : frobnicate a 0 = a := by
  simp [frobnicate_def]
```

-/

namespace Lean.Elab.Command

open Term Meta

/-- `delta% t` elaborates to a head-delta reduced version of `t`. -/
elab "delta% " t:term : term <= expectedType => do
  let t ← elabTerm t expectedType
  synthesizeSyntheticMVars
  let t ← instantiateMVars t
  let some t ← delta? t | throwError "cannot delta reduce {t}"
  pure t

/- `eta_helper f = (· + 3)` elabs to `∀ x, f x = x + 3` -/
local elab "eta_helper " t:term : term => do
  let t ← elabTerm t none
  let some (_, lhs, rhs) := t.eq? | throwError "not an equation: {t}"
  synthesizeSyntheticMVars
  let rhs ← instantiateMVars rhs
  lambdaLetTelescope rhs fun xs rhs => do
    let lhs := (mkAppN lhs xs).headBeta
    mkForallFVars xs <|← mkEq lhs rhs

/-- `value_proj x` elabs to `@x.value` -/
local elab "value_proj " e:term : term => do
  let e ← elabTerm e none
  mkProjection e `value

/--
Executes the commands,
and stops after the first error.
In short, S-A-F-E.
-/
local syntax "stop_at_first_error" command* : command
open Command in elab_rules : command
  | `(stop_at_first_error $[$cmds]*) => do
    for cmd in cmds do
      elabCommand cmd
      if (← get).messages.hasErrors then break

/--
Introduces an irreducible definition.
`irreducible_def foo := 42` generates
a constant `foo : Nat` as well as
a theorem `foo_def : foo = 42`.
-/
macro mods:declModifiers "irreducible_def" n_id:declId declSig:optDeclSig val:declVal : command => do
  let (n, us) ← match n_id with
    | `(Parser.Command.declId| $n:ident $[.{$us,*}]?) => pure (n, us)
    | _ => Macro.throwUnsupported
  let us' := us.getD (Syntax.SepArray.ofElems #[])
  let n_def := mkIdent <| (·.review) <|
    let scopes := extractMacroScopes n.getId
    { scopes with name := scopes.name.appendAfter "_def" }
  `(stop_at_first_error
    def definition$[.{$us,*}]? $declSig:optDeclSig $val
    structure Wrapper$[.{$us,*}]? where
      value : type_of% @definition.{$us',*}
      prop : Eq @value @(delta% @definition)
    constant wrapped$[.{$us,*}]? : Wrapper.{$us',*} := ⟨_, rfl⟩
    $mods:declModifiers def $n:ident$[.{$us,*}]? := value_proj @wrapped.{$us',*}
    theorem $n_def:ident $[.{$us,*}]? : eta_helper Eq @$n.{$us',*} @(delta% @definition) := by
      intros
      simp only [$n:ident]
      rw [wrapped.prop])


/-
## Interface for MLIR dialects
-/

/-
### Extended types and attributes

Dialects can define custom types and attributes that map to concrete Lean
datatypes satisfying the requirements of the type interface. These types and
attributes can be mixed together with other dialects' and used in proofs.

Because it is impossible to know in advance what set of dialects an operation
will be used with, most proofs should either (1) quantify on a dialect
interface and handle default cases, or (2) come with a lifting theorem that
performs this closure under context automatically.
-/

/-- The `DialectTypeIntf` typeclass defines the requirements for dialect-
    supplied custom types. These properties are available even on unknown
    custom types, which allows type-generic functions and proofs to be written.

    In the interface, `σ` is the type of "signatures" (parameters), while `ε`
    is the type of "extended values". For instance for `builtin.tensor`, `σ` is
    a datatype holding the dimensions and base type of the tensor, while `ε` is
    the type of tensor values (for any particular signature).

    Type families that quantify over `Type` or higher universes are not
    supported to avoid universe polymorphism in the core data structures; in
    this case, `ε` can be made to implement specific instances of the family
    instead. -/
class DialectTypeIntf (σ: Type) (ε: σ → Type): Type where

  /-- The type should be inhabited, so that SSA value accesses can return
      defaults instead of requiring a proof that every access is dominated by
      a definition (which is a very annoying statement to maintain). -/
  inhabited: forall (s: σ), ε s

  /-- The signature should have decidable equality, so we can compare types. -/
  typeEq: DecidableEq σ
  /-- The type should have decidable equality, so that rewriting tools can
      find matches involving concrete values (eg. in PDL).
      Note: This might be relaxed into a `BEq` in the future. -/
  eq: forall (s: σ), DecidableEq (ε s)

  /-- String representation of values in the type (eg. a multi-dimensional
      array for tensors).  -/
  str: (s: σ) → ε s → String
  /-- String representation of the type itself (eg. "tensor<4x4xf64>"). -/
  typeStr: σ → String

  -- TODO: DialectTypeIntf: Type signature, to match eg. "any builtin.vector"

def DialectTypeIntf.sigType {σ ε} (_: DialectTypeIntf σ ε): Type := σ
def DialectTypeIntf.extType {σ ε} (_: DialectTypeIntf σ ε) (s: σ): Type := ε s

-- Expressing typeclass instances this way helps resolution
instance {σ ε} [i: DialectTypeIntf σ ε] (s: σ): Inhabited (ε s) where
  default := i.inhabited s
instance {σ ε} [i: DialectTypeIntf σ ε] (s: σ): ToString (ε s) where
  toString := i.str s
instance _DEq1 {σ ε} [i: DialectTypeIntf σ ε]: DecidableEq σ :=
  i.typeEq
instance _DEq2 {σ ε} [i: DialectTypeIntf σ ε] (s: σ): DecidableEq (ε s) :=
  i.eq s

/-- The `DialectAttrIntf` typeclass defines the requirements for dialect-
   supplied custom attributes. -/
class DialectAttrIntf (α: Type) where

  /-- The attribute should have decidable equality so that rewriting can match
      against them (eg. in PDL). -/
  eq: DecidableEq α

  /-- String representation of attribute values. -/
  str: α → String

  -- TODO: More data on attributes

def DialectAttrIntf.type {α} (_: DialectAttrIntf α): Type := α

instance _DEq3 {α} [i: DialectAttrIntf α]: DecidableEq α := i.eq

/-
### Combinations of interfaces

As dialects can provide multiple sets of extended types and attributes (and
dialects may themselves be combined), the interface must allow for different
extensions to be combined.

The following instances allow extended types, attributes, and dialects to be
combined with `Sum`.
-/

-- Like Sum.rec, but not a recursor (hence supported for code generation)
@[reducible]
def Sum.cases {α β γ} (fα: α → γ) (fβ: β → γ): (α ⊕ β) → γ
  | .inl a => fα a
  | .inr b => fβ b

instance {σ₁ ε₁ σ₂ ε₂} [i₁: DialectTypeIntf σ₁ ε₁] [i₂: DialectTypeIntf σ₂ ε₂]:
    DialectTypeIntf (σ₁ ⊕ σ₂) (Sum.cases ε₁ ε₂) where
  inhabited s :=
    match s with
    | .inl s₁ => i₁.inhabited s₁
    | .inr s₂ => i₂.inhabited s₂
  typeEq := inferInstance
  eq s :=
    match s with
    | .inl s₁ => i₁.eq s₁
    | .inr s₂ => i₂.eq s₂
  str s :=
    match s with
    | .inl s₁ => i₁.str s₁
    | .inr s₂ => i₂.str s₂
  typeStr := Sum.cases i₁.typeStr i₂.typeStr

instance {α₁ α₂} [i₁: DialectAttrIntf α₁] [i₂: DialectAttrIntf α₂]:
    DialectAttrIntf (α₁ ⊕ α₂) where
  eq := inferInstance
  str := Sum.cases i₁.str i₂.str

/-
### Dialects
-/

-- TODO: Document and finish the Dialect interface
class Dialect (α σ) (ε: σ → Type): Type :=
  name: String
  iα: DialectAttrIntf α
  iε: DialectTypeIntf σ ε

instance {α ε} [δ: Dialect α σ ε]: DialectAttrIntf α := δ.iα
instance {α ε} [δ: Dialect α σ ε]: DialectTypeIntf σ ε := δ.iε

/-
### Empty dialect

The empty dialect is a default value to start building hierarchies from. It is
used in a couple of aliases, eg. `MLIRTy` (for `MLIRType Dialect.empty`) and
`AttrVal` (for `AttrValue Dialect.empty`).
-/

inductive Void :=
deriving DecidableEq

instance: DialectTypeIntf Void (fun _ => Unit) where
  inhabited s := nomatch s
  typeEq      := inferInstance
  eq s        := nomatch s
  str s       := nomatch s
  typeStr s   := nomatch s

instance: DialectAttrIntf Void where
  eq          := inferInstance
  str a       := nomatch a

instance Dialect.empty: Dialect Void Void (fun _ => Unit) where
  name := "Empty"
  iα := inferInstance
  iε := inferInstance

-- We write combinations of dialects with + as usual (no risk of confusion)
instance {α₁ σ₁ ε₁ α₂ σ₂ ε₂}:
  HAdd (Dialect α₁ σ₁ ε₁) (Dialect α₂ σ₂ ε₂)
       (Dialect (α₁ ⊕ α₂) (σ₁ ⊕ σ₂) (Sum.cases ε₁ ε₂)) where
  hAdd δ₁ δ₂ := {
    name := s!"({δ₁.name}+{δ₂.name})"
    iα := inferInstance
    iε := inferInstance
  }

instance {α₁ σ₁ ε₁ α₂ σ₂ ε₂} [δ₁: Dialect α₁ σ₁ ε₁] [δ₂: Dialect α₂ σ₂ ε₂]:
    Dialect (α₁ ⊕ α₂) (σ₁ ⊕ σ₂) (Sum.cases ε₁ ε₂) :=
  δ₁ + δ₂

/-
### Coercions of dialects

The `CoeDialect` ckass is used to automatically inject individual dialects into
sums of dialects, which in turn allows automatic conversion of instances of
common MLIR data such as `MLIRType`, `AttrValue` and `Op` across dialects.
-/

class CoeDialect (δ₁: Dialect α₁ σ₁ ε₁) (δ₂: Dialect α₂ σ₂ ε₂) where
  coe_α: α₁ → α₂
  coe_σ: σ₁ → σ₂
  coe_ε: forall s, ε₁ s → ε₂ (coe_σ s)

instance (δ₁: Dialect α₁ σ₁ ε₁) (δ₂: Dialect α₂ σ₂ ε₂) [c: CoeDialect δ₁ δ₂]:
  Coe α₁ α₂ where coe := c.coe_α
instance (δ₁: Dialect α₁ σ₁ ε₁) (δ₂: Dialect α₂ σ₂ ε₂) [c: CoeDialect δ₁ δ₂]:
  Coe σ₁ σ₂ where coe := c.coe_σ
instance (δ₁: Dialect α₁ σ₁ ε₁) (δ₂: Dialect α₂ σ₂ ε₂) [c: CoeDialect δ₁ δ₂] s:
  Coe (ε₁ s) (ε₂ /-coe-/s) where coe := c.coe_ε s

instance (δ: Dialect α σ ε): CoeDialect δ δ where
  coe_α := id
  coe_σ := id
  coe_ε s := id

instance (δ₁: Dialect α₁ σ₁ ε₁) (δ₂: Dialect α₂ σ₂ ε₂):
    CoeDialect δ₁ (δ₁ + δ₂) where
  coe_α := .inl
  coe_σ := .inl
  coe_ε s := id

instance (δ₁: Dialect α₁ σ₁ ε₁) (δ₂: Dialect α₂ σ₂ ε₂):
    CoeDialect δ₂ (δ₁ + δ₂) where
  coe_α := .inr
  coe_σ := .inr
  coe_ε s := id

instance (δ: Dialect α σ ε): CoeDialect Dialect.empty δ where
  coe_α a := nomatch a
  coe_σ s := nomatch s
  coe_ε s := nomatch s


/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Mario Carneiro
-/
import Std.Data.List.Basic

namespace Std

/--
`AssocList α β` is "the same as" `List (α × β)`, but flattening the structure
leads to one fewer pointer indirection (in the current code generator).
It is mainly intended as a component of `HashMap`, but it can also be used as a plain
key-value map.
-/
inductive AssocList (α : Type u) (β : Type v) where
  /-- An empty list -/
  | nil
  /-- Add a `key, value` pair to the list -/
  | cons (key : α) (value : β) (tail : AssocList α β)
  deriving Inhabited

namespace AssocList

/--
`O(n)`. Convert an `AssocList α β` into the equivalent `List (α × β)`.
This is used to give specifications for all the `AssocList` functions
in terms of corresponding list functions.
-/
@[simp] def toList : AssocList α β → List (α × β)
  | nil => []
  | cons a b es => (a, b) :: es.toList

instance : EmptyCollection (AssocList α β) := ⟨nil⟩

@[simp] theorem empty_eq : (∅ : AssocList α β) = nil := rfl

/-- `O(1)`. Is the list empty? -/
def isEmpty : AssocList α β → Bool
  | nil => true
  | _   => false

@[simp] theorem isEmpty_eq (l : AssocList α β) : isEmpty l = l.toList.isEmpty := by
  cases l <;> simp [*, isEmpty, List.isEmpty]

/-- `O(n)`. Fold a monadic function over the list, from head to tail. -/
@[specialize] def foldlM [Monad m] (f : δ → α → β → m δ) : (init : δ) → AssocList α β → m δ
  | d, nil         => pure d
  | d, cons a b es => do foldlM f (← f d a b) es

@[simp] theorem foldlM_eq [Monad m] (f : δ → α → β → m δ) (init l) :
    foldlM f init l = l.toList.foldlM (fun d (a, b) => f d a b) init := by
  induction l generalizing init <;> simp [*, foldlM]

/-- `O(n)`. Fold a function over the list, from head to tail. -/
@[inline] def foldl (f : δ → α → β → δ) (init : δ) (as : AssocList α β) : δ :=
  Id.run (foldlM f init as)

@[simp] theorem foldl_eq (f : δ → α → β → δ) (init l) :
    foldl f init l = l.toList.foldl (fun d (a, b) => f d a b) init := by
  simp [List.foldl_eq_foldlM, foldl, Id.run]

/-- Optimized version of `toList`. -/
def toListTR (as : AssocList α β) : List (α × β) :=
  as.foldl (init := #[]) (fun r a b => r.push (a, b)) |>.toList

@[csimp] theorem toList_eq_toListTR : @toList = @toListTR := by
  funext α β as; simp [toListTR]
  exact .symm <| (Array.foldl_data_eq_map (toList as) _ id).trans (List.map_id _)

/-- `O(n)`. Run monadic function `f` on all elements in the list, from head to tail. -/
@[specialize] def forM [Monad m] (f : α → β → m PUnit) : AssocList α β → m PUnit
  | nil         => pure ⟨⟩
  | cons a b es => do f a b; forM f es

@[simp] theorem forM_eq [Monad m] (f : α → β → m PUnit) (l) :
    forM f l = l.toList.forM (fun (a, b) => f a b) := by
  induction l <;> simp [*, forM]

/-- `O(n)`. Map a function `f` over the keys of the list. -/
@[simp] def mapKey (f : α → δ) : AssocList α β → AssocList δ β
  | nil        => nil
  | cons k v t => cons (f k) v (mapKey f t)

@[simp] theorem mapKey_toList (f : α → δ) (l : AssocList α β) :
    (mapKey f l).toList = l.toList.map (fun (a, b) => (f a, b)) := by
  induction l <;> simp [*]

/-- `O(n)`. Map a function `f` over the values of the list. -/
@[simp] def mapVal (f : α → β → δ) : AssocList α β → AssocList α δ
  | nil        => nil
  | cons k v t => cons k (f k v) (mapVal f t)

@[simp] theorem mapVal_toList (f : α → β → δ) (l : AssocList α β) :
    (mapVal f l).toList = l.toList.map (fun (a, b) => (a, f a b)) := by
  induction l <;> simp [*]

/-- `O(n)`. Returns the first entry in the list whose entry satisfies `p`. -/
@[specialize] def findEntryP? (p : α → β → Bool) : AssocList α β → Option (α × β)
  | nil         => none
  | cons k v es => bif p k v then some (k, v) else findEntryP? p es

@[simp] theorem findEntryP?_eq (p : α → β → Bool) (l : AssocList α β) :
    findEntryP? p l = l.toList.find? fun (a, b) => p a b := by
  induction l <;> simp [findEntryP?]; split <;> simp [*]

/-- `O(n)`. Returns the first entry in the list whose key is equal to `a`. -/
@[inline] def findEntry? [BEq α] (a : α) (l : AssocList α β) : Option (α × β) :=
  findEntryP? (fun k _ => k == a) l

@[simp] theorem findEntry?_eq [BEq α] (a : α) (l : AssocList α β) :
    findEntry? a l = l.toList.find? (·.1 == a) := findEntryP?_eq ..

/-- `O(n)`. Returns the first value in the list whose key is equal to `a`. -/
def find? [BEq α] (a : α) : AssocList α β → Option β
  | nil         => none
  | cons k v es => match k == a with
    | true  => some v
    | false => find? a es

theorem find?_eq_findEntry? [BEq α] (a : α) (l : AssocList α β) :
    find? a l = (l.findEntry? a).map (·.2) := by
  induction l <;> simp [find?]; split <;> simp [*]

@[simp] theorem find?_eq [BEq α] (a : α) (l : AssocList α β) :
    find? a l = (l.toList.find? (·.1 == a)).map (·.2) := by simp [find?_eq_findEntry?]

/-- `O(n)`. Returns true if any entry in the list satisfies `p`. -/
@[specialize] def any (p : α → β → Bool) : AssocList α β → Bool
  | nil         => false
  | cons k v es => p k v || any p es

@[simp] theorem any_eq (p : α → β → Bool) (l : AssocList α β) :
    any p l = l.toList.any fun (a, b) => p a b := by induction l <;> simp [any, *]

/-- `O(n)`. Returns true if every entry in the list satisfies `p`. -/
@[specialize] def all (p : α → β → Bool) : AssocList α β → Bool
  | nil         => true
  | cons k v es => p k v && all p es

@[simp] theorem all_eq (p : α → β → Bool) (l : AssocList α β) :
    all p l = l.toList.all fun (a, b) => p a b := by induction l <;> simp [all, *]

/-- Returns true if every entry in the list satisfies `p`. -/
def All (p : α → β → Prop) (l : AssocList α β) : Prop := ∀ a ∈ l.toList, p a.1 a.2

/-- `O(n)`. Returns true if there is an element in the list whose key is equal to `a`. -/
@[inline] def contains [BEq α] (a : α) (l : AssocList α β) : Bool := any (fun k _ => k == a) l

@[simp] theorem contains_eq [BEq α] (a : α) (l : AssocList α β) :
    contains a l = l.toList.any (·.1 == a) := by
  induction l <;> simp [*, contains]

/--
`O(n)`. Replace the first entry in the list
with key equal to `a` to have key `a` and value `b`.
-/
@[simp] def replace [BEq α] (a : α) (b : β) : AssocList α β → AssocList α β
  | nil         => nil
  | cons k v es => match k == a with
    | true  => cons a b es
    | false => cons k v (replace a b es)

@[simp] theorem replace_toList [BEq α] (a : α) (b : β) (l : AssocList α β) :
    (replace a b l).toList =
    l.toList.replaceF (bif ·.1 == a then some (a, b) else none) := by
  induction l <;> simp [replace]; split <;> simp [*]

/-- `O(n)`. Remove the first entry in the list with key equal to `a`. -/
@[specialize, simp] def eraseP (p : α → β → Bool) : AssocList α β → AssocList α β
  | nil         => nil
  | cons k v es => bif p k v then es else cons k v (eraseP p es)

@[simp] theorem eraseP_toList (p) (l : AssocList α β) :
    (eraseP p l).toList = l.toList.eraseP fun (a, b) => p a b := by
  induction l <;> simp [List.eraseP, cond]; split <;> simp [*]

/-- `O(n)`. Remove the first entry in the list with key equal to `a`. -/
@[inline] def erase [BEq α] (a : α) (l : AssocList α β) : AssocList α β :=
  eraseP (fun k _ => k == a) l

@[simp] theorem erase_toList [BEq α] (a : α) (l : AssocList α β) :
    (erase a l).toList = l.toList.eraseP (·.1 == a) := eraseP_toList ..

/-- The implementation of `ForIn`, which enables `for (k, v) in aList do ...` notation. -/
@[specialize] protected def forIn [Monad m]
    (as : AssocList α β) (init : δ) (f : (α × β) → δ → m (ForInStep δ)) : m δ :=
  match as with
  | nil => pure init
  | cons k v es => do
    match (← f (k, v) init) with
    | ForInStep.done d  => pure d
    | ForInStep.yield d => es.forIn d f

instance : ForIn m (AssocList α β) (α × β) where
  forIn := AssocList.forIn

@[simp] theorem forIn_eq [Monad m] (l : AssocList α β) (init : δ)
    (f : (α × β) → δ → m (ForInStep δ)) : forIn l init f = forIn l.toList init f := by
  simp [forIn, List.forIn]
  induction l generalizing init <;> simp [AssocList.forIn, List.forIn.loop]
  congr; funext a; split <;> simp [*]

/-- Split the list into head and tail, if possible. -/
def pop? : AssocList α β → Option ((α × β) × AssocList α β)
  | nil => none
  | cons a b l => some ((a, b), l)

instance : ToStream (AssocList α β) (AssocList α β) := ⟨fun x => x⟩
instance : Stream (AssocList α β) (α × β) := ⟨pop?⟩

/-- Converts a list into an `AssocList`. This is the inverse function to `AssocList.toList`. -/
@[simp] def _root_.List.toAssocList : List (α × β) → AssocList α β
  | []          => nil
  | (a,b) :: es => cons a b (toAssocList es)

@[simp] theorem _root_.List.toAssocList_toList (l : List (α × β)) : l.toAssocList.toList = l := by
  induction l <;> simp [*]

@[simp] theorem toList_toAssocList (l : AssocList α β) : l.toList.toAssocList = l := by
  induction l <;> simp [*]