{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PartialTypeSignatures #-}

-- | Generally useful definition used in various places in the
-- language server implementation.
module Futhark.LSP.Tool
  ( getHoverInfoFromState,
    findDefinitionRange,
    rangeFromLoc,
    posToUri,
    computeMapping,
    transformVFS,
    logWithSeverity,
    Execute,
    bindingsInRange,
    bindingsInState,
    filterByLspRange,
  )
where

import Colog.Core (LogAction, Severity, WithSeverity (WithSeverity))
import Control.Lens.Getter ((^.))
import Control.Monad.Except (ExceptT)
import Data.Foldable qualified as F
import Data.Function ((&))
import Data.Functor.Contravariant (contramap)
import Data.Map qualified as M
import Data.Text (Text)
import Data.Text qualified as T
import Futhark.Compiler.Program (lpImports)
import Futhark.LSP.PositionMapping
  ( PositionMapping,
    mappingFromDiff,
    toCurrentLoc,
    toStalePos,
  )
import Futhark.LSP.State (State (..), getStaleContent, getStaleMapping)
import Futhark.Util.Loc (Loc (Loc, NoLoc), Pos (Pos), locOf)
import Futhark.Util.Pretty (prettyText)
import Language.Futhark.Core (VName, isBuiltinLoc)
import Language.Futhark.Query
  ( AtPos (AtName),
    BoundTo (..),
    allBindings,
    atPos,
    boundLoc,
    termBindingType,
  )
import Language.LSP.Logging (logToLogMessage)
import Language.LSP.Protocol.Types
  ( ErrorCodes,
    Hover (Hover),
    LSPErrorCodes,
    Location (Location),
    MarkupContent (MarkupContent),
    MarkupKind (MarkupKind_PlainText),
    Position (Position),
    Range (Range),
    UInt,
    Uri,
    filePathToUri,
    fromNormalizedFilePath,
    toNormalizedUri,
    uriToNormalizedFilePath,
    type (|?) (InL),
  )
import Language.LSP.Server (LspM, LspT, MonadLsp, getVirtualFile)
import Language.LSP.VFS (VFS, VirtualFile, vfsMap, virtualFileText, virtualFileVersion)
import Language.LSP.VFS qualified as VFS

-- | Request Handler code usually runs in this monad
type Execute a = ExceptT (Text, LSPErrorCodes |? ErrorCodes) (LspT () IO) a

-- | Retrieve hover info for the definition referenced at the given
-- file at the given line and column number (the two 'Int's).
getHoverInfoFromState :: State -> Maybe FilePath -> Int -> Int -> Maybe Hover
getHoverInfoFromState :: State -> Maybe FilePath -> Int -> Int -> Maybe Hover
getHoverInfoFromState State
state (Just FilePath
path) Int
l Int
c = do
  AtName _ (Just def) loc <- State -> Pos -> Maybe AtPos
queryAtPos State
state (Pos -> Maybe AtPos) -> Pos -> Maybe AtPos
forall a b. (a -> b) -> a -> b
$ FilePath -> Int -> Int -> Int -> Pos
Pos FilePath
path Int
l Int
c Int
0
  let msg =
        case BoundTo
def of
          BoundTerm TermBinding
term Loc
_ -> TypeBase Size NoMode -> Text
forall a. Pretty a => a -> Text
prettyText (TypeBase Size NoMode -> Text) -> TypeBase Size NoMode -> Text
forall a b. (a -> b) -> a -> b
$ TermBinding -> TypeBase Size NoMode
termBindingType TermBinding
term
          BoundModule {} -> Text
"module"
          BoundModuleType {} -> Text
"module type"
          BoundType {} -> Text
"type"
      ms = MarkupKind -> Text -> MarkupContent
MarkupContent MarkupKind
MarkupKind_PlainText Text
msg
  Just $ Hover (InL ms) (Just (rangeFromLoc loc))
getHoverInfoFromState State
_ Maybe FilePath
_ Int
_ Int
_ = Maybe Hover
forall a. Maybe a
Nothing

-- | Find the location of the definition referenced at the given file
-- at the given line and column number (the two 'Int's).
findDefinitionRange :: State -> Maybe FilePath -> Int -> Int -> Maybe Location
findDefinitionRange :: State -> Maybe FilePath -> Int -> Int -> Maybe Location
findDefinitionRange State
state (Just FilePath
path) Int
l Int
c = do
  -- some unnessecary operations inside `queryAtPos` for this function
  -- but shouldn't affect performance much since "Go to definition" is called less frequently
  AtName _qn (Just bound) _loc <- State -> Pos -> Maybe AtPos
queryAtPos State
state (Pos -> Maybe AtPos) -> Pos -> Maybe AtPos
forall a b. (a -> b) -> a -> b
$ FilePath -> Int -> Int -> Int -> Pos
Pos FilePath
path Int
l Int
c Int
0
  let loc = BoundTo -> Loc
boundLoc BoundTo
bound
      Loc (Pos file_path _ _ _) _ = loc
  if isBuiltinLoc loc
    then Nothing
    else Just $ Location (filePathToUri file_path) (rangeFromLoc loc)
findDefinitionRange State
_ Maybe FilePath
_ Int
_ Int
_ = Maybe Location
forall a. Maybe a
Nothing

-- | Query the AST for information at certain Pos.
queryAtPos :: State -> Pos -> Maybe AtPos
queryAtPos :: State -> Pos -> Maybe AtPos
queryAtPos State
state Pos
pos = do
  let Pos FilePath
path Int
_ Int
_ Int
_ = Pos
pos
      mapping :: Maybe PositionMapping
mapping = State -> FilePath -> Maybe PositionMapping
getStaleMapping State
state FilePath
path
  loaded_prog <- State -> Maybe LoadedProg
stateProgram State
state
  stale_pos <- toStalePos mapping pos
  query_result <- atPos (lpImports loaded_prog) stale_pos
  updateAtPos mapping query_result pos state

-- Update the 'AtPos' with the current mapping.
updateAtPos :: Maybe PositionMapping -> AtPos -> Pos -> State -> Maybe AtPos
updateAtPos :: Maybe PositionMapping -> AtPos -> Pos -> State -> Maybe AtPos
updateAtPos Maybe PositionMapping
mapping (AtName QualName VName
qn (Just BoundTo
def) Loc
loc) Pos
pos State
state = do
  let def_loc :: Loc
def_loc = BoundTo -> Loc
boundLoc BoundTo
def
      Loc (Pos FilePath
def_file Int
_ Int
_ Int
_) Pos
_ = Loc
def_loc
      Pos FilePath
current_file Int
_ Int
_ Int
_ = Pos
pos
  current_loc <- Maybe PositionMapping -> Loc -> Maybe Loc
toCurrentLoc Maybe PositionMapping
mapping Loc
loc
  if def_file == current_file
    then do
      current_def_loc <- toCurrentLoc mapping def_loc
      Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc
    else do
      -- Defined in another file, get the corresponding PositionMapping.
      let def_mapping = State -> FilePath -> Maybe PositionMapping
getStaleMapping State
state FilePath
def_file
      current_def_loc <- toCurrentLoc def_mapping def_loc
      Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc
  where
    updateBoundLoc :: BoundTo -> Loc -> BoundTo
    updateBoundLoc :: BoundTo -> Loc -> BoundTo
updateBoundLoc (BoundTerm TermBinding
t Loc
_loc) Loc
current_loc = TermBinding -> Loc -> BoundTo
BoundTerm TermBinding
t Loc
current_loc
    updateBoundLoc (BoundModule Loc
_loc) Loc
current_loc = Loc -> BoundTo
BoundModule Loc
current_loc
    updateBoundLoc (BoundModuleType Loc
_loc) Loc
current_loc = Loc -> BoundTo
BoundModuleType Loc
current_loc
    updateBoundLoc (BoundType Loc
_loc) Loc
current_loc = Loc -> BoundTo
BoundType Loc
current_loc
updateAtPos Maybe PositionMapping
_ AtPos
_ Pos
_ State
_ = Maybe AtPos
forall a. Maybe a
Nothing

-- | Entry point for computing PositionMapping.
computeMapping :: State -> Maybe FilePath -> LspM () (Maybe PositionMapping)
computeMapping :: State -> Maybe FilePath -> LspM () (Maybe PositionMapping)
computeMapping State
state (Just FilePath
file_path) = do
  virtual_file <- NormalizedUri -> LspT () IO (Maybe VirtualFile)
forall config (m :: * -> *).
MonadLsp config m =>
NormalizedUri -> m (Maybe VirtualFile)
getVirtualFile (NormalizedUri -> LspT () IO (Maybe VirtualFile))
-> NormalizedUri -> LspT () IO (Maybe VirtualFile)
forall a b. (a -> b) -> a -> b
$ Uri -> NormalizedUri
toNormalizedUri (Uri -> NormalizedUri) -> Uri -> NormalizedUri
forall a b. (a -> b) -> a -> b
$ FilePath -> Uri
filePathToUri FilePath
file_path
  pure $ getMapping (getStaleContent state file_path) virtual_file
  where
    getMapping :: Maybe VirtualFile -> Maybe VirtualFile -> Maybe PositionMapping
    getMapping :: Maybe VirtualFile -> Maybe VirtualFile -> Maybe PositionMapping
getMapping (Just VirtualFile
stale_file) (Just VirtualFile
current_file) =
      if VirtualFile -> Int32
virtualFileVersion VirtualFile
stale_file Int32 -> Int32 -> Bool
forall a. Eq a => a -> a -> Bool
== VirtualFile -> Int32
virtualFileVersion VirtualFile
current_file
        then Maybe PositionMapping
forall a. Maybe a
Nothing -- Happens when other files (e.g. dependencies) fail to type-check.
        else PositionMapping -> Maybe PositionMapping
forall a. a -> Maybe a
Just (PositionMapping -> Maybe PositionMapping)
-> PositionMapping -> Maybe PositionMapping
forall a b. (a -> b) -> a -> b
$ [Text] -> [Text] -> PositionMapping
mappingFromDiff (Text -> [Text]
T.lines (Text -> [Text]) -> Text -> [Text]
forall a b. (a -> b) -> a -> b
$ VirtualFile -> Text
virtualFileText VirtualFile
stale_file) (Text -> [Text]
T.lines (Text -> [Text]) -> Text -> [Text]
forall a b. (a -> b) -> a -> b
$ VirtualFile -> Text
virtualFileText VirtualFile
current_file)
    getMapping Maybe VirtualFile
_ Maybe VirtualFile
_ = Maybe PositionMapping
forall a. Maybe a
Nothing
computeMapping State
_ Maybe FilePath
_ = Maybe PositionMapping -> LspM () (Maybe PositionMapping)
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe PositionMapping
forall a. Maybe a
Nothing

-- | Convert a Futhark 'Pos' to an LSP 'Uri'.
posToUri :: Pos -> Uri
posToUri :: Pos -> Uri
posToUri (Pos FilePath
file Int
_ Int
_ Int
_) = FilePath -> Uri
filePathToUri FilePath
file

toLspPosition :: Pos -> Position
toLspPosition :: Pos -> Position
toLspPosition (Pos FilePath
_ Int
l Int
c Int
_) =
  UInt -> UInt -> Position
Position (Int -> UInt
forall a. Enum a => Int -> a
toEnum Int
l UInt -> UInt -> UInt
forall a. Num a => a -> a -> a
- UInt
1) (Int -> UInt
forall a. Enum a => Int -> a
toEnum Int
c UInt -> UInt -> UInt
forall a. Num a => a -> a -> a
- UInt
1)

-- | Create an LSP 'Range' from a Futhark 'Loc'.
rangeFromLoc :: Loc -> Range
rangeFromLoc :: Loc -> Range
rangeFromLoc (Loc Pos
start Pos
end) = Position -> Position -> Range
Range (Pos -> Position
toLspPosition Pos
start) (Pos -> Position
toLspPosition Pos
end)
rangeFromLoc Loc
NoLoc = Position -> Position -> Range
Range (UInt -> UInt -> Position
Position UInt
0 UInt
0) (UInt -> UInt -> Position
Position UInt
0 UInt
5) -- only when file not found, throw error after moving to vfs

-- | Transform VFS to a map of file paths to file contents.
-- This is used to pass the file contents to the compiler.
transformVFS :: VFS -> M.Map FilePath T.Text
transformVFS :: VFS -> Map FilePath Text
transformVFS VFS
vfs =
  (NormalizedUri
 -> VirtualFile -> Map FilePath Text -> Map FilePath Text)
-> Map FilePath Text
-> Map NormalizedUri VirtualFile
-> Map FilePath Text
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
M.foldrWithKey
    ( \NormalizedUri
uri VirtualFile
virtual_file Map FilePath Text
acc ->
        case NormalizedUri -> Maybe NormalizedFilePath
uriToNormalizedFilePath NormalizedUri
uri of
          Maybe NormalizedFilePath
Nothing -> Map FilePath Text
acc
          Just NormalizedFilePath
file_path ->
            FilePath -> Text -> Map FilePath Text -> Map FilePath Text
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (NormalizedFilePath -> FilePath
fromNormalizedFilePath NormalizedFilePath
file_path) (VirtualFile -> Text
virtualFileText VirtualFile
virtual_file) Map FilePath Text
acc
    )
    Map FilePath Text
forall k a. Map k a
M.empty
    ((VirtualFileEntry -> Maybe VirtualFile)
-> Map NormalizedUri VirtualFileEntry
-> Map NormalizedUri VirtualFile
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe VirtualFileEntry -> Maybe VirtualFile
keepOpenFile (VFS
vfs VFS
-> Getting
     (Map NormalizedUri VirtualFileEntry)
     VFS
     (Map NormalizedUri VirtualFileEntry)
-> Map NormalizedUri VirtualFileEntry
forall s a. s -> Getting a s a -> a
^. Getting
  (Map NormalizedUri VirtualFileEntry)
  VFS
  (Map NormalizedUri VirtualFileEntry)
forall s a. HasVfsMap s a => Lens' s a
Lens' VFS (Map NormalizedUri VirtualFileEntry)
vfsMap))
  where
    keepOpenFile :: VirtualFileEntry -> Maybe VirtualFile
keepOpenFile = \case
      VFS.Open VirtualFile
file -> VirtualFile -> Maybe VirtualFile
forall a. a -> Maybe a
Just VirtualFile
file
      VFS.Closed ClosedVirtualFile
_ -> Maybe VirtualFile
forall a. Maybe a
Nothing

logWithSeverity :: (MonadLsp c m) => Severity -> LogAction m Text
logWithSeverity :: forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
severity = (Text -> WithSeverity Text)
-> LogAction m (WithSeverity Text) -> LogAction m Text
forall a' a. (a' -> a) -> LogAction m a -> LogAction m a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (Text -> Severity -> WithSeverity Text
forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
severity) LogAction m (WithSeverity Text)
forall c (m :: * -> *).
MonadLsp c m =>
LogAction m (WithSeverity Text)
logToLogMessage

bindingsInState :: State -> Maybe (M.Map VName BoundTo)
bindingsInState :: State -> Maybe (Map VName BoundTo)
bindingsInState State
state = Imports -> Map VName BoundTo
allBindings (Imports -> Map VName BoundTo)
-> (LoadedProg -> Imports) -> LoadedProg -> Map VName BoundTo
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LoadedProg -> Imports
lpImports (LoadedProg -> Map VName BoundTo)
-> Maybe LoadedProg -> Maybe (Map VName BoundTo)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State -> Maybe LoadedProg
stateProgram State
state

filterByLspRange ::
  (Foldable t) =>
  Maybe PositionMapping ->
  Range ->
  FilePath ->
  (a -> Loc) ->
  t a ->
  Maybe [a]
filterByLspRange :: forall (t :: * -> *) a.
Foldable t =>
Maybe PositionMapping
-> Range -> FilePath -> (a -> Loc) -> t a -> Maybe [a]
filterByLspRange Maybe PositionMapping
mapping Range
range FilePath
filepath a -> Loc
f t a
t = do
  let (Range (Position UInt
l1 UInt
c1) (Position UInt
l2 UInt
c2)) = Range
range
  posStart <- Maybe PositionMapping -> Pos -> Maybe Pos
toStalePos Maybe PositionMapping
mapping (Pos -> Maybe Pos) -> Pos -> Maybe Pos
forall a b. (a -> b) -> a -> b
$ UInt -> UInt -> Pos
mkPos UInt
l1 UInt
c1
  posEnd <- toStalePos mapping $ mkPos l2 c2
  pure $
    F.toList t
      & filter (isInRange posStart posEnd . f)
  where
    isInRange :: Pos -> Pos -> Loc -> Bool
isInRange Pos
range_start Pos
range_end Loc
loc = case Loc
loc of
      Loc
NoLoc -> Bool
False
      Loc Pos
start Pos
_ -> Pos
start Pos -> Pos -> Bool
forall a. Ord a => a -> a -> Bool
>= Pos
range_start Bool -> Bool -> Bool
&& Pos
start Pos -> Pos -> Bool
forall a. Ord a => a -> a -> Bool
<= Pos
range_end
    -- increment by one: Pos counts from one onwards, LSP starts at zero
    mkPos :: UInt -> UInt -> Pos
    mkPos :: UInt -> UInt -> Pos
mkPos UInt
l UInt
c = FilePath -> Int -> Int -> Int -> Pos
Pos FilePath
filepath (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
l) (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ UInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
c) (-Int
1)

bindingsInRange :: Range -> State -> FilePath -> Maybe [(VName, BoundTo)]
bindingsInRange :: Range -> State -> FilePath -> Maybe [(VName, BoundTo)]
bindingsInRange Range
range State
state FilePath
filepath = do
  let mapping :: Maybe PositionMapping
mapping = State -> FilePath -> Maybe PositionMapping
getStaleMapping State
state FilePath
filepath
  bindings <- State -> Maybe (Map VName BoundTo)
bindingsInState State
state

  bindings
    & M.assocs
    & filterByLspRange mapping range filepath (locOf . snd)