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

prelude
public import Lean.MonadEnv

public section

/-!
# Environment extension for declaration ranges
-/

namespace Lean

builtin_initialize builtinDeclRanges : IO.Ref (NameMap DeclarationRanges) ← IO.mkRef {}
builtin_initialize declRangeExt : MapDeclarationExtension DeclarationRanges ←
  mkMapDeclarationExtension (exportEntriesFn := fun _ s level =>
    if level < .server then
      #[]
    else s.toArray)

def addBuiltinDeclarationRanges (declName : Name) (declRanges : DeclarationRanges) : IO Unit :=
  builtinDeclRanges.modify (·.insert declName declRanges)

def addDeclarationRanges [Monad m] [MonadEnv m] (declName : Name) (declRanges : DeclarationRanges) : m Unit := do
  if declName.isAnonymous then
    -- This can happen on elaboration of partial syntax and would panic in `modifyState` otherwise
    return
  modifyEnv fun env => declRangeExt.insert env declName declRanges

def findDeclarationRangesCore? [Monad m] [MonadEnv m] (declName : Name) : m (Option DeclarationRanges) :=
  -- In the case of private definitions imported via `import all`, looking in `.olean.server` is not
  -- sufficient, so we look in the actual environment as well via `exported` (TODO: rethink
  -- parameter naming).
  return declRangeExt.find? (level := .exported) (← getEnv) declName <|>
    declRangeExt.find? (level := .server) (← getEnv) declName

def findDeclarationRanges? [Monad m] [MonadEnv m] [MonadLiftT BaseIO m] (declName : Name) : m (Option DeclarationRanges) := do
  let env ← getEnv
  let ranges ← if isAuxRecursor env declName || isNoConfusion env declName || (← isRec declName)  then
    findDeclarationRangesCore? declName.getPrefix
  else
    findDeclarationRangesCore? declName
  match ranges with
  | none => return (← builtinDeclRanges.get (m := BaseIO)).find? declName
  | some _ => return ranges

end Lean