InfoTree.lean

/-
Copyright (c) 2025 Marc Huisinga. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Marc Huisinga, Thomas R. Murrills
-/
module

public import Mathlib.Lean.Environment
public import Lean.Server.InfoUtils
public import Lean.Meta.TryThis
public import Batteries.Tactic.Lint.Misc
-- Import this linter explicitly to ensure that
-- this file has a valid copyright header and module docstring.
import Mathlib.Tactic.Linter.Header  --shake: keep
public import Batteries.Tactic.Lint.Basic
import Lean.Elab.Term.TermElabM

/-!
# Additions to `Lean.Elab.InfoTree.Main`
-/

@[expose] public section

namespace Lean.Elab

open Lean.Meta
open Lean.Meta.Tactic.TryThis

/--
Collects all suggestions from all `TryThisInfo`s in `trees`.
Does not require context - works with context-free trees.
-/
partial def collectTryThisSuggestions (trees : PersistentArray InfoTree) : Array Suggestion :=
  trees.foldl (init := #[]) fun acc tree => go acc tree
where
  /-- Traverses an `InfoTree` to collect `TryThisInfo` suggestions. -/
  go (acc : Array Suggestion) : InfoTree → Array Suggestion
    | .context _ t => go acc t
    | .node i children =>
      let acc := match i with
        | .ofCustomInfo ci =>
          match ci.value.get? TryThisInfo with
          | some tti => acc.push tti.suggestion
          | none => acc
        | _ => acc
      children.foldl (init := acc) fun acc tree => go acc tree
    | .hole _ => acc

namespace InfoTree

/--
Finds the first result of `← f ctx info children` which is `some a`, descending the
tree from the top. Merges and updates contexts as it descends the tree.

`f` is **only** evaluated on nodes when some context is present. An initial context should be
provided via the `ctx?` argument if invoking `findSomeM?` during a larger traversal of the
infotree. A failure to provide `ctx? := some ctx` when `t` is not the outermost `InfoTree` is thus
likely to cause `findSomeM?` to always return `none`.
-/
partial def findSomeM? {m : Type → Type} [Monad m] {α}
    (f : ContextInfo → Info → PersistentArray InfoTree → m (Option α))
    (t : InfoTree) (ctx? : Option ContextInfo := none) : m (Option α) :=
  go ctx? t
where
  /-- Accumulates contexts and visits nodes if `ctx?` is not `none`. -/
  go ctx?
  | context ctx t => go (ctx.mergeIntoOuter? ctx?) t
  | node i ts => do
    let a ← match ctx? with
      | none => pure none
      | some ctx => f ctx i ts
    match a with
    | some a => pure a
    | none => ts.findSomeM? (go <| i.updateContext? ctx?)
  | hole _ => pure none

/--
Finds the first result of `f ctx info children` which is `some a`, descending the
tree from the top. Merges and updates contexts as it descends the tree.

`f` is **only** evaluated on nodes when some context is present. An initial context should be
provided via the `ctx?` argument if invoking `findSome?` during a larger traversal of the infotree.
A failure to provide `ctx? := some ctx` when `t` is not the outermost `InfoTree` is thus likely to
cause `findSome?` to always return `none`.
-/
def findSome? {α} (f : ContextInfo → Info → PersistentArray InfoTree → Option α)
    (t : InfoTree) (ctx? : Option ContextInfo := none) : Option α :=
  Id.run <| t.findSomeM? f ctx?

/--
Returns the value of `f ctx info children` on the outermost `.node info children` which has
context, having merged and updated contexts appropriately.

If `ctx?` is `some ctx`, `ctx` is used as an initial context. A `ctx?` of `none` should **only** be
used when operating on the first node of the entire infotree. Otherwise, it is likely that no
context will be found.
-/
def onHighestNode? {α} (t : InfoTree) (ctx? : Option ContextInfo)
    (f : ContextInfo → Info → PersistentArray InfoTree → α) : Option α :=
  t.findSome? (ctx? := ctx?) fun ctx i ch => some (f ctx i ch)

/--
Returns the context and `info` on the outermost `.node info _` which has
context, having merged and updated contexts appropriately.

If `ctx?` is `some ctx`, `ctx` is used as an initial context. A `ctx?` of `none` should **only** be
used when operating on the first node of the entire infotree. Otherwise, it is likely that no
context will be found.
-/
def getHighestInfo? (t : InfoTree) (ctx? : Option ContextInfo) : Option (ContextInfo × Info) :=
  t.onHighestNode? ctx? fun ctx i _ => (ctx, i)

/--
Get the `parentDecl`s of every elaborated body.

This includes `let rec`/`where` definitions, but excludes decls without "bodies" (such as
`alias`es, `structure`s, declarations generated by attributes like `@[ext]`, and so on) as we
might find by considering every `parentDeclCtx` throughout the infotree.

Assumes that every body elaboration proceeds through `Lean.Elab.Term.BodyInfo`.
-/
def getDeclsByBody (t : InfoTree) : List Name :=
  t.collectNodesBottomUp fun ctx i _ decls =>
    match i with
    | .ofCustomInfo i =>
      if i.value.typeName == ``Lean.Elab.Term.BodyInfo then
        if let some decl := ctx.parentDecl? then
          decl :: decls
        else decls
      else decls
    | _ => decls

/-- Gets the first child info of each `Lean.Elab.BodyInfo`, which should be the only child, and
should be a `TermInfo`, `PartialTermInfo`, or `TacticInfo`. `getDeclBodyInfos` does not validate
either of these conditions. -/
def getDeclBodyInfos (t : InfoTree) : List (Syntax × ContextInfo × Info) :=
  t.foldInfoTree (init := []) fun ctx t acc =>
    match t with
    | .node (.ofCustomInfo i) body => Id.run do
      if i.value.typeName == ``Lean.Elab.Term.BodyInfo then
        if h : 0 < body.size then
          -- See through `.context`s instead of just matching on `.node`:
          let result? := body[0].getHighestInfo? ctx
          if let some result := result? then
            return (i.stx, result) :: acc
      return acc
    | _ => acc

/--
Get the declarations elaborated in the infotree `t` which are theorems according to the
environment. This includes e.g. `instance`s of `Prop` classes in addition to declarations declared
using the keyword `theorem` directly.
-/
def getTheorems (t : InfoTree) (env : Environment) : List ConstantVal :=
  t.getDeclsByBody.filterMap env.findTheoremConstVal?

end InfoTree

namespace Info

/-- Gets the local context, and the expected type of the `Info`.
Handles `TacticInfo`s (looking at the first goal), `TermInfo`s, and `PartialTermInfo`s.
Does not get the metavariable context; assumes that the caller has accumulated an ambient
`ContextInfo` at this point which is sufficient. -/
def getLCtx? : Info → Option (LocalContext × Option Expr)
  | .ofTacticInfo i => do
    let g ← i.goalsBefore.head?
    let decl ← i.mctxBefore.findDecl? g
    some (decl.lctx, decl.type)
  | .ofTermInfo i
  | .ofPartialTermInfo i => some (i.lctx, i.expectedType?)
  | _ => none

end Lean.Elab.Info