/-
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.Parser.Term
import Lean.Meta.Closure
import Lean.Meta.Check
import Lean.Elab.Command
import Lean.Elab.Match
import Lean.Elab.DefView
import Lean.Elab.PreDefinition
import Lean.Elab.DeclarationRange

namespace Lean.Elab
open Lean.Parser.Term

/- DefView after elaborating the header. -/
structure DefViewElabHeader where
  ref           : Syntax
  modifiers     : Modifiers
  kind          : DefKind
  shortDeclName : Name
  declName      : Name
  levelNames    : List Name
  binderIds     : Array Syntax
  numParams     : Nat
  type          : Expr -- including the parameters
  valueStx      : Syntax
  deriving Inhabited

namespace Term
open Meta

private def checkModifiers (m₁ m₂ : Modifiers) : TermElabM Unit := do
  unless m₁.isUnsafe == m₂.isUnsafe do
    throwError "cannot mix unsafe and safe definitions"
  unless m₁.isNoncomputable == m₂.isNoncomputable do
    throwError "cannot mix computable and non-computable definitions"
  unless m₁.isPartial == m₂.isPartial do
    throwError "cannot mix partial and non-partial definitions"

private def checkKinds (k₁ k₂ : DefKind) : TermElabM Unit := do
  unless k₁.isExample == k₂.isExample do
    throwError "cannot mix examples and definitions" -- Reason: we should discard examples
  unless k₁.isTheorem == k₂.isTheorem do
    throwError "cannot mix theorems and definitions" -- Reason: we will eventually elaborate theorems in `Task`s.

private def check (prevHeaders : Array DefViewElabHeader) (newHeader : DefViewElabHeader) : TermElabM Unit := do
  if newHeader.kind.isTheorem && newHeader.modifiers.isUnsafe then
    throwError "'unsafe' theorems are not allowed"
  if newHeader.kind.isTheorem && newHeader.modifiers.isPartial then
    throwError "'partial' theorems are not allowed, 'partial' is a code generation directive"
  if newHeader.kind.isTheorem && newHeader.modifiers.isNoncomputable then
    throwError "'theorem' subsumes 'noncomputable', code is not generated for theorems"
  if newHeader.modifiers.isNoncomputable && newHeader.modifiers.isUnsafe then
    throwError "'noncomputable unsafe' is not allowed"
  if newHeader.modifiers.isNoncomputable && newHeader.modifiers.isPartial then
    throwError "'noncomputable partial' is not allowed"
  if newHeader.modifiers.isPartial && newHeader.modifiers.isUnsafe then
    throwError "'unsafe' subsumes 'partial'"
  if h : 0 < prevHeaders.size then
    let firstHeader := prevHeaders.get ⟨0, h⟩
    try
      unless newHeader.levelNames == firstHeader.levelNames do
        throwError "universe parameters mismatch"
      checkModifiers newHeader.modifiers firstHeader.modifiers
      checkKinds newHeader.kind firstHeader.kind
    catch
       | Exception.error ref msg => throw (Exception.error ref m!"invalid mutually recursive definitions, {msg}")
       | ex => throw ex
  else
    pure ()

private def registerFailedToInferDefTypeInfo (type : Expr) (ref : Syntax) : TermElabM Unit :=
  registerCustomErrorIfMVar type ref "failed to infer definition type"

/--
  Return `some [b, c]` if the given `views` are representing a declaration of the form
  ```
  constant a b c : Nat
  ```  -/
private def isMultiConstant? (views : Array DefView) : Option (List Name) :=
  if views.size == 1 &&
     views[0].kind == DefKind.opaque &&
     views[0].binders.getArgs.size > 0 &&
     views[0].binders.getArgs.all (·.getKind == ``Parser.Term.simpleBinder) then
    some <| (views[0].binders.getArgs.toList.map (fun stx => stx[0].getArgs.toList.map (·.getId))).join
  else
    none

private def getPendindMVarErrorMessage (views : Array DefView) : String :=
  match isMultiConstant? views with
  | some ids =>
    let idsStr := ", ".intercalate <| ids.map fun id => s!"`{id}`"
    let paramsStr := ", ".intercalate <| ids.map fun id => s!"`({id} : _)`"
    s!"\nrecall that you cannot declare multiple constants in a single declaration. The identifier(s) {idsStr} are being interpreted as parameters {paramsStr}"
  | none =>
    "\nwhen the resulting type of a declaration is explicitly provided, all holes (e.g., `_`) in the header are resolved before the declaration body is processed"

private def elabHeaders (views : Array DefView) : TermElabM (Array DefViewElabHeader) := do
  let mut headers := #[]
  for view in views do
    let newHeader ← withRef view.ref do
      let ⟨shortDeclName, declName, levelNames⟩ ← Term.expandDeclId (← getCurrNamespace) (← getLevelNames) view.declId view.modifiers
      addDeclarationRanges declName view.ref
      applyAttributesAt declName view.modifiers.attrs AttributeApplicationTime.beforeElaboration
      withDeclName declName <| withAutoBoundImplicit <| withLevelNames levelNames <|
        elabBindersEx view.binders.getArgs fun xs => do
          let refForElabFunType := view.value
          let type ← match view.type? with
            | some typeStx =>
              let type ← elabType typeStx
              registerFailedToInferDefTypeInfo type typeStx
              pure type
            | none =>
              let hole := mkHole refForElabFunType
              let type ← elabType hole
              registerFailedToInferDefTypeInfo type refForElabFunType
              pure type
          Term.synthesizeSyntheticMVarsNoPostponing
          let (binderIds, xs) := xs.unzip
          let type ← mkForallFVars xs type
          let type ← mkForallFVars (← read).autoBoundImplicits.toArray type
          let type ← instantiateMVars type
          let xs ← addAutoBoundImplicits xs
          let levelNames ← getLevelNames
          if view.type?.isSome then
            let pendingMVarIds ← getMVars type
            discard <| logUnassignedUsingErrorInfos pendingMVarIds <|
              getPendindMVarErrorMessage views
          let newHeader := {
            ref           := view.ref,
            modifiers     := view.modifiers,
            kind          := view.kind,
            shortDeclName := shortDeclName,
            declName      := declName,
            levelNames    := levelNames,
            binderIds     := binderIds,
            numParams     := xs.size,
            type          := type,
            valueStx      := view.value : DefViewElabHeader }
          check headers newHeader
          pure newHeader
    headers := headers.push newHeader
  pure headers

private partial def withFunLocalDecls {α} (headers : Array DefViewElabHeader) (k : Array Expr → TermElabM α) : TermElabM α :=
  let rec loop (i : Nat) (fvars : Array Expr) := do
    if h : i < headers.size then
      let header := headers.get ⟨i, h⟩
      if header.modifiers.isNonrec then
        loop (i+1) fvars
      else
        withLocalDecl header.shortDeclName BinderInfo.auxDecl header.type fun fvar => loop (i+1) (fvars.push fvar)
    else
      k fvars
  loop 0 #[]

private def expandWhereStructInst : Macro
  | `(Parser.Command.whereStructInst|where $[$decls:letDecl$[;]?]*) => do
    let letIdDecls ← decls.mapM fun stx => match stx with
      | `(letDecl|$decl:letPatDecl)  => Macro.throwErrorAt stx "patterns are not allowed here"
      | `(letDecl|$decl:letEqnsDecl) => expandLetEqnsDecl decl
      | `(letDecl|$decl:letIdDecl)   => pure decl
      | _                               => Macro.throwUnsupported
    let structInstFields ← letIdDecls.mapM fun
      | stx@`(letIdDecl|$id:ident $[$binders]* $[: $ty?]? := $val) => withRef stx do
        let mut val := val
        if let some ty := ty? then
          val ← `(($val : $ty))
        val ← if binders.size > 0 then `(fun $[$binders]* => $val:term) else pure val
        `(structInstField|$id:ident := $val)
      | _ => Macro.throwUnsupported
    `({ $[$structInstFields,]* })
  | _ => Macro.throwUnsupported

/-
Recall that
```
def declValSimple    := leading_parser " :=\n" >> termParser >> optional Term.whereDecls
def declValEqns      := leading_parser Term.matchAltsWhereDecls
def declVal          := declValSimple <|> declValEqns <|> Term.whereDecls
```
-/
private def declValToTerm (declVal : Syntax) : MacroM Syntax := withRef declVal do
  if declVal.isOfKind ``Lean.Parser.Command.declValSimple then
    expandWhereDeclsOpt declVal[2] declVal[1]
  else if declVal.isOfKind ``Lean.Parser.Command.declValEqns then
    expandMatchAltsWhereDecls declVal[0]
  else if declVal.isOfKind ``Lean.Parser.Command.whereStructInst then
    expandWhereStructInst declVal
  else if declVal.isMissing then
    Macro.throwErrorAt declVal "declaration body is missing"
  else
    Macro.throwErrorAt declVal "unexpected declaration body"

private def elabFunValues (headers : Array DefViewElabHeader) : TermElabM (Array Expr) :=
  headers.mapM fun header => withDeclName header.declName $ withLevelNames header.levelNames do
    let valStx ← liftMacroM $ declValToTerm header.valueStx
    forallBoundedTelescope header.type header.numParams fun xs type => do
      -- Add new info nodes for new fvars. The server will detect all fvars of a binder by the binder's source location.
      for i in [0:header.binderIds.size] do
        -- skip auto-bound prefix in `xs`
        addTermInfo (isBinder := true) header.binderIds[i] xs[header.numParams - header.binderIds.size + i]
      let val ← elabTermEnsuringType valStx type
      mkLambdaFVars xs val

private def collectUsed (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
    : StateRefT CollectFVars.State MetaM Unit := do
  headers.forM fun header => collectUsedFVars header.type
  values.forM collectUsedFVars
  toLift.forM fun letRecToLift => do
    collectUsedFVars letRecToLift.type
    collectUsedFVars letRecToLift.val

private def removeUnusedVars (vars : Array Expr) (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
    : TermElabM (LocalContext × LocalInstances × Array Expr) := do
  let (_, used) ← (collectUsed headers values toLift).run {}
  removeUnused vars used

private def withUsed {α} (vars : Array Expr) (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
    (k : Array Expr → TermElabM α) : TermElabM α := do
  let (lctx, localInsts, vars) ← removeUnusedVars vars headers values toLift
  withLCtx lctx localInsts $ k vars

private def isExample (views : Array DefView) : Bool :=
  views.any (·.kind.isExample)

private def isTheorem (views : Array DefView) : Bool :=
  views.any (·.kind.isTheorem)

private def instantiateMVarsAtHeader (header : DefViewElabHeader) : TermElabM DefViewElabHeader := do
  let type ← instantiateMVars header.type
  pure { header with type := type }

private def instantiateMVarsAtLetRecToLift (toLift : LetRecToLift) : TermElabM LetRecToLift := do
  let type ← instantiateMVars toLift.type
  let val ← instantiateMVars toLift.val
  pure { toLift with type := type, val := val }

private def typeHasRecFun (type : Expr) (funFVars : Array Expr) (letRecsToLift : List LetRecToLift) : Option FVarId :=
  let occ? := type.find? fun e => match e with
    | Expr.fvar fvarId _ => funFVars.contains e || letRecsToLift.any fun toLift => toLift.fvarId == fvarId
    | _ => false
  match occ? with
  | some (Expr.fvar fvarId _) => some fvarId
  | _ => none

private def getFunName (fvarId : FVarId) (letRecsToLift : List LetRecToLift) : TermElabM Name := do
  match (← findLocalDecl? fvarId) with
  | some decl => pure decl.userName
  | none =>
    /- Recall that the FVarId of nested let-recs are not in the current local context. -/
    match letRecsToLift.findSome? fun toLift => if toLift.fvarId == fvarId then some toLift.shortDeclName else none with
    | none   => throwError "unknown function"
    | some n => pure n

/-
Ensures that the of let-rec definition types do not contain functions being defined.
In principle, this test can be improved. We could perform it after we separate the set of functions is strongly connected components.
However, this extra complication doesn't seem worth it.
-/
private def checkLetRecsToLiftTypes (funVars : Array Expr) (letRecsToLift : List LetRecToLift) : TermElabM Unit :=
  letRecsToLift.forM fun toLift =>
    match typeHasRecFun toLift.type funVars letRecsToLift with
    | none        => pure ()
    | some fvarId => do
      let fnName ← getFunName fvarId letRecsToLift
      throwErrorAt toLift.ref "invalid type in 'let rec', it uses '{fnName}' which is being defined simultaneously"

namespace MutualClosure

/- A mapping from FVarId to Set of FVarIds. -/
abbrev UsedFVarsMap := FVarIdMap FVarIdSet

/-
Create the `UsedFVarsMap` mapping that takes the variable id for the mutually recursive functions being defined to the set of
free variables in its definition.

For `mainFVars`, this is just the set of section variables `sectionVars` used.
For nested let-rec functions, we collect their free variables.

Recall that a `let rec` expressions are encoded as follows in the elaborator.
```lean
let rec
  f : A := t,
  g : B := s;
body
```
is encoded as
```lean
let f : A := ?m₁;
let g : B := ?m₂;
body
```
where `?m₁` and `?m₂` are synthetic opaque metavariables. That are assigned by this module.
We may have nested `let rec`s.
```lean
let rec f : A :=
    let rec g : B := t;
    s;
body
```
is encoded as
```lean
let f : A := ?m₁;
body
```
and the body of `f` is stored the field `val` of a `LetRecToLift`. For the example above,
we would have a `LetRecToLift` containing:
```
{
  mvarId := m₁,
  val    := `(let g : B := ?m₂; body)
  ...
}
```
Note that `g` is not a free variable at `(let g : B := ?m₂; body)`. We recover the fact that
`f` depends on `g` because it contains `m₂`
-/
private def mkInitialUsedFVarsMap (mctx : MetavarContext) (sectionVars : Array Expr) (mainFVarIds : Array FVarId) (letRecsToLift : List LetRecToLift)
    : UsedFVarsMap := Id.run <| do
  let mut sectionVarSet := {}
  for var in sectionVars do
    sectionVarSet := sectionVarSet.insert var.fvarId!
  let mut usedFVarMap := {}
  for mainFVarId in mainFVarIds do
    usedFVarMap := usedFVarMap.insert mainFVarId sectionVarSet
  for toLift in letRecsToLift do
    let state := Lean.collectFVars {} toLift.val
    let state := Lean.collectFVars state toLift.type
    let mut set := state.fvarSet
    /- toLift.val may contain metavariables that are placeholders for nested let-recs. We should collect the fvarId
       for the associated let-rec because we need this information to compute the fixpoint later. -/
    let mvarIds := (toLift.val.collectMVars {}).result
    for mvarId in mvarIds do
      match letRecsToLift.findSome? fun (toLift : LetRecToLift) => if toLift.mvarId == mctx.getDelayedRoot mvarId then some toLift.fvarId else none with
      | some fvarId => set := set.insert fvarId
      | none        => pure ()
    usedFVarMap := usedFVarMap.insert toLift.fvarId set
  pure usedFVarMap

/-
The let-recs may invoke each other. Example:
```
let rec
  f (x : Nat) := g x + y
  g : Nat → Nat
    | 0   => 1
    | x+1 => f x + z
```
`y` is free variable in `f`, and `z` is a free variable in `g`.
To close `f` and `g`, `y` and `z` must be in the closure of both.
That is, we need to generate the top-level definitions.
```
def f (y z x : Nat) := g y z x + y
def g (y z : Nat) : Nat → Nat
  | 0 => 1
  | x+1 => f y z x + z
```
-/
namespace FixPoint

structure State where
  usedFVarsMap : UsedFVarsMap := {}
  modified     : Bool         := false

abbrev M := ReaderT (List FVarId) $ StateM State

private def isModified : M Bool := do pure (← get).modified
private def resetModified : M Unit := modify fun s => { s with modified := false }
private def markModified : M Unit := modify fun s => { s with modified := true }
private def getUsedFVarsMap : M UsedFVarsMap := do pure (← get).usedFVarsMap
private def modifyUsedFVars (f : UsedFVarsMap → UsedFVarsMap) : M Unit := modify fun s => { s with usedFVarsMap := f s.usedFVarsMap }

-- merge s₂ into s₁
private def merge (s₁ s₂ : FVarIdSet) : M FVarIdSet :=
  s₂.foldM (init := s₁) fun s₁ k => do
    if s₁.contains k then
      pure s₁
    else
      markModified
      pure $ s₁.insert k

private def updateUsedVarsOf (fvarId : FVarId) : M Unit := do
  let usedFVarsMap ← getUsedFVarsMap
  match usedFVarsMap.find? fvarId with
  | none         => pure ()
  | some fvarIds =>
    let fvarIdsNew ← fvarIds.foldM (init := fvarIds) fun fvarIdsNew fvarId' =>
      if fvarId == fvarId' then
        pure fvarIdsNew
      else
        match usedFVarsMap.find? fvarId' with
        | none => pure fvarIdsNew
          /- We are being sloppy here `otherFVarIds` may contain free variables that are
             not in the context of the let-rec associated with fvarId.
             We filter these out-of-context free variables later. -/
        | some otherFVarIds => merge fvarIdsNew otherFVarIds
    modifyUsedFVars fun usedFVars => usedFVars.insert fvarId fvarIdsNew

private partial def fixpoint : Unit → M Unit
  | _ => do
    resetModified
    let letRecFVarIds ← read
    letRecFVarIds.forM updateUsedVarsOf
    if (← isModified) then
      fixpoint ()

def run (letRecFVarIds : List FVarId) (usedFVarsMap : UsedFVarsMap) : UsedFVarsMap :=
  let (_, s) := ((fixpoint ()).run letRecFVarIds).run { usedFVarsMap := usedFVarsMap }
  s.usedFVarsMap

end FixPoint

abbrev FreeVarMap := FVarIdMap (Array FVarId)

private def mkFreeVarMap
    (mctx : MetavarContext) (sectionVars : Array Expr) (mainFVarIds : Array FVarId)
    (recFVarIds : Array FVarId) (letRecsToLift : List LetRecToLift) : FreeVarMap := Id.run <| do
  let usedFVarsMap  := mkInitialUsedFVarsMap mctx sectionVars mainFVarIds letRecsToLift
  let letRecFVarIds := letRecsToLift.map fun toLift => toLift.fvarId
  let usedFVarsMap  := FixPoint.run letRecFVarIds usedFVarsMap
  let mut freeVarMap := {}
  for toLift in letRecsToLift do
    let lctx       := toLift.lctx
    let fvarIdsSet := (usedFVarsMap.find? toLift.fvarId).get!
    let fvarIds    := fvarIdsSet.fold (init := #[]) fun fvarIds fvarId =>
      if lctx.contains fvarId && !recFVarIds.contains fvarId then
        fvarIds.push fvarId
      else
        fvarIds
    freeVarMap := freeVarMap.insert toLift.fvarId fvarIds
  pure freeVarMap

structure ClosureState where
  newLocalDecls : Array LocalDecl := #[]
  localDecls    : Array LocalDecl := #[]
  newLetDecls   : Array LocalDecl := #[]
  exprArgs      : Array Expr      := #[]

private def pickMaxFVar? (lctx : LocalContext) (fvarIds : Array FVarId) : Option FVarId :=
  fvarIds.getMax? fun fvarId₁ fvarId₂ => (lctx.get! fvarId₁).index < (lctx.get! fvarId₂).index

private def preprocess (e : Expr) : TermElabM Expr := do
  let e ← instantiateMVars e
  -- which let-decls are dependent. We say a let-decl is dependent if its lambda abstraction is type incorrect.
  Meta.check e
  pure e

/- Push free variables in `s` to `toProcess` if they are not already there. -/
private def pushNewVars (toProcess : Array FVarId) (s : CollectFVars.State) : Array FVarId :=
  s.fvarSet.fold (init := toProcess) fun toProcess fvarId =>
    if toProcess.contains fvarId then toProcess else toProcess.push fvarId

private def pushLocalDecl (toProcess : Array FVarId) (fvarId : FVarId) (userName : Name) (type : Expr) (bi := BinderInfo.default)
    : StateRefT ClosureState TermElabM (Array FVarId) := do
  let type ← preprocess type
  modify fun s => { s with
    newLocalDecls := s.newLocalDecls.push $ LocalDecl.cdecl default fvarId userName type bi,
    exprArgs      := s.exprArgs.push (mkFVar fvarId)
  }
  pure $ pushNewVars toProcess (collectFVars {} type)

private partial def mkClosureForAux (toProcess : Array FVarId) : StateRefT ClosureState TermElabM Unit := do
  let lctx ← getLCtx
  match pickMaxFVar? lctx toProcess with
  | none        => pure ()
  | some fvarId =>
    trace[Elab.definition.mkClosure] "toProcess: {toProcess.map mkFVar}, maxVar: {mkFVar fvarId}"
    let toProcess := toProcess.erase fvarId
    let localDecl ← getLocalDecl fvarId
    match localDecl with
    | LocalDecl.cdecl _ _ userName type bi =>
      let toProcess ← pushLocalDecl toProcess fvarId userName type bi
      mkClosureForAux toProcess
    | LocalDecl.ldecl _ _ userName type val _ =>
      let zetaFVarIds ← getZetaFVarIds
      if !zetaFVarIds.contains fvarId then
        /- Non-dependent let-decl. See comment at src/Lean/Meta/Closure.lean -/
        let toProcess ← pushLocalDecl toProcess fvarId userName type
        mkClosureForAux toProcess
      else
        /- Dependent let-decl. -/
        let type ← preprocess type
        let val  ← preprocess val
        modify fun s => { s with
          newLetDecls   := s.newLetDecls.push $ LocalDecl.ldecl default fvarId userName type val false,
          /- We don't want to interleave let and lambda declarations in our closure. So, we expand any occurrences of fvarId
             at `newLocalDecls` and `localDecls` -/
          newLocalDecls := s.newLocalDecls.map (replaceFVarIdAtLocalDecl fvarId val),
          localDecls := s.localDecls.map (replaceFVarIdAtLocalDecl fvarId val)
        }
        mkClosureForAux (pushNewVars toProcess (collectFVars (collectFVars {} type) val))

private partial def mkClosureFor (freeVars : Array FVarId) (localDecls : Array LocalDecl) : TermElabM ClosureState := do
  let (_, s) ← (mkClosureForAux freeVars).run { localDecls := localDecls }
  pure { s with
    newLocalDecls := s.newLocalDecls.reverse,
    newLetDecls   := s.newLetDecls.reverse,
    exprArgs      := s.exprArgs.reverse
  }

structure LetRecClosure where
  ref        : Syntax
  localDecls : Array LocalDecl
  closed     : Expr -- expression used to replace occurrences of the let-rec FVarId
  toLift     : LetRecToLift

private def mkLetRecClosureFor (toLift : LetRecToLift) (freeVars : Array FVarId) : TermElabM LetRecClosure := do
  let lctx := toLift.lctx
  withLCtx lctx toLift.localInstances do
  lambdaTelescope toLift.val fun xs val => do
    let type ← instantiateForall toLift.type xs
    let lctx ← getLCtx
    let s ← mkClosureFor freeVars $ xs.map fun x => lctx.get! x.fvarId!
    let type := Closure.mkForall s.localDecls $ Closure.mkForall s.newLetDecls type
    let val  := Closure.mkLambda s.localDecls $ Closure.mkLambda s.newLetDecls val
    let c    := mkAppN (Lean.mkConst toLift.declName) s.exprArgs
    assignExprMVar toLift.mvarId c
    return {
      ref        := toLift.ref
      localDecls := s.newLocalDecls
      closed     := c
      toLift     := { toLift with val := val, type := type }
    }

private def mkLetRecClosures (letRecsToLift : List LetRecToLift) (freeVarMap : FreeVarMap) : TermElabM (List LetRecClosure) :=
  letRecsToLift.mapM fun toLift => mkLetRecClosureFor toLift (freeVarMap.find? toLift.fvarId).get!

/- Mapping from FVarId of mutually recursive functions being defined to "closure" expression. -/
abbrev Replacement := FVarIdMap Expr

def insertReplacementForMainFns (r : Replacement) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainFVars : Array Expr) : Replacement :=
  mainFVars.size.fold (init := r) fun i r =>
    r.insert mainFVars[i].fvarId! (mkAppN (Lean.mkConst mainHeaders[i].declName) sectionVars)

def insertReplacementForLetRecs (r : Replacement) (letRecClosures : List LetRecClosure) : Replacement :=
  letRecClosures.foldl (init := r) fun r c =>
    r.insert c.toLift.fvarId c.closed

def Replacement.apply (r : Replacement) (e : Expr) : Expr :=
  e.replace fun e => match e with
    | Expr.fvar fvarId _ => match r.find? fvarId with
      | some c => some c
      | _      => none
    | _ => none

def pushMain (preDefs : Array PreDefinition) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainVals : Array Expr)
    : TermElabM (Array PreDefinition) :=
  mainHeaders.size.foldM (init := preDefs) fun i preDefs => do
    let header := mainHeaders[i]
    let val  ← mkLambdaFVars sectionVars mainVals[i]
    let type ← mkForallFVars sectionVars header.type
    return preDefs.push {
      ref         := getDeclarationSelectionRef header.ref
      kind        := header.kind
      declName    := header.declName
      levelParams := [], -- we set it later
      modifiers   := header.modifiers
      type        := type
      value       := val
    }

def pushLetRecs (preDefs : Array PreDefinition) (letRecClosures : List LetRecClosure) (kind : DefKind) (modifiers : Modifiers) : Array PreDefinition :=
  letRecClosures.foldl (init := preDefs) fun preDefs c =>
    let type := Closure.mkForall c.localDecls c.toLift.type
    let val  := Closure.mkLambda c.localDecls c.toLift.val
    preDefs.push {
      ref         := c.ref
      kind        := kind
      declName    := c.toLift.declName
      levelParams := [] -- we set it later
      modifiers   := { modifiers with attrs := c.toLift.attrs }
      type        := type
      value       := val
    }

def getKindForLetRecs (mainHeaders : Array DefViewElabHeader) : DefKind :=
  if mainHeaders.any fun h => h.kind.isTheorem then DefKind.«theorem»
  else DefKind.«def»

def getModifiersForLetRecs (mainHeaders : Array DefViewElabHeader) : Modifiers := {
  isNoncomputable := mainHeaders.any fun h => h.modifiers.isNoncomputable
  recKind         := if mainHeaders.any fun h => h.modifiers.isPartial then RecKind.partial else RecKind.default
  isUnsafe        := mainHeaders.any fun h => h.modifiers.isUnsafe
}

/-
- `sectionVars`:   The section variables used in the `mutual` block.
- `mainHeaders`:   The elaborated header of the top-level definitions being defined by the mutual block.
- `mainFVars`:     The auxiliary variables used to represent the top-level definitions being defined by the mutual block.
- `mainVals`:      The elaborated value for the top-level definitions
- `letRecsToLift`: The let-rec's definitions that need to be lifted
-/
def main (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainFVars : Array Expr) (mainVals : Array Expr) (letRecsToLift : List LetRecToLift)
    : TermElabM (Array PreDefinition) := do
  -- Store in recFVarIds the fvarId of every function being defined by the mutual block.
  let mainFVarIds := mainFVars.map Expr.fvarId!
  let recFVarIds  := (letRecsToLift.toArray.map fun toLift => toLift.fvarId) ++ mainFVarIds
  -- Compute the set of free variables (excluding `recFVarIds`) for each let-rec.
  let mctx ← getMCtx
  let freeVarMap := mkFreeVarMap mctx sectionVars mainFVarIds recFVarIds letRecsToLift
  resetZetaFVarIds
  withTrackingZeta do
    -- By checking `toLift.type` and `toLift.val` we populate `zetaFVarIds`. See comments at `src/Lean/Meta/Closure.lean`.
    letRecsToLift.forM fun toLift => withLCtx toLift.lctx toLift.localInstances do Meta.check toLift.type; Meta.check toLift.val
    let letRecClosures ← mkLetRecClosures letRecsToLift freeVarMap
    -- mkLetRecClosures assign metavariables that were placeholders for the lifted declarations.
    let mainVals    ← mainVals.mapM (instantiateMVars ·)
    let mainHeaders ← mainHeaders.mapM instantiateMVarsAtHeader
    let letRecClosures ← letRecClosures.mapM fun closure => do pure { closure with toLift := (← instantiateMVarsAtLetRecToLift closure.toLift) }
    -- Replace fvarIds for functions being defined with closed terms
    let r              := insertReplacementForMainFns {} sectionVars mainHeaders mainFVars
    let r              := insertReplacementForLetRecs r letRecClosures
    let mainVals       := mainVals.map r.apply
    let mainHeaders    := mainHeaders.map fun h => { h with type := r.apply h.type }
    let letRecClosures := letRecClosures.map fun c => { c with toLift := { c.toLift with type := r.apply c.toLift.type, val := r.apply c.toLift.val } }
    let letRecKind     := getKindForLetRecs mainHeaders
    let letRecMods     := getModifiersForLetRecs mainHeaders
    pushMain (pushLetRecs #[] letRecClosures letRecKind letRecMods) sectionVars mainHeaders mainVals

end MutualClosure

private def getAllUserLevelNames (headers : Array DefViewElabHeader) : List Name :=
  if h : 0 < headers.size then
    -- Recall that all top-level functions must have the same levels. See `check` method above
    (headers.get ⟨0, h⟩).levelNames
  else
    []

/-- Eagerly convert universe metavariables occurring in theorem headers to universe parameters. -/
private def levelMVarToParamHeaders (views : Array DefView) (headers : Array DefViewElabHeader) : TermElabM (Array DefViewElabHeader) := do
  let rec process : StateRefT Nat TermElabM (Array DefViewElabHeader) := do
    let mut newHeaders := #[]
    for view in views, header in headers do
      if view.kind.isTheorem then
        newHeaders := newHeaders.push { header with type := (← levelMVarToParam' header.type) }
      else
        newHeaders := newHeaders.push header
    return newHeaders
  let newHeaders ← (process).run' 1
  newHeaders.mapM fun header => return { header with type := (← instantiateMVars header.type) }

/-- Result for `mkInst?` -/
structure MkInstResult where
  instVal   : Expr
  instType  : Expr
  outParams : Array Expr := #[]

/--
  Construct an instance for `className out₁ ... outₙ type`.
  The method support classes with a prefix of `outParam`s (e.g. `MonadReader`). -/
private partial def mkInst? (className : Name) (type : Expr) : MetaM (Option MkInstResult) := do
  let rec go? (instType instTypeType : Expr) (outParams : Array Expr) : MetaM (Option MkInstResult) := do
    let instTypeType ← whnfD instTypeType
    unless instTypeType.isForall do
      return none
    let d := instTypeType.bindingDomain!
    if isOutParam d then
      let mvar ← mkFreshExprMVar d
      go? (mkApp instType mvar) (instTypeType.bindingBody!.instantiate1 mvar) (outParams.push mvar)
    else
      unless (← isDefEqGuarded (← inferType type) d) do
        return none
      let instType ← instantiateMVars (mkApp instType type)
      let instVal ← synthInstance instType
      return some { instVal, instType, outParams }
  let instType ← mkConstWithFreshMVarLevels className
  go? instType (← inferType instType) #[]

def processDefDeriving (className : Name) (declName : Name) : TermElabM Bool := do
  try
    let ConstantInfo.defnInfo info ← getConstInfo declName | return false
    let some result ← mkInst? className info.value | return false
    let instTypeNew := mkApp result.instType.appFn! (Lean.mkConst declName (info.levelParams.map mkLevelParam))
    Meta.check instTypeNew
    let instName ← liftMacroM <| mkUnusedBaseName (declName.appendBefore "inst" |>.appendAfter className.getString!)
    addAndCompile <| Declaration.defnDecl {
      name        := instName
      levelParams := info.levelParams
      type        := (← instantiateMVars instTypeNew)
      value       := (← instantiateMVars result.instVal)
      hints       := info.hints
      safety      := info.safety
    }
    addInstance instName AttributeKind.global (eval_prio default)
    return true
  catch ex =>
    return false

/-- Remove auxiliary match discriminant let-declarations. -/
def eraseAuxDiscr (e : Expr) : CoreM Expr := do
  Core.transform e fun e => match e with
    | Expr.letE n _ v b .. =>
      if isAuxDiscrName n then
        return TransformStep.visit (b.instantiate1 v)
      else
        return TransformStep.visit e
    | e => return TransformStep.visit e

def elabMutualDef (vars : Array Expr) (views : Array DefView) (hints : TerminationHints) : TermElabM Unit :=
  if isExample views then
    withoutModifyingEnv go
  else
    go
where
  go := do
    let scopeLevelNames ← getLevelNames
    let headers ← elabHeaders views
    let headers ← levelMVarToParamHeaders views headers
    let allUserLevelNames := getAllUserLevelNames headers
    withFunLocalDecls headers fun funFVars => do
      let values ← elabFunValues headers
      Term.synthesizeSyntheticMVarsNoPostponing
      let values ← values.mapM (instantiateMVars ·)
      let headers ← headers.mapM instantiateMVarsAtHeader
      let letRecsToLift ← getLetRecsToLift
      let letRecsToLift ← letRecsToLift.mapM instantiateMVarsAtLetRecToLift
      checkLetRecsToLiftTypes funFVars letRecsToLift
      withUsed vars headers values letRecsToLift fun vars => do
        let preDefs ← MutualClosure.main vars headers funFVars values letRecsToLift
        for preDef in preDefs do
          trace[Elab.definition] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
        let preDefs ← levelMVarToParamPreDecls preDefs
        let preDefs ← instantiateMVarsAtPreDecls preDefs
        let preDefs ← fixLevelParams preDefs scopeLevelNames allUserLevelNames
        let preDefs ← preDefs.mapM fun preDef =>
          if preDef.kind.isTheorem || preDef.kind.isExample then
            return preDef
          else
            return { preDef with value := (← eraseAuxDiscr preDef.value) }
        addPreDefinitions preDefs hints
        processDeriving headers

  processDeriving (headers : Array DefViewElabHeader) := do
    for header in headers, view in views do
      if let some classNamesStx := view.deriving? then
        for classNameStx in classNamesStx do
          let className ← resolveGlobalConstNoOverload classNameStx
          withRef classNameStx do
            unless (← processDefDeriving className header.declName) do
              throwError "failed to synthesize instance '{className}' for '{header.declName}'"

end Term
namespace Command

def elabMutualDef (ds : Array Syntax) (hints : TerminationHints) : CommandElabM Unit := do
  let views ← ds.mapM fun d => do
    let modifiers ← elabModifiers d[0]
    if ds.size > 1 && modifiers.isNonrec then
      throwErrorAt d "invalid use of 'nonrec' modifier in 'mutual' block"
    mkDefView modifiers d[1]
  runTermElabM none fun vars => Term.elabMutualDef vars views hints

end Command
end Lean.Elab


import Lean
import Mathlib.Tactic.Simps.Basic
import Mathlib.Tactic.PermuteGoals
import Mathlib.Tactic.Classical

open Lean Elab Parser Term Meta Tactic

syntax (name := seq) "seq" tacticSeq : tactic

section Source
  -- modified from `Lean.Elab.Tactic.Simp`
  def traceSimpCall' (stx : Syntax) (usedSimps : Simp.UsedSimps) : MetaM Syntax := do
    let mut stx := stx
    if stx[3].isNone then
      stx := stx.setArg 3 (mkNullNode #[mkAtom "only"])
    let mut args := #[]
    let mut localsOrStar := some #[]
    let lctx ← getLCtx
    let env ← getEnv
    for (thm, _) in usedSimps.toArray.qsort (·.2 < ·.2) do
      match thm with
      | .decl declName => -- global definitions in the environment
        if env.contains declName && !simpOnlyBuiltins.contains declName then
          args := args.push (← `(Parser.Tactic.simpLemma| $(mkIdent (← unresolveNameGlobal declName)):ident))
      | .fvar fvarId => -- local hypotheses in the context
        if let some ldecl := lctx.find? fvarId then
          localsOrStar := localsOrStar.bind fun locals =>
            if !ldecl.userName.isInaccessibleUserName &&
                (lctx.findFromUserName? ldecl.userName).get!.fvarId == ldecl.fvarId then
              some (locals.push ldecl.userName)
            else
              none
        -- Note: the `if let` can fail for `simp (config := {contextual := true})` when
        -- rewriting with a variable that was introduced in a scope. In that case we just ignore.
      | .stx _ thmStx => -- simp theorems provided in the local invocation
        args := args.push ⟨thmStx⟩
      | .other _ => -- Ignore "special" simp lemmas such as constructed by `simp_all`.
        pure ()     -- We can't display them anyway.
    if let some locals := localsOrStar then
      args := args ++ (← locals.mapM fun id => `(Parser.Tactic.simpLemma| $(mkIdent id):ident))
    else
      args := args.push ⟨(← `(Parser.Tactic.simpStar| *))⟩
    let argsStx := if args.isEmpty then #[] else #[mkAtom "[", (mkAtom ",").mkSep args, mkAtom "]"]
    stx := stx.setArg 4 (mkNullNode argsStx)
    return stx

  def dsimpLocation' (ctx : Simp.Context) (loc : Location) : TacticM Syntax := do
    match loc with
    | Location.targets hyps simplifyTarget =>
      withMainContext do
        let fvarIds ← getFVarIds hyps
        go fvarIds simplifyTarget
    | Location.wildcard =>
      withMainContext do
        go (← (← getMainGoal).getNondepPropHyps) (simplifyTarget := true)
  where
    go (fvarIdsToSimp : Array FVarId) (simplifyTarget : Bool) : TacticM Syntax := do
      let mvarId ← getMainGoal
      let (result?, usedSimps) ← dsimpGoal mvarId ctx (simplifyTarget := simplifyTarget) (fvarIdsToSimp := fvarIdsToSimp)
      match result? with
      | none => replaceMainGoal []
      | some mvarId => replaceMainGoal [mvarId]
      traceSimpCall' (← getRef) usedSimps

  def getMainGoal' : TacticM (MVarId × List MVarId) := do
  loop (← getGoals)
    where
  loop : List MVarId → TacticM (MVarId × List MVarId)
    | [] => throwNoGoalsToBeSolved
    | mvarId :: mvarIds => do
      if (← mvarId.isAssigned) then
        loop mvarIds
      else
        setGoals (mvarId :: mvarIds)
        return (mvarId, mvarIds)

/--
  Searches for a metavariable `g` s.t. `tag` is its exact name.
  If none then searches for a metavariable `g` s.t. `tag` is a suffix of its name.
  If none, then it searches for a metavariable `g` s.t. `tag` is a prefix of its name. -/
def findTag? (mvarIds : List MVarId) (tag : Name) : TacticM (Option MVarId) := do
  match (← mvarIds.findM? fun mvarId => return tag == (← mvarId.getDecl).userName) with
  | some mvarId => return mvarId
  | none =>
  match (← mvarIds.findM? fun mvarId => return tag.isSuffixOf (← mvarId.getDecl).userName) with
  | some mvarId => return mvarId
  | none => mvarIds.findM? fun mvarId => return tag.isPrefixOf (← mvarId.getDecl).userName

def getCaseGoals (tag : TSyntax `Lean.binderIdent) : TacticM (MVarId × List MVarId) := do
  let gs ← getUnsolvedGoals
  let g ← if let `(Lean.binderIdent| $tag:ident) := tag then
    let tag := tag.getId
    let some g ← findTag? gs tag | throwError "tag not found"
    pure g
  else
    getMainGoal
  return (g, gs.erase g)

def matchAltTac := Term.matchAlt (rhsParser := matchRhs)

end Source

def traceGoalsAt (stx : TSyntax `tactic) : TacticM Unit := do
  let gs ← getUnsolvedGoals
  withRef stx <| addRawTrace (goalsToMessageData gs)

def traceTacticCallAt (stx : TSyntax `tactic) (tac : TSyntax `tactic) : TacticM Unit := do
  withRef stx <| addRawTrace m!"[TACTIC] {tac}"

#check Split.applyMatchSplitter

partial def evalTacticWithTrace : TSyntax `tactic → TacticM Unit
  /- Dealing with bracketing -/
  | `(tactic| { $[$tacs]* }) => do
    for tac in tacs do
      evalTacticWithTrace tac
  | `(tactic| ( $[$tacs]* )) => do
    for tac in tacs do
      evalTacticWithTrace tac
  /- Dealing with focused goals -/
  | `(tactic| · $[$tacs]*) => do
    let (mainGoal, otherGoals) ← getMainGoal'
    setGoals [mainGoal]
    for tac in tacs do
      evalTacticWithTrace tac
    setGoals otherGoals
  | `(tactic| focus $[$tacs]*) => do
    let (mainGoal, otherGoals) ← getMainGoal'
    setGoals [mainGoal]
    for tac in tacs do
      evalTacticWithTrace tac
    setGoals otherGoals
  /- Handle trace for the `classical` tactic -/
  | `(tactic| classical $[$tacs]*) => do
      modifyEnv Meta.instanceExtension.pushScope
      Meta.addInstance ``Classical.propDecidable .local 10
      try for tac in tacs do evalTacticWithTrace tac
      finally modifyEnv Meta.instanceExtension.popScope
   /- Trace `simp` calls with the complete list of theorems used -/
  | stx@`(tactic| simp%$tk $(config)? $(discharger)? $[only%$o]? $[[$args,*]]? $(loc)?) => do
    traceGoalsAt stx
    let { ctx, dischargeWrapper } ← withMainContext <| mkSimpContext stx (eraseLocal := false)
    let usedSimps ← dischargeWrapper.with fun discharge? =>
      simpLocation ctx discharge? (expandOptLocation stx.raw[5])
    traceTacticCallAt stx ⟨← traceSimpCall' stx usedSimps⟩
    traceGoalsAt stx
  | stx@`(tactic| simp_all%$tk $(config)? $(discharger)? $[only%$o]? $[[$args,*]]?) => do
    traceGoalsAt stx
    let { ctx, .. } ← mkSimpContext stx (eraseLocal := true) (kind := .simpAll) (ignoreStarArg := true)
    let (result?, usedSimps) ← simpAll (← getMainGoal) ctx
    match result? with
    | none => replaceMainGoal []
    | some mvarId => replaceMainGoal [mvarId]
    traceTacticCallAt stx ⟨← traceSimpCall' stx usedSimps⟩
    traceGoalsAt stx
  | stx@`(tactic| dsimp%$tk $(config)? $[only%$o]? $[[$args,*]]? $(loc)?) => do
    traceGoalsAt stx
    let { ctx, .. } ← withMainContext <| mkSimpContext stx (eraseLocal := false) (kind := .dsimp)
    traceTacticCallAt stx ⟨← dsimpLocation' ctx (expandOptLocation stx.raw[5])⟩
    traceGoalsAt stx
  /- Treat a rewrite sequence as a sequence of individual rewrites, with a trace provided at each step -/
  | `(tactic| rw $[$cfg]? [$rs,*] $[$loc]?) => do
    for r in (rs : TSyntaxArray `Lean.Parser.Tactic.rwRule) do
      traceGoalsAt ⟨r.raw⟩
      let rtac ← `(tactic| rw $[$cfg]? [$r] $[$loc]?)
      evalTactic rtac
      traceTacticCallAt ⟨r.raw⟩ rtac
      traceGoalsAt ⟨r.raw⟩
  | `(tactic| erw [$rs,*] $[$loc]?) => do
    for r in (rs : TSyntaxArray `Lean.Parser.Tactic.rwRule) do
      traceGoalsAt ⟨r.raw⟩
      let rtac ← `(tactic| erw [$r] $[$loc]?)
      evalTactic rtac
      traceTacticCallAt ⟨r.raw⟩ rtac
      traceGoalsAt ⟨r.raw⟩
  | `(tactic| rwa [$rs,*] $[$loc]?) => do
    for r in (rs : TSyntaxArray `Lean.Parser.Tactic.rwRule) do
      traceGoalsAt ⟨r.raw⟩
      let rtac ← `(tactic| rw [$r] $[$loc]?)
      evalTactic rtac
      traceTacticCallAt ⟨r.raw⟩ rtac
      traceGoalsAt ⟨r.raw⟩
    `(tactic| assumption) >>= evalTactic ∘ TSyntax.raw
  /- Annotate `apply` and `exact` applications with the type information -/
  | stx@`(tactic| apply $v) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm v none
    let typ ← inferType trm
    let (mainGoal, otherGoals) ← getMainGoal'
    let newGoals ← mainGoal.apply trm
    setGoals <| newGoals ++ otherGoals
    let typ ← instantiateMVars typ
    let typStx ← PrettyPrinter.delab typ
    `(tactic| apply ($v : $typStx)) >>= traceTacticCallAt stx
    traceGoalsAt stx
  | stx@`(tactic| exact $v) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm v none
    let typ ← inferType trm
    let typStx ← PrettyPrinter.delab typ
    `(tactic| exact ($v : $typStx)) >>= traceTacticCallAt stx
    traceGoalsAt stx
  /- Handling `match`, `induction` and `cases` -/
  | stx@`(tactic| case $[$tag $hs*]|* =>%$arr $tac:tacticSeq) => do
    for tag in tag, h in hs do
      traceGoalsAt ⟨arr⟩
      traceTacticCallAt ⟨arr⟩ stx -- TODO (unimportant) remove the `tacticSeq` from `stx`
      let (g, _) ← getCaseGoals tag
      withRef arr <| addRawTrace (goalsToMessageData [g])
    let stx ← `(tactic| case $[$tag $hs*]|* =>%$arr seq $tac:tacticSeq)
    evalTactic stx.raw
  | stx@`(tactic| case' $[$tag $hs*]|* =>%$arr $tac:tacticSeq) => do
    for tag in tag, h in hs do
      traceGoalsAt ⟨arr⟩
      traceTacticCallAt ⟨arr⟩ stx -- TODO (unimportant) remove the `tacticSeq` from `stx`
      let (g, _) ← getCaseGoals tag
      withRef arr <| addRawTrace (goalsToMessageData [g])
    let stx ← `(tactic| case' $[$tag $hs*]|* =>%$arr seq $tac:tacticSeq)
  | `(tactic| induction $[$ts],* $[using $id:ident]?  $[generalizing $gs*]? with $[$tac]? $is*) => do
    let is' : TSyntaxArray ``inductionAlt ←
      is.mapM <|
        fun
          | `(inductionAlt| $il* => $ts:tacticSeq) => `(inductionAlt| $il* => seq $ts)
          | i => return ⟨i⟩
    let stx' ← `(tactic| induction $[$ts],* $[using $id:ident]?  $[generalizing $gs*]? with $[$tac]? $is'*)
    evalTactic stx'
  | `(tactic| cases $[$cs],* $[using $id:ident]? with $[$tac]? $is*) => do
    let is' : TSyntaxArray ``inductionAlt ←
      is.mapM <|
        fun
          | `(inductionAlt| $il* => $ts:tacticSeq) => `(inductionAlt| $il* => seq $ts)
          | i => return ⟨i⟩
    let stx' ← `(tactic| cases $[$cs],* $[using $id:ident]? with $[$tac]? $is'*)
    evalTactic stx'
  | `(tactic| match $[$gen]? $[$motive]? $discrs,* with $alts:matchAlt*) => do
    let alts' : TSyntaxArray ``matchAlt ←
      alts.mapM <|
        fun
          | `(matchAltTac| | $[$pats,*]|* => $rhs:tacticSeq) => do
              let alt ← `(matchAltTac| | $[$pats,*]|* => seq $rhs)
              return ⟨alt⟩
          | alt =>  return ⟨alt⟩
      let stx' ← `(tactic| match $[$gen]? $[$motive]? $discrs,* with $alts':matchAlt*)
      evalTactic stx'
  /- Display the expected type in `have` and `let` statements -/
  | stx@`(tactic| have $[$x:ident]? := $prf) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm prf none
    let typ ← inferType trm
    let typStx ← PrettyPrinter.delab typ
    `(tactic| have $[$x:ident]? : $typStx := $prf) >>= traceTacticCallAt stx
    traceGoalsAt stx
  | stx@`(tactic| let $x:ident := $val) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm val none
    let typ ← inferType trm
    let typStx ← PrettyPrinter.delab typ
    `(tactic| let $x:ident : $typStx := $val) >>= traceTacticCallAt stx
    traceGoalsAt stx
  /- Otherwise, evaluate the tactic normally -/
  | stx@`(tactic| $tac) => do
    traceGoalsAt stx
    evalTactic tac
    traceTacticCallAt stx tac
    traceGoalsAt stx

@[tactic seq]
def traceSequence : Tactic := fun s => do
  -- Leonardo de Moura's code for extracting the list of tactics
  match s with
  | `(tactic| seq $[$tacs]*) =>
    for tac in tacs do
      evalTacticWithTrace tac
  | _ => evalTactic <| ← `(tactic.sorry)

-- an example of the `seq` tactic
example (h : x = y) : 0 + x = y ∧ 1 = 1 := by
  seq
    rw [Nat.zero_add, (h : x = y)]
  refine' ⟨_, _⟩
  focus
    rw [←h, h]
  · rfl

-- a deep copy of Lean's `by` tactic, called `by'`
syntax (name := byTactic') "by' " tacticSeq : term

@[term_elab byTactic'] def elabByTactic' : TermElab := fun stx expectedType? => do
  match expectedType? with
  | some expectedType =>
    let mvar ← mkFreshExprMVar expectedType MetavarKind.syntheticOpaque
    let mvarId := mvar.mvarId!
    let ref ← getRef
    registerSyntheticMVar ref mvarId <| SyntheticMVarKind.tactic stx (← saveContext)
    return mvar
  | none =>
    tryPostpone
    throwError ("invalid 'by\'' tactic, expected type has not been provided")

example : 1 + 1 = 2 := by' -- the new `by'` syntax can be used to replace `by`
  focus
  rfl

-- intercepting the `by` tactic to output intermediate trace data
-- the `by'` clone is needed here to avoid infinite recursion
macro_rules
  | `(by $t) => `(by' seq $t)

section Test

set_option linter.unreachableTactic false

-- the `by` tactic now generates trace data by default
example (h : x = y) : x + 0 + x = x + y ∧ 1 = 1 := by
  have := (rfl : 1 = 1)
  let a := 5
  simp at this
  refine' ⟨_, _⟩
  · apply Eq.symm
    apply Eq.symm
    erw [h, ← h, h, ← h]
    simp_all
  · apply Eq.symm
    apply Eq.symm
    subst h
    rfl
  done

example : ∀ n : Nat, n = n := by
  intro n
  let x : ∀ m : ℕ, m = m := by
    intro a
    rfl
  match n with
  | .zero => rfl
  | .succ _ => rfl

example : P ∧ Q ↔ Q ∧ P := by
  constructor
  · intro ⟨x, y⟩
    constructor
    · assumption
    · assumption
  · intro ⟨_, _⟩
    constructor
    · assumption
    · assumption

example : ∀ n : Nat, n = n := by
  intro n
  induction n
  case zero => rfl
  case succ _ ih => simp

example : ∀ n : Nat, n + n = n + n := by
  intro n
  induction n with
    | zero => rfl
    | succ _ _ => rfl

example : ∀ n : Nat, n + n = n + n := by
  intro n
  cases n with
    | zero =>
      let h := (rfl : 1 = 1)
      rfl
    | succ _ => rfl

example : ∀ n : Nat, n + n = n + n := by
  intro n
  match n with
    | .zero => rfl
    | .succ _ => rfl

end Test


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

Notation for operators defined at Prelude.lean
-/
prelude
import Init.Prelude

-- DSL for specifying parser precedences and priorities

namespace Lean.Parser.Syntax

syntax:65 (name := addPrec) prec " + " prec:66 : prec
syntax:65 (name := subPrec) prec " - " prec:66 : prec

syntax:65 (name := addPrio) prio " + " prio:66 : prio
syntax:65 (name := subPrio) prio " - " prio:66 : prio

end Lean.Parser.Syntax

macro "max"  : prec => `(1024) -- maximum precedence used in term parsers, in particular for terms in function position (`ident`, `paren`, ...)
macro "arg"  : prec => `(1023) -- precedence used for application arguments (`do`, `by`, ...)
macro "lead" : prec => `(1022) -- precedence used for terms not supposed to be used as arguments (`let`, `have`, ...)
macro "(" p:prec ")" : prec => p
macro "min"  : prec => `(10)   -- minimum precedence used in term parsers
macro "min1" : prec => `(11)   -- `(min+1) we can only `min+1` after `Meta.lean`
/-
  `max:prec` as a term. It is equivalent to `eval_prec max` for `eval_prec` defined at `Meta.lean`.
  We use `max_prec` to workaround bootstrapping issues. -/
macro "max_prec" : term => `(1024)

macro "default" : prio => `(1000)
macro "low"     : prio => `(100)
macro "mid"     : prio => `(1000)
macro "high"    : prio => `(10000)
macro "(" p:prio ")" : prio => p

-- Basic notation for defining parsers
-- NOTE: precedence must be at least `arg` to be used in `macro` without parentheses
syntax:arg stx:max "+" : stx
syntax:arg stx:max "*" : stx
syntax:arg stx:max "?" : stx
syntax:2 stx:2 " <|> " stx:1 : stx

macro_rules
  | `(stx| $p +) => `(stx| many1($p))
  | `(stx| $p *) => `(stx| many($p))
  | `(stx| $p ?) => `(stx| optional($p))
  | `(stx| $p₁ <|> $p₂) => `(stx| orelse($p₁, $p₂))

/- Comma-separated sequence. -/
macro:arg x:stx:max ",*"   : stx => `(stx| sepBy($x, ",", ", "))
macro:arg x:stx:max ",+"   : stx => `(stx| sepBy1($x, ",", ", "))
/- Comma-separated sequence with optional trailing comma. -/
macro:arg x:stx:max ",*,?" : stx => `(stx| sepBy($x, ",", ", ", allowTrailingSep))
macro:arg x:stx:max ",+,?" : stx => `(stx| sepBy1($x, ",", ", ", allowTrailingSep))

macro:arg "!" x:stx:max : stx => `(stx| notFollowedBy($x))

syntax (name := rawNatLit) "nat_lit " num : term

infixr:90 " ∘ "  => Function.comp
infixr:35 " × "  => Prod

infixl:55 " ||| " => HOr.hOr
infixl:58 " ^^^ " => HXor.hXor
infixl:60 " &&& " => HAnd.hAnd
infixl:65 " + "   => HAdd.hAdd
infixl:65 " - "   => HSub.hSub
infixl:70 " * "   => HMul.hMul
infixl:70 " / "   => HDiv.hDiv
infixl:70 " % "   => HMod.hMod
infixl:75 " <<< " => HShiftLeft.hShiftLeft
infixl:75 " >>> " => HShiftRight.hShiftRight
infixr:80 " ^ "   => HPow.hPow
infixl:65 " ++ "  => HAppend.hAppend
prefix:100 "-"    => Neg.neg
prefix:100 "~~~"  => Complement.complement
/-
  Remark: the infix commands above ensure a delaborator is generated for each relations.
  We redefine the macros below to be able to use the auxiliary `binop%` elaboration helper for binary operators.
  It addresses issue #382. -/
macro_rules | `($x ||| $y) => `(binop% HOr.hOr $x $y)
macro_rules | `($x ^^^ $y) => `(binop% HXor.hXor $x $y)
macro_rules | `($x &&& $y) => `(binop% HAnd.hAnd $x $y)
macro_rules | `($x + $y)   => `(binop% HAdd.hAdd $x $y)
macro_rules | `($x - $y)   => `(binop% HSub.hSub $x $y)
macro_rules | `($x * $y)   => `(binop% HMul.hMul $x $y)
macro_rules | `($x / $y)   => `(binop% HDiv.hDiv $x $y)
macro_rules | `($x ++ $y)  => `(binop% HAppend.hAppend $x $y)

-- declare ASCII alternatives first so that the latter Unicode unexpander wins
infix:50 " <= " => LE.le
infix:50 " ≤ "  => LE.le
infix:50 " < "  => LT.lt
infix:50 " >= " => GE.ge
infix:50 " ≥ "  => GE.ge
infix:50 " > "  => GT.gt
infix:50 " = "  => Eq
infix:50 " == " => BEq.beq
/-
  Remark: the infix commands above ensure a delaborator is generated for each relations.
  We redefine the macros below to be able to use the auxiliary `binrel%` elaboration helper for binary relations.
  It has better support for applying coercions. For example, suppose we have `binrel% Eq n i` where `n : Nat` and
  `i : Int`. The default elaborator fails because we don't have a coercion from `Int` to `Nat`, but
  `binrel%` succeeds because it also tries a coercion from `Nat` to `Int` even when the nat occurs before the int. -/
macro_rules | `($x <= $y) => `(binrel% LE.le $x $y)
macro_rules | `($x ≤ $y)  => `(binrel% LE.le $x $y)
macro_rules | `($x < $y)  => `(binrel% LT.lt $x $y)
macro_rules | `($x > $y)  => `(binrel% GT.gt $x $y)
macro_rules | `($x >= $y) => `(binrel% GE.ge $x $y)
macro_rules | `($x ≥ $y)  => `(binrel% GE.ge $x $y)
macro_rules | `($x = $y)  => `(binrel% Eq $x $y)
macro_rules | `($x == $y) => `(binrel% BEq.beq $x $y)

infixr:35 " /\\ " => And
infixr:35 " ∧ "   => And
infixr:30 " \\/ " => Or
infixr:30 " ∨  "  => Or
notation:max "¬" p:40 => Not p

infixl:35 " && " => and
infixl:30 " || " => or
notation:max "!" b:40 => not b

infixr:67 " :: " => List.cons
syntax:20 term:21 " <|> " term:20 : term
syntax:60 term:61 " >> " term:60 : term
infixl:55  " >>= " => Bind.bind
notation:60 a:60 " <*> " b:61 => Seq.seq a fun _ : Unit => b
notation:60 a:60 " <* " b:61 => SeqLeft.seqLeft a fun _ : Unit => b
notation:60 a:60 " *> " b:61 => SeqRight.seqRight a fun _ : Unit => b
infixr:100 " <$> " => Functor.map

macro_rules | `($x <|> $y) => `(binop_lazy% HOrElse.hOrElse $x $y)
macro_rules | `($x >> $y)  => `(binop_lazy% HAndThen.hAndThen $x $y)

syntax (name := termDepIfThenElse) ppGroup(ppDedent("if " ident " : " term " then" ppSpace term ppDedent(ppSpace "else") ppSpace term)) : term

macro_rules
  | `(if $h:ident : $c then $t:term else $e:term) => ``(dite $c (fun $h:ident => $t) (fun $h:ident => $e))

syntax (name := termIfThenElse) ppGroup(ppDedent("if " term " then" ppSpace term ppDedent(ppSpace "else") ppSpace term)) : term

macro_rules
  | `(if $c then $t:term else $e:term) => ``(ite $c $t $e)

macro "if " "let " pat:term " := " d:term " then " t:term " else " e:term : term =>
  `(match $d:term with | $pat:term => $t | _ => $e)

syntax:min term "<|" term:min : term

macro_rules
  | `($f $args* <| $a) => let args := args.push a; `($f $args*)
  | `($f <| $a) => `($f $a)

syntax:min term "|>" term:min1 : term

macro_rules
  | `($a |> $f $args*) => let args := args.push a; `($f $args*)
  | `($a |> $f)        => `($f $a)

-- Haskell-like pipe <|
-- Note that we have a whitespace after `$` to avoid an ambiguity with the antiquotations.
syntax:min term atomic("$" ws) term:min : term

macro_rules
  | `($f $args* $ $a) => let args := args.push a; `($f $args*)
  | `($f $ $a) => `($f $a)

syntax "{ " ident (" : " term)? " // " term " }" : term

macro_rules
  | `({ $x : $type // $p }) => ``(Subtype (fun ($x:ident : $type) => $p))
  | `({ $x // $p })         => ``(Subtype (fun ($x:ident : _) => $p))

/-
  `without_expected_type t` instructs Lean to elaborate `t` without an expected type.
  Recall that terms such as `match ... with ...` and `⟨...⟩` will postpone elaboration until
  expected type is known. So, `without_expected_type` is not effective in this case. -/
macro "without_expected_type " x:term : term => `(let aux := $x; aux)

syntax "[" term,* "]"  : term
syntax "%[" term,* "|" term "]" : term -- auxiliary notation for creating big list literals

namespace Lean

macro_rules
  | `([ $elems,* ]) => do
    let rec expandListLit (i : Nat) (skip : Bool) (result : Syntax) : MacroM Syntax := do
      match i, skip with
      | 0,   _     => pure result
      | i+1, true  => expandListLit i false result
      | i+1, false => expandListLit i true  (← ``(List.cons $(elems.elemsAndSeps[i]) $result))
    if elems.elemsAndSeps.size < 64 then
      expandListLit elems.elemsAndSeps.size false (← ``(List.nil))
    else
      `(%[ $elems,* | List.nil ])

notation:50 e:51 " matches " p:51 => match e with | p => true | _ => false

namespace Parser.Tactic
/--
Introduce one or more hypotheses, optionally naming and/or pattern-matching them.
For each hypothesis to be introduced, the remaining main goal's target type must be a `let` or function type.
* `intro` by itself introduces one anonymous hypothesis, which can be accessed by e.g. `assumption`.
* `intro x y` introduces two hypotheses and names them. Individual hypotheses can be anonymized via `_`,
  or matched against a pattern:
  ```lean
  -- ... ⊢ α × β → ...
  intro (a, b)
  -- ..., a : α, b : β ⊢ ...
  ```
* Alternatively, `intro` can be combined with pattern matching much like `fun`:
  ```lean
  intro
  | n + 1, 0 => tac
  | ...
  ```
-/
syntax (name := intro) "intro " notFollowedBy("|") (colGt term:max)* : tactic
/-- `intros x...` behaves like `intro x...`, but then keeps introducing (anonymous) hypotheses until goal is not of a function type. -/
syntax (name := intros) "intros " (colGt (ident <|> "_"))* : tactic
/--
`rename t => x` renames the most recent hypothesis whose type matches `t` (which may contain placeholders) to `x`,
or fails if no such hypothesis could be found. -/
syntax (name := rename) "rename " term " => " ident : tactic
/-- `revert x...` is the inverse of `intro x...`: it moves the given hypotheses into the main goal's target type. -/
syntax (name := revert) "revert " (colGt ident)+ : tactic
/-- `clear x...` removes the given hypotheses, or fails if there are remaining references to a hypothesis. -/
syntax (name := clear) "clear " (colGt ident)+ : tactic
/--
`subst x...` substitutes each `x` with `e` in the goal if there is a hypothesis of type `x = e` or `e = x`.
If `x` is itself a hypothesis of type `y = e` or `e = y`, `y` is substituted instead. -/
syntax (name := subst) "subst " (colGt ident)+ : tactic
/--
`assumption` tries to solve the main goal using a hypothesis of compatible type, or else fails.
Note also the `‹t›` term notation, which is a shorthand for `show t by assumption`. -/
syntax (name := assumption) "assumption" : tactic
/--
`contradiction` closes the main goal if its hypotheses are "trivially contradictory".
```lean
example (h : False) : p := by contradiction  -- inductive type/family with no applicable constructors
example (h : none = some true) : p := by contradiction  -- injectivity of constructors
example (h : 2 + 2 = 3) : p := by contradiction  -- decidable false proposition
example (h : p) (h' : ¬ p) : q := by contradiction
example (x : Nat) (h : x ≠ x) : p := by contradiction
```
-/
syntax (name := contradiction) "contradiction" : tactic
/--
`apply e` tries to match the current goal against the conclusion of `e`'s type.
If it succeeds, then the tactic returns as many subgoals as the number of premises that
have not been fixed by type inference or type class resolution.
Non-dependent premises are added before dependent ones.

The `apply` tactic uses higher-order pattern matching, type class resolution, and first-order unification with dependent types.
-/
syntax (name := apply) "apply " term : tactic
/--
`exact e` closes the main goal if its target type matches that of `e`.
-/
syntax (name := exact) "exact " term : tactic
/--
`refine e` behaves like `exact e`, except that named (`?x`) or unnamed (`?_`) holes in `e` that are not solved
by unification with the main goal's target type are converted into new goals, using the hole's name, if any, as the goal case name.
-/
syntax (name := refine) "refine " term : tactic
/-- `refine' e` behaves like `refine e`, except that unsolved placeholders (`_`) and implicit parameters are also converted into new goals. -/
syntax (name := refine') "refine' " term : tactic
/-- If the main goal's target type is an inductive type, `constructor` solves it with the first matching constructor, or else fails. -/
syntax (name := constructor) "constructor" : tactic
/--
`case tag => tac` focuses on the goal with case name `tag` and solves it using `tac`, or else fails.
`case tag x₁ ... xₙ => tac` additionally renames the `n` most recent hypotheses with inaccessible names to the given names. -/
syntax (name := case) "case " (ident <|> "_") (ident <|> "_")* " => " tacticSeq : tactic
/--
`next => tac` focuses on the next goal solves it using `tac`, or else fails.
`next x₁ ... xₙ => tac` additionally renames the `n` most recent hypotheses with inaccessible names to the given names. -/
macro "next " args:(ident <|> "_")* " => " tac:tacticSeq : tactic => `(tactic| case _ $(args.getArgs)* => $tac)

/-- `allGoals tac` runs `tac` on each goal, concatenating the resulting goals, if any. -/
syntax (name := allGoals) "all_goals " tacticSeq : tactic
/-- `anyGoals tac` applies the tactic `tac` to every goal, and succeeds if at least one application succeeds.  -/
syntax (name := anyGoals) "any_goals " tacticSeq : tactic
/--
`focus tac` focuses on the main goal, suppressing all other goals, and runs `tac` on it.
Usually `· tac`, which enforces that the goal is closed by `tac`, should be preferred. -/
syntax (name := focus) "focus " tacticSeq : tactic
/-- `skip` does nothing. -/
syntax (name := skip) "skip" : tactic
/-- `done` succeeds iff there are no remaining goals. -/
syntax (name := done) "done" : tactic
syntax (name := traceState) "trace_state" : tactic
syntax (name := failIfSuccess) "fail_if_success " tacticSeq : tactic
syntax (name := paren) "(" tacticSeq ")" : tactic
syntax (name := withReducible) "with_reducible " tacticSeq : tactic
syntax (name := withReducibleAndInstances) "with_reducible_and_instances " tacticSeq : tactic
/-- `first | tac | ...` runs each `tac` until one succeeds, or else fails. -/
syntax (name := first) "first " withPosition((group(colGe "|" tacticSeq))+) : tactic
syntax (name := rotateLeft) "rotate_left" (num)? : tactic
syntax (name := rotateRight) "rotate_right" (num)? : tactic
/-- `try tac` runs `tac` and succeeds even if `tac` failed. -/
macro "try " t:tacticSeq : tactic => `(first | $t | skip)
/-- `tac <;> tac'` runs `tac` on the main goal and `tac'` on each produced goal, concatenating all goals produced by `tac'`. -/
macro:1 x:tactic " <;> " y:tactic:0 : tactic => `(tactic| focus ($x:tactic; all_goals $y:tactic))

/-- `· tac` focuses on the main goal and tries to solve it using `tac`, or else fails. -/
macro dot:("·" <|> ".") ts:tacticSeq : tactic => `(tactic| {%$dot ($ts:tacticSeq) })

/-- `rfl` is a shorthand for `exact rfl`. -/
macro "rfl" : tactic => `(exact rfl)
/-- `admit` is a shorthand for `exact sorry`. -/
macro "admit" : tactic => `(exact sorry)
/-- The `sorry` tactic isnxo a shorthand for `exact sorry`. -/
macro "sorry" : tactic => `(exact sorry)
macro "infer_instance" : tactic => `(exact inferInstance)

/-- Optional configuration option for tactics -/
syntax config := atomic("(" &"config") " := " term ")"

syntax locationWildcard := "*"
syntax locationHyp      := (colGt ident)+ ("⊢" <|> "|-")? -- TODO: delete
syntax locationTargets  := (colGt ident)+ ("⊢" <|> "|-")?
syntax location         := withPosition(" at " (locationWildcard <|> locationHyp))

syntax (name := change) "change " term (location)? : tactic
syntax (name := changeWith) "change " term " with " term (location)? : tactic

syntax rwRule    := ("←" <|> "<-")? term
syntax rwRuleSeq := "[" rwRule,+,? "]"

syntax (name := rewriteSeq) "rewrite " (config)? rwRuleSeq (location)? : tactic

syntax (name := rwSeq) "rw " (config)? rwRuleSeq (location)? : tactic

def rwWithRfl (kind : SyntaxNodeKind) (atom : String) (stx : Syntax) : MacroM Syntax := do
  -- We show the `rfl` state on `]`
  let seq   := stx[2]
  let rbrak := seq[2]
  -- Replace `]` token with one without position information in the expanded tactic
  let seq   := seq.setArg 2 (mkAtom "]")
  let tac   := stx.setKind kind |>.setArg 0 (mkAtomFrom stx atom) |>.setArg 2 seq
  `(tactic| $tac; try (with_reducible rfl%$rbrak))

@[macro rwSeq] def expandRwSeq : Macro :=
  rwWithRfl ``Lean.Parser.Tactic.rewriteSeq "rewrite"

syntax (name := injection) "injection " term (" with " (colGt (ident <|> "_"))+)? : tactic

syntax (name := injections) "injections" : tactic

syntax discharger := atomic("(" (&"discharger" <|> &"disch")) " := " tacticSeq ")"

syntax simpPre   := "↓"
syntax simpPost  := "↑"
syntax simpLemma := (simpPre <|> simpPost)? ("←" <|> "<-")? term
syntax simpErase := "-" ident
syntax simpStar  := "*"
syntax (name := simp) "simp " (config)? (discharger)? (&"only ")? ("[" (simpStar <|> simpErase <|> simpLemma),* "]")? (location)? : tactic
syntax (name := simpAll) "simp_all " (config)? (discharger)? (&"only ")? ("[" (simpErase <|> simpLemma),* "]")? : tactic

/--
  Delta expand the given definition.
  This is a low-level tactic, it will expose how recursive definitions have been compiled by Lean. -/
syntax (name := delta) "delta " ident (location)? : tactic

-- Auxiliary macro for lifting have/suffices/let/...
-- It makes sure the "continuation" `?_` is the main goal after refining
macro "refine_lift " e:term : tactic => `(focus (refine noImplicitLambda% $e; rotate_right))

macro "have " d:haveDecl : tactic => `(refine_lift have $d:haveDecl; ?_)
/- We use a priority > default, to avoid ambiguity with previous `have` notation -/
macro (priority := high) "have" x:ident " := " p:term : tactic => `(have $x:ident : _ := $p)
macro "suffices " d:sufficesDecl : tactic => `(refine_lift suffices $d:sufficesDecl; ?_)
macro "let " d:letDecl : tactic => `(refine_lift let $d:letDecl; ?_)
macro "show " e:term : tactic => `(refine_lift show $e:term from ?_)
syntax (name := letrec) withPosition(atomic(group("let " &"rec ")) letRecDecls) : tactic
macro_rules
  | `(tactic| let rec $d:letRecDecls) => `(tactic| refine_lift let rec $d:letRecDecls; ?_)

-- Similar to `refineLift`, but using `refine'`
macro "refine_lift' " e:term : tactic => `(focus (refine' noImplicitLambda% $e; rotate_right))
macro "have' " d:haveDecl : tactic => `(refine_lift' have $d:haveDecl; ?_)
macro (priority := high) "have'" x:ident " := " p:term : tactic => `(have' $x:ident : _ := $p)
macro "let' " d:letDecl : tactic => `(refine_lift' let $d:letDecl; ?_)

syntax inductionAlt  := "| " (group("@"? ident) <|> "_") (ident <|> "_")* " => " (hole <|> syntheticHole <|> tacticSeq)
syntax inductionAlts := "with " (tactic)? withPosition( (colGe inductionAlt)+)
syntax (name := induction) "induction " term,+ (" using " ident)?  ("generalizing " ident+)? (inductionAlts)? : tactic

syntax generalizeArg := atomic(ident " : ")? term:51 " = " ident
/--
`generalize ([h :] e = x),+` replaces all occurrences `e`s in the main goal with a fresh hypothesis `x`s.
If `h` is given, `h : e = x` is introduced as well. -/
syntax (name := generalize) "generalize " generalizeArg,+ : tactic

syntax casesTarget := atomic(ident " : ")? term
syntax (name := cases) "cases " casesTarget,+ (" using " ident)? (inductionAlts)? : tactic

syntax (name := existsIntro) "exists " term : tactic

/-- `rename_i x_1 ... x_n` renames the last `n` inaccessible names using the given names. -/
syntax (name := renameI) "rename_i " (colGt (ident <|> "_"))+ : tactic

syntax "repeat " tacticSeq : tactic
macro_rules
  | `(tactic| repeat $seq) => `(tactic| first | ($seq); repeat $seq | skip)

syntax "trivial" : tactic

syntax (name := split) "split " (colGt term)? (location)? : tactic

/--
The tactic `specialize h a₁ ... aₙ` works on local hypothesis `h`.
The premises of this hypothesis, either universal quantifications or non-dependent implications,
are instantiated by concrete terms coming either from arguments `a₁` ... `aₙ`.
The tactic adds a new hypothesis with the same name `h := h a₁ ... aₙ` and tries to clear the previous one.
-/
syntax (name := specialize) "specialize " term : tactic

macro_rules | `(tactic| trivial) => `(tactic| assumption)
macro_rules | `(tactic| trivial) => `(tactic| rfl)
macro_rules | `(tactic| trivial) => `(tactic| contradiction)
macro_rules | `(tactic| trivial) => `(tactic| apply True.intro)
macro_rules | `(tactic| trivial) => `(tactic| apply And.intro <;> trivial)

macro "unhygienic " t:tacticSeq : tactic => `(set_option tactic.hygienic false in $t:tacticSeq)

end Tactic

namespace Attr
-- simp attribute syntax
syntax (name := simp) "simp" (Tactic.simpPre <|> Tactic.simpPost)? (prio)? : attr
end Attr

end Parser
end Lean

macro "‹" type:term "›" : term => `((by assumption : $type))


import Lean
import Mathlib.Control.Writer
import PremiseSelection.Utils
open Lean
/-!

# Theorem feature extraction

Input: the goal state
Ouput: the theorem statement as an expr

 -/
open Std

def Std.RBMap.modify' (k : κ) (fn : Option α → Option α) (r : RBMap κ α cmp) :=
  match fn <| r.find? k with
  | none => r.erase k
  | some v => r.insert k v

def Std.RBMap.mergeBy (fn : κ → α → α → α) (r1 r2 : RBMap κ α cmp) :  RBMap κ α cmp :=
  r2.foldl (fun r1 k v2 => r1.modify' k (fun | none => some v2 | some v1 => some (fn k v1 v2))) r1

namespace PremiseSelection

def Multiset (α : Type) [Ord α] := Std.RBMap α Nat compare

variable {α : Type} [Ord α]

def Multiset.empty : Multiset α := mkRBMap _ _ _

instance : EmptyCollection  (Multiset α) :=  ⟨Multiset.empty⟩

instance : Append  (Multiset α) where
  append x y := x.mergeBy (fun _ => (·+·)) y

def Multiset.add : Multiset α → α → Multiset α
  | m, a => m.modify' a (fun | none => some 1 | some v => some (v + 1))

def Multiset.singleton : α → Multiset α
  | a => Multiset.empty |>.add a

instance : Ord Name := ⟨Name.quickCmp⟩

structure Bigram where
  fst : Name
  snd : Name
  deriving Ord

structure Trigram where
  fst : Name
  snd : Name
  trd : Name
  deriving Ord

instance : ToJson Bigram where
  toJson b := s!"{b.fst}/{b.snd}"

instance : ToString Bigram where
  toString b := s!"{b.fst}/{b.snd}"

instance : ToString Trigram where
  toString t := s!"{t.fst}/{t.snd}/{t.trd}"

instance : ToJson Trigram where
  toJson t := s!"{t.fst}/{t.snd}/{t.trd}"

structure StatementFeatures where
  /-- Just the constant's names and how frequently they arise. -/
  nameCounts : Multiset Name := ∅
  bigramCounts : Multiset Bigram := ∅
  trigramCounts : Multiset Trigram := ∅

instance : ForIn M (Multiset α) (α × Nat) :=
  show ForIn _ (Std.RBMap _ _ _) _ by infer_instance

def Multiset.toList (m : Multiset α) : List α :=
  m.foldl (fun l x _ => x :: l) []

def Multiset.toHFeatures [ToString α] (m : Multiset α) : Array String :=
  Array.mk <| m.toList.map (s!"H:{·}")

def Multiset.toTFeatures [ToString α] (m : Multiset α) : Array String :=
  Array.mk <| m.toList.map (s!"T:{·}")

instance [ToJson α] : ToJson (Multiset α) where
  toJson m := Json.arr (Array.mk (m.toList.map toJson))

instance : EmptyCollection StatementFeatures := ⟨{}⟩
instance : Append StatementFeatures where
  append x y := {
    nameCounts := x.nameCounts ++ y.nameCounts
    bigramCounts := x.bigramCounts ++ y.bigramCounts
    trigramCounts := x.trigramCounts ++ y.trigramCounts
  }

instance : ToJson StatementFeatures where
  toJson f := Json.mkObj [
    ("nameCounts", toJson f.nameCounts),
    ("bigramCounts", toJson f.bigramCounts),
    ("trigramCounts", toJson f.trigramCounts)
  ]

def StatementFeatures.mkName : Name → StatementFeatures
  | n => {nameCounts := Multiset.singleton n}

def StatementFeatures.mkBigram : Name → Name → StatementFeatures
  | n1, n2 => {bigramCounts := Multiset.singleton ⟨n1, n2⟩}

def StatementFeatures.mkTrigram : Name → Name → Name → StatementFeatures
  | n1, n2, n3 => {trigramCounts := Multiset.singleton ⟨n1, n2, n3⟩}

def StatementFeatures.toHFeatures (f : StatementFeatures) : Array String :=
  f.nameCounts.toHFeatures ++
  f.bigramCounts.toHFeatures ++
  f.trigramCounts.toHFeatures

def StatementFeatures.toTFeatures (f : StatementFeatures) : Array String :=
  f.nameCounts.toTFeatures ++
  f.bigramCounts.toTFeatures ++
  f.trigramCounts.toTFeatures

def immediateName (e : Expr) : Option Name :=
  if let .const n _ := e then
    some n
  else if let some n := e.natLit? then
    some <| toString n
  else
    none

def getHeadName? (e : Expr) : Option Name := do
  immediateName <| e.getAppFn

def visitFeature (e : Expr) : WriterT StatementFeatures MetaM Unit  := do
  --let ppe ← Lean.PrettyPrinter.ppExpr e
  if let some n := immediateName e then
    tell <| StatementFeatures.mkName n
  if e.isApp then
    e.withApp (fun f args => do
      if let some n1 := immediateName f then
        for arg in args do
          if let some n2 := getHeadName? arg then
            tell <| StatementFeatures.mkBigram n1 n2
            if arg.isApp then
              arg.withApp (fun f args => do
                if let some n2 := immediateName f then
                  for arg in args do
                    if let some n3 := getHeadName? arg then
                      tell <| StatementFeatures.mkTrigram n1 n2 n3
              )
        for p in args.toList.allPairs do
          if let some n2 := getHeadName? p.1 then
            if let some n3 := getHeadName? p.2 then
              tell <| StatementFeatures.mkTrigram n1 n2 n3
    )
  return ()

def getStatementFeatures (e : Expr) : MetaM StatementFeatures := do
  let ((), features) ← WriterT.run <| forEachExpr visitFeature e
  return features

open Lean.Meta

def getArgsFeatures (args : List Expr) : MetaM (Array StatementFeatures) := do
  let mut argsFeats := #[]
  for arg in args do
    let argType ← inferType arg
    if (← inferType argType).isProp then
      let argFeats ← getStatementFeatures argType
      if ! argFeats.nameCounts.isEmpty then
        argsFeats := argsFeats ++ #[argFeats]
  return argsFeats

def getThmAndArgsFeatures (e : Expr)
  : MetaM (StatementFeatures × Array StatementFeatures) := do
  forallTelescope e <| fun args thm => do
      let thmFeats ← getStatementFeatures thm
      let argsFeats ← getArgsFeatures args.data
      return (thmFeats, argsFeats)

end PremiseSelection