/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Sebastian Ullrich
-/
import Lean.Parser.Extension
-- necessary for auto-generation
import Lean.PrettyPrinter.Parenthesizer
import Lean.PrettyPrinter.Formatter

namespace Lean
namespace Parser

-- synthesize pretty printers for parsers declared prior to `Lean.PrettyPrinter`
-- (because `Parser.Extension` depends on them)
attribute [run_builtin_parser_attribute_hooks]
  leadingNode termParser commandParser mkAntiquot nodeWithAntiquot sepBy sepBy1
  unicodeSymbol nonReservedSymbol
  withCache withResetCache withPosition withPositionAfterLinebreak withoutPosition withForbidden withoutForbidden setExpected
  incQuotDepth decQuotDepth suppressInsideQuot evalInsideQuot
  withOpen withOpenDecl
  dbgTraceState

@[run_builtin_parser_attribute_hooks] def optional (p : Parser) : Parser :=
  optionalNoAntiquot (withAntiquotSpliceAndSuffix `optional p (symbol "?"))

@[run_builtin_parser_attribute_hooks] def many (p : Parser) : Parser :=
  manyNoAntiquot (withAntiquotSpliceAndSuffix `many p (symbol "*"))

@[run_builtin_parser_attribute_hooks] def many1 (p : Parser) : Parser :=
  many1NoAntiquot (withAntiquotSpliceAndSuffix `many p (symbol "*"))

@[run_builtin_parser_attribute_hooks] def ident : Parser :=
  withAntiquot (mkAntiquot "ident" identKind) identNoAntiquot

-- `ident` and `rawIdent` produce the same syntax tree, so we reuse the antiquotation kind name
@[run_builtin_parser_attribute_hooks] def rawIdent : Parser :=
  withAntiquot (mkAntiquot "ident" identKind) rawIdentNoAntiquot

@[run_builtin_parser_attribute_hooks] def numLit : Parser :=
  withAntiquot (mkAntiquot "num" numLitKind) numLitNoAntiquot

@[run_builtin_parser_attribute_hooks] def scientificLit : Parser :=
  withAntiquot (mkAntiquot "scientific" scientificLitKind) scientificLitNoAntiquot

@[run_builtin_parser_attribute_hooks] def strLit : Parser :=
  withAntiquot (mkAntiquot "str" strLitKind) strLitNoAntiquot

@[run_builtin_parser_attribute_hooks] def charLit : Parser :=
  withAntiquot (mkAntiquot "char" charLitKind) charLitNoAntiquot

@[run_builtin_parser_attribute_hooks] def nameLit : Parser :=
  withAntiquot (mkAntiquot "name" nameLitKind) nameLitNoAntiquot

@[run_builtin_parser_attribute_hooks, inline] def group (p : Parser) : Parser :=
  node groupKind p

@[run_builtin_parser_attribute_hooks, inline] def many1Indent (p : Parser) : Parser :=
  withPosition $ many1 (checkColGe "irrelevant" >> p)

@[run_builtin_parser_attribute_hooks, inline] def manyIndent (p : Parser) : Parser :=
  withPosition $ many (checkColGe "irrelevant" >> p)

@[inline] def sepByIndent (p : Parser) (sep : String) (psep : Parser := symbol sep) (allowTrailingSep : Bool := false) : Parser :=
  let p := withAntiquotSpliceAndSuffix `sepBy p (symbol "*")
  withPosition $ sepBy (checkColGe "irrelevant" >> p) sep (psep <|> checkColEq "irrelevant" >> checkLinebreakBefore >> pushNone) allowTrailingSep

@[inline] def sepBy1Indent (p : Parser) (sep : String) (psep : Parser := symbol sep) (allowTrailingSep : Bool := false) : Parser :=
  let p := withAntiquotSpliceAndSuffix `sepBy p (symbol "*")
  withPosition $ sepBy1 (checkColGe "irrelevant" >> p) sep (psep <|> checkColEq "irrelevant" >> checkLinebreakBefore >> pushNone) allowTrailingSep

open PrettyPrinter Syntax.MonadTraverser Formatter in
@[combinator_formatter sepByIndent]
def sepByIndent.formatter (p : Formatter) (_sep : String) (pSep : Formatter) : Formatter := do
  let stx ← getCur
  let hasNewlineSep := stx.getArgs.mapIdx (fun ⟨i, _⟩ n =>
    i % 2 == 1 && n.matchesNull 0 && i != stx.getArgs.size - 1) |>.any id
  visitArgs do
    for i in (List.range stx.getArgs.size).reverse do
      if i % 2 == 0 then p else pSep <|>
        -- If the final separator is a newline, skip it.
        ((if i == stx.getArgs.size - 1 then pure () else pushWhitespace "\n") *> goLeft)
  -- If there is any newline separator, then we add an `align` at the start
  -- so that `withPosition` will pick up the right column.
  if hasNewlineSep then
    pushAlign (force := true)

@[combinator_formatter sepBy1Indent] def sepBy1Indent.formatter := sepByIndent.formatter

attribute [run_builtin_parser_attribute_hooks] sepByIndent sepBy1Indent

@[run_builtin_parser_attribute_hooks] abbrev notSymbol (s : String) : Parser :=
  notFollowedBy (symbol s) s

/-- No-op parser combinator that annotates subtrees to be ignored in syntax patterns. -/
@[inline, run_builtin_parser_attribute_hooks] def patternIgnore : Parser → Parser := node `patternIgnore

/-- No-op parser that advises the pretty printer to emit a non-breaking space. -/
@[inline] def ppHardSpace : Parser := skip
/-- No-op parser that advises the pretty printer to emit a space/soft line break. -/
@[inline] def ppSpace : Parser := skip
/-- No-op parser that advises the pretty printer to emit a hard line break. -/
@[inline] def ppLine : Parser := skip
/-- No-op parser combinator that advises the pretty printer to emit a `Format.fill` node. -/
@[inline] def ppRealFill : Parser → Parser := id
/-- No-op parser combinator that advises the pretty printer to emit a `Format.group` node. -/
@[inline] def ppRealGroup : Parser → Parser := id
/-- No-op parser combinator that advises the pretty printer to indent the given syntax without grouping it. -/
@[inline] def ppIndent : Parser → Parser := id
/--
  No-op parser combinator that advises the pretty printer to group and indent the given syntax.
  By default, only syntax categories are grouped. -/
@[inline] def ppGroup (p : Parser) : Parser := ppRealFill (ppIndent p)
/--
  No-op parser combinator that advises the pretty printer to dedent the given syntax.
  Dedenting can in particular be used to counteract automatic indentation. -/
@[inline] def ppDedent : Parser → Parser := id

/--
  No-op parser combinator that allows the pretty printer to omit the group and
  indent operation in the enclosing category parser.
  ```
  syntax ppAllowUngrouped "by " tacticSeq : term
  -- allows a `by` after `:=` without linebreak in between:
  theorem foo : True := by
    trivial
  ```
-/
@[inline] def ppAllowUngrouped : Parser := skip

/--
  No-op parser combinator that advises the pretty printer to dedent the given syntax,
  if it was grouped by the category parser.
  Dedenting can in particular be used to counteract automatic indentation. -/
@[inline] def ppDedentIfGrouped : Parser → Parser := id

/--
  No-op parser combinator that prints a line break.
  The line break is soft if the combinator is followed
  by an ungrouped parser (see ppAllowUngrouped), otherwise hard. -/
@[inline] def ppHardLineUnlessUngrouped : Parser := skip

end Parser

section
open PrettyPrinter Parser

@[combinator_formatter ppHardSpace] def ppHardSpace.formatter : Formatter := Formatter.pushWhitespace " "
@[combinator_formatter ppSpace] def ppSpace.formatter : Formatter := Formatter.pushLine
@[combinator_formatter ppLine] def ppLine.formatter : Formatter := Formatter.pushWhitespace "\n"
@[combinator_formatter ppRealFill] def ppRealFill.formatter (p : Formatter) : Formatter := Formatter.fill p
@[combinator_formatter ppRealGroup] def ppRealGroup.formatter (p : Formatter) : Formatter := Formatter.group p
@[combinator_formatter ppIndent] def ppIndent.formatter (p : Formatter) : Formatter := Formatter.indent p
@[combinator_formatter ppDedent] def ppDedent.formatter (p : Formatter) : Formatter := do
  let opts ← getOptions
  Formatter.indent p (some ((0:Int) - Std.Format.getIndent opts))

@[combinator_formatter ppAllowUngrouped] def ppAllowUngrouped.formatter : Formatter := do
  modify ({ · with mustBeGrouped := false })
@[combinator_formatter ppDedentIfGrouped] def ppDedentIfGrouped.formatter (p : Formatter) : Formatter := do
  Formatter.concat p
  let indent := Std.Format.getIndent (← getOptions)
  unless (← get).isUngrouped do
    modify fun st => { st with stack := st.stack.modify (st.stack.size - 1) (·.nest (0 - indent)) }
@[combinator_formatter ppHardLineUnlessUngrouped] def ppHardLineUnlessUngrouped.formatter : Formatter := do
  if (← get).isUngrouped then
    Formatter.pushLine
  else
    ppLine.formatter

end

namespace Parser

-- now synthesize parenthesizers
attribute [run_builtin_parser_attribute_hooks]
  ppHardSpace ppSpace ppLine ppGroup ppRealGroup ppRealFill ppIndent ppDedent
  ppAllowUngrouped ppDedentIfGrouped ppHardLineUnlessUngrouped

syntax "register_parser_alias" group("(" &"kind" " := " term ")")? (strLit)? ident (colGt term)? : term
macro_rules
  | `(register_parser_alias $[(kind := $kind?)]? $(aliasName?)? $declName $(info?)?) => do
    let [(fullDeclName, [])] ← Macro.resolveGlobalName declName.getId |
      Macro.throwError "expected non-overloaded constant name"
    let aliasName := aliasName?.getD (Syntax.mkStrLit declName.getId.toString)
    `(do Parser.registerAlias $aliasName ``$declName $declName $(info?.getD (Unhygienic.run `({}))) (kind? := some $(kind?.getD (quote fullDeclName)))
         PrettyPrinter.Formatter.registerAlias $aliasName $(mkIdentFrom declName (declName.getId ++ `formatter))
         PrettyPrinter.Parenthesizer.registerAlias $aliasName $(mkIdentFrom declName (declName.getId ++ `parenthesizer)))

builtin_initialize
  register_parser_alias patternIgnore { autoGroupArgs := false }

  register_parser_alias group { autoGroupArgs := false }
  register_parser_alias ppHardSpace { stackSz? := some 0 }
  register_parser_alias ppSpace { stackSz? := some 0 }
  register_parser_alias ppLine { stackSz? := some 0 }
  register_parser_alias ppGroup { stackSz? := none }
  register_parser_alias ppRealGroup { stackSz? := none }
  register_parser_alias ppRealFill { stackSz? := none }
  register_parser_alias ppIndent { stackSz? := none }
  register_parser_alias ppDedent { stackSz? := none }
  register_parser_alias ppDedentIfGrouped { stackSz? := none }
  register_parser_alias ppAllowUngrouped { stackSz? := some 0 }
  register_parser_alias ppHardLineUnlessUngrouped { stackSz? := some 0 }

end Parser

end Lean


/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Sebastian Ullrich, Leonardo de Moura
-/
import Lean.Data.Name
import Lean.Data.Format

/--
A position range inside a string. This type is mostly in combination with syntax trees,
as there might not be a single underlying string in this case that could be used for a `Substring`.
-/
protected structure String.Range where
  start : String.Pos
  stop  : String.Pos
  deriving Inhabited, Repr, BEq, Hashable

def String.Range.contains (r : String.Range) (pos : String.Pos) (includeStop := false) : Bool :=
  r.start <= pos && (if includeStop then pos <= r.stop else pos < r.stop)

def String.Range.includes (super sub : String.Range) : Bool :=
  super.start <= sub.start && super.stop >= sub.stop

namespace Lean

def SourceInfo.updateTrailing (trailing : Substring) : SourceInfo → SourceInfo
  | SourceInfo.original leading pos _ endPos => SourceInfo.original leading pos trailing endPos
  | info                                     => info

/-! # Syntax AST -/

inductive IsNode : Syntax → Prop where
  | mk (info : SourceInfo) (kind : SyntaxNodeKind) (args : Array Syntax) : IsNode (Syntax.node info kind args)

def SyntaxNode : Type := {s : Syntax // IsNode s }

def unreachIsNodeMissing {β} (h : IsNode Syntax.missing) : β := False.elim (nomatch h)
def unreachIsNodeAtom {β} {info val} (h : IsNode (Syntax.atom info val)) : β := False.elim (nomatch h)
def unreachIsNodeIdent {β info rawVal val preresolved} (h : IsNode (Syntax.ident info rawVal val preresolved)) : β := False.elim (nomatch h)

def isLitKind (k : SyntaxNodeKind) : Bool :=
  k == strLitKind || k == numLitKind || k == charLitKind || k == nameLitKind || k == scientificLitKind

namespace SyntaxNode

@[inline] def getKind (n : SyntaxNode) : SyntaxNodeKind :=
  match n with
  | ⟨Syntax.node _ k _, _⟩  => k
  | ⟨Syntax.missing, h⟩     => unreachIsNodeMissing h
  | ⟨Syntax.atom .., h⟩     => unreachIsNodeAtom h
  | ⟨Syntax.ident .., h⟩    => unreachIsNodeIdent h

@[inline] def withArgs {β} (n : SyntaxNode) (fn : Array Syntax → β) : β :=
  match n with
  | ⟨Syntax.node _ _ args, _⟩   => fn args
  | ⟨Syntax.missing, h⟩       => unreachIsNodeMissing h
  | ⟨Syntax.atom _ _, h⟩      => unreachIsNodeAtom h
  | ⟨Syntax.ident _ _ _ _, h⟩ => unreachIsNodeIdent h

@[inline] def getNumArgs (n : SyntaxNode) : Nat :=
  withArgs n fun args => args.size

@[inline] def getArg (n : SyntaxNode) (i : Nat) : Syntax :=
  withArgs n fun args => args.get! i

@[inline] def getArgs (n : SyntaxNode) : Array Syntax :=
  withArgs n fun args => args

@[inline] def modifyArgs (n : SyntaxNode) (fn : Array Syntax → Array Syntax) : Syntax :=
  match n with
  | ⟨Syntax.node i k args, _⟩  => Syntax.node i k (fn args)
  | ⟨Syntax.missing, h⟩        => unreachIsNodeMissing h
  | ⟨Syntax.atom _ _, h⟩       => unreachIsNodeAtom h
  | ⟨Syntax.ident _ _ _ _,  h⟩ => unreachIsNodeIdent h

end SyntaxNode

namespace Syntax

def getAtomVal : Syntax → String
  | atom _ val => val
  | _          => ""

def setAtomVal : Syntax → String → Syntax
  | atom info _, v => (atom info v)
  | stx,         _ => stx

@[inline] def ifNode {β} (stx : Syntax) (hyes : SyntaxNode → β) (hno : Unit → β) : β :=
  match stx with
  | Syntax.node i k args => hyes ⟨Syntax.node i k args, IsNode.mk i k args⟩
  | _                    => hno ()

@[inline] def ifNodeKind {β} (stx : Syntax) (kind : SyntaxNodeKind) (hyes : SyntaxNode → β) (hno : Unit → β) : β :=
  match stx with
  | Syntax.node i k args => if k == kind then hyes ⟨Syntax.node i k args, IsNode.mk i k args⟩ else hno ()
  | _                    => hno ()

def asNode : Syntax → SyntaxNode
  | Syntax.node info kind args => ⟨Syntax.node info kind args, IsNode.mk info kind args⟩
  | _                          => ⟨mkNullNode, IsNode.mk _ _ _⟩

def getIdAt (stx : Syntax) (i : Nat) : Name :=
  (stx.getArg i).getId

@[inline] def modifyArgs (stx : Syntax) (fn : Array Syntax → Array Syntax) : Syntax :=
  match stx with
  | node i k args => node i k (fn args)
  | stx           => stx

@[inline] def modifyArg (stx : Syntax) (i : Nat) (fn : Syntax → Syntax) : Syntax :=
  match stx with
  | node info k args => node info k (args.modify i fn)
  | stx              => stx

@[specialize] partial def replaceM {m : Type → Type} [Monad m] (fn : Syntax → m (Option Syntax)) : Syntax → m (Syntax)
  | stx@(node info kind args) => do
    match (← fn stx) with
    | some stx => return stx
    | none     => return node info kind (← args.mapM (replaceM fn))
  | stx => do
    let o ← fn stx
    return o.getD stx

@[specialize] partial def rewriteBottomUpM {m : Type → Type} [Monad m] (fn : Syntax → m (Syntax)) : Syntax → m (Syntax)
  | node info kind args   => do
    let args ← args.mapM (rewriteBottomUpM fn)
    fn (node info kind args)
  | stx => fn stx

@[inline] def rewriteBottomUp (fn : Syntax → Syntax) (stx : Syntax) : Syntax :=
  Id.run <| stx.rewriteBottomUpM fn

private def updateInfo : SourceInfo → String.Pos → String.Pos → SourceInfo
  | SourceInfo.original lead pos trail endPos, leadStart, trailStop =>
    SourceInfo.original { lead with startPos := leadStart } pos { trail with stopPos := trailStop } endPos
  | info, _, _ => info

private def chooseNiceTrailStop (trail : Substring) : String.Pos :=
trail.startPos + trail.posOf '\n'

/-- Remark: the State `String.Pos` is the `SourceInfo.trailing.stopPos` of the previous token,
   or the beginning of the String. -/
@[inline]
private def updateLeadingAux : Syntax → StateM String.Pos (Option Syntax)
  | atom info@(SourceInfo.original _ _ trail _) val => do
    let trailStop := chooseNiceTrailStop trail
    let newInfo := updateInfo info (← get) trailStop
    set trailStop
    return some (atom newInfo val)
  | ident info@(SourceInfo.original _ _ trail _) rawVal val pre => do
    let trailStop := chooseNiceTrailStop trail
    let newInfo := updateInfo info (← get) trailStop
    set trailStop
    return some (ident newInfo rawVal val pre)
  | _ => pure none

/-- Set `SourceInfo.leading` according to the trailing stop of the preceding token.
    The result is a round-tripping syntax tree IF, in the input syntax tree,
    * all leading stops, atom contents, and trailing starts are correct
    * trailing stops are between the trailing start and the next leading stop.

    Remark: after parsing, all `SourceInfo.leading` fields are empty.
    The `Syntax` argument is the output produced by the parser for `source`.
    This function "fixes" the `source.leading` field.

    Additionally, we try to choose "nicer" splits between leading and trailing stops
    according to some heuristics so that e.g. comments are associated to the (intuitively)
    correct token.

    Note that the `SourceInfo.trailing` fields must be correct.
    The implementation of this Function relies on this property. -/
def updateLeading : Syntax → Syntax :=
  fun stx => (replaceM updateLeadingAux stx).run' 0

partial def updateTrailing (trailing : Substring) : Syntax → Syntax
  | Syntax.atom info val               => Syntax.atom (info.updateTrailing trailing) val
  | Syntax.ident info rawVal val pre   => Syntax.ident (info.updateTrailing trailing) rawVal val pre
  | n@(Syntax.node info k args)        =>
    if args.size == 0 then n
    else
     let i    := args.size - 1
     let last := updateTrailing trailing args[i]!
     let args := args.set! i last;
     Syntax.node info k args
  | s => s

partial def getTailWithPos : Syntax → Option Syntax
  | stx@(atom info _)   => info.getPos?.map fun _ => stx
  | stx@(ident info ..) => info.getPos?.map fun _ => stx
  | node SourceInfo.none _ args => args.findSomeRev? getTailWithPos
  | stx@(node ..) => stx
  | _ => none

open SourceInfo in
/-- Split an `ident` into its dot-separated components while preserving source info.
Macro scopes are first erased.  For example, `` `foo.bla.boo._@._hyg.4 `` ↦ `` [`foo, `bla, `boo] ``.
If `nFields` is set, we take that many fields from the end and keep the remaining components
as one name. For example, `` `foo.bla.boo `` with `(nFields := 1)` ↦ `` [`foo.bla, `boo] ``. -/
def identComponents (stx : Syntax) (nFields? : Option Nat := none) : List Syntax :=
  match stx with
  | ident (SourceInfo.original lead pos trail _) rawStr val _ =>
    let val := val.eraseMacroScopes
    -- With original info, we assume that `rawStr` represents `val`.
    let nameComps := nameComps val nFields?
    let rawComps := splitNameLit rawStr
    let rawComps :=
      if let some nFields := nFields? then
        let nPrefix := rawComps.length - nFields
        let prefixSz := rawComps.take nPrefix |>.foldl (init := 0) fun acc (ss : Substring) => acc + ss.bsize + 1
        let prefixSz := prefixSz - 1 -- The last component has no dot
        rawStr.extract 0 ⟨prefixSz⟩ :: rawComps.drop nPrefix
      else
        rawComps
    assert! nameComps.length == rawComps.length
    nameComps.zip rawComps |>.map fun (id, ss) =>
      let off := ss.startPos - rawStr.startPos
      let lead := if off == 0 then lead else "".toSubstring
      let trail := if ss.stopPos == rawStr.stopPos then trail else "".toSubstring
      let info := original lead (pos + off) trail (pos + off + ⟨ss.bsize⟩)
      ident info ss id []
  | ident si _ val _ =>
    let val := val.eraseMacroScopes
    /- With non-original info:
     - `rawStr` can take all kinds of forms so we only use `val`.
     - there is no source extent to offset, so we pass it as-is. -/
    nameComps val nFields? |>.map fun n => ident si n.toString.toSubstring n []
  | _ => unreachable!
  where
    nameComps (n : Name) (nFields? : Option Nat) : List Name :=
      if let some nFields := nFields? then
        let nameComps := n.components
        let nPrefix := nameComps.length - nFields
        let namePrefix := nameComps.take nPrefix |>.foldl (init := Name.anonymous) fun acc n => acc ++ n
        namePrefix :: nameComps.drop nPrefix
      else
        n.components

structure TopDown where
  firstChoiceOnly : Bool
  stx : Syntax

/--
`for _ in stx.topDown` iterates through each node and leaf in `stx` top-down, left-to-right.
If `firstChoiceOnly` is `true`, only visit the first argument of each choice node.
-/
def topDown (stx : Syntax) (firstChoiceOnly := false) : TopDown := ⟨firstChoiceOnly, stx⟩

partial instance : ForIn m TopDown Syntax where
  forIn := fun ⟨firstChoiceOnly, stx⟩ init f => do
    let rec @[specialize] loop stx b [Inhabited (type_of% b)] := do
      match (← f stx b) with
      | ForInStep.yield b' =>
        let mut b := b'
        if let Syntax.node _ k args := stx then
          if firstChoiceOnly && k == choiceKind then
            return ← loop args[0]! b
          else
            for arg in args do
              match (← loop arg b) with
              | ForInStep.yield b' => b := b'
              | ForInStep.done b'  => return ForInStep.done b'
        return ForInStep.yield b
      | ForInStep.done b => return ForInStep.done b
    match (← @loop stx init ⟨init⟩) with
    | ForInStep.yield b => return b
    | ForInStep.done b  => return b

partial def reprint (stx : Syntax) : Option String := do
  let mut s := ""
  for stx in stx.topDown (firstChoiceOnly := true) do
    match stx with
    | atom info val           => s := s ++ reprintLeaf info val
    | ident info rawVal _ _   => s := s ++ reprintLeaf info rawVal.toString
    | node _    kind args     =>
      if kind == choiceKind then
        -- this visit the first arg twice, but that should hardly be a problem
        -- given that choice nodes are quite rare and small
        let s0 ← reprint args[0]!
        for arg in args[1:] do
          let s' ← reprint arg
          guard (s0 == s')
    | _ => pure ()
  return s
where
  reprintLeaf (info : SourceInfo) (val : String) : String :=
    match info with
    | SourceInfo.original lead _ trail _ => s!"{lead}{val}{trail}"
    -- no source info => add gracious amounts of whitespace to definitely separate tokens
    -- Note that the proper pretty printer does not use this function.
    -- The parser as well always produces source info, so round-tripping is still
    -- guaranteed.
    | _                                => s!" {val} "

def hasMissing (stx : Syntax) : Bool := Id.run do
  for stx in stx.topDown do
    if stx.isMissing then
      return true
  return false

def getRange? (stx : Syntax) (canonicalOnly := false) : Option String.Range :=
  match stx.getPos? canonicalOnly, stx.getTailPos? canonicalOnly with
  | some start, some stop => some { start, stop }
  | _,          _         => none

/--
Represents a cursor into a syntax tree that can be read, written, and advanced down/up/left/right.
Indices are allowed to be out-of-bound, in which case `cur` is `Syntax.missing`.
If the `Traverser` is used linearly, updates are linear in the `Syntax` object as well.
-/
structure Traverser where
  cur     : Syntax
  parents : Array Syntax
  idxs    : Array Nat

namespace Traverser

def fromSyntax (stx : Syntax) : Traverser :=
  ⟨stx, #[], #[]⟩

def setCur (t : Traverser) (stx : Syntax) : Traverser :=
  { t with cur := stx }

/-- Advance to the `idx`-th child of the current node. -/
def down (t : Traverser) (idx : Nat) : Traverser :=
  if idx < t.cur.getNumArgs then
    { cur := t.cur.getArg idx, parents := t.parents.push <| t.cur.setArg idx default, idxs := t.idxs.push idx }
  else
    { cur := Syntax.missing, parents := t.parents.push t.cur, idxs := t.idxs.push idx }

/-- Advance to the parent of the current node, if any. -/
def up (t : Traverser) : Traverser :=
  if t.parents.size > 0 then
    let cur := if t.idxs.back < t.parents.back.getNumArgs then t.parents.back.setArg t.idxs.back t.cur else t.parents.back
    { cur := cur, parents := t.parents.pop, idxs := t.idxs.pop }
  else
    t

/-- Advance to the left sibling of the current node, if any. -/
def left (t : Traverser) : Traverser :=
  if t.parents.size > 0 then
    t.up.down (t.idxs.back - 1)
  else
    t

/-- Advance to the right sibling of the current node, if any. -/
def right (t : Traverser) : Traverser :=
  if t.parents.size > 0 then
    t.up.down (t.idxs.back + 1)
  else
    t

end Traverser

/-- Monad class that gives read/write access to a `Traverser`. -/
class MonadTraverser (m : Type → Type) where
  st : MonadState Traverser m

namespace MonadTraverser

variable {m : Type → Type} [Monad m] [t : MonadTraverser m]

def getCur : m Syntax := Traverser.cur <$> t.st.get
def setCur (stx : Syntax) : m Unit := @modify _ _ t.st (fun t => t.setCur stx)
def goDown (idx : Nat)    : m Unit := @modify _ _ t.st (fun t => t.down idx)
def goUp                  : m Unit := @modify _ _ t.st (fun t => t.up)
def goLeft                : m Unit := @modify _ _ t.st (fun t => t.left)
def goRight               : m Unit := @modify _ _ t.st (fun t => t.right)

def getIdx : m Nat := do
  let st ← t.st.get
  return st.idxs.back?.getD 0

end MonadTraverser
end Syntax

namespace SyntaxNode

@[inline] def getIdAt (n : SyntaxNode) (i : Nat) : Name :=
  (n.getArg i).getId

end SyntaxNode

def mkListNode (args : Array Syntax) : Syntax :=
  mkNullNode args

namespace Syntax

-- quotation node kinds are formed from a unique quotation name plus "quot"
def isQuot : Syntax → Bool
  | Syntax.node _ (Name.str _ "quot")           _ => true
  | Syntax.node _ `Lean.Parser.Term.dynamicQuot _ => true
  | _                                             => false

def getQuotContent (stx : Syntax) : Syntax :=
  let stx := if stx.getNumArgs == 1 then stx[0] else stx
  if stx.isOfKind `Lean.Parser.Term.dynamicQuot then
    stx[3]
  else
    stx[1]

-- antiquotation node kinds are formed from the original node kind (if any) plus "antiquot"
def isAntiquot : Syntax → Bool
  | .node _ (.str _ "antiquot") _ => true
  | _                             => false

def isAntiquots (stx : Syntax) : Bool :=
  stx.isAntiquot || (stx.isOfKind choiceKind && stx.getNumArgs > 0 && stx.getArgs.all isAntiquot)

def getCanonicalAntiquot (stx : Syntax) : Syntax :=
  if stx.isOfKind choiceKind then
    stx[0]
  else
    stx

def mkAntiquotNode (kind : Name) (term : Syntax) (nesting := 0) (name : Option String := none) (isPseudoKind := false) : Syntax :=
  let nesting := mkNullNode (mkArray nesting (mkAtom "$"))
  let term :=
    if term.isIdent then term
    else if term.isOfKind `Lean.Parser.Term.hole then term[0]
    else mkNode `antiquotNestedExpr #[mkAtom "(", term, mkAtom ")"]
  let name := match name with
    | some name => mkNode `antiquotName #[mkAtom ":", mkAtom name]
    | none      => mkNullNode
  mkNode (kind ++ (if isPseudoKind then `pseudo else Name.anonymous) ++ `antiquot) #[mkAtom "$", nesting, term, name]

-- Antiquotations can be escaped as in `$$x`, which is useful for nesting macros. Also works for antiquotation splices.
def isEscapedAntiquot (stx : Syntax) : Bool :=
  !stx[1].getArgs.isEmpty

-- Also works for antiquotation splices.
def unescapeAntiquot (stx : Syntax) : Syntax :=
  if isAntiquot stx then
    stx.setArg 1 <| mkNullNode stx[1].getArgs.pop
  else
    stx

-- Also works for token antiquotations.
def getAntiquotTerm (stx : Syntax) : Syntax :=
  let e := if stx.isAntiquot then stx[2] else stx[3]
  if e.isIdent then e
  else if e.isAtom then mkNode `Lean.Parser.Term.hole #[e]
  else
    -- `e` is from `"(" >> termParser >> ")"`
    e[1]

/-- Return kind of parser expected at this antiquotation, and whether it is a "pseudo" kind (see `mkAntiquot`). -/
def antiquotKind? : Syntax → Option (SyntaxNodeKind × Bool)
  | .node _ (.str (.str k "pseudo") "antiquot") _ => (k, true)
  | .node _ (.str k                 "antiquot") _ => (k, false)
  | _                                             => none

def antiquotKinds (stx : Syntax) : List (SyntaxNodeKind × Bool) :=
  if stx.isOfKind choiceKind then
    stx.getArgs.filterMap antiquotKind? |>.toList
  else
    match antiquotKind? stx with
    | some stx => [stx]
    | none     => []

-- An "antiquotation splice" is something like `$[...]?` or `$[...]*`.
def antiquotSpliceKind? : Syntax → Option SyntaxNodeKind
  | .node _ (.str k "antiquot_scope") _ => some k
  | _ => none

def isAntiquotSplice (stx : Syntax) : Bool :=
  antiquotSpliceKind? stx |>.isSome

def getAntiquotSpliceContents (stx : Syntax) : Array Syntax :=
  stx[3].getArgs

-- `$[..],*` or `$x,*` ~> `,*`
def getAntiquotSpliceSuffix (stx : Syntax) : Syntax :=
  if stx.isAntiquotSplice then
    stx[5]
  else
    stx[1]

def mkAntiquotSpliceNode (kind : SyntaxNodeKind) (contents : Array Syntax) (suffix : String) (nesting := 0) : Syntax :=
  let nesting := mkNullNode (mkArray nesting (mkAtom "$"))
  mkNode (kind ++ `antiquot_splice) #[mkAtom "$", nesting, mkAtom "[", mkNullNode contents, mkAtom "]", mkAtom suffix]

-- `$x,*` etc.
def antiquotSuffixSplice? : Syntax → Option SyntaxNodeKind
  | .node _ (.str k "antiquot_suffix_splice") _ => some k
  | _ => none

def isAntiquotSuffixSplice (stx : Syntax) : Bool :=
  antiquotSuffixSplice? stx |>.isSome

-- `$x` in the example above
def getAntiquotSuffixSpliceInner (stx : Syntax) : Syntax :=
  stx[0]

def mkAntiquotSuffixSpliceNode (kind : SyntaxNodeKind) (inner : Syntax) (suffix : String) : Syntax :=
  mkNode (kind ++ `antiquot_suffix_splice) #[inner, mkAtom suffix]

def isTokenAntiquot (stx : Syntax) : Bool :=
  stx.isOfKind `token_antiquot

def isAnyAntiquot (stx : Syntax) : Bool :=
  stx.isAntiquot || stx.isAntiquotSplice || stx.isAntiquotSuffixSplice || stx.isTokenAntiquot

/-- List of `Syntax` nodes in which each succeeding element is the parent of
the current. The associated index is the index of the preceding element in the
list of children of the current element. -/
protected abbrev Stack := List (Syntax × Nat)

/-- Return stack of syntax nodes satisfying `visit`, starting with such a node that also fulfills `accept` (default "is leaf"), and ending with the root. -/
partial def findStack? (root : Syntax) (visit : Syntax → Bool) (accept : Syntax → Bool := fun stx => !stx.hasArgs) : Option Syntax.Stack :=
  if visit root then go [] root else none
where
  go (stack : Syntax.Stack) (stx : Syntax) : Option Syntax.Stack := Id.run do
    if accept stx then
      return (stx, 0) :: stack  -- the first index is arbitrary as there is no preceding element
    for i in [0:stx.getNumArgs] do
      if visit stx[i] then
        if let some stack := go ((stx, i) :: stack) stx[i] then
          return stack
    return none

/-- Compare the `SyntaxNodeKind`s in `pattern` to those of the `Syntax`
elements in `stack`. Return `false` if `stack` is shorter than `pattern`. -/
def Stack.matches (stack : Syntax.Stack) (pattern : List $ Option SyntaxNodeKind) : Bool :=
  stack.length >= pattern.length &&
  (stack
    |>.zipWith (fun (s, _) p => p |>.map (s.isOfKind ·) |>.getD true) pattern
    |>.all id)

end Syntax
end Lean


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.

Authors: Wojciech Nawrocki
-/
import Lean.Linter.UnusedVariables
import Lean.Server.Utils
import Lean.Widget.InteractiveGoal

namespace Lean.Widget
open Lsp Server

inductive StrictOrLazy (α β : Type) : Type
  | strict : α → StrictOrLazy α β
  | lazy : β → StrictOrLazy α β
  deriving Inhabited, RpcEncodable

structure LazyTraceChildren where
  indent : Nat
  children : Array (WithRpcRef MessageData)
  deriving TypeName

inductive MsgEmbed where
  /-- A piece of Lean code with elaboration/typing data.
  Note: does not necessarily correspond to an `Expr`, the name is for RPC API compatibility. -/
  | expr : CodeWithInfos → MsgEmbed
  /-- An interactive goal display. -/
  | goal : InteractiveGoal → MsgEmbed
  /-- Some messages (in particular, traces) are too costly to print eagerly. Instead, we allow
  the user to expand sub-traces interactively. -/
  | trace (indent : Nat) (cls : Name) (msg : TaggedText MsgEmbed) (collapsed : Bool)
      (children : StrictOrLazy (Array (TaggedText MsgEmbed)) (WithRpcRef LazyTraceChildren))
  deriving Inhabited, RpcEncodable

/-- The `message` field is the text of a message possibly containing interactive *embeds* of type
`MsgEmbed`. We maintain the invariant that embeds are stored in `.tag`s with empty `.text` subtrees,
i.e. `.tag embed (.text "")`, because a `MsgEmbed` display involve more than just text. -/
abbrev InteractiveDiagnostic := Lsp.DiagnosticWith (TaggedText MsgEmbed)

deriving instance RpcEncodable for Lsp.DiagnosticWith

namespace InteractiveDiagnostic
open MsgEmbed

def toDiagnostic (diag : InteractiveDiagnostic) : Lsp.Diagnostic :=
  { diag with message := prettyTt diag.message }
where
  prettyTt (tt : TaggedText MsgEmbed) : String :=
    let tt : TaggedText MsgEmbed := tt.rewrite fun
      | .expr tt,  _ => .text tt.stripTags
      | .goal g,   _ => .text (toString g.pretty)
      | .trace .., _ => .text "(trace)"
    tt.stripTags

end InteractiveDiagnostic

private def mkPPContext (nCtx : NamingContext) (ctx : MessageDataContext) : PPContext := {
  env := ctx.env, mctx := ctx.mctx, lctx := ctx.lctx, opts := ctx.opts,
  currNamespace := nCtx.currNamespace, openDecls := nCtx.openDecls
}

/-! The `msgToInteractive` algorithm turns a `MessageData` into `TaggedText MsgEmbed` in two stages.

First, in `msgToInteractiveAux` we produce a `Format` object whose `.tag` nodes refer to `EmbedFmt`
objects stored in an auxiliary array. Only the most shallow `.tag` in every branch through the
`Format` corresponds to an `EmbedFmt`. The kind of this tag determines how the nested `Format`
object (possibly including further `.tag`s), is processed. For example, if the output is
`.tag (.expr ctx infos) fmt` then tags in the nested `fmt` object refer to elements of `infos`.

In the second stage, we recursively transform such a `Format` into `TaggedText MsgEmbed` according
to the rule above by first pretty-printing it and then grabbing data referenced by the tags from
all the nested arrays (such as the `infos` array in the example above).

We cannot easily do the translation in a single `MessageData → TaggedText MsgEmbed` step because
that would effectively require reimplementing the (stateful, to keep track of indentation)
`Format.prettyM` algorithm.
-/

private inductive EmbedFmt
  /-- Nested tags denote `Info` objects in `infos`. -/
  | code (ctx : Elab.ContextInfo) (infos : RBMap Nat Elab.Info compare)
  /-- Nested text is ignored. -/
  | goal (ctx : Elab.ContextInfo) (lctx : LocalContext) (g : MVarId)
  /-- Nested text is ignored. -/
  | trace (cls : Name) (msg : Format) (collapsed : Bool)
    (children : StrictOrLazy (Array Format) (Array MessageData))
  /-- Nested tags are ignored, show nested text as-is. -/
  | ignoreTags
  deriving Inhabited

private abbrev MsgFmtM := StateT (Array EmbedFmt) IO

open MessageData in
private partial def msgToInteractiveAux (msgData : MessageData) : IO (Format × Array EmbedFmt) :=
  go { currNamespace := Name.anonymous, openDecls := [] } none msgData #[]
where
  pushEmbed (e : EmbedFmt) : MsgFmtM Nat :=
    modifyGet fun es => (es.size, es.push e)

  withIgnoreTags (fmt : Format) : MsgFmtM Format := do
    let t ← pushEmbed EmbedFmt.ignoreTags
    return Format.tag t fmt

  mkContextInfo (nCtx : NamingContext) (ctx : MessageDataContext) : Elab.ContextInfo := {
    env           := ctx.env
    mctx          := ctx.mctx
    fileMap       := default
    options       := ctx.opts
    currNamespace := nCtx.currNamespace
    openDecls     := nCtx.openDecls
    -- Hack: to make sure unique ids created at `ppExprWithInfos` do not collide with ones in `ctx.mctx`
    ngen          := { namePrefix := `_diag }
  }

  go (nCtx : NamingContext) : Option MessageDataContext → MessageData → MsgFmtM Format
  | _,         ofFormat fmt             => withIgnoreTags fmt
  | none,      ofPPFormat fmt           => (·.fmt) <$> fmt.pp none
  | some ctx,  ofPPFormat fmt           => do
    let ⟨fmt, infos⟩ ← fmt.pp (mkPPContext nCtx ctx)
    let t ← pushEmbed <| EmbedFmt.code (mkContextInfo nCtx ctx) infos
    return Format.tag t fmt
  | none,      ofGoal mvarId            => pure $ "goal " ++ format (mkMVar mvarId)
  | some ctx,  ofGoal mvarId            =>
    return .tag (← pushEmbed (.goal (mkContextInfo nCtx ctx) ctx.lctx mvarId)) "\n"
  | _,         withContext ctx d        => go nCtx ctx d
  | ctx,       withNamingContext nCtx d => go nCtx ctx d
  | ctx,       tagged _ d               => go nCtx ctx d
  | ctx,       nest n d                 => Format.nest n <$> go nCtx ctx d
  | ctx,       compose d₁ d₂            => do let d₁ ← go nCtx ctx d₁; let d₂ ← go nCtx ctx d₂; pure $ d₁ ++ d₂
  | ctx,       group d                  => Format.group <$> go nCtx ctx d
  | ctx,       .trace cls header children collapsed => do
    let header := (← go nCtx ctx header).nest 4
    let nodes ←
      if collapsed && !children.isEmpty then
        let children := children.map fun child =>
          MessageData.withNamingContext nCtx <|
            match ctx with
            | some ctx => MessageData.withContext ctx child
            | none     => child
        pure (.lazy children)
      else
        pure (.strict (← children.mapM (go nCtx ctx)))
    let e := .trace cls header collapsed nodes
    return .tag (← pushEmbed e) ".\n"

partial def msgToInteractive (msgData : MessageData) (hasWidgets : Bool) (indent : Nat := 0) : IO (TaggedText MsgEmbed) := do
  if !hasWidgets then
    return (TaggedText.prettyTagged (← msgData.format)).rewrite fun _ tt => .text tt.stripTags
  let (fmt, embeds) ← msgToInteractiveAux msgData
  let rec fmtToTT (fmt : Format) (indent : Nat) : IO (TaggedText MsgEmbed) :=
    (TaggedText.prettyTagged fmt indent).rewriteM fun (n, col) tt =>
      match embeds[n]! with
        | .code ctx infos =>
          return .tag (.expr (tagCodeInfos ctx infos tt)) default
        | .goal ctx lctx g =>
          ctx.runMetaM lctx do
            return .tag (.goal (← goalToInteractive g)) default
        | .trace cls msg collapsed children => do
          let col := col + tt.stripTags.length - 2
          let children ←
            match children with
              | .lazy children => pure <| .lazy ⟨{indent := col+2, children := children.map .mk}⟩
              | .strict children => pure <| .strict (← children.mapM (fmtToTT · (col+2)))
          return .tag (.trace indent cls (← fmtToTT msg col) collapsed children) default
        | .ignoreTags => return .text tt.stripTags
  fmtToTT fmt indent

/-- Transform a Lean Message concerning the given text into an LSP Diagnostic. -/
def msgToInteractiveDiagnostic (text : FileMap) (m : Message) (hasWidgets : Bool) : IO InteractiveDiagnostic := do
  let low : Lsp.Position := text.leanPosToLspPos m.pos
  let fullHigh := text.leanPosToLspPos <| m.endPos.getD m.pos
  let high : Lsp.Position := match m.endPos with
    | some endPos =>
      /-
        Truncate messages that are more than one line long.
        This is a workaround to avoid big blocks of "red squiggly lines" on VS Code.
        TODO: should it be a parameter?
      -/
      let endPos := if endPos.line > m.pos.line then { line := m.pos.line + 1, column := 0 } else endPos
      text.leanPosToLspPos endPos
    | none        => low
  let range : Range := ⟨low, high⟩
  let fullRange : Range := ⟨low, fullHigh⟩
  let severity? := some <| match m.severity with
    | .information => .information
    | .warning     => .warning
    | .error       => .error
  let source? := some "Lean 4"
  let tags? :=
    if m.data.isDeprecationWarning then some #[.deprecated]
    else if m.data.isUnusedVariableWarning then some #[.unnecessary]
    else none
  let message ← try
      msgToInteractive m.data hasWidgets
    catch ex =>
      pure <| TaggedText.text s!"[error when printing message: {ex.toString}]"
  pure { range, fullRange? := some fullRange, severity?, source?, message, tags? }

end Lean.Widget


import data.option.basic

/-!

# Model of the Prisma migration engine `devDiagnostic` command

The `devDiagnostic` RPC command acts as a wrapper around `diagnoseMigrationHistory`. Its
role is to interpret the diagnostic output, and translate it to a concrete
action to be performed by the CLI.

The corresponding control flow in the CLI should be:

1. Call `RPC devDiagnostic`. Check the output:
  - Error / BrokenMigration -> display the error (regular user-facing error, no
    CLI code should be needed)
  - Reset -> Prompt the user to reset with the provided reason. Call
    `RPC reset`, then proceed with 2.
  - CreateMigration -> proceed with 2.
2. Call `RPC applyMigrations`
3. If we have no migration name, prompt for it.
4. Check for the `--create-only` flag
  - If it was passed, call `RPC evaluateDataLoss`, show the warnings,
    `RPC createMigration`. Done.
  - Otherwise, call `RPC evaluateDataLoss`, potentially ask for confirmation,
    `RPC createMigration`, `RPC applyMigrations`. Generate the client. Done.

Implemented JSON-RPC API:

```typescript
interface DevDiagnosticInput {}

interface DevDiagnosticOutput {
  action: DevAction
}

type DevAction =
  { tag: "reset", reason: string }
  | { tag: "createMigration" }

```

-/

open except (error ok)

variables { α : Type }
universes u v

/-- The top-level RPC input type. -/
inductive DevInput : Type
| mk : DevInput

inductive ResetReason
| Drifted
| Unspecified

/-- The top-level RPC output type. -/
inductive DevOutput
| CreateMigration
| Reset : ResetReason → DevOutput
-- This manifests itself as a user-facing error output, it does need any special
-- handling in the CLI.
| BrokenMigration : string -> DevOutput

open DevOutput

inductive DriftDiagnostic : Type
| DriftDetected : string -> DriftDiagnostic
| MigrationFailedToApply : string -> DriftDiagnostic

inductive HistoryDiagnostic : Type
| DatabaseIsBehind
| MigrationDirectoryIsBehind
| HistoriesDiverge

structure DiagnoseMigrationHistoryOutput :=
mk ::
  ( drift : option DriftDiagnostic )
  ( history : option HistoryDiagnostic )
  ( failedMigrationNames : list string )
  ( editedMigrationNames : list string )
  ( errorInUnappliedMigrations : option string )
  ( hasMigrationsTable : bool )

def DiagnoseMigrationHistoryOutput.resetReason : DiagnoseMigrationHistoryOutput → option ResetReason :=
λ projectState,
if (
  ¬projectState.failedMigrationNames.is_nil ||
  ¬projectState.editedMigrationNames.is_nil
) then
  some ResetReason.Unspecified
else if ¬projectState.drift.is_none then
  some ResetReason.Drifted
else
  match projectState.history with
  | some HistoryDiagnostic.MigrationDirectoryIsBehind := some ResetReason.Unspecified
  | some HistoryDiagnostic.HistoriesDiverge := some ResetReason.Unspecified
  | _ := none
  end

def DiagnoseMigrationHistoryOutput.brokenMigration : DiagnoseMigrationHistoryOutput → option string :=
λ o,
match (o.drift, o.errorInUnappliedMigrations) with
| ⟨some (DriftDiagnostic.MigrationFailedToApply name), _⟩ := some name
| ⟨_, some name⟩ := some name
| _ := none
end

example : monad id := by apply_instance

/-- Machinery to define early returns. -/
def devState : Type → Type := except_t DevOutput id

instance devStateMonad : monad devState := by { unfold devState, apply_instance }
instance devStateMonadError : monad_except DevOutput devState := by { unfold devState, apply_instance }
instance devStateMonadRun : monad_run (except DevOutput) devState := by { unfold devState, apply_instance }

/-- Check that no migration (applied or unapplied) is broken. -/
def checkBrokenMigration : DiagnoseMigrationHistoryOutput → devState punit :=
λ state, match state.brokenMigration with
| some name := throw $ BrokenMigration name
| none := pure ()
end

/-- Check whether we have a ground for a reset. -/
def checkReset : DiagnoseMigrationHistoryOutput → devState punit :=
λ state, match state.resetReason with
| some reason := throw $ Reset reason
| none := pure ()
end

/-- The model implementation of `dev`. -/
def dev : DevInput → DiagnoseMigrationHistoryOutput → devState DevOutput :=
λ input projectState,
checkBrokenMigration projectState >>
  checkReset projectState >>
  pure CreateMigration

/-- Convenience wrapper around `dev` to make proof types more readable. -/
def runDev : DevInput → DiagnoseMigrationHistoryOutput → DevOutput :=
λ input diagnostics, match run (dev input diagnostics) with
| (error output) := output
| (ok output) := output
end

-- -- ---- ---- ---- ---- ---- ---- ---- ---- -
-- Proofs about `dev`'s model defined above. --
-- ---- ---- ---- ---- ---- ---- ---- ---- ----

/--
If the migrations are working and we should reset, we will always return
`Reset`. -/
theorem devReset :
  ∀ (input : DevInput) (projectState : DiagnoseMigrationHistoryOutput),
  projectState.brokenMigration = none →
  projectState.resetReason.is_some →
  ∃ r, runDev input projectState = Reset r :=
begin
  intros input projectState hBroken hReset,
  delta runDev dev checkBrokenMigration checkReset,
  obtain ⟨r, hSome⟩ : ∃ r, projectState.resetReason = some r, from option.is_some_iff_exists.mp hReset,
  existsi r,
  simp [hReset],
  rw [hSome, hBroken],
  refl
end

/--
Whenever we are not in a reset situation and no migration is broken, we will
return `CreateMigration`. -/
theorem devCreateMigration :
  ∀ (input : DevInput) (projectState : DiagnoseMigrationHistoryOutput),
  projectState.resetReason = none →
  projectState.brokenMigration = none →
  runDev input projectState = CreateMigration :=
begin
  intros input projectState hNoReset hNoBrokenMigration,
  delta runDev dev checkBrokenMigration checkReset,
  rw [hNoReset, hNoBrokenMigration],
  refl
end

/--
`dev` will always return an error before asking for a reset in case a
migration doesn't apply cleanly to the dev database. It never prematurely asks
for a reset. -/
theorem devBrokenMigration :
  ∀ (input : DevInput) (projectState : DiagnoseMigrationHistoryOutput) (brokenMigrationName : string),
  projectState.brokenMigration = some brokenMigrationName →
  runDev input projectState = BrokenMigration brokenMigrationName :=
begin
  rintros input projectState brokenMigrationName hBroken,
  unfold runDev dev checkBrokenMigration,
  simpa [hBroken]
end


namespace Lean
syntax "foo " binderIdent : term
example : Syntax → MacroM Syntax
  | `(foo _) => `(_)
  | `(foo $x:ident) => `($x:ident)
  | _ => `(_)


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Daniel Selsam
-/
import Lean.Meta
import Lean.Util.FindMVar
import Lean.Util.FindLevelMVar
import Lean.Util.CollectLevelParams
import Lean.Util.ReplaceLevel
import Lean.PrettyPrinter.Delaborator.Options
import Lean.PrettyPrinter.Delaborator.SubExpr
import Std.Data.RBMap

/-!
The top-down analyzer is an optional preprocessor to the delaborator that aims
to determine the minimal annotations necessary to ensure that the delaborated
expression can be re-elaborated correctly. Currently, the top-down analyzer
is neither sound nor complete: there may be edge-cases in which the expression
can still not be re-elaborated correctly, and it may also add many annotations
that are not strictly necessary.
-/

namespace Lean

open Lean.Meta
open Std (RBMap)

register_builtin_option pp.analyze : Bool := {
  defValue := false
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) determine annotations sufficient to ensure round-tripping"
}

register_builtin_option pp.analyze.checkInstances : Bool := {
  -- TODO: It would be great to make this default to `true`, but currently, `MessageData` does not
  -- include the `LocalInstances`, so this will be very over-aggressive in inserting instances
  -- that would otherwise be easy to synthesize. We may consider threading the instances in the future,
  -- or at least tracking a bool for whether the instances have been lost.
  defValue := false
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) confirm that instances can be re-synthesized"
}

register_builtin_option pp.analyze.typeAscriptions : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) add type ascriptions when deemed necessary"
}

register_builtin_option pp.analyze.trustSubst : Bool := {
  defValue := false
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) always 'pretend' applications that can delab to ▸ are 'regular'"
}

register_builtin_option pp.analyze.trustOfNat : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) always 'pretend' `OfNat.ofNat` applications can elab bottom-up"
}

register_builtin_option pp.analyze.trustOfScientific : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) always 'pretend' `OfScientific.ofScientific` applications can elab bottom-up"
}

register_builtin_option pp.analyze.trustCoe : Bool := {
  defValue := false
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) always assume a coercion can be correctly inserted"
}

-- TODO: this is an arbitrary special case of a more general principle.
register_builtin_option pp.analyze.trustSubtypeMk : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) assume the implicit arguments of Subtype.mk can be inferred"
}

register_builtin_option pp.analyze.trustId : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) always assume an implicit `fun x => x` can be inferred"
}

register_builtin_option pp.analyze.trustKnownFOType2TypeHOFuns : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) omit higher-order functions whose values seem to be knownType2Type"
}

register_builtin_option pp.analyze.omitMax : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) omit universe `max` annotations (these constraints can actually hurt)"
}

register_builtin_option pp.analyze.knowsType : Bool := {
  defValue := true
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) assume the type of the original expression is known"
}

register_builtin_option pp.analyze.explicitHoles : Bool := {
  defValue := false
  group    := "pp.analyze"
  descr    := "(pretty printer analyzer) use `_` for explicit arguments that can be inferred"
}

def getPPAnalyze                            (o : Options) : Bool := o.get pp.analyze.name pp.analyze.defValue
def getPPAnalyzeCheckInstances              (o : Options) : Bool := o.get pp.analyze.checkInstances.name pp.analyze.checkInstances.defValue
def getPPAnalyzeTypeAscriptions             (o : Options) : Bool := o.get pp.analyze.typeAscriptions.name pp.analyze.typeAscriptions.defValue
def getPPAnalyzeTrustSubst                  (o : Options) : Bool := o.get pp.analyze.trustSubst.name pp.analyze.trustSubst.defValue
def getPPAnalyzeTrustOfNat                  (o : Options) : Bool := o.get pp.analyze.trustOfNat.name pp.analyze.trustOfNat.defValue
def getPPAnalyzeTrustOfScientific           (o : Options) : Bool := o.get pp.analyze.trustOfScientific.name pp.analyze.trustOfScientific.defValue
def getPPAnalyzeTrustId                     (o : Options) : Bool := o.get pp.analyze.trustId.name pp.analyze.trustId.defValue
def getPPAnalyzeTrustCoe                    (o : Options) : Bool := o.get pp.analyze.trustCoe.name pp.analyze.trustCoe.defValue
def getPPAnalyzeTrustSubtypeMk              (o : Options) : Bool := o.get pp.analyze.trustSubtypeMk.name pp.analyze.trustSubtypeMk.defValue
def getPPAnalyzeTrustKnownFOType2TypeHOFuns (o : Options) : Bool := o.get pp.analyze.trustKnownFOType2TypeHOFuns.name pp.analyze.trustKnownFOType2TypeHOFuns.defValue
def getPPAnalyzeOmitMax                     (o : Options) : Bool := o.get pp.analyze.omitMax.name pp.analyze.omitMax.defValue
def getPPAnalyzeKnowsType                   (o : Options) : Bool := o.get pp.analyze.knowsType.name pp.analyze.knowsType.defValue
def getPPAnalyzeExplicitHoles               (o : Options) : Bool := o.get pp.analyze.explicitHoles.name pp.analyze.explicitHoles.defValue

def getPPAnalysisSkip            (o : Options) : Bool := o.get `pp.analysis.skip false
def getPPAnalysisHole            (o : Options) : Bool := o.get `pp.analysis.hole false
def getPPAnalysisNamedArg        (o : Options) : Bool := o.get `pp.analysis.namedArg false
def getPPAnalysisLetVarType      (o : Options) : Bool := o.get `pp.analysis.letVarType false
def getPPAnalysisNeedsType       (o : Options) : Bool := o.get `pp.analysis.needsType false
def getPPAnalysisBlockImplicit   (o : Options) : Bool := o.get `pp.analysis.blockImplicit false

namespace PrettyPrinter.Delaborator

def returnsPi (motive : Expr) : MetaM Bool := do
  lambdaTelescope motive fun xs b => b.isForall

def isNonConstFun (motive : Expr) : MetaM Bool := do
  match motive with
  | Expr.lam name d b _ => isNonConstFun b
  | _ => motive.hasLooseBVars

def isSimpleHOFun (motive : Expr) : MetaM Bool := do
  not (← returnsPi motive) && not (← isNonConstFun motive)

def isType2Type (motive : Expr) : MetaM Bool := do
  match ← inferType motive with
  | Expr.forallE _ (Expr.sort ..) (Expr.sort ..) .. => true
  | _ => false

def isFOLike (motive : Expr) : MetaM Bool := do
  let f := motive.getAppFn
  f.isFVar || f.isConst

def isIdLike (arg : Expr) : Bool := do
  -- TODO: allow `id` constant as well?
  match arg with
  | Expr.lam _ _ (Expr.bvar ..) .. => true
  | _ => false

def isCoe (e : Expr) : Bool :=
  -- TODO: `coeSort? Builtins doesn't seem to render them anyway
  e.isAppOfArity `coe 4
  || (e.isAppOf `coeFun && e.getAppNumArgs >= 4)
  || e.isAppOfArity `coeSort 4

def isStructureInstance (e : Expr) : MetaM Bool := do
  match e.isConstructorApp? (← getEnv) with
  | some s => isStructure (← getEnv) s.induct
  | none   => false

namespace TopDownAnalyze

partial def hasMVarAtCurrDepth (e : Expr) : MetaM Bool := do
  let mctx ← getMCtx
  Option.isSome $ e.findMVar? fun mvarId =>
    match mctx.findDecl? mvarId with
    | some mdecl => mdecl.depth == mctx.depth
    | _ => false

partial def hasLevelMVarAtCurrDepth (e : Expr) : MetaM Bool := do
  let mctx ← getMCtx
  Option.isSome $ e.findLevelMVar? fun mvarId =>
    mctx.findLevelDepth? mvarId == some mctx.depth

private def valUnknown (e : Expr) : MetaM Bool := do
  hasMVarAtCurrDepth (← instantiateMVars e)

private def typeUnknown (e : Expr) : MetaM Bool := do
  valUnknown (← inferType e)

def isHBinOp (e : Expr) : Bool := do
  -- TODO: instead of tracking these explicitly,
  -- consider a more general solution that checks for defaultInstances
  if e.getAppNumArgs != 6 then return false
  let f := e.getAppFn
  if !f.isConst then return false

  -- Note: we leave out `HPow.hPow because we expect its homogeneous
  -- version will change soon
  let ops := #[
    `HOr.hOr, `HXor.hXor, `HAnd.hAnd,
    `HAppend.hAppend, `HOrElse.hOrElse, `HAndThen.hAndThen,
    `HAdd.hAdd, `HSub.hSub, `HMul.hMul, `HDiv.hDiv, `HMod.hMod,
    `HShiftLeft.hShiftLeft, `HShiftRight]
  ops.any fun op => op == f.constName!

def replaceLPsWithVars (e : Expr) : MetaM Expr := do
  if !e.hasLevelParam then return e
  let lps := collectLevelParams {} e |>.params
  let mut replaceMap : Std.HashMap Name Level := {}
  for lp in lps do replaceMap := replaceMap.insert lp (← mkFreshLevelMVar)
  return e.replaceLevel fun
    | Level.param n .. => replaceMap.find! n
    | l => if !l.hasParam then some l else none

def isDefEqAssigning (t s : Expr) : MetaM Bool := do
  withReader (fun ctx => { ctx with config := { ctx.config with assignSyntheticOpaque := true }}) $
    Meta.isDefEq t s

def checkpointDefEq (t s : Expr) : MetaM Bool := do
  Meta.checkpointDefEq (mayPostpone := false) do
    isDefEqAssigning t s

def isHigherOrder (type : Expr) : MetaM Bool := do
  forallTelescopeReducing type fun xs b => xs.size > 0 && b.isSort

def isFunLike (e : Expr) : MetaM Bool := do
  forallTelescopeReducing (← inferType e) fun xs b => xs.size > 0

def isSubstLike (e : Expr) : Bool :=
  e.isAppOfArity `Eq.ndrec 6 || e.isAppOfArity `Eq.rec 6

def nameNotRoundtrippable (n : Name) : Bool :=
  n.hasMacroScopes || isPrivateName n || containsNum n
where
  containsNum
    | Name.str p .. => containsNum p
    | Name.num ..   => true
    | Name.anonymous => false

def mvarName (mvar : Expr) : MetaM Name := do
  (← getMVarDecl mvar.mvarId!).userName

def containsBadMax : Level → Bool
  | Level.succ u ..   => containsBadMax u
  | Level.max u v ..  => (u.hasParam && v.hasParam) || containsBadMax u || containsBadMax v
  | Level.imax u v .. => (u.hasParam && v.hasParam) || containsBadMax u || containsBadMax v
  | _                 => false

open SubExpr

structure Context where
  knowsType   : Bool
  knowsLevel  : Bool -- only constants look at this
  inBottomUp  : Bool := false
  parentIsApp : Bool := false
  subExpr     : SubExpr
  deriving Inhabited

structure State where
  annotations : RBMap Pos Options compare := {}
  postponed   : Array (Expr × Expr) := #[] -- not currently used

abbrev AnalyzeM := ReaderT Context (StateRefT State MetaM)

instance (priority := low) : MonadReaderOf SubExpr AnalyzeM where
  read := Context.subExpr <$> read

instance (priority := low) : MonadWithReaderOf SubExpr AnalyzeM where
  withReader f x := fun ctx => x { ctx with subExpr := f ctx.subExpr }

def tryUnify (e₁ e₂ : Expr) : AnalyzeM Unit := do
  try
    let r ← isDefEqAssigning e₁ e₂
    if !r then modify fun s => { s with postponed := s.postponed.push (e₁, e₂) }
    pure ()
  catch ex =>
    modify fun s => { s with postponed := s.postponed.push (e₁, e₂) }

partial def inspectOutParams (arg mvar : Expr) : AnalyzeM Unit := do
  let argType  ← inferType arg -- HAdd α α α
  let mvarType ← inferType mvar
  let fType ← inferType argType.getAppFn -- Type → Type → outParam Type
  let mType ← inferType mvarType.getAppFn
  inspectAux fType mType 0 argType.getAppArgs mvarType.getAppArgs
where
  inspectAux (fType mType : Expr) (i : Nat) (args mvars : Array Expr) := do
    let fType ← whnf fType
    let mType ← whnf mType
    if not (i < args.size) then return ()
    match fType, mType with
    | Expr.forallE _ fd fb _, Expr.forallE _ md mb _ => do
      -- TODO: do I need to check (← okBottomUp? args[i] mvars[i] fuel).isSafe here?
      -- if so, I'll need to take a callback
      if isOutParam fd then
        tryUnify (args[i]) (mvars[i])
      inspectAux (fb.instantiate1 args[i]) (mb.instantiate1 mvars[i]) (i+1) args mvars
    | _, _ => return ()

partial def isTrivialBottomUp (e : Expr) : AnalyzeM Bool := do
  let opts ← getOptions
  return e.isFVar
         || e.isConst || e.isMVar || e.isNatLit || e.isStringLit || e.isSort
         || (getPPAnalyzeTrustOfNat opts && e.isAppOfArity `OfNat.ofNat 3)
         || (getPPAnalyzeTrustOfScientific opts && e.isAppOfArity `OfScientific.ofScientific 5)

partial def canBottomUp (e : Expr) (mvar? : Option Expr := none) (fuel : Nat := 10) : AnalyzeM Bool := do
  -- Here we check if `e` can be safely elaborated without its expected type.
  -- These are incomplete (and possibly unsound) heuristics.
  -- TODO: do I need to snapshot the state before calling this?
  match fuel with
  | 0 => false
  | fuel + 1 =>
    if ← isTrivialBottomUp e then return true
    let f := e.getAppFn
    if !f.isConst && !f.isFVar then return false
    let args := e.getAppArgs
    let fType ← replaceLPsWithVars (← inferType e.getAppFn)
    let (mvars, bInfos, resultType) ← forallMetaBoundedTelescope fType e.getAppArgs.size
    for i in [:mvars.size] do
      if bInfos[i] == BinderInfo.instImplicit then
        inspectOutParams args[i] mvars[i]
      else if ← bInfos[i] == BinderInfo.default then
        if ← isTrivialBottomUp args[i] then tryUnify args[i] mvars[i]
        else if ← typeUnknown mvars[i] <&&> canBottomUp args[i] mvars[i] fuel then tryUnify args[i] mvars[i]
    if ← (isHBinOp e <&&> (valUnknown mvars[0] <||> valUnknown mvars[1])) then tryUnify mvars[0] mvars[1]
    if mvar?.isSome then tryUnify resultType (← inferType mvar?.get!)
    return !(← valUnknown resultType)

def withKnowing (knowsType knowsLevel : Bool) (x : AnalyzeM α) : AnalyzeM α := do
  withReader (fun ctx => { ctx with knowsType := knowsType, knowsLevel := knowsLevel }) x

builtin_initialize analyzeFailureId : InternalExceptionId ← registerInternalExceptionId `analyzeFailure

def checkKnowsType : AnalyzeM Unit := do
  if not (← read).knowsType then
    throw $ Exception.internal analyzeFailureId

def annotateBoolAt (n : Name) (pos : Pos) : AnalyzeM Unit := do
  let opts := (← get).annotations.findD pos {} |>.setBool n true
  trace[pp.analyze.annotate] "{pos} {n}"
  modify fun s => { s with annotations := s.annotations.insert pos opts }

def annotateBool (n : Name) : AnalyzeM Unit := do
  annotateBoolAt n (← getPos)

structure App.Context where
  f               : Expr
  fType           : Expr
  args            : Array Expr
  mvars           : Array Expr
  bInfos          : Array BinderInfo
  forceRegularApp : Bool

structure App.State where
  bottomUps       : Array Bool
  higherOrders    : Array Bool
  funBinders      : Array Bool
  provideds       : Array Bool
  namedArgs       : Array Name := #[]

abbrev AnalyzeAppM := ReaderT App.Context (StateT App.State AnalyzeM)

mutual

  partial def analyze (parentIsApp : Bool := false) : AnalyzeM Unit := do
    checkMaxHeartbeats "Delaborator.topDownAnalyze"
    trace[pp.analyze] "{(← read).knowsType}.{(← read).knowsLevel}"
    let e ← getExpr
    let opts ← getOptions
    if ← !e.isAtomic <&&> !(getPPProofs opts) <&&> (try Meta.isProof e catch ex => false) then
      if getPPProofsWithType opts then
        withType $ withKnowing true true $ analyze
      return ()
    else
      withReader (fun ctx => { ctx with parentIsApp := parentIsApp }) do
        match (← getExpr) with
        | Expr.app ..     => analyzeApp
        | Expr.forallE .. => analyzePi
        | Expr.lam ..     => analyzeLam
        | Expr.const ..   => analyzeConst
        | Expr.sort ..    => analyzeSort
        | Expr.proj ..    => analyzeProj
        | Expr.fvar ..    => analyzeFVar
        | Expr.mdata ..   => analyzeMData
        | Expr.letE ..    => analyzeLet
        | Expr.lit ..     => pure ()
        | Expr.mvar ..    => pure ()
        | Expr.bvar ..    => pure ()
  where
    analyzeApp := do
      let mut willKnowType := (← read).knowsType
      if !(← read).knowsType && !(← canBottomUp (← getExpr)) then
        annotateBool `pp.analysis.needsType
        withType $ withKnowing true false $ analyze
        willKnowType := true

      else if ← (!(← read).knowsType <||> (← read).inBottomUp) <&&> isStructureInstance (← getExpr) then
        withType do
          annotateBool `pp.structureInstanceTypes
          withKnowing true false $ analyze
        willKnowType := true

      withKnowing willKnowType true $ analyzeAppStaged (← getExpr).getAppFn (← getExpr).getAppArgs

    analyzeAppStaged (f : Expr) (args : Array Expr) : AnalyzeM Unit := do
      let fType ← replaceLPsWithVars (← inferType f)
      let (mvars, bInfos, resultType) ← forallMetaBoundedTelescope fType args.size
      let rest := args.extract mvars.size args.size
      let args := args.shrink mvars.size

      -- Unify with the expected type
      if (← read).knowsType then tryUnify (← inferType (mkAppN f args)) resultType

      let forceRegularApp : Bool :=
        (getPPAnalyzeTrustSubst (← getOptions) && isSubstLike (← getExpr))
        || (getPPAnalyzeTrustCoe (← getOptions) && isCoe (← getExpr))
        || (getPPAnalyzeTrustSubtypeMk (← getOptions) && (← getExpr).isAppOfArity `Subtype.mk 4)

      analyzeAppStagedCore { f, fType, args, mvars, bInfos, forceRegularApp } |>.run' {
        bottomUps    := mkArray args.size false,
        higherOrders := mkArray args.size false,
        provideds    := mkArray args.size false,
        funBinders   := mkArray args.size false
      }

      if not rest.isEmpty then
        -- Note: this shouldn't happen for type-correct terms
        if !args.isEmpty then
          analyzeAppStaged (mkAppN f args) rest

    maybeAddBlockImplicit : AnalyzeM Unit := do
      -- See `MonadLift.noConfusion for an example where this is necessary.
      if !(← read).parentIsApp then
        let type ← inferType (← getExpr)
        if type.isForall && type.bindingInfo! == BinderInfo.implicit then
          annotateBool `pp.analysis.blockImplicit

    analyzeConst : AnalyzeM Unit := do
      let Expr.const n ls .. ← getExpr | unreachable!
      if !(← read).knowsLevel && !ls.isEmpty then
        -- TODO: this is a very crude heuristic, motivated by https://github.com/leanprover/lean4/issues/590
        unless getPPAnalyzeOmitMax (← getOptions) && ls.any containsBadMax do
        annotateBool `pp.universes
      maybeAddBlockImplicit

    analyzePi : AnalyzeM Unit := do
      withBindingDomain $ withKnowing true false analyze
      withBindingBody Name.anonymous analyze

    analyzeLam : AnalyzeM Unit := do
      if !(← read).knowsType then annotateBool `pp.funBinderTypes
      withBindingDomain $ withKnowing true false analyze
      withBindingBody Name.anonymous analyze

    analyzeLet : AnalyzeM Unit := do
      let Expr.letE n t v body .. ← getExpr | unreachable!
      if !(← canBottomUp v) then
        annotateBool `pp.analysis.letVarType
        withLetVarType $ withKnowing true false analyze
        withLetValue $ withKnowing true true analyze
      else
        withReader (fun ctx => { ctx with inBottomUp := true }) do
          withLetValue $ withKnowing true true analyze

      withLetBody analyze

    analyzeSort  : AnalyzeM Unit := pure ()
    analyzeProj  : AnalyzeM Unit := withProj analyze
    analyzeFVar  : AnalyzeM Unit := maybeAddBlockImplicit
    analyzeMData : AnalyzeM Unit := withMDataExpr analyze

  partial def analyzeAppStagedCore : AnalyzeAppM Unit := do
    collectBottomUps
    checkOutParams
    collectHigherOrders
    hBinOpHeuristic
    collectTrivialBottomUps
    discard <| processPostponed (mayPostpone := true)
    applyFunBinderHeuristic
    analyzeFn
    for i in [:(← read).args.size] do analyzeArg i
    maybeSetExplicit

  where
    collectBottomUps := do
      let { args, mvars, bInfos, ..} ← read
      for target in [fun _ => none, fun i => some mvars[i]] do
        for i in [:args.size] do
          if bInfos[i] == BinderInfo.default then
            if ← typeUnknown mvars[i] <&&> canBottomUp args[i] (target i) then
              tryUnify args[i] mvars[i]
              modify fun s => { s with bottomUps := s.bottomUps.set! i true }

    checkOutParams := do
      let { args, mvars, bInfos, ..} ← read
      for i in [:args.size] do
        if bInfos[i] == BinderInfo.instImplicit then inspectOutParams args[i] mvars[i]

    collectHigherOrders := do
      let { args, mvars, bInfos, ..} ← read
      for i in [:args.size] do
        if not (bInfos[i] == BinderInfo.implicit || bInfos[i] == BinderInfo.strictImplicit) then continue
        if not (← isHigherOrder (← inferType args[i])) then continue
        if getPPAnalyzeTrustId (← getOptions) && isIdLike args[i] then continue

        if getPPAnalyzeTrustKnownFOType2TypeHOFuns (← getOptions) && not (← valUnknown mvars[i])
          && (← isType2Type (args[i])) && (← isFOLike (args[i])) then continue

        tryUnify args[i] mvars[i]
        modify fun s => { s with higherOrders := s.higherOrders.set! i true }

    hBinOpHeuristic := do
      let { args, mvars, bInfos, ..} ← read
      if ← (isHBinOp (← getExpr) <&&> (valUnknown mvars[0] <||> valUnknown mvars[1])) then
        tryUnify mvars[0] mvars[1]

    collectTrivialBottomUps := do
      -- motivation: prevent levels from printing in
      -- Boo.mk : {α : Type u_1} → {β : Type u_2} → α → β → Boo.{u_1, u_2} α β
      let { args, mvars, bInfos, ..} ← read
      for i in [:args.size] do
        if bInfos[i] == BinderInfo.default then
          if ← valUnknown mvars[i] <&&> isTrivialBottomUp args[i] then
            tryUnify args[i] mvars[i]
            modify fun s => { s with bottomUps := s.bottomUps.set! i true }

    applyFunBinderHeuristic := do
      let { f, args, mvars, bInfos, .. } ← read

      let rec core (argIdx : Nat) (mvarType : Expr) : AnalyzeAppM Bool := do
        match ← getExpr, mvarType with
        | Expr.lam .., Expr.forallE n t b .. =>
          let mut annotated := false
          for i in [:argIdx] do
            if ← bInfos[i] == BinderInfo.implicit <&&> valUnknown mvars[i] <&&> withNewMCtxDepth (checkpointDefEq t mvars[i]) then
              annotateBool `pp.funBinderTypes
              tryUnify args[i] mvars[i]
              -- Note: currently we always analyze the lambda binding domains in `analyzeLam`
              -- (so we don't need to analyze it again here)
              annotated := true
              break
          let annotatedBody ← withBindingBody Name.anonymous (core argIdx b)
          return annotated || annotatedBody

        | _, _ => return false

      for i in [:args.size] do
        if ← bInfos[i] == BinderInfo.default then
          let b ← withNaryArg i (core i (← inferType mvars[i]))
          if b then modify fun s => { s with funBinders := s.funBinders.set! i true }

    analyzeFn := do
      -- Now, if this is the first staging, analyze the n-ary function without expected type
      let {f, fType, forceRegularApp ..} ← read
      if !f.isApp then withKnowing false (forceRegularApp || !(← hasLevelMVarAtCurrDepth (← instantiateMVars fType))) $ withNaryFn (analyze (parentIsApp := true))

    annotateNamedArg (n : Name) : AnalyzeAppM Unit := do
      annotateBool `pp.analysis.namedArg
      modify fun s => { s with namedArgs := s.namedArgs.push n }

    analyzeArg (i : Nat) := do
      let { f, args, mvars, bInfos, forceRegularApp ..} ← read
      let { bottomUps, higherOrders, funBinders, ..} ← get
      let arg := args[i]
      let argType ← inferType arg

      let processNaturalImplicit : AnalyzeAppM Unit := do
        if (← valUnknown mvars[i] <||> higherOrders[i]) && !forceRegularApp then
          annotateNamedArg (← mvarName mvars[i])
          modify fun s => { s with provideds := s.provideds.set! i true }
        else
          annotateBool `pp.analysis.skip

      withNaryArg (f.getAppNumArgs + i) do
        withTheReader Context (fun ctx => { ctx with inBottomUp := ctx.inBottomUp || bottomUps[i] }) do

          match bInfos[i] with
          | BinderInfo.default =>
            if ← getPPAnalyzeExplicitHoles (← getOptions) <&&> !(← valUnknown mvars[i]) <&&> !(← readThe Context).inBottomUp <&&> !(← isFunLike arg) <&&> !funBinders[i] <&&> checkpointDefEq mvars[i] arg then
              annotateBool `pp.analysis.hole
            else
              modify fun s => { s with provideds := s.provideds.set! i true }

          | BinderInfo.implicit => processNaturalImplicit
          | BinderInfo.strictImplicit => processNaturalImplicit

          | BinderInfo.instImplicit =>
            -- Note: apparently checking valUnknown here is not sound, because the elaborator
            -- will not happily assign instImplicits that it cannot synthesize
            let mut provided := true
            if !getPPInstances (← getOptions) then
              annotateBool `pp.analysis.skip
              provided := false
            else if getPPAnalyzeCheckInstances (← getOptions) then
              let instResult ← try trySynthInstance argType catch _ => LOption.undef
              match instResult with
              | LOption.some inst =>
                if ← checkpointDefEq inst arg then annotateBool `pp.analysis.skip; provided := false
                else annotateNamedArg (← mvarName mvars[i])
              | _                 => annotateNamedArg (← mvarName mvars[i])
            else annotateBool `pp.analysis.skip; provided := false
            modify fun s => { s with provideds := s.provideds.set! i provided }
          | BinderInfo.auxDecl => pure ()
          if (← get).provideds[i] then withKnowing (not (← typeUnknown mvars[i])) true analyze
          tryUnify mvars[i] args[i]

    maybeSetExplicit := do
      let { f, args, mvars, bInfos, forceRegularApp, ..} ← read
      if (← get).namedArgs.any nameNotRoundtrippable then
        annotateBool `pp.explicit
        for i in [:args.size] do
          if !(← get).provideds[i] then
            withNaryArg (f.getAppNumArgs + i) do annotateBool `pp.analysis.hole
          if bInfos[i] == BinderInfo.instImplicit && getPPInstanceTypes (← getOptions) then
            withType (withKnowing true false analyze)

end

end TopDownAnalyze

open TopDownAnalyze SubExpr

def topDownAnalyze (e : Expr) : MetaM OptionsPerPos := do
  let s₀ ← get
  traceCtx `pp.analyze do
    withReader (fun ctx => { ctx with config := Lean.Elab.Term.setElabConfig ctx.config }) do
      let ϕ : AnalyzeM OptionsPerPos := do withNewMCtxDepth analyze; (← get).annotations
      try
        let knowsType := getPPAnalyzeKnowsType (← getOptions)
        ϕ { knowsType := knowsType, knowsLevel := knowsType, subExpr := mkRoot e }
          |>.run' { : TopDownAnalyze.State }
      catch ex =>
        trace[pp.analyze.error] "failed"
        pure {}
      finally set s₀

builtin_initialize
  registerTraceClass `pp.analyze
  registerTraceClass `pp.analyze.annotate
  registerTraceClass `pp.analyze.tryUnify
  registerTraceClass `pp.analyze.error

end Lean.PrettyPrinter.Delaborator