/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Elab.App

/-
Auxiliary elaboration functions: AKA custom elaborators
-/

namespace Lean.Elab.Term
open Meta

def elabBinRelCore (noProp : Bool) (stx : Syntax) (expectedType? : Option Expr) : TermElabM Expr :=  do
  match (← resolveId? stx[1]) with
  | some f =>
    let s ← saveState
    let (lhs, rhs) ← withSynthesize (mayPostpone := true) do
      let mut lhs ← elabTerm stx[2] none
      let mut rhs ← elabTerm stx[3] none
      if lhs.isAppOfArity ``OfNat.ofNat 3 then
        lhs ← ensureHasType (← inferType rhs) lhs
      else if rhs.isAppOfArity ``OfNat.ofNat 3 then
        rhs ← ensureHasType (← inferType lhs) rhs
      return (lhs, rhs)
    let lhs ← toBoolIfNecessary lhs
    let rhs ← toBoolIfNecessary rhs
    let lhsType ← inferType lhs
    let rhsType ← inferType rhs

    let (lhs, rhs) ←
      try
        pure (lhs, ← withRef stx[3] do ensureHasType lhsType rhs)
      catch _ =>
        try
          pure (← withRef stx[2] do ensureHasType rhsType lhs, rhs)
        catch _ =>
          s.restore
          -- Use default approach
          let lhs ← elabTerm stx[2] none
          let rhs ← elabTerm stx[3] none
          let lhsType ← inferType lhs
          let rhsType ← inferType rhs
          pure (lhs, ← withRef stx[3] do ensureHasType lhsType rhs)
    elabAppArgs f #[] #[Arg.expr lhs, Arg.expr rhs] expectedType? (explicit := false) (ellipsis := false)
  | none   => throwUnknownConstant stx[1].getId
where
  /-- If `noProp == true` and `e` has type `Prop`, then coerce it to `Bool`. -/
  toBoolIfNecessary (e : Expr) : TermElabM Expr := do
    if noProp then
      -- We use `withNewMCtxDepth` to make sure metavariables are not assigned
      if (← withNewMCtxDepth <| isDefEq (← inferType e) (mkSort levelZero)) then
        return (← ensureHasType (Lean.mkConst ``Bool) e)
    return e

@[builtinTermElab binrel] def elabBinRel : TermElab := elabBinRelCore false

@[builtinTermElab binrel_no_prop] def elabBinRelNoProp : TermElab := elabBinRelCore true

@[builtinTermElab forInMacro] def elabForIn : TermElab :=  fun stx expectedType? => do
  match stx with
  | `(for_in% $col $init $body) =>
      match (← isLocalIdent? col) with
      | none   => elabTerm (← `(let col := $col; for_in% col $init $body)) expectedType?
      | some colFVar =>
        tryPostponeIfNoneOrMVar expectedType?
        let m ← getMonad expectedType?
        let colType ← inferType colFVar
        let elemType ← mkFreshExprMVar (mkSort (mkLevelSucc (← mkFreshLevelMVar)))
        let forInInstance ←
          try
            mkAppM ``ForIn #[m, colType, elemType]
          catch
            ex => tryPostpone; throwError "failed to construct 'ForIn' instance for collection{indentExpr colType}\nand monad{indentExpr m}"
        match (← trySynthInstance forInInstance) with
        | LOption.some val =>
          let ref ← getRef
          let forInFn ← mkConst ``forIn
          elabAppArgs forInFn #[] #[Arg.stx col, Arg.stx init, Arg.stx body] expectedType? (explicit := false) (ellipsis := false)
        | LOption.undef    => tryPostpone; throwFailure forInInstance
        | LOption.none     => throwFailure forInInstance
  | _ => throwUnsupportedSyntax
where
  getMonad (expectedType? : Option Expr) : TermElabM Expr := do
    match expectedType? with
    | none => throwError "invalid 'for_in%' notation, expected type is not available"
    | some expectedType =>
      match (← isTypeApp? expectedType) with
      | some (m, _) => return m
      | none => throwError "invalid 'for_in%' notation, expected type is not of of the form `M α`{indentExpr expectedType}"
  throwFailure (forInInstance : Expr) : TermElabM Expr :=
    throwError "failed to synthesize instance for 'for_in%' notation{indentExpr forInInstance}"

namespace BinOp
/-

The elaborator for `binop%` terms works as follows:

1- Expand macros.
2- Convert `Syntax` object corresponding to the `binop%` term into a `Tree`.
   The `toTree` method visits nested `binop%` terms and parentheses.
3- Synthesize pending metavariables without applying default instances and using the
   `(mayPostpone := true)`.
4- Tries to compute a maximal type for the tree computed at step 2.
   We say a type α is smaller than type β if there is a (nondependent) coercion from α to β.
   We are currently ignoring the case we may have cycles in the coercion graph.
   If there are "uncomparable" types α and β in the tree, we skip the next step.
   We say two types are "uncomparable" if there isn't a coercion between them.
   Note that two types may be "uncomparable" because some typing information may still be missing.
5- We traverse the tree and inject coercions to the "maximal" type when needed.

Recall that the coercions are expanded eagerly by the elaborator.

Properties:

a) Given `n : Nat` and `i : Nat`, it can successfully elaborate `n + i` and `i + n`. Recall that Lean 3
   fails on the former.

b) The coercions are inserted in the "leaves" like in Lean 3.

c) There are no coercions "hidden" inside instances, and we can elaborate
```
axiom Int.add_comm (i j : Int) : i + j = j + i

example (n : Nat) (i : Int) : n + i = i + n := by
  rw [Int.add_comm]
```
Recall that the `rw` tactic used to fail because our old `binop%` elaborator would hide
coercions inside of a `HAdd` instance.

Remarks:

In the new `binop%` elaborator the decision whether a coercion will be inserted or not
is made at `binop%` elaboration time. This was not the case in the old elaborator.
For example, an instance, such as `HAdd Int ?m ?n`, could be created when executing
the `binop%` elaborator, and only resolved much later. We try to minimize this problem
by synthesizing pending metavariables at step 3.

For types containing heterogeneous operators (e.g., matrix multiplication), step 4 will fail
and we will skip coercion insertion. For example, `x : Matrix Real 5 4` and `y : Matrix Real 4 8`,
there is no coercion `Matrix Real 5 4` from `Matrix Real 4 8` and vice-versa, but
`x * y` is elaborated successfully and has type `Matrix Real 5 8`.
-/

private inductive Tree where
  | term  (ref : Syntax) (val : Expr)
  | op    (ref : Syntax) (lazy : Bool) (f : Expr) (lhs rhs : Tree)

private partial def toTree (s : Syntax) : TermElabM Tree := do
  let result ← go (← liftMacroM <| expandMacros s)
  synthesizeSyntheticMVars (mayPostpone := true)
  return result
where
  go (s : Syntax) := do
    match s with
    | `(binop% $f $lhs $rhs) => processOp (lazy := false) f lhs rhs
    | `(binop_lazy% $f $lhs $rhs) => processOp (lazy := true) f lhs rhs
    | `(($e)) => go e
    | _ =>
       return Tree.term s (← elabTerm s none)

  processOp (f lhs rhs : Syntax) (lazy : Bool) := do
    let some f ← resolveId? f | throwUnknownConstant f.getId
    return Tree.op s (lazy := lazy) f (← go lhs) (← go rhs)

-- Auxiliary function used at `analyze`
private def hasCoe (fromType toType : Expr) : TermElabM Bool := do
  if (← getEnv).contains ``CoeHTCT then
    let u ← getLevel fromType
    let v ← getLevel toType
    let coeInstType := mkAppN (Lean.mkConst ``CoeHTCT [u, v]) #[fromType, toType]
    match ← trySynthInstance coeInstType (some (maxCoeSize.get (← getOptions))) with
    | LOption.some _ => return true
    | LOption.none   => return false
    | LOption.undef  => return false -- TODO: should we do something smarter here?
  else
    return false

private structure AnalyzeResult where
  max?            : Option Expr := none
  hasUncomparable : Bool := false -- `true` if there are two types `α` and `β` where we don't have coercions in any direction.

private def isUnknow : Expr → Bool
  | Expr.mvar ..        => true
  | Expr.app f ..       => isUnknow f
  | Expr.letE _ _ _ b _ => isUnknow b
  | Expr.mdata _ b _    => isUnknow b
  | _                   => false

private def analyze (t : Tree) (expectedType? : Option Expr) : TermElabM AnalyzeResult := do
  let max? ←
    match expectedType? with
    | none => pure none
    | some expectedType =>
      let expectedType ← instantiateMVars expectedType
      if isUnknow expectedType then pure none else pure (some expectedType)
  (go t *> get).run' { max? }
where
   go (t : Tree) : StateRefT AnalyzeResult TermElabM Unit := do
     unless (← get).hasUncomparable do
       match t with
       | Tree.op _ _ _ lhs rhs => go lhs; go rhs
       | Tree.term _ val =>
         let type ← instantiateMVars (← inferType val)
         unless isUnknow type do
           match (← get).max? with
           | none     => modify fun s => { s with max? := type }
           | some max =>
             unless (← withNewMCtxDepth <| isDefEqGuarded max type) do
               if (← hasCoe type max) then
                 return ()
               else if (← hasCoe max type) then
                 modify fun s => { s with max? := type }
               else
                 trace[Elab.binop] "uncomparable types: {max}, {type}"
                 modify fun s => { s with hasUncomparable := true }

private def mkOp (f : Expr) (lhs rhs : Expr) : TermElabM Expr :=
  elabAppArgs f #[] #[Arg.expr lhs, Arg.expr rhs] (expectedType? := none) (explicit := false) (ellipsis := false)

private def toExpr (t : Tree) : TermElabM Expr := do
  match t with
  | Tree.term _ e               => return e
  | Tree.op ref true f lhs rhs  => withRef ref <| mkOp f (← toExpr lhs) (← mkFunUnit (← toExpr rhs))
  | Tree.op ref false f lhs rhs => withRef ref <| mkOp f (← toExpr lhs) (← toExpr rhs)

private def applyCoe (t : Tree) (maxType : Expr) : TermElabM Tree := do
  go t
where
  go (t : Tree) : TermElabM Tree := do
    match t with
    | Tree.op ref lazy f lhs rhs => return Tree.op ref lazy f (← go lhs) (← go rhs)
    | Tree.term ref e       =>
      let type ← inferType e
      trace[Elab.binop] "visiting {e} : {type} =?= {maxType}"
      if (← isDefEqGuarded maxType type) then
        return t
      else
        trace[Elab.binop] "added coercion: {e} : {type} => {maxType}"
        withRef ref <| return Tree.term ref (← mkCoe maxType type e)

@[builtinTermElab binop]
def elabBinOp : TermElab :=  fun stx expectedType? => do
  let tree ← toTree stx
  let r    ← analyze tree expectedType?
  trace[Elab.binop] "hasUncomparable: {r.hasUncomparable}, maxType: {r.max?}"
  if r.hasUncomparable || r.max?.isNone then
    let result ← toExpr tree
    ensureHasType expectedType? result
  else
    let result ← toExpr (← applyCoe tree r.max?.get!)
    trace[Elab.binop] "result: {result}"
    ensureHasType expectedType? result

@[builtinTermElab binop_lazy]
def elabBinOpLazy : TermElab := elabBinOp

/--
  Decompose `e` into `(r, a, b)`.

  Remark: it assumes the last two arguments are explicit. -/
private def relation? (e : Expr) : MetaM (Option (Expr × Expr × Expr)) :=
  if e.getAppNumArgs < 2 then
    return none
  else
    return some (e.appFn!.appFn!, e.appFn!.appArg!, e.appArg!)

/-- Step-wise reasoning over transitive relations.
```
calc
  a = b := pab
  b = c := pbc
  ...
  y = z := pyz
```
proves `a = z` from the given step-wise proofs. `=` can be replaced with any
relation implementing the typeclass `Trans`. Instead of repeating the right-
hand sides, subsequent left-hand sides can be replaced with `_`. -/
@[builtinTermElab «calc»]
def elabBinCalc : TermElab :=  fun stx expectedType? => do
  let stepStxs := stx[1].getArgs
  let mut proofs := #[]
  let mut types  := #[]
  for stepStx in stepStxs do
    let type  ← elabType stepStx[0]
    let some (_, lhs, _) ← relation? type |
      throwErrorAt stepStx[0] "invalid 'calc' step, relation expected{indentExpr type}"
    if types.size > 0 then
      let some (_, _, prevRhs) ← relation? types.back | unreachable!
      unless (← isDefEqGuarded lhs prevRhs) do
        throwErrorAt stepStx[0] "invalid 'calc' step, left-hand-side is {indentD m!"{lhs} : {← inferType lhs}"}\nprevious right-hand-side is{indentD m!"{prevRhs} : {← inferType prevRhs}"}"
    types := types.push type
    let proof ← elabTermEnsuringType stepStx[2] type
    synthesizeSyntheticMVars
    proofs := proofs.push proof
  let mut result := proofs[0]
  let mut resultType := types[0]
  for i in [1:proofs.size] do
    let some (r, a, b) ← relation? resultType | unreachable!
    let some (s, _, c) ← relation? (← instantiateMVars types[i]) | unreachable!
    let (α, β, γ)       := (← inferType a, ← inferType b, ← inferType c)
    let (u_1, u_2, u_3) := (← getLevel α, ← getLevel β, ← getLevel γ)
    let t ← mkFreshExprMVar (← mkArrow α (← mkArrow γ (mkSort levelZero)))
    let selfType := mkAppN (Lean.mkConst ``Trans [u_1, u_2, u_3]) #[α, β, γ, r, s, t]
    match (← trySynthInstance selfType) with
    | LOption.some self =>
      result := mkAppN (Lean.mkConst ``Trans.trans [u_1, u_2, u_3]) #[α, β, γ, r, s, t, self, a, b, c, result, proofs[i]]
      resultType := (← instantiateMVars (← inferType result)).headBeta
      unless (← relation? resultType).isSome do
        throwErrorAt stepStxs[i] "invalid 'calc' step, step result is not a relation{indentExpr resultType}"
    | _ => throwErrorAt stepStxs[i] "invalid 'calc' step, failed to synthesize `Trans` instance{indentExpr selfType}"
    pure ()
  ensureHasType expectedType? result

@[builtinTermElab defaultOrOfNonempty]
def elabDefaultOrNonempty : TermElab :=  fun stx expectedType? => do
  tryPostponeIfNoneOrMVar expectedType?
  match expectedType? with
  | none => throwError "invalid 'default_or_ofNonempty%', expected type is not known"
  | some expectedType =>
    try
      mkDefault expectedType
    catch ex => try
      mkOfNonempty expectedType
    catch _ =>
      throw ex

builtin_initialize
  registerTraceClass `Elab.binop

end BinOp

end Lean.Elab.Term


/-
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) 2018 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Mario Carneiro, Minchao Wu
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.tactic.core
import Mathlib.PostPort

universes l

namespace Mathlib

/-!
# `#explode` command

Displays a proof term in a line by line format somewhat akin to a Fitch style
proof or the Metamath proof style.
-/

namespace tactic

namespace explode

inductive status where
| reg : status
| intro : status
| lam : status
| sintro : status

/--
A type to distinguish introduction or elimination rules represented as
strings from theorems referred to by their names.
-/
/--
Turn a thm into a string.
-/
end explode

/--
`#explode decl_name` displays a proof term in a line-by-line format somewhat akin to a Fitch-style
proof or the Metamath proof style.
`#explode_widget decl_name` renders a widget that displays an `#explode` proof.

`#explode iff_true_intro` produces

```lean
iff_true_intro : ∀ {a : Prop}, a → (a ↔ true)
0│   │ a         ├ Prop
1│   │ h         ├ a
2│   │ hl        │ ┌ a
3│   │ trivial   │ │ true
4│2,3│ ∀I        │ a → true
5│   │ hr        │ ┌ true
6│5,1│ ∀I        │ true → a
7│4,6│ iff.intro │ a ↔ true
8│1,7│ ∀I        │ a → (a ↔ true)
9│0,8│ ∀I        │ ∀ {a : Prop}, a → (a ↔ true)
```

In more detail:

The output of `#explode` is a Fitch-style proof in a four-column diagram modeled after Metamath
proof displays like [this](http://us.metamath.org/mpeuni/ru.html). The headers of the columns are
"Step", "Hyp", "Ref", "Type" (or "Expression" in the case of Metamath):
* Step: An increasing sequence of numbers to number each step in the proof, used in the Hyp field.
* Hyp: The direct children of the current step. Most theorems are implications like `A -> B -> C`,
  and so on the step proving `C` the Hyp field will refer to the steps that prove `A` and `B`.
* Ref: The name of the theorem being applied. This is well-defined in Metamath, but in Lean there
  are some special steps that may have long names because the structure of proof terms doesn't
  exactly match this mold.
  * If the theorem is `foo (x y : Z) : A x -> B y -> C x y`:
    * the Ref field will contain `foo`,
    * `x` and `y` will be suppressed, because term construction is not interesting, and
    * the Hyp field will reference steps proving `A x` and `B y`. This corresponds to a proof term
      like `@foo x y pA pB` where `pA` and `pB` are subproofs.
  * If the head of the proof term is a local constant or lambda, then in this case the Ref will
    say `∀E` for forall-elimination. This happens when you have for example `h : A -> B` and
    `ha : A` and prove `b` by `h ha`; we reinterpret this as if it said `∀E h ha` where `∀E` is
    (n-ary) modus ponens.
  * If the proof term is a lambda, we will also use `∀I` for forall-introduction, referencing the
    body of the lambda. The indentation level will increase, and a bracket will surround the proof
    of the body of the lambda, starting at a proof step labeled with the name of the lambda variable
    and its type, and ending with the `∀I` step. Metamath doesn't have steps like this, but the
    style is based on Fitch proofs in first-order logic.
* Type: This contains the type of the proof term, the theorem being proven at the current step.
  This proof layout differs from `#print` in using lots of intermediate step displays so that you
  can follow along and don't have to see term construction steps because they are implicitly in the
  intermediate step displays.

Also, it is common for a Lean theorem to begin with a sequence of lambdas introducing local
constants of the theorem. In order to minimize the indentation level, the `∀I` steps at the end of
the proof will be introduced in a group and the indentation will stay fixed. (The indentation
brackets are only needed in order to delimit the scope of assumptions, and these assumptions
have global scope anyway so detailed tracking is not necessary.)
-/
end Mathlib


/-
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) 2023 Heather Macbeth. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Heather Macbeth
-/
import Mathlib.Tactic.LabelAttr

/-! # The @[mono] attribute -/

namespace Mathlib.Tactic.Monotonicity

syntax mono.side := &"left" <|> &"right" <|> &"both"

namespace Attr

/-- A lemma stating the monotonicity of some function, with respect to appropriate relations on its
domain and range, and possibly with side conditions. -/
syntax (name := mono) "mono" (ppSpace mono.side)? : attr

-- The following is inlined from `register_label_attr`.
/- TODO: currently `left`/`right`/`both` is ignored, and e.g. `@[mono left]` means the same as
`@[mono]`. No error is thrown by e.g. `@[mono left]`. -/
-- TODO: possibly extend `register_label_attr` to handle trailing syntax

open LabelAttr in
@[inherit_doc mono]
initialize ext : LabelExtension ← (
  let descr := "A lemma stating the monotonicity of some function, with respect to appropriate
relations on its domain and range, and possibly with side conditions."
  let mono := `mono
  registerLabelAttr mono descr mono)


/-
Copyright 2022 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Authors: Moritz Firsching
-/
import tactic
/-!
# Pigeon-hole and double counting

## TODO
  - statement
  - 1. Numbers
    - Claim
    - Claim
  - 2. Sequences
    - Claim
      - proof
  - 3. Sums
    - Claim
  - Double Counting
  - 4. Numbers again
  - 5. Graphs
    - Theorem
      - proof
    - Claim
  - 6. Sperner's Lemma
    - proof
    - Proof of Brouwer's fixed point theorem (for $n = 2$)
-/


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Leonardo de Moura
-/
prelude
import Init.Control.Except
import Init.Data.ByteArray
import Init.SimpLemmas
import Init.Data.Nat.Linear
import Init.Util
import Init.WFTactics

namespace String

/-- Interpret the string as the decimal representation of a natural number.

Panics if the string is not a string of digits. -/
def toNat! (s : String) : Nat :=
  if s.isNat then
    s.foldl (fun n c => n*10 + (c.toNat - '0'.toNat)) 0
  else
    panic! "Nat expected"

/--
  Convert a [UTF-8](https://en.wikipedia.org/wiki/UTF-8) encoded `ByteArray` string to `String`.
  The result is unspecified if `a` is not properly UTF-8 encoded.
-/
@[extern "lean_string_from_utf8_unchecked"]
opaque fromUTF8Unchecked (a : @& ByteArray) : String

/-- Convert the given `String` to a [UTF-8](https://en.wikipedia.org/wiki/UTF-8) encoded byte array. -/
@[extern "lean_string_to_utf8"]
opaque toUTF8 (a : @& String) : ByteArray

theorem one_le_csize (c : Char) : 1 ≤ csize c := by
  simp [csize, Char.utf8Size]
  repeat (first | split | decide)

@[simp] theorem pos_lt_eq (p₁ p₂ : Pos) : (p₁ < p₂) = (p₁.1 < p₂.1) := rfl

@[simp] theorem pos_add_char (p : Pos) (c : Char) : (p + c).byteIdx = p.byteIdx + csize c := rfl

theorem eq_empty_of_bsize_eq_zero (h : s.endPos = {}) : s = "" := by
  match s with
  | ⟨[]⟩   => rfl
  | ⟨c::cs⟩ =>
    injection h with h
    simp [endPos, utf8ByteSize, utf8ByteSize.go] at h
    have : utf8ByteSize.go cs + 1 ≤ utf8ByteSize.go cs + csize c := Nat.add_le_add_left (one_le_csize c) _
    simp_arith [h] at this

theorem lt_next (s : String) (i : String.Pos) : i.1 < (s.next i).1 := by
  simp_arith [next]; apply one_le_csize

theorem Iterator.sizeOf_next_lt_of_hasNext (i : String.Iterator) (h : i.hasNext) : sizeOf i.next < sizeOf i := by
  cases i; rename_i s pos; simp [Iterator.next, Iterator.sizeOf_eq]; simp [Iterator.hasNext] at h
  have := String.lt_next s pos
  apply Nat.sub.elim (motive := fun k => k < _) (utf8ByteSize s) (String.next s pos).1
  . intro _ k he
    simp [he]; rw [Nat.add_comm, Nat.add_sub_assoc (Nat.le_of_lt this)]
    have := Nat.zero_lt_sub_of_lt this
    simp_all_arith
  . intro; apply Nat.zero_lt_sub_of_lt h

macro_rules | `(tactic| decreasing_trivial) => `(tactic| apply String.Iterator.sizeOf_next_lt_of_hasNext; assumption)

theorem Iterator.sizeOf_next_lt_of_atEnd (i : String.Iterator) (h : ¬ i.atEnd = true) : sizeOf i.next < sizeOf i :=
  have h : i.hasNext = true := by simp_arith [atEnd] at h; simp_arith [hasNext, h]
  sizeOf_next_lt_of_hasNext i h

macro_rules | `(tactic| decreasing_trivial) => `(tactic| apply String.Iterator.sizeOf_next_lt_of_atEnd; assumption)

namespace Iterator

/-- Advance the given iterator until the predicate returns true or the end of the string is reached. -/
@[specialize] def find (it : Iterator) (p : Char → Bool) : Iterator :=
  if it.atEnd then it
  else if p it.curr then it
  else find it.next p

@[specialize] def foldUntil (it : Iterator) (init : α) (f : α → Char → Option α) : α × Iterator :=
  if it.atEnd then
    (init, it)
  else if let some a := f init it.curr then
    foldUntil it.next a f
  else
    (init, it)

end Iterator

end String


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


example : 2 + 3 = 5 :=
begin
  generalize : 3 = x,
  sorry
end


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

import Aesop

set_option aesop.check.all true

-- These test cases test the builtin subst tactic.

example (h₁ : x = 5) (h₂ : y = 5) : x = y := by
  fail_if_success
    aesop
      (erase Aesop.BuiltinRules.subst)
      (simp_options := { useHyps := false })
      (options := { terminal := true })
  aesop (simp_options := { useHyps := false })

example (h₁ : x = y) (h₂ : y = z) : x = z := by
  fail_if_success
    aesop
      (erase Aesop.BuiltinRules.subst)
      (simp_options := { useHyps := false })
      (options := { terminal := true })
  aesop (simp_options := { useHyps := false })

example (P : ∀ x y, x = y → Prop) (h₁ : x = y) (h₂ : P x y h₁) : x = y := by
  fail_if_success
    aesop
      (erase Aesop.BuiltinRules.subst,
             Aesop.BuiltinRules.assumption,
             Aesop.BuiltinRules.applyHyps)
      (simp_options := { useHyps := false })
      (options := { terminal := true })
  aesop
    (erase Aesop.BuiltinRules.assumption,
           Aesop.BuiltinRules.applyHyps)
    (simp_options := { useHyps := false })

-- Subst also works for bi-implications.
example (h₁ : P ↔ Q) (h₂ : Q ↔ R) (h₃ : P) : R  := by
  fail_if_success
    aesop
      (erase Aesop.BuiltinRules.subst)
      (simp_options := { useHyps := false })
      (options := { terminal := true })
  aesop (simp_options := { useHyps := false })

-- Subst also works for morally-homogeneous heterogeneous equalities (using a
-- builtin simp rule which turns these into actual homogeneous equalities).
example {P : α → Prop} {x y z : α} (h₁ : HEq x y) (h₂ : HEq y z) (h₃ : P x) :
    P z  := by
  fail_if_success
    aesop
      (erase Aesop.BuiltinRules.subst)
      (simp_options := { useHyps := false })
      (options := { terminal := true })
  aesop (simp_options := { useHyps := false })


/-
Copyright (c) 2018 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import Init.Data.Nat.Basic
import Init.Data.Fin.Basic
import Init.Data.UInt
import Init.Data.Repr
import Init.Data.ToString.Basic
import Init.Util
universes u v w

namespace Array
variable {α : Type u}

@[extern "lean_mk_array"]
def mkArray {α : Type u} (n : Nat) (v : α) : Array α := {
  data := List.replicate n v
}

@[simp] theorem size_mkArray (n : Nat) (v : α) : (mkArray n v).size = n :=
  List.length_replicate ..

instance : EmptyCollection (Array α) := ⟨Array.empty⟩
instance : Inhabited (Array α) where
  default := Array.empty

def isEmpty (a : Array α) : Bool :=
  a.size = 0

def singleton (v : α) : Array α :=
  mkArray 1 v

/- Low-level version of `fget` which is as fast as a C array read.
   `Fin` values are represented as tag pointers in the Lean runtime. Thus,
   `fget` may be slightly slower than `uget`. -/
@[extern "lean_array_uget"]
def uget (a : @& Array α) (i : USize) (h : i.toNat < a.size) : α :=
  a.get ⟨i.toNat, h⟩

def back [Inhabited α] (a : Array α) : α :=
  a.get! (a.size - 1)

def get? (a : Array α) (i : Nat) : Option α :=
  if h : i < a.size then some (a.get ⟨i, h⟩) else none

def back? (a : Array α) : Option α :=
  a.get? (a.size - 1)

-- auxiliary declaration used in the equation compiler when pattern matching array literals.
abbrev getLit {α : Type u} {n : Nat} (a : Array α) (i : Nat) (h₁ : a.size = n) (h₂ : i < n) : α :=
  a.get ⟨i, h₁.symm ▸ h₂⟩

@[simp] theorem size_set (a : Array α) (i : Fin a.size) (v : α) : (set a i v).size = a.size :=
  List.length_set ..

@[simp] theorem size_push (a : Array α) (v : α) : (push a v).size = a.size + 1 :=
  List.length_concat ..

/- Low-level version of `fset` which is as fast as a C array fset.
   `Fin` values are represented as tag pointers in the Lean runtime. Thus,
   `fset` may be slightly slower than `uset`. -/
@[extern "lean_array_uset"]
def uset (a : Array α) (i : USize) (v : α) (h : i.toNat < a.size) : Array α :=
  a.set ⟨i.toNat, h⟩ v

@[extern "lean_array_fswap"]
def swap (a : Array α) (i j : @& Fin a.size) : Array α :=
  let v₁ := a.get i
  let v₂ := a.get j
  let a'  := a.set i v₂
  a'.set (size_set a i v₂ ▸ j) v₁

@[extern "lean_array_swap"]
def swap! (a : Array α) (i j : @& Nat) : Array α :=
  if h₁ : i < a.size then
  if h₂ : j < a.size then swap a ⟨i, h₁⟩ ⟨j, h₂⟩
  else panic! "index out of bounds"
  else panic! "index out of bounds"

@[inline] def swapAt (a : Array α) (i : Fin a.size) (v : α) : α × Array α :=
  let e := a.get i
  let a := a.set i v
  (e, a)

@[inline]
def swapAt! (a : Array α) (i : Nat) (v : α) : α × Array α :=
  if h : i < a.size then
    swapAt a ⟨i, h⟩ v
  else
    have : Inhabited α := ⟨v⟩
    panic! ("index " ++ toString i ++ " out of bounds")

@[extern "lean_array_pop"]
def pop (a : Array α) : Array α := {
  data := a.data.dropLast
}

def shrink (a : Array α) (n : Nat) : Array α :=
  let rec loop
    | 0,   a => a
    | n+1, a => loop n a.pop
  loop (a.size - n) a

@[inline]
def modifyM [Monad m] [Inhabited α] (a : Array α) (i : Nat) (f : α → m α) : m (Array α) := do
  if h : i < a.size then
    let idx : Fin a.size := ⟨i, h⟩
    let v                := a.get idx
    let a'               := a.set idx arbitrary
    let v ← f v
    pure <| a'.set (size_set a .. ▸ idx) v
  else
    pure a

@[inline]
def modify [Inhabited α] (a : Array α) (i : Nat) (f : α → α) : Array α :=
  Id.run <| a.modifyM i f

@[inline]
def modifyOp [Inhabited α] (self : Array α) (idx : Nat) (f : α → α) : Array α :=
  self.modify idx f

/-
  We claim this unsafe implementation is correct because an array cannot have more than `usizeSz` elements in our runtime.

  This kind of low level trick can be removed with a little bit of compiler support. For example, if the compiler simplifies `as.size < usizeSz` to true. -/
@[inline] unsafe def forInUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (b : β) (f : α → β → m (ForInStep β)) : m β :=
  let sz := USize.ofNat as.size
  let rec @[specialize] loop (i : USize) (b : β) : m β := do
    if i < sz then
      let a := as.uget i lcProof
      match (← f a b) with
      | ForInStep.done  b => pure b
      | ForInStep.yield b => loop (i+1) b
    else
      pure b
  loop 0 b

-- Move?
private theorem zeroLtOfLt : {a b : Nat} → a < b → 0 < b
  | 0,   _, h => h
  | a+1, b, h =>
    have : a < b := Nat.ltTrans (Nat.ltSuccSelf _) h
    zeroLtOfLt this

/- Reference implementation for `forIn` -/
@[implementedBy Array.forInUnsafe]
protected def forIn {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (b : β) (f : α → β → m (ForInStep β)) : m β :=
  let rec loop (i : Nat) (h : i ≤ as.size) (b : β) : m β := do
    match i, h with
    | 0,   _ => pure b
    | i+1, h =>
      have h' : i < as.size            := Nat.ltOfLtOfLe (Nat.ltSuccSelf i) h
      have : as.size - 1 < as.size     := Nat.subLt (zeroLtOfLt h') (by decide)
      have : as.size - 1 - i < as.size := Nat.ltOfLeOfLt (Nat.subLe (as.size - 1) i) this
      match (← f (as.get ⟨as.size - 1 - i, this⟩) b) with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop i (Nat.leOfLt h') b
  loop as.size (Nat.leRefl _) b

instance : ForIn m (Array α) α where
  forIn := Array.forIn

/- See comment at forInUnsafe -/
@[inline]
unsafe def foldlMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : β → α → m β) (init : β) (as : Array α) (start := 0) (stop := as.size) : m β :=
  let rec @[specialize] fold (i : USize) (stop : USize) (b : β) : m β := do
    if i == stop then
      pure b
    else
      fold (i+1) stop (← f b (as.uget i lcProof))
  if start < stop then
    if stop ≤ as.size then
      fold (USize.ofNat start) (USize.ofNat stop) init
    else
      pure init
  else
    pure init

/- Reference implementation for `foldlM` -/
@[implementedBy foldlMUnsafe]
def foldlM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : β → α → m β) (init : β) (as : Array α) (start := 0) (stop := as.size) : m β :=
  let fold (stop : Nat) (h : stop ≤ as.size) :=
    let rec loop (i : Nat) (j : Nat) (b : β) : m β := do
      if hlt : j < stop then
        match i with
        | 0    => pure b
        | i'+1 =>
          loop i' (j+1) (← f b (as.get ⟨j, Nat.ltOfLtOfLe hlt h⟩))
      else
        pure b
    loop (stop - start) start init
  if h : stop ≤ as.size then
    fold stop h
  else
    fold as.size (Nat.leRefl _)

/- See comment at forInUnsafe -/
@[inline]
unsafe def foldrMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → β → m β) (init : β) (as : Array α) (start := as.size) (stop := 0) : m β :=
  let rec @[specialize] fold (i : USize) (stop : USize) (b : β) : m β := do
    if i == stop then
      pure b
    else
      fold (i-1) stop (← f (as.uget (i-1) lcProof) b)
  if start ≤ as.size then
    if stop < start then
      fold (USize.ofNat start) (USize.ofNat stop) init
    else
      pure init
  else if stop < as.size then
    fold (USize.ofNat as.size) (USize.ofNat stop) init
  else
    pure init

/- Reference implementation for `foldrM` -/
@[implementedBy foldrMUnsafe]
def foldrM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → β → m β) (init : β) (as : Array α) (start := as.size) (stop := 0) : m β :=
  let rec fold (i : Nat) (h : i ≤ as.size) (b : β) : m β := do
    if i == stop then
      pure b
    else match i, h with
      | 0, _   => pure b
      | i+1, h =>
        have : i < as.size := Nat.ltOfLtOfLe (Nat.ltSuccSelf _) h
        fold i (Nat.leOfLt this) (← f (as.get ⟨i, this⟩) b)
  if h : start ≤ as.size then
    if stop < start then
      fold start h init
    else
      pure init
  else if stop < as.size then
    fold as.size (Nat.leRefl _) init
  else
    pure init

/- See comment at forInUnsafe -/
@[inline]
unsafe def mapMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m β) (as : Array α) : m (Array β) :=
  let sz := USize.ofNat as.size
  let rec @[specialize] map (i : USize) (r : Array NonScalar) : m (Array PNonScalar.{v}) := do
    if i < sz then
     let v    := r.uget i lcProof
     let r    := r.uset i arbitrary lcProof
     let vNew ← f (unsafeCast v)
     map (i+1) (r.uset i (unsafeCast vNew) lcProof)
    else
     pure (unsafeCast r)
  unsafeCast <| map 0 (unsafeCast as)

/- Reference implementation for `mapM` -/
@[implementedBy mapMUnsafe]
def mapM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m β) (as : Array α) : m (Array β) :=
  as.foldlM (fun bs a => do let b ← f a; pure (bs.push b)) (mkEmpty as.size)

@[inline]
def mapIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : Fin as.size → α → m β) : m (Array β) :=
  let rec @[specialize] map (i : Nat) (j : Nat) (inv : i + j = as.size) (bs : Array β) : m (Array β) := do
    match i, inv with
    | 0,    _  => pure bs
    | i+1, inv =>
      have : j < as.size := by rw [← inv, Nat.add_assoc, Nat.add_comm 1 j, Nat.add_left_comm]; apply Nat.leAddRight
      let idx : Fin as.size := ⟨j, this⟩
      have : i + (j + 1) = as.size := by rw [← inv, Nat.add_comm j 1, Nat.add_assoc]
      map i (j+1) this (bs.push (← f idx (as.get idx)))
  map as.size 0 rfl (mkEmpty as.size)

@[inline]
def findSomeM? {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : α → m (Option β)) : m (Option β) := do
  for a in as do
    match (← f a) with
    | some b => return b
    | _      => pure ⟨⟩
  return none

@[inline]
def findM? {α : Type} {m : Type → Type} [Monad m] (as : Array α) (p : α → m Bool) : m (Option α) := do
  for a in as do
    if (← p a) then
      return a
  return none

@[inline]
def findIdxM? [Monad m] (as : Array α) (p : α → m Bool) : m (Option Nat) := do
  let mut i := 0
  for a in as do
    if (← p a) then
      return some i
    i := i + 1
  return none

@[inline]
unsafe def anyMUnsafe {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m Bool :=
  let rec @[specialize] any (i : USize) (stop : USize) : m Bool := do
    if i == stop then
      pure false
    else
      if (← p (as.uget i lcProof)) then
        pure true
      else
        any (i+1) stop
  if start < stop then
    if stop ≤ as.size then
      any (USize.ofNat start) (USize.ofNat stop)
    else
      pure false
  else
    pure false

@[implementedBy anyMUnsafe]
def anyM {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m Bool :=
  let any (stop : Nat) (h : stop ≤ as.size) :=
    let rec loop (i : Nat) (j : Nat) : m Bool := do
      if hlt : j < stop then
        match i with
        | 0    => pure false
        | i'+1 =>
          if (← p (as.get ⟨j, Nat.ltOfLtOfLe hlt h⟩)) then
            pure true
          else
            loop i' (j+1)
      else
        pure false
    loop (stop - start) start
  if h : stop ≤ as.size then
    any stop h
  else
    any as.size (Nat.leRefl _)

@[inline]
def allM {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m Bool :=
  return !(← as.anyM fun v => return !(← p v))

@[inline]
def findSomeRevM? {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : α → m (Option β)) : m (Option β) :=
  let rec @[specialize] find : (i : Nat) → i ≤ as.size → m (Option β)
    | 0,   h => pure none
    | i+1, h => do
      have : i < as.size := Nat.ltOfLtOfLe (Nat.ltSuccSelf _) h
      let r ← f (as.get ⟨i, this⟩)
      match r with
      | some v => pure r
      | none   =>
        have : i ≤ as.size := Nat.leOfLt this
        find i this
  find as.size (Nat.leRefl _)

@[inline]
def findRevM? {α : Type} {m : Type → Type w} [Monad m] (as : Array α) (p : α → m Bool) : m (Option α) :=
  as.findSomeRevM? fun a => return if (← p a) then some a else none

@[inline]
def forM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Array α) (start := 0) (stop := as.size) : m PUnit :=
  as.foldlM (fun _ => f) ⟨⟩ start stop

@[inline]
def forRevM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Array α) (start := as.size) (stop := 0) : m PUnit :=
  as.foldrM (fun a _ => f a) ⟨⟩ start stop

@[inline]
def foldl {α : Type u} {β : Type v} (f : β → α → β) (init : β) (as : Array α) (start := 0) (stop := as.size) : β :=
  Id.run <| as.foldlM f init start stop

@[inline]
def foldr {α : Type u} {β : Type v} (f : α → β → β) (init : β) (as : Array α) (start := as.size) (stop := 0) : β :=
  Id.run <| as.foldrM f init start stop

@[inline]
def map {α : Type u} {β : Type v} (f : α → β) (as : Array α) : Array β :=
  Id.run <| as.mapM f

@[inline]
def mapIdx {α : Type u} {β : Type v} (as : Array α) (f : Fin as.size → α → β) : Array β :=
  Id.run <| as.mapIdxM f

@[inline]
def find? {α : Type} (as : Array α) (p : α → Bool) : Option α :=
  Id.run <| as.findM? p

@[inline]
def findSome? {α : Type u} {β : Type v} (as : Array α) (f : α → Option β) : Option β :=
  Id.run <| as.findSomeM? f

@[inline]
def findSome! {α : Type u} {β : Type v} [Inhabited β] (a : Array α) (f : α → Option β) : β :=
  match findSome? a f with
  | some b => b
  | none   => panic! "failed to find element"

@[inline]
def findSomeRev? {α : Type u} {β : Type v} (as : Array α) (f : α → Option β) : Option β :=
  Id.run <| as.findSomeRevM? f

@[inline]
def findRev? {α : Type} (as : Array α) (p : α → Bool) : Option α :=
  Id.run <| as.findRevM? p

@[inline]
def findIdx? {α : Type u} (as : Array α) (p : α → Bool) : Option Nat :=
  let rec loop (i : Nat) (j : Nat) (inv : i + j = as.size) : Option Nat :=
    if hlt : j < as.size then
      match i, inv with
      | 0, inv => by
        apply False.elim
        rw [Nat.zero_add] at inv
        rw [inv] at hlt
        exact absurd hlt (Nat.ltIrrefl _)
      | i+1, inv =>
        if p (as.get ⟨j, hlt⟩) then
          some j
        else
          have : i + (j+1) = as.size := by
            rw [← inv, Nat.add_comm j 1, Nat.add_assoc]
          loop i (j+1) this
    else
      none
  loop as.size 0 rfl

def getIdx? [BEq α] (a : Array α) (v : α) : Option Nat :=
a.findIdx? fun a => a == v

@[inline]
def any (as : Array α) (p : α → Bool) (start := 0) (stop := as.size) : Bool :=
  Id.run <| as.anyM p start stop

@[inline]
def all (as : Array α) (p : α → Bool) (start := 0) (stop := as.size) : Bool :=
  Id.run <| as.allM p start stop

def contains [BEq α] (as : Array α) (a : α) : Bool :=
  as.any fun b => a == b

def elem [BEq α] (a : α) (as : Array α) : Bool :=
  as.contains a

-- TODO(Leo): justify termination using wf-rec, and use `swap`
partial def reverse (as : Array α) : Array α :=
  let n   := as.size
  let mid := n / 2
  let rec rev (as : Array α) (i : Nat) :=
    if i < mid then
      rev (as.swap! i (n - i - 1)) (i+1)
    else
      as
  rev as 0

@[inline] def getEvenElems (as : Array α) : Array α :=
  (·.2) <| as.foldl (init := (true, Array.empty)) fun (even, r) a =>
    if even then
      (false, r.push a)
    else
      (true, r)

@[export lean_array_to_list]
def toList (as : Array α) : List α :=
  as.foldr List.cons []

instance {α : Type u} [Repr α] : Repr (Array α) where
  reprPrec a _ :=
    if a.size == 0 then
      "#[]"
    else
      Std.Format.bracketFill "#[" (@Std.Format.joinSep _ ⟨repr⟩ (toList a) ("," ++ Std.Format.line)) "]"

instance [ToString α] : ToString (Array α) where
  toString a := "#" ++ toString a.toList

protected def append (as : Array α) (bs : Array α) : Array α :=
  bs.foldl (init := as) fun r v => r.push v

instance : Append (Array α) := ⟨Array.append⟩

protected def appendList (as : Array α) (bs : List α) : Array α :=
  bs.foldl (init := as) fun r v => r.push v

instance : HAppend (Array α) (List α) (Array α) := ⟨Array.appendList⟩

@[inline]
def concatMapM [Monad m] (f : α → m (Array β)) (as : Array α) : m (Array β) :=
  as.foldlM (init := empty) fun bs a => do return bs ++ (← f a)

@[inline]
def concatMap (f : α → Array β) (as : Array α) : Array β :=
  as.foldl (init := empty) fun bs a => bs ++ f a

end Array

export Array (mkArray)

syntax "#[" sepBy(term, ", ") "]" : term

macro_rules
  | `(#[ $elems,* ]) => `(List.toArray [ $elems,* ])

namespace Array

-- TODO(Leo): cleanup
@[specialize]
partial def isEqvAux (a b : Array α) (hsz : a.size = b.size) (p : α → α → Bool) (i : Nat) : Bool :=
  if h : i < a.size then
     let aidx : Fin a.size := ⟨i, h⟩;
     let bidx : Fin b.size := ⟨i, hsz ▸ h⟩;
     match p (a.get aidx) (b.get bidx) with
     | true  => isEqvAux a b hsz p (i+1)
     | false => false
  else
    true

@[inline] def isEqv (a b : Array α) (p : α → α → Bool) : Bool :=
  if h : a.size = b.size then
    isEqvAux a b h p 0
  else
    false

instance [BEq α] : BEq (Array α) :=
  ⟨fun a b => isEqv a b BEq.beq⟩

@[inline]
def filter (p : α → Bool) (as : Array α) (start := 0) (stop := as.size) : Array α :=
  as.foldl (init := #[]) (start := start) (stop := stop) fun r a =>
    if p a then r.push a else r

@[inline]
def filterM [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m (Array α) :=
  as.foldlM (init := #[]) (start := start) (stop := stop) fun r a => do
    if (← p a) then r.push a else r

@[specialize]
def filterMapM [Monad m] (f : α → m (Option β)) (as : Array α) (start := 0) (stop := as.size) : m (Array β) :=
  as.foldlM (init := #[]) (start := start) (stop := stop) fun bs a => do
    match (← f a) with
    | some b => pure (bs.push b)
    | none   => pure bs

@[inline]
def filterMap (f : α → Option β) (as : Array α) (start := 0) (stop := as.size) : Array β :=
  Id.run <| as.filterMapM f (start := start) (stop := stop)

@[specialize]
def getMax? (as : Array α) (lt : α → α → Bool) : Option α :=
  if h : 0 < as.size then
    let a0 := as.get ⟨0, h⟩
    some <| as.foldl (init := a0) (start := 1) fun best a =>
      if lt best a then a else best
  else
    none

@[inline]
def partition (p : α → Bool) (as : Array α) : Array α × Array α := do
  let mut bs := #[]
  let mut cs := #[]
  for a in as do
    if p a then
      bs := bs.push a
    else
      cs := cs.push a
  return (bs, cs)

theorem ext (a b : Array α)
    (h₁ : a.size = b.size)
    (h₂ : (i : Nat) → (hi₁ : i < a.size) → (hi₂ : i < b.size) → a.get ⟨i, hi₁⟩ = b.get ⟨i, hi₂⟩)
    : a = b := by
  let rec extAux (a b : List α)
      (h₁ : a.length = b.length)
      (h₂ : (i : Nat) → (hi₁ : i < a.length) → (hi₂ : i < b.length) → a.get i hi₁ = b.get i hi₂)
      : a = b := by
    induction a generalizing b with
    | nil =>
      cases b with
      | nil       => rfl
      | cons b bs => rw [List.length_cons] at h₁; injection h₁
    | cons a as ih =>
      cases b with
      | nil => rw [List.length_cons] at h₁; injection h₁
      | cons b bs =>
        have hz₁ : 0 < (a::as).length := by rw [List.length_cons]; apply Nat.zeroLtSucc
        have hz₂ : 0 < (b::bs).length := by rw [List.length_cons]; apply Nat.zeroLtSucc
        have headEq : a = b := h₂ 0 hz₁ hz₂
        have h₁' : as.length = bs.length := by rw [List.length_cons, List.length_cons] at h₁; injection h₁; assumption
        have h₂' : (i : Nat) → (hi₁ : i < as.length) → (hi₂ : i < bs.length) → as.get i hi₁ = bs.get i hi₂ := by
          intro i hi₁ hi₂
          have hi₁' : i+1 < (a::as).length := by rw [List.length_cons]; apply Nat.succ_lt_succ; assumption
          have hi₂' : i+1 < (b::bs).length := by rw [List.length_cons]; apply Nat.succ_lt_succ; assumption
          have : (a::as).get (i+1) hi₁' = (b::bs).get (i+1) hi₂' := h₂ (i+1) hi₁' hi₂'
          apply this
        have tailEq : as = bs := ih bs h₁' h₂'
        rw [headEq, tailEq]
  cases a; cases b
  apply congrArg
  apply extAux
  assumption
  assumption

theorem extLit {n : Nat}
    (a b : Array α)
    (hsz₁ : a.size = n) (hsz₂ : b.size = n)
    (h : (i : Nat) → (hi : i < n) → a.getLit i hsz₁ hi = b.getLit i hsz₂ hi) : a = b :=
  Array.ext a b (hsz₁.trans hsz₂.symm) fun i hi₁ hi₂ => h i (hsz₁ ▸ hi₁)

end Array

-- CLEANUP the following code
namespace Array

partial def indexOfAux [BEq α] (a : Array α) (v : α) : Nat → Option (Fin a.size)
  | i =>
    if h : i < a.size then
      let idx : Fin a.size := ⟨i, h⟩;
      if a.get idx == v then some idx
      else indexOfAux a v (i+1)
    else none

def indexOf? [BEq α] (a : Array α) (v : α) : Option (Fin a.size) :=
  indexOfAux a v 0

partial def eraseIdxAux : Nat → Array α → Array α
  | i, a =>
    if h : i < a.size then
      let idx  : Fin a.size := ⟨i, h⟩;
      let idx1 : Fin a.size := ⟨i - 1, by exact Nat.ltOfLeOfLt (Nat.predLe i) h⟩;
      eraseIdxAux (i+1) (a.swap idx idx1)
    else
      a.pop

def feraseIdx (a : Array α) (i : Fin a.size) : Array α :=
  eraseIdxAux (i.val + 1) a

def eraseIdx (a : Array α) (i : Nat) : Array α :=
  if i < a.size then eraseIdxAux (i+1) a else a

@[simp] theorem size_swap (a : Array α) (i j : Fin a.size) : (a.swap i j).size = a.size := by
  show ((a.set i (a.get j)).set (size_set a i _ ▸ j) (a.get i)).size = a.size
  rw [size_set, size_set]

@[simp] theorem size_pop (a : Array α) : a.pop.size = a.size - 1 :=
  List.length_dropLast ..

section
/- Instance for justifying `partial` declaration.
   We should be able to delete it as soon as we restore support for well-founded recursion. -/
instance eraseIdxSzAuxInstance (a : Array α) : Inhabited { r : Array α // r.size = a.size - 1 } where
  default := ⟨a.pop, size_pop a⟩

partial def eraseIdxSzAux (a : Array α) : ∀ (i : Nat) (r : Array α), r.size = a.size → { r : Array α // r.size = a.size - 1 }
  | i, r, heq =>
    if h : i < r.size then
      let idx  : Fin r.size := ⟨i, h⟩;
      let idx1 : Fin r.size := ⟨i - 1, by exact Nat.ltOfLeOfLt (Nat.predLe i) h⟩;
      eraseIdxSzAux a (i+1) (r.swap idx idx1) ((size_swap r idx idx1).trans heq)
    else
      ⟨r.pop, (size_pop r).trans (heq ▸ rfl)⟩
end

def eraseIdx' (a : Array α) (i : Fin a.size) : { r : Array α // r.size = a.size - 1 } :=
  eraseIdxSzAux a (i.val + 1) a rfl

def erase [BEq α] (as : Array α) (a : α) : Array α :=
  match as.indexOf? a with
  | none   => as
  | some i => as.feraseIdx i

partial def insertAtAux (i : Nat) : Array α → Nat → Array α
  | as, j =>
    if i == j then as
    else
      let as := as.swap! (j-1) j;
      insertAtAux i as (j-1)

/--
  Insert element `a` at position `i`.
  Pre: `i < as.size` -/
def insertAt (as : Array α) (i : Nat) (a : α) : Array α :=
  if i > as.size then panic! "invalid index"
  else
    let as := as.push a;
    as.insertAtAux i as.size

def toListLitAux (a : Array α) (n : Nat) (hsz : a.size = n) : ∀ (i : Nat), i ≤ a.size → List α → List α
  | 0,     hi, acc => acc
  | (i+1), hi, acc => toListLitAux a n hsz i (Nat.leOfSuccLe hi) (a.getLit i hsz (Nat.ltOfLtOfEq (Nat.ltOfLtOfLe (Nat.ltSuccSelf i) hi) hsz) :: acc)

def toArrayLit (a : Array α) (n : Nat) (hsz : a.size = n) : Array α :=
  List.toArray <| toListLitAux a n hsz n (hsz ▸ Nat.leRefl _) []

theorem toArrayLitEq (a : Array α) (n : Nat) (hsz : a.size = n) : a = toArrayLit a n hsz :=
  -- TODO: this is painful to prove without proper automation
  sorry
  /-
  First, we need to prove
  ∀ i j acc, i ≤ a.size → (toListLitAux a n hsz (i+1) hi acc).index j = if j < i then a.getLit j hsz _ else acc.index (j - i)
  by induction

  Base case is trivial
  (j : Nat) (acc : List α) (hi : 0 ≤ a.size)
       |- (toListLitAux a n hsz 0 hi acc).index j = if j < 0 then a.getLit j hsz _ else acc.index (j - 0)
  ...  |- acc.index j = acc.index j

  Induction

  (j : Nat) (acc : List α) (hi : i+1 ≤ a.size)
        |- (toListLitAux a n hsz (i+1) hi acc).index j = if j < i + 1 then a.getLit j hsz _ else acc.index (j - (i + 1))
    ... |- (toListLitAux a n hsz i hi' (a.getLit i hsz _ :: acc)).index j = if j < i + 1 then a.getLit j hsz _ else acc.index (j - (i + 1))  * by def
    ... |- if j < i     then a.getLit j hsz _ else (a.getLit i hsz _ :: acc).index (j-i)    * by induction hypothesis
           =
           if j < i + 1 then a.getLit j hsz _ else acc.index (j - (i + 1))
  If j < i, then both are a.getLit j hsz _
  If j = i, then lhs reduces else-branch to (a.getLit i hsz _) and rhs is then-brachn (a.getLit i hsz _)
  If j >= i + 1, we use
     - j - i >= 1 > 0
     - (a::as).index k = as.index (k-1) If k > 0
     - j - (i + 1) = (j - i) - 1
     Then lhs = (a.getLit i hsz _ :: acc).index (j-i) = acc.index (j-i-1) = acc.index (j-(i+1)) = rhs

  With this proof, we have

  ∀ j, j < n → (toListLitAux a n hsz n _ []).index j = a.getLit j hsz _

  We also need

  - (toListLitAux a n hsz n _ []).length = n
  - j < n -> (List.toArray as).getLit j _ _ = as.index j

  Then using Array.extLit, we have that a = List.toArray <| toListLitAux a n hsz n _ []
  -/

partial def isPrefixOfAux [BEq α] (as bs : Array α) (hle : as.size ≤ bs.size) : Nat → Bool
  | i =>
    if h : i < as.size then
      let a := as.get ⟨i, h⟩;
      let b := bs.get ⟨i, Nat.ltOfLtOfLe h hle⟩;
      if a == b then
        isPrefixOfAux as bs hle (i+1)
      else
        false
    else
      true

/- Return true iff `as` is a prefix of `bs` -/
def isPrefixOf [BEq α] (as bs : Array α) : Bool :=
  if h : as.size ≤ bs.size then
    isPrefixOfAux as bs h 0
  else
    false

private def allDiffAuxAux [BEq α] (as : Array α) (a : α) : forall (i : Nat), i < as.size → Bool
  | 0,   h => true
  | i+1, h =>
    have : i < as.size := Nat.ltTrans (Nat.ltSuccSelf _) h;
    a != as.get ⟨i, this⟩ && allDiffAuxAux as a i this

private partial def allDiffAux [BEq α] (as : Array α) : Nat → Bool
  | i =>
    if h : i < as.size then
      allDiffAuxAux as (as.get ⟨i, h⟩) i h && allDiffAux as (i+1)
    else
      true

def allDiff [BEq α] (as : Array α) : Bool :=
  allDiffAux as 0

@[specialize] partial def zipWithAux (f : α → β → γ) (as : Array α) (bs : Array β) : Nat → Array γ → Array γ
  | i, cs =>
    if h : i < as.size then
      let a := as.get ⟨i, h⟩;
      if h : i < bs.size then
        let b := bs.get ⟨i, h⟩;
        zipWithAux f as bs (i+1) <| cs.push <| f a b
      else
        cs
    else
      cs

@[inline] def zipWith (as : Array α) (bs : Array β) (f : α → β → γ) : Array γ :=
  zipWithAux f as bs 0 #[]

def zip (as : Array α) (bs : Array β) : Array (α × β) :=
  zipWith as bs Prod.mk

def unzip (as : Array (α × β)) : Array α × Array β :=
  as.foldl (init := (#[], #[])) fun (as, bs) (a, b) => (as.push a, bs.push b)

def split (as : Array α) (p : α → Bool) : Array α × Array α :=
  as.foldl (init := (#[], #[])) fun (as, bs) a =>
    if p a then (as.push a, bs) else (as, bs.push a)

end Array


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