/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.Check
import Lean.Meta.Closure
import Lean.Meta.Tactic.Cases
import Lean.Meta.Tactic.Contradiction
import Lean.Meta.GeneralizeTelescope
import Lean.Meta.Match.Basic

namespace Lean.Meta.Match

/-- The number of patterns in each AltLHS must be equal to the number of discriminants. -/
private def checkNumPatterns (numDiscrs : Nat) (lhss : List AltLHS) : MetaM Unit := do
  if lhss.any fun lhs => lhs.patterns.length != numDiscrs then
    throwError "incorrect number of patterns"

/--
  Execute `k hs` where `hs` contains new equalities `h : lhs[i] = rhs[i]` for each `discrInfos[i] = some h`.
  Assume `lhs.size == rhs.size == discrInfos.size`
-/
private partial def withEqs (lhs rhs : Array Expr) (discrInfos : Array DiscrInfo) (k : Array Expr → MetaM α) : MetaM α := do
  go 0 #[]
where
  go (i : Nat) (hs : Array Expr) : MetaM α := do
    if i < lhs.size then
      if let some hName := discrInfos[i]!.hName? then
        withLocalDeclD hName (← mkEqHEq lhs[i]! rhs[i]!) fun h =>
          go (i+1) (hs.push h)
      else
        go (i+1) hs
    else
      k hs

/-- Given a list of `AltLHS`, create a minor premise for each one, convert them into `Alt`, and then execute `k` -/
private def withAlts {α} (motive : Expr) (discrs : Array Expr) (discrInfos : Array DiscrInfo) (lhss : List AltLHS) (k : List Alt → Array (Expr × Nat) → MetaM α) : MetaM α :=
  loop lhss [] #[]
where
  mkMinorType (xs : Array Expr) (lhs : AltLHS) : MetaM Expr :=
    withExistingLocalDecls lhs.fvarDecls do
      let args ← lhs.patterns.toArray.mapM (Pattern.toExpr · (annotate := true))
      let minorType := mkAppN motive args
      withEqs discrs args discrInfos fun eqs => do
        mkForallFVars (xs ++ eqs) minorType

  loop (lhss : List AltLHS) (alts : List Alt) (minors : Array (Expr × Nat)) : MetaM α := do
    match lhss with
    | [] => k alts.reverse minors
    | lhs::lhss =>
      let xs := lhs.fvarDecls.toArray.map LocalDecl.toExpr
      let minorType ← mkMinorType xs lhs
      let hasParams := !xs.isEmpty || discrInfos.any fun info => info.hName?.isSome
      let (minorType, minorNumParams) := if hasParams then (minorType, xs.size) else (mkSimpleThunkType minorType, 1)
      let idx       := alts.length
      let minorName := (`h).appendIndexAfter (idx+1)
      trace[Meta.Match.debug] "minor premise {minorName} : {minorType}"
      withLocalDeclD minorName minorType fun minor => do
        let rhs    := if hasParams then mkAppN minor xs else mkApp minor (mkConst `Unit.unit)
        let minors := minors.push (minor, minorNumParams)
        let fvarDecls ← lhs.fvarDecls.mapM instantiateLocalDeclMVars
        let alts   := { ref := lhs.ref, idx := idx, rhs := rhs, fvarDecls := fvarDecls, patterns := lhs.patterns, cnstrs := [] } :: alts
        loop lhss alts minors

structure State where
  used            : HashSet Nat := {} -- used alternatives
  counterExamples : List (List Example) := []

/-- Return true if the given (sub-)problem has been solved. -/
private def isDone (p : Problem) : Bool :=
  p.vars.isEmpty

/-- Return true if the next element on the `p.vars` list is a variable. -/
private def isNextVar (p : Problem) : Bool :=
  match p.vars with
  | .fvar _ :: _ => true
  | _            => false

private def hasAsPattern (p : Problem) : Bool :=
  p.alts.any fun alt => match alt.patterns with
    | .as .. :: _ => true
    | _           => false

private def hasCtorPattern (p : Problem) : Bool :=
  p.alts.any fun alt => match alt.patterns with
    | .ctor .. :: _ => true
    | _             => false

private def hasValPattern (p : Problem) : Bool :=
  p.alts.any fun alt => match alt.patterns with
    | .val _ :: _ => true
    | _           => false

private def hasNatValPattern (p : Problem) : Bool :=
  p.alts.any fun alt => match alt.patterns with
    | .val v :: _ => v.isNatLit
    | _           => false

private def hasVarPattern (p : Problem) : Bool :=
  p.alts.any fun alt => match alt.patterns with
    | .var _ :: _ => true
    | _           => false

private def hasArrayLitPattern (p : Problem) : Bool :=
  p.alts.any fun alt => match alt.patterns with
    | .arrayLit .. :: _ => true
    | _                 => false

private def isVariableTransition (p : Problem) : Bool :=
  p.alts.all fun alt => match alt.patterns with
    | .inaccessible _ :: _ => true
    | .var _ :: _          => true
    | _                    => false

private def isConstructorTransition (p : Problem) : Bool :=
  (hasCtorPattern p || p.alts.isEmpty)
  && p.alts.all fun alt => match alt.patterns with
     | .ctor .. :: _        => true
     | .var _ :: _          => true
     | .inaccessible _ :: _ => true
     | _                    => false

private def isValueTransition (p : Problem) : Bool :=
  hasVarPattern p && hasValPattern p
  && p.alts.all fun alt => match alt.patterns with
     | .val _ :: _ => true
     | .var _ :: _ => true
     | _           => false

private def isArrayLitTransition (p : Problem) : Bool :=
  hasArrayLitPattern p && hasVarPattern p
  && p.alts.all fun alt => match alt.patterns with
     | .arrayLit .. :: _ => true
     | .var _ :: _       => true
     | _                 => false

private def isNatValueTransition (p : Problem) : Bool :=
  hasNatValPattern p
  && (!isNextVar p ||
      p.alts.any fun alt => match alt.patterns with
      | .ctor .. :: _        => true
      | .inaccessible _ :: _ => true
      | _                    => false)

private def processSkipInaccessible (p : Problem) : Problem := Id.run do
  let x :: xs := p.vars | unreachable!
  let alts := p.alts.map fun alt => Id.run do
    let .inaccessible e :: ps := alt.patterns | unreachable!
    { alt with patterns := ps, cnstrs := (x, e) :: alt.cnstrs }
  { p with alts := alts, vars := xs }

/--
If contraint is of the form `e ≋ x` where `x` is a free variable, reorient it
as `x ≋ e` If
- `x` is an `alt`-local declaration
- `e` is not a free variable.
-/
private def reorientCnstrs (alt : Alt) : Alt :=
  let cnstrs := alt.cnstrs.map fun (lhs, rhs) =>
    if rhs.isFVar && alt.isLocalDecl rhs.fvarId! then
      (rhs, lhs)
    else if !lhs.isFVar && rhs.isFVar then
      (rhs, lhs)
    else
      (lhs, rhs)
  { alt with cnstrs }

/--
Remove constraints of the form `lhs ≋ rhs` where `lhs` and `rhs` are definitionally equal,
or `lhs` is a free variable.
-/
private def filterTrivialCnstrs (alt : Alt) : MetaM Alt := do
   let cnstrs ← withExistingLocalDecls alt.fvarDecls do
     alt.cnstrs.filterM fun (lhs, rhs) => do
       if (← isDefEqGuarded lhs rhs) then
         return false
       else if lhs.isFVar then
         return false
       else
         return true
   return { alt with cnstrs }

/--
Find an alternative constraint of the form `x ≋ e` where `x` is an alternative
local declarations, and `x` and `e` have definitionally equal types.
Then, replace `x` with `e` in the alternative, and return it.
Return `none` if the alternative does not contain a constraint of this form.
-/
private def solveSomeLocalFVarIdCnstr? (alt : Alt) : MetaM (Option Alt) :=
  withExistingLocalDecls alt.fvarDecls do
    let (some (fvarId, val), cnstrs) ← go alt.cnstrs | return none
    trace[Meta.Match.match] "found cnstr to solve {mkFVar fvarId} ↦ {val}"
    return some <| { alt with cnstrs }.replaceFVarId fvarId val
where
  go (cnstrs : List (Expr × Expr)) := do
    match cnstrs with
    | [] => return (none, [])
    | (lhs, rhs) :: cnstrs =>
      if lhs.isFVar && alt.isLocalDecl lhs.fvarId! then
        if !(← dependsOn rhs lhs.fvarId!) && (← isDefEqGuarded (← inferType lhs) (← inferType rhs)) then
          return (some (lhs.fvarId!, rhs), cnstrs)
      let (p, cnstrs) ← go cnstrs
      return (p, (lhs, rhs) :: cnstrs)

/--
Solve pending alternative constraints. If all constraints can be solved perform assignment
`mvarId := alt.rhs`, and return true.
-/
private partial def solveCnstrs (mvarId : MVarId) (alt : Alt) : StateRefT State MetaM Bool := do
  go (reorientCnstrs alt)
where
  go (alt : Alt) : StateRefT State MetaM Bool := do
    match (← solveSomeLocalFVarIdCnstr? alt) with
    | some alt => go alt
    | none =>
      let alt ← filterTrivialCnstrs alt
      if alt.cnstrs.isEmpty then
        let eType ← inferType alt.rhs
        let targetType ← mvarId.getType
        unless (← isDefEqGuarded targetType eType) do
          trace[Meta.Match.match] "assignGoalOf failed {eType} =?= {targetType}"
          throwError "dependent elimination failed, type mismatch when solving alternative with type{indentExpr eType}\nbut expected{indentExpr targetType}"
        mvarId.assign alt.rhs
        modify fun s => { s with used := s.used.insert alt.idx }
        return true
      else
        trace[Meta.Match.match] "alt has unsolved cnstrs:\n{← alt.toMessageData}"
        return false

/--
Try to solve the problem by using the first alternative whose pending constraints can be resolved.
-/
private def processLeaf (p : Problem) : StateRefT State MetaM Unit :=
  p.mvarId.withContext do
    trace[Meta.Match.match] "local context at processLeaf:\n{(← mkFreshTypeMVar).mvarId!}"
    go p.alts
where
  go (alts : List Alt) : StateRefT State MetaM Unit := do
    match alts with
    | [] =>
      /- TODO: allow users to configure which tactic is used to close leaves. -/
      unless (← p.mvarId.contradictionCore {}) do
        trace[Meta.Match.match] "missing alternative"
        p.mvarId.admit
        modify fun s => { s with counterExamples := p.examples :: s.counterExamples }
    | alt :: alts =>
      unless (← solveCnstrs p.mvarId alt) do
        go alts

private def processAsPattern (p : Problem) : MetaM Problem := withGoalOf p do
  let x :: _ := p.vars | unreachable!
  let alts ← p.alts.mapM fun alt => do
    match alt.patterns with
    | .as fvarId p h :: ps =>
      /- We used to use `checkAndReplaceFVarId` here, but `x` and `fvarId` may have different types
        when dependent types are beind used. Let's consider the repro for issue #471
        ```
        inductive vec : Nat → Type
        | nil : vec 0
        | cons : Int → vec n → vec n.succ

        def vec_len : vec n → Nat
        | vec.nil => 0
        | x@(vec.cons h t) => vec_len t + 1

        ```
        we reach the state
        ```
          [Meta.Match.match] remaining variables: [x✝:(vec n✝)]
          alternatives:
            [n:(Nat), x:(vec (Nat.succ n)), h:(Int), t:(vec n)] |- [x@(vec.cons n h t)] => h_1 n x h t
            [x✝:(vec n✝)] |- [x✝] => h_2 n✝ x✝
        ```
        The variables `x✝:(vec n✝)` and `x:(vec (Nat.succ n))` have different types, but we perform the substitution anyway,
        because we claim the "discrepancy" will be corrected after we process the pattern `(vec.cons n h t)`.
        The right-hand-side is temporarily type incorrect, but we claim this is fine because it will be type correct again after
        we the pattern `(vec.cons n h t)`. TODO: try to find a cleaner solution.
       -/
      let r ← mkEqRefl x
      return { alt with patterns := p :: ps }.replaceFVarId fvarId x |>.replaceFVarId h r
    | _ => return alt
  return { p with alts := alts }

private def processVariable (p : Problem) : MetaM Problem := withGoalOf p do
  let x :: xs := p.vars | unreachable!
  let alts ← p.alts.mapM fun alt => do
    match alt.patterns with
    | .inaccessible e :: ps => return { alt with patterns := ps, cnstrs := (x, e) :: alt.cnstrs }
    | .var fvarId :: ps     =>
      withExistingLocalDecls alt.fvarDecls do
        if (← isDefEqGuarded (← fvarId.getType) (← inferType x)) then
          return { alt with patterns := ps }.replaceFVarId fvarId x
        else
          return { alt with patterns := ps, cnstrs := (mkFVar fvarId, x) :: alt.cnstrs }
    | _  => unreachable!
  return { p with alts := alts, vars := xs }

/-!
Note that we decided to store pending constraints to address issues exposed by #1279 and #1361.
Here is a simplified version of the example on this issue (see test: `1279_simplified.lean`)
```lean
inductive Arrow : Type → Type → Type 1
  | id   : Arrow a a
  | unit : Arrow Unit Unit
  | comp : Arrow β γ → Arrow α β → Arrow α γ
deriving Repr

def Arrow.compose (f : Arrow β γ) (g : Arrow α β) : Arrow α γ :=
  match f, g with
  | id, g => g
  | f, id => f
  | f, g => comp f g
```
The initial state for the `match`-expression above is
```lean
[Meta.Match.match] remaining variables: [β✝:(Type), γ✝:(Type), f✝:(Arrow β✝ γ✝), g✝:(Arrow α β✝)]
alternatives:
  [β:(Type), g:(Arrow α β)] |- [β, .(β), (Arrow.id .(β)), g] => h_1 β g
  [γ:(Type), f:(Arrow α γ)] |- [.(α), γ, f, (Arrow.id .(α))] => h_2 γ f
  [β:(Type), γ:(Type), f:(Arrow β γ), g:(Arrow α β)] |- [β, γ, f, g] => h_3 β γ f g
```
The first step is a variable-transition which replaces `β` with `β✝` in the first and third alternatives.
The constraint `β✝ ≋ α` in the second alternative used to be discarded. We now store it at the
alternative `cnstrs` field.
-/

private def inLocalDecls (localDecls : List LocalDecl) (fvarId : FVarId) : Bool :=
  localDecls.any fun d => d.fvarId == fvarId

private def expandVarIntoCtor? (alt : Alt) (fvarId : FVarId) (ctorName : Name) : MetaM (Option Alt) :=
  withExistingLocalDecls alt.fvarDecls do
    trace[Meta.Match.unify] "expandVarIntoCtor? fvarId: {mkFVar fvarId}, ctorName: {ctorName}, alt:\n{← alt.toMessageData}"
    let expectedType ← inferType (mkFVar fvarId)
    let expectedType ← whnfD expectedType
    let (ctorLevels, ctorParams) ← getInductiveUniverseAndParams expectedType
    let ctor := mkAppN (mkConst ctorName ctorLevels) ctorParams
    let ctorType ← inferType ctor
    forallTelescopeReducing ctorType fun ctorFields resultType => do
      let ctor := mkAppN ctor ctorFields
      let alt  := alt.replaceFVarId fvarId ctor
      let ctorFieldDecls ← ctorFields.mapM fun ctorField => ctorField.fvarId!.getDecl
      let newAltDecls := ctorFieldDecls.toList ++ alt.fvarDecls
      let mut cnstrs := alt.cnstrs
      unless (← isDefEqGuarded resultType expectedType) do
         cnstrs := (resultType, expectedType) :: cnstrs
      trace[Meta.Match.unify] "expandVarIntoCtor? {mkFVar fvarId} : {expectedType}, ctor: {ctor}"
      let ctorFieldPatterns := ctorFieldDecls.toList.map fun decl => Pattern.var decl.fvarId
      return some { alt with fvarDecls := newAltDecls, patterns := ctorFieldPatterns ++ alt.patterns, cnstrs }

private def getInductiveVal? (x : Expr) : MetaM (Option InductiveVal) := do
  let xType ← inferType x
  let xType ← whnfD xType
  match xType.getAppFn with
  | Expr.const constName _ =>
    let cinfo ← getConstInfo constName
    match cinfo with
    | ConstantInfo.inductInfo val => return some val
    | _ => return none
  | _ => return none

private def hasRecursiveType (x : Expr) : MetaM Bool := do
  match (← getInductiveVal? x) with
  | some val => return val.isRec
  | _        => return false

/-- Given `alt` s.t. the next pattern is an inaccessible pattern `e`,
   try to normalize `e` into a constructor application.
   If it is not a constructor, throw an error.
   Otherwise, if it is a constructor application of `ctorName`,
   update the next patterns with the fields of the constructor.
   Otherwise, return none. -/
def processInaccessibleAsCtor (alt : Alt) (ctorName : Name) : MetaM (Option Alt) := do
  let env ← getEnv
  match alt.patterns with
  | p@(.inaccessible e) :: ps =>
    trace[Meta.Match.match] "inaccessible in ctor step {e}"
    withExistingLocalDecls alt.fvarDecls do
      -- Try to push inaccessible annotations.
      let e ← whnfD e
      match e.constructorApp? env with
      | some (ctorVal, ctorArgs) =>
        if ctorVal.name == ctorName then
          let fields := ctorArgs.extract ctorVal.numParams ctorArgs.size
          let fields := fields.toList.map .inaccessible
          return some { alt with patterns := fields ++ ps }
        else
          return none
      | _ => throwErrorAt alt.ref "dependent match elimination failed, inaccessible pattern found{indentD p.toMessageData}\nconstructor expected"
  | _ => unreachable!

private def hasNonTrivialExample (p : Problem) : Bool :=
  p.examples.any fun | Example.underscore => false | _ => true

private def throwCasesException (p : Problem) (ex : Exception) : MetaM α := do
  match ex with
  | .error ref msg =>
    let exampleMsg :=
      if hasNonTrivialExample p then m!" after processing{indentD <| examplesToMessageData p.examples}" else ""
    throw <| Exception.error ref <| m!"{msg}{exampleMsg}\n" ++
              "the dependent pattern matcher can solve the following kinds of equations\n" ++
              "- <var> = <term> and <term> = <var>\n" ++
              "- <term> = <term> where the terms are definitionally equal\n" ++
              "- <constructor> = <constructor>, examples: List.cons x xs = List.cons y ys, and List.cons x xs = List.nil"
  | _ => throw ex

private def processConstructor (p : Problem) : MetaM (Array Problem) := do
  trace[Meta.Match.match] "constructor step"
  let x :: xs := p.vars | unreachable!
  let subgoals? ← commitWhenSome? do
     let subgoals ←
       try
         p.mvarId.cases x.fvarId!
       catch ex =>
         if p.alts.isEmpty then
           /- If we have no alternatives and dependent pattern matching fails, then a "missing cases" error is bettern than a "stuck" error message. -/
           return none
         else
           throwCasesException p ex
     if subgoals.isEmpty then
       /- Easy case: we have solved problem `p` since there are no subgoals -/
       return some #[]
     else if !p.alts.isEmpty then
       return some subgoals
     else do
       let isRec ← withGoalOf p <| hasRecursiveType x
        /- If there are no alternatives and the type of the current variable is recursive, we do NOT consider
          a constructor-transition to avoid nontermination.
          TODO: implement a more general approach if this is not sufficient in practice -/
       if isRec then
         return none
       else
         return some subgoals
  let some subgoals := subgoals? | return #[{ p with vars := xs }]
  subgoals.mapM fun subgoal => subgoal.mvarId.withContext do
    let subst    := subgoal.subst
    let fields   := subgoal.fields.toList
    let newVars  := fields ++ xs
    let newVars  := newVars.map fun x => x.applyFVarSubst subst
    let subex    := Example.ctor subgoal.ctorName <| fields.map fun field => match field with
      | .fvar fvarId => Example.var fvarId
      | _            => Example.underscore -- This case can happen due to dependent elimination
    let examples := p.examples.map <| Example.replaceFVarId x.fvarId! subex
    let examples := examples.map <| Example.applyFVarSubst subst
    let newAlts  := p.alts.filter fun alt => match alt.patterns with
      | .ctor n .. :: _       => n == subgoal.ctorName
      | .var _ :: _           => true
      | .inaccessible _ :: _  => true
      | _                     => false
    let newAlts  := newAlts.map fun alt => alt.applyFVarSubst subst
    let newAlts ← newAlts.filterMapM fun alt => do
      match alt.patterns with
      | .ctor _ _ _ fields :: ps  => return some { alt with patterns := fields ++ ps }
      | .var fvarId :: ps         => expandVarIntoCtor? { alt with patterns := ps } fvarId subgoal.ctorName
      | .inaccessible _ :: _      => processInaccessibleAsCtor alt subgoal.ctorName
      | _                         => unreachable!
    return { mvarId := subgoal.mvarId, vars := newVars, alts := newAlts, examples := examples }

private def altsAreCtorLike (p : Problem) : MetaM Bool := withGoalOf p do
  p.alts.allM fun alt => do match alt.patterns with
    | .ctor .. :: _ => return true
    | .inaccessible e :: _ => return (← whnfD e).isConstructorApp (← getEnv)
    | _ => return false

private def processNonVariable (p : Problem) : MetaM Problem := withGoalOf p do
  let x :: xs := p.vars | unreachable!
  if let some (ctorVal, xArgs) := (← whnfD x).constructorApp? (← getEnv) then
    if (← altsAreCtorLike p) then
      let alts ← p.alts.filterMapM fun alt => do
        match alt.patterns with
        | .ctor ctorName _ _ fields :: ps   =>
          if ctorName != ctorVal.name then
            return none
          else
            return some { alt with patterns := fields ++ ps }
        | .inaccessible _ :: _ => processInaccessibleAsCtor alt ctorVal.name
        | _ => unreachable!
      let xFields := xArgs.extract ctorVal.numParams xArgs.size
      return { p with alts := alts, vars := xFields.toList ++ xs }
  let alts ← p.alts.mapM fun alt => do
    match alt.patterns with
    | p :: ps => return { alt with patterns := ps, cnstrs := (x, ← p.toExpr) :: alt.cnstrs }
    | _      => unreachable!
  return { p with alts := alts, vars := xs }

private def collectValues (p : Problem) : Array Expr :=
  p.alts.foldl (init := #[]) fun values alt =>
    match alt.patterns with
    | .val v :: _ => if values.contains v then values else values.push v
    | _           => values

private def isFirstPatternVar (alt : Alt) : Bool :=
  match alt.patterns with
  | .var _ :: _ => true
  | _           => false

private def processValue (p : Problem) : MetaM (Array Problem) := do
  trace[Meta.Match.match] "value step"
  let x :: xs := p.vars | unreachable!
  let values := collectValues p
  let subgoals ← caseValues p.mvarId x.fvarId! values (substNewEqs := true)
  subgoals.mapIdxM fun i subgoal => do
    trace[Meta.Match.match] "processValue subgoal\n{MessageData.ofGoal subgoal.mvarId}"
    if h : i.val < values.size then
      let value := values.get ⟨i, h⟩
      -- (x = value) branch
      let subst := subgoal.subst
      trace[Meta.Match.match] "processValue subst: {subst.map.toList.map fun p => mkFVar p.1}, {subst.map.toList.map fun p => p.2}"
        let examples := p.examples.map <| Example.replaceFVarId x.fvarId! (Example.val value)
      let examples := examples.map <| Example.applyFVarSubst subst
      let newAlts  := p.alts.filter fun alt => match alt.patterns with
        | .val v :: _ => v == value
        | .var _ :: _ => true
        | _           => false
      let newAlts := newAlts.map fun alt => alt.applyFVarSubst subst
      let newAlts := newAlts.map fun alt => match alt.patterns with
        | .val _ :: ps      => { alt with patterns := ps }
        | .var fvarId :: ps =>
          let alt := { alt with patterns := ps }
          alt.replaceFVarId fvarId value
        | _  => unreachable!
      let newVars := xs.map fun x => x.applyFVarSubst subst
      return { mvarId := subgoal.mvarId, vars := newVars, alts := newAlts, examples := examples }
    else
      -- else branch for value
      let newAlts := p.alts.filter isFirstPatternVar
      return { p with mvarId := subgoal.mvarId, alts := newAlts, vars := x::xs }

private def collectArraySizes (p : Problem) : Array Nat :=
  p.alts.foldl (init := #[]) fun sizes alt =>
    match alt.patterns with
    | .arrayLit _ ps :: _ => let sz := ps.length; if sizes.contains sz then sizes else sizes.push sz
    | _                   => sizes

private def expandVarIntoArrayLit (alt : Alt) (fvarId : FVarId) (arrayElemType : Expr) (arraySize : Nat) : MetaM Alt :=
  withExistingLocalDecls alt.fvarDecls do
    let fvarDecl ← fvarId.getDecl
    let varNamePrefix := fvarDecl.userName
    let rec loop (n : Nat) (newVars : Array Expr) := do
      match n with
      | n+1 =>
        withLocalDeclD (varNamePrefix.appendIndexAfter (n+1)) arrayElemType fun x =>
          loop n (newVars.push x)
      | 0 =>
        let arrayLit ← mkArrayLit arrayElemType newVars.toList
        let alt := alt.replaceFVarId fvarId arrayLit
        let newDecls ← newVars.toList.mapM fun newVar => newVar.fvarId!.getDecl
        let newPatterns := newVars.toList.map fun newVar => .var newVar.fvarId!
        return { alt with fvarDecls := newDecls ++ alt.fvarDecls, patterns := newPatterns ++ alt.patterns }
    loop arraySize #[]

private def processArrayLit (p : Problem) : MetaM (Array Problem) := do
  trace[Meta.Match.match] "array literal step"
  let x :: xs := p.vars | unreachable!
  let sizes := collectArraySizes p
  let subgoals ← caseArraySizes p.mvarId x.fvarId! sizes
  subgoals.mapIdxM fun i subgoal => do
    if i.val < sizes.size then
      let size     := sizes.get! i
      let subst    := subgoal.subst
      let elems    := subgoal.elems.toList
      let newVars  := elems.map mkFVar ++ xs
      let newVars  := newVars.map fun x => x.applyFVarSubst subst
      let subex    := Example.arrayLit <| elems.map Example.var
      let examples := p.examples.map <| Example.replaceFVarId x.fvarId! subex
      let examples := examples.map <| Example.applyFVarSubst subst
      let newAlts  := p.alts.filter fun alt => match alt.patterns with
        | .arrayLit _ ps :: _ => ps.length == size
        | .var _ :: _         => true
        | _                          => false
      let newAlts := newAlts.map fun alt => alt.applyFVarSubst subst
      let newAlts ← newAlts.mapM fun alt => do
        match alt.patterns with
        | .arrayLit _ pats :: ps => return { alt with patterns := pats ++ ps }
        | .var fvarId :: ps      =>
          let α ← getArrayArgType <| subst.apply x
          expandVarIntoArrayLit { alt with patterns := ps } fvarId α size
        | _  => unreachable!
      return { mvarId := subgoal.mvarId, vars := newVars, alts := newAlts, examples := examples }
    else
      -- else branch
      let newAlts := p.alts.filter isFirstPatternVar
      return { p with mvarId := subgoal.mvarId, alts := newAlts, vars := x::xs }

private def expandNatValuePattern (p : Problem) : Problem :=
  let alts := p.alts.map fun alt => match alt.patterns with
    | .val (.lit (.natVal 0)) :: ps     => { alt with patterns := .ctor ``Nat.zero [] [] [] :: ps }
    | .val (.lit (.natVal (n+1))) :: ps => { alt with patterns := .ctor ``Nat.succ [] [] [.val (mkRawNatLit n)] :: ps }
    | _                                 => alt
  { p with alts := alts }

private def traceStep (msg : String) : StateRefT State MetaM Unit := do
  trace[Meta.Match.match] "{msg} step"

private def traceState (p : Problem) : MetaM Unit :=
  withGoalOf p (traceM `Meta.Match.match p.toMessageData)

private def throwNonSupported (p : Problem) : MetaM Unit :=
  withGoalOf p do
    let msg ← p.toMessageData
    throwError "failed to compile pattern matching, stuck at{indentD msg}"

def isCurrVarInductive (p : Problem) : MetaM Bool := do
  match p.vars with
  | []   => return false
  | x::_ => withGoalOf p do
    let val? ← getInductiveVal? x
    return val?.isSome

private def checkNextPatternTypes (p : Problem) : MetaM Unit := do
  match p.vars with
  | []   => return ()
  | x::_ => withGoalOf p do
    for alt in p.alts do
      withRef alt.ref do
        match alt.patterns with
        | []   => return ()
        | p::_ =>
          let e ← p.toExpr
          let xType ← inferType x
          let eType ← inferType e
          unless (← isDefEq xType eType) do
            throwError "pattern{indentExpr e}\n{← mkHasTypeButIsExpectedMsg eType xType}"

private partial def process (p : Problem) : StateRefT State MetaM Unit := do
  traceState p
  let isInductive ← isCurrVarInductive p
  if isDone p then
    traceStep ("leaf")
    processLeaf p
  else if hasAsPattern p then
    traceStep ("as-pattern")
    let p ← processAsPattern p
    process p
  else if isNatValueTransition p then
    traceStep ("nat value to constructor")
    process (expandNatValuePattern p)
  else if !isNextVar p then
    traceStep ("non variable")
    let p ← processNonVariable p
    process p
  else if isInductive && isConstructorTransition p then
    let ps ← processConstructor p
    ps.forM process
  else if isVariableTransition p then
    traceStep ("variable")
    let p ← processVariable p
    process p
  else if isValueTransition p then
    let ps ← processValue p
    ps.forM process
  else if isArrayLitTransition p then
    let ps ← processArrayLit p
    ps.forM process
  else if hasNatValPattern p then
    -- This branch is reachable when `p`, for example, is just values without an else-alternative.
    -- We added it just to get better error messages.
    traceStep ("nat value to constructor")
    process (expandNatValuePattern p)
  else
    checkNextPatternTypes p
    throwNonSupported p

private def getUElimPos? (matcherLevels : List Level) (uElim : Level) : MetaM (Option Nat) :=
  if uElim == levelZero then
    return none
  else match matcherLevels.toArray.indexOf? uElim with
    | none => throwError "dependent match elimination failed, universe level not found"
    | some pos => return some pos.val

/- See comment at `mkMatcher` before `mkAuxDefinition` -/
register_builtin_option bootstrap.genMatcherCode : Bool := {
  defValue := true
  group := "bootstrap"
  descr := "disable code generation for auxiliary matcher function"
}

builtin_initialize matcherExt : EnvExtension (PHashMap (Expr × Bool) Name) ← registerEnvExtension (pure {})

/-- Similar to `mkAuxDefinition`, but uses the cache `matcherExt`.
   It also returns an Boolean that indicates whether a new matcher function was added to the environment or not. -/
def mkMatcherAuxDefinition (name : Name) (type : Expr) (value : Expr) : MetaM (Expr × Option (MatcherInfo → MetaM Unit)) := do
  trace[Meta.Match.debug] "{name} : {type} := {value}"
  let compile := bootstrap.genMatcherCode.get (← getOptions)
  let result ← Closure.mkValueTypeClosure type value (zeta := false)
  let env ← getEnv
  let mkMatcherConst name :=
    mkAppN (mkConst name result.levelArgs.toList) result.exprArgs
  match (matcherExt.getState env).find? (result.value, compile) with
  | some nameNew => return (mkMatcherConst nameNew, none)
  | none =>
    let decl := Declaration.defnDecl {
      name
      levelParams := result.levelParams.toList
      type        := result.type
      value       := result.value
      hints       := ReducibilityHints.abbrev
      safety      := if env.hasUnsafe result.type || env.hasUnsafe result.value then DefinitionSafety.unsafe else DefinitionSafety.safe
    }
    trace[Meta.Match.debug] "{name} : {result.type} := {result.value}"
    let addMatcher : MatcherInfo → MetaM Unit := fun mi => do
      addDecl decl
      modifyEnv fun env => matcherExt.modifyState env fun s => s.insert (result.value, compile) name
      addMatcherInfo name mi
      setInlineAttribute name
      if compile then
        compileDecl decl
    return (mkMatcherConst name, some addMatcher)

structure MkMatcherInput where
  matcherName : Name
  matchType   : Expr
  discrInfos  : Array DiscrInfo
  lhss        : List AltLHS

def MkMatcherInput.numDiscrs (m : MkMatcherInput) :=
  m.discrInfos.size

def MkMatcherInput.collectFVars (m : MkMatcherInput) : StateRefT CollectFVars.State MetaM Unit := do
  m.matchType.collectFVars
  m.lhss.forM fun alt => alt.collectFVars

def MkMatcherInput.collectDependencies (m : MkMatcherInput) : MetaM FVarIdSet := do
  let (_, s) ← m.collectFVars |>.run {}
  let s ← s.addDependencies
  return s.fvarSet

/--
Auxiliary method used at `mkMatcher`. It executes `k` in a local context that contains only
the local declarations `m` depends on. This is important because otherwise dependent elimination
may "refine" the types of unnecessary declarations and accidentally introduce unnecessary dependencies
in the auto-generated auxiliary declaration. Note that this is not just an optimization because the
unnecessary dependencies may prevent the termination checker from succeeding. For an example,
see issue #1237.
-/
def withCleanLCtxFor (m : MkMatcherInput) (k : MetaM α) : MetaM α := do
  let s ← m.collectDependencies
  let lctx ← getLCtx
  let lctx := lctx.foldr (init := lctx) fun localDecl lctx =>
    if s.contains localDecl.fvarId then lctx else lctx.erase localDecl.fvarId
  let localInstances := (← getLocalInstances).filter fun localInst => s.contains localInst.fvar.fvarId!
  withLCtx lctx localInstances k

/--
Create a dependent matcher for `matchType` where `matchType` is of the form
`(a_1 : A_1) -> (a_2 : A_2[a_1]) -> ... -> (a_n : A_n[a_1, a_2, ... a_{n-1}]) -> B[a_1, ..., a_n]`
where `n = numDiscrs`, and the `lhss` are the left-hand-sides of the `match`-expression alternatives.
Each `AltLHS` has a list of local declarations and a list of patterns.
The number of patterns must be the same in each `AltLHS`.
The generated matcher has the structure described at `MatcherInfo`. The motive argument is of the form
`(motive : (a_1 : A_1) -> (a_2 : A_2[a_1]) -> ... -> (a_n : A_n[a_1, a_2, ... a_{n-1}]) -> Sort v)`
where `v` is a universe parameter or 0 if `B[a_1, ..., a_n]` is a proposition. -/
def mkMatcher (input : MkMatcherInput) : MetaM MatcherResult := withCleanLCtxFor input do
  let ⟨matcherName, matchType, discrInfos, lhss⟩ := input
  let numDiscrs := discrInfos.size
  let numEqs := getNumEqsFromDiscrInfos discrInfos
  checkNumPatterns numDiscrs lhss
  forallBoundedTelescope matchType numDiscrs fun discrs matchTypeBody => do
  /- We generate an matcher that can eliminate using different motives with different universe levels.
     `uElim` is the universe level the caller wants to eliminate to.
     If it is not levelZero, we create a matcher that can eliminate in any universe level.
     This is useful for implementing `MatcherApp.addArg` because it may have to change the universe level. -/
  let uElim ← getLevel matchTypeBody
  let uElimGen ← if uElim == levelZero then pure levelZero else mkFreshLevelMVar
  let mkMatcher (type val : Expr) (minors : Array (Expr × Nat)) (s : State) : MetaM MatcherResult := do
    trace[Meta.Match.debug] "matcher value: {val}\ntype: {type}"
    trace[Meta.Match.debug] "minors num params: {minors.map (·.2)}"
    /- The option `bootstrap.gen_matcher_code` is a helper hack. It is useful, for example,
       for compiling `src/Init/Data/Int`. It is needed because the compiler uses `Int.decLt`
       for generating code for `Int.casesOn` applications, but `Int.casesOn` is used to
       give the reference implementation for
       ```
       @[extern "lean_int_neg"] def neg (n : @& Int) : Int :=
       match n with
       | ofNat n   => negOfNat n
       | negSucc n => succ n
       ```
       which is defined **before** `Int.decLt` -/

    let (matcher, addMatcher) ← mkMatcherAuxDefinition matcherName type val
    trace[Meta.Match.debug] "matcher levels: {matcher.getAppFn.constLevels!}, uElim: {uElimGen}"
    let uElimPos? ← getUElimPos? matcher.getAppFn.constLevels! uElimGen
    discard <| isLevelDefEq uElimGen uElim
    let addMatcher :=
      match addMatcher with
      | some addMatcher => addMatcher <|
        { numParams := matcher.getAppNumArgs
          altNumParams := minors.map fun minor => minor.2 + numEqs
          discrInfos
          numDiscrs
          uElimPos?
          }
      | none => pure ()

    trace[Meta.Match.debug] "matcher: {matcher}"
    let unusedAltIdxs := lhss.length.fold (init := []) fun i r =>
      if s.used.contains i then r else i::r
    return {
      matcher,
      counterExamples := s.counterExamples,
      unusedAltIdxs := unusedAltIdxs.reverse,
      addMatcher
    }

  let motiveType ← mkForallFVars discrs (mkSort uElimGen)
  trace[Meta.Match.debug] "motiveType: {motiveType}"
  withLocalDeclD `motive motiveType fun motive => do
  if discrInfos.any fun info => info.hName?.isSome then
    forallBoundedTelescope matchType numDiscrs fun discrs' _ => do
    let (mvarType, isEqMask) ← withEqs discrs discrs' discrInfos fun eqs => do
      let mvarType ← mkForallFVars eqs (mkAppN motive discrs')
      let isEqMask ← eqs.mapM fun eq => return (← inferType eq).isEq
      return (mvarType, isEqMask)
    trace[Meta.Match.debug] "target: {mvarType}"
    withAlts motive discrs discrInfos lhss fun alts minors => do
      let mvar ← mkFreshExprMVar mvarType
      trace[Meta.Match.debug] "goal\n{mvar.mvarId!}"
      let examples := discrs'.toList.map fun discr => Example.var discr.fvarId!
      let (_, s) ← (process { mvarId := mvar.mvarId!, vars := discrs'.toList, alts := alts, examples := examples }).run {}
      let val ← mkLambdaFVars discrs' mvar
      trace[Meta.Match.debug] "matcher\nvalue: {val}\ntype: {← inferType val}"
      let mut rfls := #[]
      let mut isEqMaskIdx := 0
      for discr in discrs, info in discrInfos do
        if info.hName?.isSome then
          if isEqMask[isEqMaskIdx]! then
            rfls := rfls.push (← mkEqRefl discr)
          else
            rfls := rfls.push (← mkHEqRefl discr)
          isEqMaskIdx := isEqMaskIdx + 1
      let val := mkAppN (mkAppN val discrs) rfls
      let args := #[motive] ++ discrs ++ minors.map Prod.fst
      let val ← mkLambdaFVars args val
      let type ← mkForallFVars args (mkAppN motive discrs)
      mkMatcher type val minors s
  else
    let mvarType  := mkAppN motive discrs
    trace[Meta.Match.debug] "target: {mvarType}"
    withAlts motive discrs discrInfos lhss fun alts minors => do
      let mvar ← mkFreshExprMVar mvarType
      let examples := discrs.toList.map fun discr => Example.var discr.fvarId!
      let (_, s) ← (process { mvarId := mvar.mvarId!, vars := discrs.toList, alts := alts, examples := examples }).run {}
      let args := #[motive] ++ discrs ++ minors.map Prod.fst
      let type ← mkForallFVars args mvarType
      let val  ← mkLambdaFVars args mvar
      mkMatcher type val minors s

def getMkMatcherInputInContext (matcherApp : MatcherApp) : MetaM MkMatcherInput := do
  let matcherName := matcherApp.matcherName
  let some matcherInfo ← getMatcherInfo? matcherName | throwError "not a matcher: {matcherName}"
  let matcherConst ← getConstInfo matcherName
  let matcherType ← instantiateForall matcherConst.type <| matcherApp.params ++ #[matcherApp.motive]
  let matchType ← do
    let u :=
      if let some idx := matcherInfo.uElimPos?
      then mkLevelParam matcherConst.levelParams.toArray[idx]!
      else levelZero
    forallBoundedTelescope matcherType (some matcherInfo.numDiscrs) fun discrs _ => do
    mkForallFVars discrs (mkConst ``PUnit [u])

  let matcherType ← instantiateForall matcherType matcherApp.discrs
  let lhss ← forallBoundedTelescope matcherType (some matcherApp.alts.size) fun alts _ =>
    alts.mapM fun alt => do
    let ty ← inferType alt
    forallTelescope ty fun xs body => do
    let xs ← xs.filterM fun x => dependsOn body x.fvarId!
    body.withApp fun _ args => do
    let ctx ← getLCtx
    let localDecls := xs.map ctx.getFVar!
    let patterns ← args.mapM Match.toPattern
    return {
      ref := Syntax.missing
      fvarDecls := localDecls.toList
      patterns := patterns.toList : Match.AltLHS }

  return { matcherName, matchType, discrInfos := matcherInfo.discrInfos, lhss := lhss.toList }

/-- This function is only used for testing purposes -/
def withMkMatcherInput (matcherName : Name) (k : MkMatcherInput → MetaM α) : MetaM α := do
  let some matcherInfo ← getMatcherInfo? matcherName | throwError "not a matcher: {matcherName}"
  let matcherConst ← getConstInfo matcherName
  forallBoundedTelescope matcherConst.type (some matcherInfo.arity) fun xs _ => do
  let matcherApp ← mkConstWithLevelParams matcherConst.name
  let matcherApp := mkAppN matcherApp xs
  let some matcherApp ← matchMatcherApp? matcherApp | throwError "not a matcher app: {matcherApp}"
  let mkMatcherInput ← getMkMatcherInputInContext matcherApp
  k mkMatcherInput

end Match

/-- Auxiliary function for MatcherApp.addArg -/
private partial def updateAlts (typeNew : Expr) (altNumParams : Array Nat) (alts : Array Expr) (i : Nat) : MetaM (Array Nat × Array Expr) := do
  if h : i < alts.size then
    let alt       := alts.get ⟨i, h⟩
    let numParams := altNumParams[i]!
    let typeNew ← whnfD typeNew
    match typeNew with
    | Expr.forallE _ d b _ =>
      let alt ← forallBoundedTelescope d (some numParams) fun xs d => do
        let alt ← try instantiateLambda alt xs catch _ => throwError "unexpected matcher application, insufficient number of parameters in alternative"
        forallBoundedTelescope d (some 1) fun x _ => do
          let alt ← mkLambdaFVars x alt -- x is the new argument we are adding to the alternative
          mkLambdaFVars xs alt
      updateAlts (b.instantiate1 alt) (altNumParams.set! i (numParams+1)) (alts.set ⟨i, h⟩ alt) (i+1)
    | _ => throwError "unexpected type at MatcherApp.addArg"
  else
    return (altNumParams, alts)

/-- Given
  - matcherApp `match_i As (fun xs => motive[xs]) discrs (fun ys_1 => (alt_1 : motive (C_1[ys_1])) ... (fun ys_n => (alt_n : motive (C_n[ys_n]) remaining`, and
  - expression `e : B[discrs]`,
  Construct the term
  `match_i As (fun xs => B[xs] -> motive[xs]) discrs (fun ys_1 (y : B[C_1[ys_1]]) => alt_1) ... (fun ys_n (y : B[C_n[ys_n]]) => alt_n) e remaining`, and
  We use `kabstract` to abstract the discriminants from `B[discrs]`.
  This method assumes
  - the `matcherApp.motive` is a lambda abstraction where `xs.size == discrs.size`
  - each alternative is a lambda abstraction where `ys_i.size == matcherApp.altNumParams[i]`
-/
def MatcherApp.addArg (matcherApp : MatcherApp) (e : Expr) : MetaM MatcherApp :=
  lambdaTelescope matcherApp.motive fun motiveArgs motiveBody => do
    unless motiveArgs.size == matcherApp.discrs.size do
      -- This error can only happen if someone implemented a transformation that rewrites the motive created by `mkMatcher`.
      throwError "unexpected matcher application, motive must be lambda expression with #{matcherApp.discrs.size} arguments"
    let eType ← inferType e
    let eTypeAbst ← matcherApp.discrs.size.foldRevM (init := eType) fun i eTypeAbst => do
      let motiveArg := motiveArgs[i]!
      let discr     := matcherApp.discrs[i]!
      let eTypeAbst ← kabstract eTypeAbst discr
      return eTypeAbst.instantiate1 motiveArg
    let motiveBody ← mkArrow eTypeAbst motiveBody
    let matcherLevels ← match matcherApp.uElimPos? with
      | none     => pure matcherApp.matcherLevels
      | some pos =>
        let uElim ← getLevel motiveBody
        pure <| matcherApp.matcherLevels.set! pos uElim
    let motive ← mkLambdaFVars motiveArgs motiveBody
    -- Construct `aux` `match_i As (fun xs => B[xs] → motive[xs]) discrs`, and infer its type `auxType`.
    -- We use `auxType` to infer the type `B[C_i[ys_i]]` of the new argument in each alternative.
    let aux := mkAppN (mkConst matcherApp.matcherName matcherLevels.toList) matcherApp.params
    let aux := mkApp aux motive
    let aux := mkAppN aux matcherApp.discrs
    unless (← isTypeCorrect aux) do
      throwError "failed to add argument to matcher application, type error when constructing the new motive"
    let auxType ← inferType aux
    let (altNumParams, alts) ← updateAlts auxType matcherApp.altNumParams matcherApp.alts 0
    return { matcherApp with
      matcherLevels := matcherLevels,
      motive        := motive,
      alts          := alts,
      altNumParams  := altNumParams,
      remaining     := #[e] ++ matcherApp.remaining
    }

/-- Similar `MatcherApp.addArg?`, but returns `none` on failure. -/
def MatcherApp.addArg? (matcherApp : MatcherApp) (e : Expr) : MetaM (Option MatcherApp) :=
  try
    return some (← matcherApp.addArg e)
  catch _ =>
    return none

builtin_initialize
  registerTraceClass `Meta.Match.match
  registerTraceClass `Meta.Match.debug
  registerTraceClass `Meta.Match.unify

end Lean.Meta


theorem ex
     (h₁ : α = β)
     (as : List α)
     (bs : List β)
     (h₂ : (h ▸ as) = bs)
     : True :=
  True.intro


import Std.Data.PersistentHashMap

open Std
def m : PersistentHashMap Nat Nat :=
let m : PersistentHashMap Nat Nat := {};
m.insert 1 1

def natDiffHash : Hashable Nat :=
⟨fun n => USize.ofNat $ n+10⟩

-- The following example should fail since the `Hashable` instance used to create `m` is not `natDiffHash`
#eval @PersistentHashMap.find? Nat Nat _ natDiffHash m 1


/-
Copyright (c) 2020 Gabriel Ebner. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Gabriel Ebner
-/
import Mathlib.Tactic.Lint.Basic
import Mathlib.Tactic.OpenPrivate
import Mathlib.Util.LibraryNote
open Lean Meta

namespace Mathlib.Tactic.Lint

/-!
# Linter for simplification lemmas

This files defines several linters that prevent common mistakes when declaring simp lemmas:

 * `simpNF` checks that the left-hand side of a simp lemma is not simplified by a different lemma.
 * `simpVarHead` checks that the head symbol of the left-hand side is not a variable.
 * `simpComm` checks that commutativity lemmas are not marked as simplification lemmas.
-/

structure SimpTheoremInfo where
  hyps : Array Expr
  isConditional : Bool
  lhs : Expr
  rhs : Expr

def isConditionalHyps (eq : Expr) : List Expr → MetaM Bool
  | [] => pure false
  | h :: hs => do
    let ldecl ← getFVarLocalDecl h
    if !ldecl.binderInfo.isInstImplicit
        && !(← hs.anyM fun h' => do pure $ (← inferType h').containsFVar h.fvarId!)
        && !eq.containsFVar h.fvarId! then
      return true
    isConditionalHyps eq hs

open private preprocess from Lean.Meta.Tactic.Simp.SimpTheorems in
def withSimpTheoremInfos (ty : Expr) (k : SimpTheoremInfo → MetaM α) : MetaM (Array α) := withReducible do
  (← preprocess (← mkSorry ty true) ty (inv := false) (isGlobal := true))
      |>.toArray.mapM fun (_, ty') => do
    forallTelescopeReducing ty' fun hyps eq => do
      let some (_, lhs, rhs) := eq.eq? | throwError "not an equality {eq}"
      k {
        hyps, lhs, rhs
        isConditional := ← isConditionalHyps eq hyps.toList
      }

/-- Checks whether two expressions are equal for the simplifier. That is,
they are reducibly-definitional equal, and they have the same head symbol. -/
def isSimpEq (a b : Expr) (whnfFirst := true) : MetaM Bool := withReducible do
  let a ← if whnfFirst then whnf a else pure a
  let b ← if whnfFirst then whnf b else pure b
  if a.getAppFn.constName? != b.getAppFn.constName? then return false
  isDefEq a b

def checkAllSimpTheoremInfos (ty : Expr) (k : SimpTheoremInfo → MetaM (Option MessageData)) : MetaM (Option MessageData) := do
  let errors := (← withSimpTheoremInfos ty fun i => do (← k i).mapM addMessageContextFull).filterMap id
  if errors.isEmpty then
    return none
  else
    return MessageData.joinSep errors.toList Format.line

def isSimpTheorem (declName : Name) : MetaM Bool := do
  pure $ (← getSimpTheorems).lemmaNames.contains declName

open Lean.Meta.DiscrTree
partial def trieElements : Trie α → StateT (Array α) Id Unit
  | Trie.node vs children => do
    modify (· ++ vs)
    for (_, child) in children do
      trieElements child

def elements (d : DiscrTree α) : StateT (Array α) Id Unit := do
  for (_, child) in d.root.toList do
    trieElements child

open Std

-- In Lean 4, the simplifier adds auxiliary lemmas if the declaration is not an equation.
-- For example, for `decl : a ↔ b` it generates `decl._auxLemma.1 : a = b`.
-- This function computes the map ``{`decl._auxLemma.1 ↦ `decl}``
def constToSimpDeclMap (ctx : Simp.Context) : HashMap Name Name := Id.run do
  let mut map : HashMap Name Name := {}
  for sls' in ctx.simpTheorems do
    for sls in [sls'.pre, sls'.post] do
      for sl in ((elements sls).run #[]).2 do
        if let some declName := sl.name? then
          if let some auxDeclName := sl.proof.getAppFn.constName? then
            map := map.insert auxDeclName declName
  return map

def isEqnLemma? (n : Name) : Option Name :=
  if n.isStr && n.getString!.startsWith "_eq_" then
    n.getPrefix
  else
    none

def heuristicallyExtractSimpTheoremsCore (ctx : Simp.Context) (constToSimpDecl : HashMap Name Name) (prf : Expr) : Array Name := Id.run do
  let mut cnsts : HashSet Name := {}
  for c in prf.getUsedConstants do
    if ctx.simpTheorems.isDeclToUnfold c then
      cnsts := cnsts.insert c
    else if ctx.congrTheorems.lemmas.contains c then
      cnsts := cnsts.insert c
    else if let some c' := constToSimpDecl.find? c then
      cnsts := cnsts.insert c'
    else if let some c' := isEqnLemma? c then
      cnsts := cnsts.insert c'
  return cnsts.toArray

@[inline] def heuristicallyExtractSimpTheorems (ctx : Simp.Context) (prf : Expr) : Array Name :=
  heuristicallyExtractSimpTheoremsCore ctx (constToSimpDeclMap ctx) prf

def decorateError (msg : MessageData) (k : MetaM α) : MetaM α := do
  try k catch e => throw e

def formatLemmas (lems : Array Name) : CoreM MessageData := do
  toMessageData <$> lems.mapM mkConstWithLevelParams

/-- A linter for simp lemmas whose lhs is not in simp-normal form, and which hence never fire. -/
@[mathlibLinter] def simpNF : Linter where
  noErrorsFound := "All left-hand sides of simp lemmas are in simp-normal form."
  errorsFound := "SOME SIMP LEMMAS ARE NOT IN SIMP-NORMAL FORM.
see note [simp-normal form] for tips how to debug this.
https://leanprover-community.github.io/mathlib_docs/notes.html#simp-normal%20form"

  test := fun declName => do
    unless ← isSimpTheorem declName do return none
    -- TODO: equation lemmas
    let ctx ← Simp.Context.mkDefault
    checkAllSimpTheoremInfos (← getConstInfo declName).type fun {lhs, rhs, isConditional, ..} => do
    let ⟨lhs', prf1⟩ ← decorateError "simplify fails on left-hand side:" <| simp lhs ctx
    let prf1_lems := heuristicallyExtractSimpTheorems ctx (prf1.getD (mkBVar 0))
    if prf1_lems.contains declName then return none
    let ⟨rhs', prf2⟩ ← decorateError "simplify fails on right-hand side:" <| simp rhs ctx
    let lhs'_eq_rhs' ← isSimpEq lhs' rhs' (whnfFirst := false)
    let lhs_in_nf ← isSimpEq lhs' lhs
    if lhs'_eq_rhs' then do
      if prf1.isNone then return none -- TODO: cannot detect used rfl-lemmas
      let used_lemmas := heuristicallyExtractSimpTheorems ctx <|
        mkApp (prf1.getD (mkBVar 0)) (prf2.getD (mkBVar 0))
      return m!"simp can prove this:
  by simp only {← formatLemmas used_lemmas}
One of the lemmas above could be a duplicate.
If that's not the case try reordering lemmas or adding @[priority].
"
    else if ¬ lhs_in_nf then do
      return m!"Left-hand side simplifies from
  {lhs}
to
  {lhs'}
using
  {← formatLemmas prf1_lems}
Try to change the left-hand side to the simplified term!
"
    else if !isConditional && lhs == lhs' then
      return m!"Left-hand side does not simplify.
You need to debug this yourself using `set_option trace.Meta.Tactic.simp.rewrite true`"
    else
      return none

library_note "simp-normal form" /--
This note gives you some tips to debug any errors that the simp-normal form linter raises.

The reason that a lemma was considered faulty is because its left-hand side is not in simp-normal
form.
These lemmas are hence never used by the simplifier.

This linter gives you a list of other simp lemmas: look at them!

Here are some tips depending on the error raised by the linter:

  1. 'the left-hand side reduces to XYZ':
     you should probably use XYZ as the left-hand side.

  2. 'simp can prove this':
     This typically means that lemma is a duplicate, or is shadowed by another lemma:

     2a. Always put more general lemmas after specific ones:
      ```
      @[simp] lemma zero_add_zero : 0 + 0 = 0 := rfl
      @[simp] lemma add_zero : x + 0 = x := rfl
      ```

      And not the other way around!  The simplifier always picks the last matching lemma.

     2b. You can also use `@[priority]` instead of moving simp-lemmas around in the file.

      Tip: the default priority is 1000.
      Use `@[priority 1100]` instead of moving a lemma down,
      and `@[priority 900]` instead of moving a lemma up.

     2c. Conditional simp lemmas are tried last. If they are shadowed
         just remove the `simp` attribute.

     2d. If two lemmas are duplicates, the linter will complain about the first one.
         Try to fix the second one instead!
         (You can find it among the other simp lemmas the linter prints out!)

  3. 'try_for tactic failed, timeout':
     This typically means that there is a loop of simp lemmas.
     Try to apply squeeze_simp to the right-hand side (removing this lemma from the simp set) to see
     what lemmas might be causing the loop.

     Another trick is to `set_option trace.simplify.rewrite true` and
     then apply `try_for 10000 { simp }` to the right-hand side.  You will
     see a periodic sequence of lemma applications in the trace message.
-/

/--
A linter for simp lemmas whose lhs has a variable as head symbol,
and which hence never fire.
-/
@[mathlibLinter] def simpVarHead : Linter where
  noErrorsFound :=
    "No left-hand sides of a simp lemma has a variable as head symbol."
  errorsFound := "LEFT-HAND SIDE HAS VARIABLE AS HEAD SYMBOL.
Some simp lemmas have a variable as head symbol of the left-hand side:"
  test := fun declName => do
    unless ← isSimpTheorem declName do return none
    checkAllSimpTheoremInfos (← getConstInfo declName).type fun {lhs, ..} => do
    let headSym := lhs.getAppFn
    unless headSym.isFVar do return none
    return m!"Left-hand side has variable as head symbol: {headSym}"

/-- A linter for commutativity lemmas that are marked simp. -/
@[mathlibLinter] def simpComm : Linter where
  noErrorsFound := "No commutativity lemma is marked simp."
  errorsFound := "COMMUTATIVITY LEMMA IS SIMP.
Some commutativity lemmas are simp lemmas:"
  test := fun declName => withReducible do
    unless ← isSimpTheorem declName do return none
    let ty := (← getConstInfo declName).type
    forallTelescopeReducing ty fun xs ty => do
    let some (_, lhs, rhs) := ty.eq? | return none
    unless lhs.getAppFn.constName? == rhs.getAppFn.constName? do return none
    let (mvars, bis, ty') ← forallMetaTelescopeReducing ty
    let some (_, lhs', rhs') := ty'.eq? | return none
    unless ← isDefEq rhs lhs' do return none
    unless ← withNewMCtxDepth (isDefEq rhs lhs') do return none
    -- ensure that the second application makes progress:
    if ← isDefEq lhs' rhs' then return none
    pure m!"should not be marked simp"


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

      http://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.
 -/
import measure_theory.measurable_space

import measure_theory.measure_space
import measure_theory.outer_measure
import measure_theory.lebesgue_measure
import measure_theory.integration

import measure_theory.borel_space
import data.set.countable
import formal_ml.nnreal
import formal_ml.sum
import formal_ml.lattice
import formal_ml.measurable_space
import formal_ml.classical
import data.equiv.list
import formal_ml.prod_measure
import formal_ml.finite_pi_measure
import formal_ml.probability_space
import formal_ml.random_variable_identical
import formal_ml.prod_probability_space

/-
  This file proves that two IID distributions are identical.

  This result is subtle (i.e., why isn't it just by definition).
  IID variables can be constructed as IID, at which point they are
  obviously IID. However, they can be observed to be IID.

  Notice that, given X_1 and X_2, and Y_1 and Y_2, all coin flips.

  if X_1 is identical to Y_1, and X_2 is identical to Y_2, it is
  not enough to show that every event has the same probability. For
  example, maybe `Pr[ X_1 =ᵣ X_2] ≠ Pr[Y_1 =ᵣ Y_2]`. This kind of
  reasoning gives us a deeper understanding of the proof of VCD, which
  are based upon constructing an IID distribution through shuffling
  and subsampling.

-/

def pi_base {δ:Type*} {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (S:set (Π (a:δ), π a)):Prop := ∃ (f:Π (a:δ), set (π a)), (∀ i, measurable_set (f i)) ∧
  set.pi set.univ f = S

lemma pi_base.closed_inter
  {δ:Type*} {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (S T:set (Π (a:δ), π a)):pi_base S → pi_base T → pi_base (S ∩ T) :=
begin
  intros hS hT,
  cases hS with fS hS,
  cases hT with fT hT,
  apply exists.intro (λ i, (fS i) ∩ (fT i)),
  split,
  { intros i, apply measurable_set.inter, apply hS.left, apply hT.left },
  rw ← hS.right,
  rw ← hT.right,
  ext ω, split; intros h1; simp at h1; simp [h1],
end

lemma pi_base.empty
  {δ:Type*} [nonempty δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]:
  pi_base (∅:set (Π (a : δ), π a)) := begin
  apply exists.intro (λ (d:δ), (∅:set (π d))),
  simp [set.pi],
end

lemma pi_base.univ
  {δ:Type*} [nonempty δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]:
  pi_base (set.univ:set (Π (a : δ), π a)) := begin
  apply exists.intro (λ (d:δ), (set.univ:set (π d))),
  simp [set.pi],
end

lemma pi_base.closed_compl
  {δ:Type*} [fintype δ] [nonempty δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
    (s : set (Π (a : δ), π a)):
    pi_base s →
    (set.disjoint_union_closure pi_base) sᶜ  := begin
  classical,
  intros h1,
  simp [pi_base] at h1,
  cases h1 with f h1,
  let f':(δ → bool) → set (Π (a : δ), π a) :=
    (λ x, if (∀ (d:δ), x d) then ∅  else
   (set.pi (@set.univ δ) (λ (d:δ), if (x d) then (f d) else (f d)ᶜ))),
  begin
    have h2:sᶜ = set.Union f',
    { rw ← h1.right, simp [f'], ext; split; intros h2_1;
      simp at h2_1; simp [h2_1],
      { let i:δ → bool := (λ (d:δ), x d ∈ (f d)),
        apply exists.intro i,
        rw if_neg,
        simp [set.pi, i],
        intros d,
        cases classical.em (x d ∈ f d) with h2_2 h2_2,
        rw if_pos h2_2, apply h2_2,
        rw if_neg h2_2, apply h2_2,
        intros contra,
        cases h2_1 with d h2_1,
        apply h2_1,
        have contra_2 := contra d,
        simp [i] at contra_2,
        apply contra_2 },
      { cases h2_1 with i h2_1,
        cases classical.em (∀ (d : δ), (i d)) with h2_3 h2_3,
        { exfalso, rw if_pos h2_3 at h2_1, apply h2_1 },
        rw if_neg h2_3 at h2_1,
        rw classical.not_forall_iff_exists_not at h2_3,
        cases h2_3 with d h2_3,
        apply exists.intro d,
        simp [set.pi] at h2_1,
        have h2_4 := h2_1 d,
        rw if_neg h2_3 at h2_4,
        apply h2_4 } },
    rw h2,
    apply set.disjoint_union_closure_intro,
    intros b,
    { simp [f',pi_base],
      cases classical.em (∀ (d : δ), (b d)) with h3 h3,
      rw if_pos,
      apply pi_base.empty,
      apply h3,
      rw if_neg,
      { apply exists.intro (λ (d : δ), ite ↥(b d) (f d) (f d)ᶜ),
        rw and.comm,
        split,
        refl,
        intros d,
        simp,
        cases classical.em (b d) with h4 h4,
        rw if_pos,
        apply h1.left,
        apply h4,
        rw if_neg,
        apply measurable_set.compl,
        apply h1.left, apply h4 },
      apply h3 },
    intros b b' h_ne,
    simp [function.on_fun], rw disjoint_iff,
    simp,
    rw ← set.subset_empty_iff,
    rw set.subset_def,
    intros ω h_contra,
    exfalso,
    simp [f'] at h_contra,
    cases h_contra with h_contra_1 h_contra_2,
    cases classical.em (∀ (d : δ), (b d)) with h5 h5,
    { rw if_pos at h_contra_1,
      apply h_contra_1,
      apply h5 },
    rw if_neg h5 at h_contra_1,
    cases classical.em (∀ (d : δ), (b' d)) with h6 h6,
    { rw if_pos h6 at h_contra_2,
      apply h_contra_2 },
    rw if_neg h6 at h_contra_2,
    apply h_ne,
    simp [set.pi] at h_contra_1,
    simp [set.pi] at h_contra_2,
    ext i,
    have h_contra_1_i := h_contra_1 i,
    have h_contra_2_i := h_contra_2 i,
    destruct (b i); destruct (b' i);
    intros h_b'_i h_b_i;
    simp [h_b_i, h_b'_i];
    simp [h_b_i] at h_contra_1_i;
    simp [h_b'_i] at h_contra_2_i,
    apply h_contra_1_i h_contra_2_i,
    apply h_contra_2_i h_contra_1_i,
  end
end

lemma set.pi_as_preimage {δ:Type*} {π:δ → Type*} {f:Π (a:δ), set (π a)}:
  (set.pi set.univ f) = (⋂ (i:δ), (λ (d:Π (a:δ), π a), d i) ⁻¹' (f i)) := begin
  ext,split;intros A1; simp at A1; simp [A1],
end

lemma pi_base.measurable_set {δ:Type*} [encodable δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]
  (S:set (Π (a:δ), π a)):pi_base S → measurable_set S :=
begin
  intros hS,
  cases hS with fS hS,
  cases hS with h_meas h_def,
  rw ← h_def,
  rw set.pi_as_preimage,
  apply measurable_set.Inter,
  intros b,
  have h_proj_meas:measurable (λ (d : Π (a : δ), π a), d b),
  { apply measurable_pi_apply },
  apply h_proj_meas,
  apply h_meas,
end

lemma pi_base.covers_generate {δ:Type*} [encodable δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]
  (S:set (Π (a:δ), π a)):
  S ∈
     (⋃ (i : δ),
       (measurable_space.comap (λ (b : Π (a : δ), (λ (a : δ), π a) a), b i)
          ((λ (a : δ), M a) i)).measurable_set') →
  pi_base S := begin
  intros h1,
  simp [pi_base],
  simp at h1,
  cases h1 with d h1,
  simp [measurable_space.comap] at h1,
  cases h1 with s h1,
  cases h1 with h1 h3,
  have h2:∀ (d':δ), ∃ (T:set (π d')), measurable_set T ∧
     (∀ (s':(Π (a:δ), π a)), s' ∈ S → s' d' ∈ T) ∧ (d' = d → T == s),
  { intros d',
    cases classical.em (d = d') with h2_1 h2_1,
    { subst d', apply exists.intro s,
      split,
      apply h1,
      split,
      intros s' h4, rw ← h3 at h4,
      simp at h4, apply h4, intros h5, refl },
    { apply exists.intro set.univ,
      split,
      apply measurable_set.univ,
      split,
      intros s' h6,
      simp,
      intros h2_1_not, exfalso, apply h2_1, rw h2_1_not, },
   },
   have h7 := classical.axiom_of_choice h2,
   cases h7 with f h7,
   apply exists.intro f,
   split,
   { intros d',
     apply (h7 d').left },
   { ext ω,
     simp [set.pi],split; intros h8,
     { rw ← h3,
       have h9:d = d := rfl,
       have h10 := (h7 d).right.right h9,
       simp at h10, rw ← h10,
       simp, apply h8 },
     { intros d',
       have h11 := (h7 d').right.left ω h8, apply h11 } },
end

lemma pi_base_eq_measurable_space_pi {δ:Type*} [encodable δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]:
 @measurable_space.pi δ π M  = measurable_space.generate_from pi_base :=
begin
  apply le_antisymm;
  apply measurable_space.generate_from_le;
  intros a h_a,
  { simp [measurable_space.generate_from],
    apply measurable_space.generate_measurable.basic,
    apply pi_base.covers_generate,
    simp at h_a, simp, cases h_a with i h_a, apply exists.intro i, apply h_a },
  { apply pi_base.measurable_set, apply h_a },
end

lemma pi_base_eq {δ:Type*} [fintype δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (S:measurable_setB (@measurable_space.pi δ π M)):pi_base S.val →
  ∃ (f:Π (b:δ), measurable_setB (M b)), (set.pi_measurable set.univ f) = S := begin
  intros h1,
  unfold pi_base at h1,
  cases h1 with g h1,
  cases h1 with h1 h2,
  let f := (λ (b:δ), @measurable_setB.mk _ _ (g b) (h1 b)),
  begin
     apply exists.intro f,
     apply subtype.eq,
     rw ← h2,
     simp [set.pi, set.pi_measurable, f],
     ext ω,split;intros h1; simp at h1; simp [h1]; intros i; apply h1,
  end
end

lemma pi_base_is_semialgebra {δ:Type*} [fintype δ] [nonempty δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]:set.is_semialgebra (@pi_base δ π M) := {
  univ := pi_base.univ,
  empty := pi_base.empty,
  compl := pi_base.closed_compl,
  inter := pi_base.closed_inter,
}

lemma pi_union_all {Ω δ:Type*} [fintype δ]
  {P:probability_space Ω}
  {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (f:Π (b:δ), measurable_setB (M b))
  {X:Π (b:δ), P →ᵣ (M b)}:
  (pi.random_variable_combine X ∈ᵣ (set.pi_measurable set.univ f)) =
  (∀ᵣ (b:δ), (X b) ∈ᵣ (f b)) := begin
  apply event.eq,
  simp [pi.random_variable_combine, set.pi_measurable, pi.measurable_fun],
  ext ω, split; intros h1; simp at h1; simp [h1],
end

lemma random_variable_independent_all {Ω β:Type*} {P:probability_space Ω} [fintype β]
  {γ:β → Type*} {M:Π (b:β), measurable_space (γ b)}
   (S:Π (b:β), measurable_setB (M b))
   {X:Π (b:β), P →ᵣ (M b)}:
   random_variable_independent X →
   Pr[∀ᵣ (b : β), X b ∈ᵣ S b] =  finset.univ.prod (λ (b:β), Pr[X b ∈ᵣ S b]) := begin
  intros h1,
  simp [random_variable_independent] at h1,
  have h2 := h1 S,
  simp [independent_events] at h2,
  have h3 := h2 finset.univ,
  rw h3,
  refl,
end

lemma IID_identical_joint_base' {Ω₁ Ω₂ β γ:Type*} [fintype β]
  {P₁:probability_space Ω₁}
  {P₂:probability_space Ω₂}
  {M:measurable_space γ}
  {X₁:β → P₁ →ᵣ M}  {X₂:β  → P₂ →ᵣ M}
  (S:measurable_setB (@measurable_space.pi β (λ _, γ) (λ _, M))):
  pi_base S.val →
  (∀ (b:β), random_variable_identical (X₁ b) (X₂ b)) →
  (random_variables_IID X₁) →
  (random_variables_IID X₂) →
  (Pr[pi.random_variable_combine X₁ ∈ᵣ S] =
   Pr[pi.random_variable_combine X₂ ∈ᵣ S]) :=
begin
  intros h1 h2 h3 h4,
  have h5:=pi_base_eq S h1,
  cases h5 with f h5,
  rw ← h5,
  rw pi_union_all,
  rw pi_union_all,
  rw random_variable_independent_all,
  rw random_variable_independent_all,
  have h6:(λ (b : β), Pr[X₁ b ∈ᵣ f b]) = (λ (b : β), Pr[X₂ b ∈ᵣ f b]),
  { ext1 b, apply h2, },
  rw h6,
  apply h4.left,
  apply h3.left,
end

lemma IID_identical_joint' {Ω₁ Ω₂ β γ:Type*} [fintype β] [nonempty β]
  {P₁:probability_space Ω₁}
  {P₂:probability_space Ω₂}
  {M:measurable_space γ}
  {X₁:β → P₁ →ᵣ M}  {X₂:β  → P₂ →ᵣ M}:
  (∀ (b:β), random_variable_identical (X₁ b) (X₂ b)) →
  (random_variables_IID X₁) →
  (random_variables_IID X₂) →
  random_variable_identical (pi.random_variable_combine X₁)
                            (pi.random_variable_combine X₂)
  :=
begin
  intros h1 h2 h3,
  haveI:encodable β := fintype.encodable β,
  apply random_variable_identical_on_semialgebra' pi_base,
  --apply pi_base
  apply pi_base.closed_inter,
  apply pi_base.closed_compl,
  apply pi_base.empty,
  apply pi_base.univ,
  {
    apply pi_base_eq_measurable_space_pi },
  intros T h_T,
  apply IID_identical_joint_base',
  apply h_T,
  apply h1,
  apply h2,
  apply h3,
end

lemma IID_identical_joint {Ω₁ Ω₂ β γ:Type*} [fintype β] [nonempty β]
  {P₁:probability_space Ω₁}
  {P₂:probability_space Ω₂}
  {M:measurable_space γ}
  {X₁:β → P₁ →ᵣ M}  {X₂:β  → P₂ →ᵣ M}:
  (∃ (b₁ b₂:β), random_variable_identical (X₁ b₁) (X₂ b₂)) →
  (random_variables_IID X₁) →
  (random_variables_IID X₂) →
  random_variable_identical (pi.random_variable_combine X₁)
                            (pi.random_variable_combine X₂) :=
begin
  intros h1 h2 h3,
  apply IID_identical_joint',
  intros b,
  cases h1 with b₁ h1,
  cases h1 with b₂ h1,
  have h4:random_variable_identical (X₁ b) (X₁ b₁),
  { apply h2.right },
  apply random_variable_identical.trans h4,
  apply random_variable_identical.trans h1,
  apply h3.right,
  apply h2,
  apply h3,
end


/-
Copyright © 2022 François G. Dorais, Kyrill Serdyuk, Emma Shroyer. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
-/

import Std
import UnicodeBasic

structure ByteSubarray extends ByteArray where
  start : Nat
  size  : Nat
  valid : start + size ≤ toByteArray.size := by simp_arith [*]

abbrev ByteSubarray.stop (bs : ByteSubarray) : Nat := bs.start + bs.size

def ByteArray.toByteSubarray (bs : ByteArray) : ByteSubarray where
  toByteArray := bs
  start := 0
  size := bs.size

namespace ByteSubarray
variable (bs : ByteSubarray)

@[inline] def extract : ByteArray := bs.toByteArray.extract bs.start bs.stop

@[inline] def get (i : Fin bs.size) : UInt8 :=
  have : bs.start + i.val < bs.toByteArray.size := calc
    _ < bs.start + bs.size := Nat.add_lt_add_left i.isLt bs.start
    _ ≤ bs.toByteArray.size := bs.valid
  bs.toByteArray[bs.start + i.val]'this

instance : GetElem ByteSubarray Nat UInt8 fun bs i => i < bs.size where
  getElem xs i h := xs.get ⟨i, h⟩

@[inline] def get? (i : Nat) : Option UInt8 :=
  if h : i < bs.size then some (bs.get ⟨i, h⟩) else none

@[inline] def getD (i : Nat) (default : UInt8) : UInt8 :=
  if h : i < bs.size then bs.get ⟨i, h⟩ else default

abbrev get! (i : Nat) : UInt8 := getD bs i default

def popFront : ByteSubarray :=
  if h : bs.size ≥ 1 then
    have : (bs.start+1) + (bs.size-1) = bs.start + bs.size := by
      rw [Nat.add_assoc, Nat.add_sub_cancel' h]
    {bs with start := bs.start+1, size := bs.size-1, valid := by rw [this]; exact bs.valid}
  else bs

def slice (start stop : Nat) (h : start ≤ stop ∧ stop ≤ bs.size := by simp_arith [*]) : ByteSubarray where
  toByteArray := bs.toByteArray
  start := bs.start + start
  size := stop - start
  valid := by
    rw [Nat.add_assoc]
    rw [Nat.add_sub_cancel' h.1]
    apply Nat.le_trans _ bs.valid
    apply Nat.add_le_add_left h.2

@[simp] theorem size_slice (start stop : Nat) (h : start ≤ stop ∧ stop ≤ bs.size := by simp_arith [*]) : (bs.slice start stop h).size = stop - start := rfl

@[inline] unsafe def forInUnsafe {α m} [Monad m] (bs : ByteSubarray) (a : α) (f : UInt8 → α → m (ForInStep α)) : m α :=
  let sz := USize.ofNat bs.stop
  let rec @[specialize] loop (i : USize) (a : α) : m α := do
    if i < sz then
      let b := bs.uget i lcProof
      match (← f b a) with
      | ForInStep.done  a => pure a
      | ForInStep.yield a => loop (i+1) a
    else
      pure a
  loop (USize.ofNat bs.start) a

@[implemented_by ByteSubarray.forInUnsafe]
protected def forIn {α m} [Monad m] (bs : ByteSubarray) (a : α) (f : UInt8 → α → m (ForInStep α)) : m α :=
  let rec @[specialize] loop (i : Nat) (a : α) : m α := do
    if h : i < bs.size then
      let b := bs.get ⟨i, h⟩
      match (← f b a) with
      | ForInStep.done  a => pure a
      | ForInStep.yield a => loop (i+1) a
    else
      pure a
  loop bs.start a
termination_by loop => bs.size - i

instance : ForIn m ByteSubarray UInt8 where
  forIn := ByteSubarray.forIn

instance : Stream ByteSubarray UInt8 where
  next? s :=
    if h : s.size > 0 then
      some (s.get ⟨0, h⟩, s.popFront)
    else
      none

end ByteSubarray


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

! This file was ported from Lean 3 source module init.meta.environment
! leanprover-community/mathlib commit 1340477dccb7fbe0cf2146aa1f1995022c13cd30
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
prelude
import Leanbin.Init.Meta.Declaration
import Leanbin.Init.Meta.Exceptional
import Leanbin.Init.Data.Option.Basic
import Leanbin.Init.Meta.RbMap

/--
An __environment__ contains all of the declarations and notation that have been defined so far.   -/
unsafe axiom environment : Type
#align environment environment

namespace Environment

/--
Consider a type `ψ` which is an inductive datatype using a single constructor `mk (a : α) (b : β) : ψ`.
Lean will automatically make two projection functions `a : ψ → α`, `b : ψ → β`.
Lean tags these declarations as __projections__.
This helps the simplifier / rewriter not have to expand projectors.
Eg `a (mk x y)` will automatically reduce to `x`.
If you `extend` a structure, all of the projections on the parent will also be created for the child.
Projections are also treated differently in the VM for efficiency.

Note that projections have nothing to do with the dot `mylist.map` syntax.

You can find out if a declaration is a projection using `environment.is_projection` which returns `projection_info`.

Data for a projection declaration:
- `cname`    is the name of the constructor associated with the projection.
- `nparams`  is the number of constructor parameters. Eg `and.intro` has two type parameters.
- `idx`      is the parameter being projected by this projection.
- `is_class` is tt iff this is a typeclass projection.

### Examples:

- `and.right` is a projection with ``{cname := `and.intro, nparams := 2, idx := 1, is_class := ff}``
- `ordered_ring.neg` is a projection with ``{cname := `ordered_ring.mk, nparams := 1, idx := 5, is_class := tt}``.

-/
structure ProjectionInfo where
  cname : Name
  nparams : Nat
  idx : Nat
  isClass : Bool
#align environment.projection_info Environment.ProjectionInfo

/-- A marking on the binders of structures and inductives indicating
   how this constructor should mark its parameters.

       inductive foo
       | one {} : foo -> foo   -- relaxed_implicit
       | two ( ) : foo -> foo  -- explicit
       | two [] : foo -> foo   -- implicit
       | three : foo -> foo    -- relaxed implicit (default)
-/
inductive ImplicitInferKind
  | implicit
  | relaxed_implicit
  | none
#align environment.implicit_infer_kind Environment.ImplicitInferKind

instance ImplicitInferKind.inhabited : Inhabited ImplicitInferKind :=
  ⟨ImplicitInferKind.implicit⟩
#align environment.implicit_infer_kind.inhabited Environment.ImplicitInferKind.inhabited

/-- One introduction rule in an inductive declaration -/
unsafe structure intro_rule where
  constr : Name
  type : expr
  infer : ImplicitInferKind := ImplicitInferKind.implicit
#align environment.intro_rule environment.intro_rule

/-- Create a standard environment using the given trust level -/
unsafe axiom mk_std : Nat → environment
#align environment.mk_std environment.mk_std

/-- Return the trust level of the given environment -/
unsafe axiom trust_lvl : environment → Nat
#align environment.trust_lvl environment.trust_lvl

/-- Add a new declaration to the environment -/
unsafe axiom add : environment → declaration → exceptional environment
#align environment.add environment.add

/-- make declaration `n` protected -/
unsafe axiom mk_protected : environment → Name → environment
#align environment.mk_protected environment.mk_protected

/-- add declaration `d` and make it protected -/
unsafe def add_protected (env : environment) (d : declaration) : exceptional environment := do
  let env ← env.add d
  pure <| env d
#align environment.add_protected environment.add_protected

/-- check if `n` is the name of a protected declaration -/
unsafe axiom is_protected : environment → Name → Bool
#align environment.is_protected environment.is_protected

/-- Retrieve a declaration from the environment -/
unsafe axiom get : environment → Name → exceptional declaration
#align environment.get environment.get

unsafe def contains (env : environment) (d : Name) : Bool :=
  match env.get d with
  | exceptional.success _ => true
  | exceptional.exception _ => false
#align environment.contains environment.contains

unsafe axiom add_defn_eqns (env : environment) (opt : options) (lp_params : List Name)
    (params : List expr) (sig : expr) (eqns : List (List (expr false) × expr)) (is_meta : Bool) :
    exceptional environment
#align environment.add_defn_eqns environment.add_defn_eqns

/-- Register the given name as a namespace, making it available to the `open` command -/
unsafe axiom add_namespace : environment → Name → environment
#align environment.add_namespace environment.add_namespace

/-- Mark a namespace as open -/
unsafe axiom mark_namespace_as_open : environment → Name → environment
#align environment.mark_namespace_as_open environment.mark_namespace_as_open

/-- Modify the environment as if `open %%name` had been parsed -/
unsafe axiom execute_open : environment → Name → environment
#align environment.execute_open environment.execute_open

/-- Retrieve all registered namespaces -/
unsafe axiom get_namespaces : environment → List Name
#align environment.get_namespaces environment.get_namespaces

/-- Return tt iff the given name is a namespace -/
unsafe axiom is_namespace : environment → Name → Bool
#align environment.is_namespace environment.is_namespace

/-- Add a new inductive datatype to the environment
   name, universe parameters, number of parameters, type, constructors (name and type), is_meta -/
unsafe axiom add_inductive (env : environment) (n : Name) (levels : List Name) (num_params : Nat)
    (type : expr) (intros : List (Name × expr)) (is_meta : Bool) : exceptional environment
#align environment.add_inductive environment.add_inductive

/-- Add a new general inductive declaration to the environment.
  This has the same effect as a `inductive` in the file, including generating
  all the auxiliary definitions, as well as triggering mutual/nested inductive
  compilation, by contrast to `environment.add_inductive` which only adds the
  core axioms supported by the kernel.

  The `inds` argument should be a list of inductives in the mutual family.
  The first argument is a pair of the name of the type being constructed
  and the type of this inductive family (not including the params).
  The second argument is a list of intro rules, specified by a name, an
  `implicit_infer_kind` giving the implicitness of the params for this constructor,
  and an expression with the type of the constructor (not including the params).
-/
unsafe axiom add_ginductive (env : environment) (opt : options) (levels : List Name)
    (params : List expr) (inds : List ((Name × expr) × List intro_rule)) (is_meta : Bool) :
    exceptional environment
#align environment.add_ginductive environment.add_ginductive

/-- Return tt iff the given name is an inductive datatype -/
unsafe axiom is_inductive : environment → Name → Bool
#align environment.is_inductive environment.is_inductive

/-- Return tt iff the given name is a constructor -/
unsafe axiom is_constructor : environment → Name → Bool
#align environment.is_constructor environment.is_constructor

/-- Return tt iff the given name is a recursor -/
unsafe axiom is_recursor : environment → Name → Bool
#align environment.is_recursor environment.is_recursor

/-- Return tt iff the given name is a recursive inductive datatype -/
unsafe axiom is_recursive : environment → Name → Bool
#align environment.is_recursive environment.is_recursive

/-- Return the name of the inductive datatype of the given constructor. -/
unsafe axiom inductive_type_of : environment → Name → Option Name
#align environment.inductive_type_of environment.inductive_type_of

/-- Return the constructors of the inductive datatype with the given name -/
unsafe axiom constructors_of : environment → Name → List Name
#align environment.constructors_of environment.constructors_of

/-- Return the recursor of the given inductive datatype -/
unsafe axiom recursor_of : environment → Name → Option Name
#align environment.recursor_of environment.recursor_of

/-- Return the number of parameters of the inductive datatype -/
unsafe axiom inductive_num_params : environment → Name → Nat
#align environment.inductive_num_params environment.inductive_num_params

/-- Return the number of indices of the inductive datatype -/
unsafe axiom inductive_num_indices : environment → Name → Nat
#align environment.inductive_num_indices environment.inductive_num_indices

/-- Return tt iff the inductive datatype recursor supports dependent elimination -/
unsafe axiom inductive_dep_elim : environment → Name → Bool
#align environment.inductive_dep_elim environment.inductive_dep_elim

/-- Functionally equivalent to `is_inductive`.

Technically, this works by checking if the name is in the ginductive environment
extension which is outside the kernel, whereas `is_inductive` works by looking at the kernel extension.
But there are no `is_inductive`s which are not `is_ginductive`.
 -/
unsafe axiom is_ginductive : environment → Name → Bool
#align environment.is_ginductive environment.is_ginductive

/-- See the docstring for `projection_info`. -/
unsafe axiom is_projection : environment → Name → Option ProjectionInfo
#align environment.is_projection environment.is_projection

/-- Fold over declarations in the environment. -/
unsafe axiom fold {α : Type} : environment → α → (declaration → α → α) → α
#align environment.fold environment.fold

/-- `relation_info env n` returns some value if n is marked as a relation in the given environment.
   the tuple contains: total number of arguments of the relation, lhs position and rhs position. -/
unsafe axiom relation_info : environment → Name → Option (Nat × Nat × Nat)
#align environment.relation_info environment.relation_info

/-- `refl_for env R` returns the name of the reflexivity theorem for the relation R -/
unsafe axiom refl_for : environment → Name → Option Name
#align environment.refl_for environment.refl_for

/-- `symm_for env R` returns the name of the symmetry theorem for the relation R -/
unsafe axiom symm_for : environment → Name → Option Name
#align environment.symm_for environment.symm_for

/-- `trans_for env R` returns the name of the transitivity theorem for the relation R -/
unsafe axiom trans_for : environment → Name → Option Name
#align environment.trans_for environment.trans_for

/-- `decl_olean env d` returns the name of the .olean file where d was defined.
   The result is none if d was not defined in an imported file. -/
unsafe axiom decl_olean : environment → Name → Option String
#align environment.decl_olean environment.decl_olean

/-- `decl_pos env d` returns the source location of d if available. -/
unsafe axiom decl_pos : environment → Name → Option Pos
#align environment.decl_pos environment.decl_pos

/-- `decl_pos env d` returns the name of a declaration that d inherits
noncomputability from, or `none` if it is computable.

Note that this also returns `none` on `axiom`s and `constant`s. These can be detected by using
`environment.get_decl` and `declaration.is_axiom` and `declaration.is_constant`. -/
unsafe axiom decl_noncomputable_reason : environment → Name → Option Name
#align environment.decl_noncomputable_reason environment.decl_noncomputable_reason

/-- Return the fields of the structure with the given name, or `none` if it is not a structure -/
unsafe axiom structure_fields : environment → Name → Option (List Name)
#align environment.structure_fields environment.structure_fields

/-- `get_class_attribute_symbols env attr_name` return symbols
   occurring in instances of type classes tagged with the attribute `attr_name`.
   Example: [algebra] -/
unsafe axiom get_class_attribute_symbols : environment → Name → name_set
#align environment.get_class_attribute_symbols environment.get_class_attribute_symbols

/--
The fingerprint of the environment is a hash formed from all of the declarations in the environment. -/
unsafe axiom fingerprint : environment → Nat
#align environment.fingerprint environment.fingerprint

/-- Gets the equation lemmas for the declaration `n`. -/
unsafe axiom get_eqn_lemmas_for (env : environment) (n : Name) : List Name
#align environment.get_eqn_lemmas_for environment.get_eqn_lemmas_for

/-- Gets the equation lemmas for the declaration `n`, including lemmas for match statements, etc. -/
unsafe axiom get_ext_eqn_lemmas_for (env : environment) (n : Name) : List Name
#align environment.get_ext_eqn_lemmas_for environment.get_ext_eqn_lemmas_for

/-- Adds the equation lemma `n`.
It is added for the declaration `t.pi_codomain.get_app_fn.const_name` where `t` is the type of the equation lemma.
-/
unsafe axiom add_eqn_lemma (env : environment) (n : Name) : environment
#align environment.add_eqn_lemma environment.add_eqn_lemma

open Expr

unsafe axiom unfold_untrusted_macros : environment → expr → expr
#align environment.unfold_untrusted_macros environment.unfold_untrusted_macros

unsafe axiom unfold_all_macros : environment → expr → expr
#align environment.unfold_all_macros environment.unfold_all_macros

unsafe def is_constructor_app (env : environment) (e : expr) : Bool :=
  is_constant (get_app_fn e) && is_constructor env (const_name (get_app_fn e))
#align environment.is_constructor_app environment.is_constructor_app

unsafe def is_refl_app (env : environment) (e : expr) : Option (Name × expr × expr) :=
  match refl_for env (const_name (get_app_fn e)) with
  | some n => if get_app_num_args e ≥ 2 then some (n, app_arg (app_fn e), app_arg e) else none
  | none => none
#align environment.is_refl_app environment.is_refl_app

/-- Return true if 'n' has been declared in the current file -/
unsafe def in_current_file (env : environment) (n : Name) : Bool :=
  (env.decl_olean n).isNone && env.contains n &&
    n ∉ [`` Quot, `` Quot.mk, `` Quot.lift, `` Quot.ind]
#align environment.in_current_file environment.in_current_file

unsafe def is_definition (env : environment) (n : Name) : Bool :=
  match env.get n with
  | exceptional.success (declaration.defn _ _ _ _ _ _) => true
  | _ => false
#align environment.is_definition environment.is_definition

end Environment

unsafe instance : Repr environment :=
  ⟨fun e => "[environment]"⟩

unsafe instance : Inhabited environment :=
  ⟨environment.mk_std 0⟩


/-
Copyright (c) 2014 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import Init.Core
import Init.Control.Basic
import Init.Coe

namespace Option

def toMonad [Monad m] [Alternative m] : Option α → m α
  | none     => failure
  | some a   => pure a

@[inline] def toBool : Option α → Bool
  | some _ => true
  | none   => false

@[inline] def isSome : Option α → Bool
  | some _ => true
  | none   => false

@[inline] def isNone : Option α → Bool
  | some _ => false
  | none   => true

@[inline] def isEqSome [BEq α] : Option α → α → Bool
  | some a, b => a == b
  | none,   _ => false

@[inline] protected def bind : Option α → (α → Option β) → Option β
  | none,   b => none
  | some a, b => b a

@[inline] protected def map (f : α → β) (o : Option α) : Option β :=
  Option.bind o (some ∘ f)

theorem map_id : (Option.map id : Option α → Option α) = id :=
  funext (fun o => match o with | none => rfl | some x => rfl)

instance : Functor Option where
  map := Option.map

@[inline] protected def filter (p : α → Bool) : Option α → Option α
  | some a => if p a then some a else none
  | none   => none

@[inline] protected def all (p : α → Bool) : Option α → Bool
  | some a => p a
  | none   => true

@[inline] protected def any (p : α → Bool) : Option α → Bool
  | some a => p a
  | none   => false

@[macroInline] protected def orElse : Option α → (Unit → Option α) → Option α
  | some a, _ => some a
  | none,   b => b ()

instance : OrElse (Option α) where
  orElse := Option.orElse

@[inline] protected def lt (r : α → α → Prop) : Option α → Option α → Prop
  | none, some x     => True
  | some x,   some y => r x y
  | _, _             => False

instance (r : α → α → Prop) [s : DecidableRel r] : DecidableRel (Option.lt r)
  | none,   some y => isTrue  trivial
  | some x, some y => s x y
  | some x, none   => isFalse not_false
  | none,   none   => isFalse not_false

end Option

deriving instance DecidableEq for Option
deriving instance BEq for Option

instance [LT α] : LT (Option α) where
  lt := Option.lt (· < ·)


example (P Q R : Type) : (P → (Q → R)) → ((P → Q) → (P → R)) := by
  intro p_qr
  intro p_q
  intro p
  apply p_qr
  case a =>
    exact p
  case a =>
    apply p_q
    exact p

example (P Q R : Type) : (P → (Q → R)) → ((P → Q) → (P → R)) := by
  intro p_qr
  intro p_q
  intro p
  apply p_qr
  . exact p
  apply p_q
  exact p