{-# LANGUAGE DataKinds #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}

module Unison.LSP.VFS where

import Colog.Core qualified as Colog
import Control.Lens
import Control.Monad.Reader
import Control.Monad.State
import Data.Map qualified as Map
import Data.Set qualified as Set
import Data.Set.Lens (setOf)
import Data.Text qualified as Text
import Data.Text.Utf16.Rope.Mixed qualified as Rope
import Language.LSP.Logging qualified as LSP
import Language.LSP.Protocol.Lens (HasCharacter (character), HasParams (params), HasTextDocument (textDocument), HasUri (uri))
import Language.LSP.Protocol.Lens qualified as LSP
import Language.LSP.Protocol.Message qualified as Msg
import Language.LSP.Protocol.Types
import Language.LSP.VFS as VFS hiding (character)
import Unison.LSP.Orphans ()
import Unison.LSP.Types
import Unison.LSP.VFS.CompletionPrefix as VFS
import Unison.Prelude
import Unison.Syntax.Lexer qualified as Lexer
import UnliftIO

-- | Some VFS combinators require Monad State, this provides it in a transactionally safe
-- manner so we're sure we don't edit the same file in two different actions at the same time.
usingVFS :: forall a. StateT VFS Lsp a -> Lsp a
usingVFS :: forall a. StateT VFS Lsp a -> Lsp a
usingVFS StateT VFS Lsp a
m = do
  vfsVar' <- (Env -> MVar VFS) -> Lsp (MVar VFS)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> MVar VFS
vfsVar
  modifyMVar vfsVar' $ \VFS
vfs -> do
    (!a, !vfs) <- StateT VFS Lsp a -> VFS -> Lsp (a, VFS)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT VFS Lsp a
m VFS
vfs
    pure (vfs, a)

getVirtualFile :: (Lspish m) => Uri -> MaybeT m VirtualFile
getVirtualFile :: forall (m :: * -> *). Lspish m => Uri -> MaybeT m VirtualFile
getVirtualFile Uri
fileUri = do
  vfs <- (Env -> MVar VFS) -> MaybeT m (MVar VFS)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> MVar VFS
vfsVar MaybeT m (MVar VFS) -> (MVar VFS -> MaybeT m VFS) -> MaybeT m VFS
forall a b. MaybeT m a -> (a -> MaybeT m b) -> MaybeT m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar VFS -> MaybeT m VFS
forall (m :: * -> *) a. MonadIO m => MVar a -> m a
readMVar
  MaybeT . pure $ vfs ^. vfsMap . at (toNormalizedUri $ fileUri)

getFileContents :: (Lspish m) => Uri -> MaybeT m (FileVersion, Text)
getFileContents :: forall (m :: * -> *).
Lspish m =>
Uri -> MaybeT m (FileVersion, Text)
getFileContents Uri
fileUri = do
  vf <- Uri -> MaybeT m VirtualFile
forall (m :: * -> *). Lspish m => Uri -> MaybeT m VirtualFile
getVirtualFile Uri
fileUri
  pure (vf ^. lsp_version, Rope.toText $ vf ^. file_text)

vfsLogger :: Colog.LogAction (StateT VFS Lsp) (Colog.WithSeverity VfsLog)
vfsLogger :: LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
vfsLogger = (WithSeverity VfsLog -> WithSeverity Text)
-> LogAction (StateT VFS Lsp) (WithSeverity Text)
-> LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
Colog.cmap ((VfsLog -> Text) -> WithSeverity VfsLog -> WithSeverity Text
forall a b. (a -> b) -> WithSeverity a -> WithSeverity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VfsLog -> Text
forall a. Show a => a -> Text
tShow) ((forall x. Lsp x -> StateT VFS Lsp x)
-> LogAction Lsp (WithSeverity Text)
-> LogAction (StateT VFS Lsp) (WithSeverity Text)
forall (m :: * -> *) (n :: * -> *) a.
(forall x. m x -> n x) -> LogAction m a -> LogAction n a
Colog.hoistLogAction Lsp x -> StateT VFS Lsp x
forall x. Lsp x -> StateT VFS Lsp x
forall (m :: * -> *) a. Monad m => m a -> StateT VFS m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift LogAction Lsp (WithSeverity Text)
forall c (m :: * -> *).
MonadLsp c m =>
LogAction m (WithSeverity Text)
LSP.defaultClientLogger)

-- | Mark some files as needing to be checked.
markFilesDirty :: (Foldable f, HasUri doc Uri) => f doc -> Lsp ()
markFilesDirty :: forall (f :: * -> *) doc.
(Foldable f, HasUri doc Uri) =>
f doc -> Lsp ()
markFilesDirty f doc
docs = do
  dirtyFilesV <- (Env -> TVar (Set Uri)) -> Lsp (TVar (Set Uri))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TVar (Set Uri)
dirtyFilesVar
  checkedFilesV <- asks checkedFilesVar
  let dirtyUris = Getting (Set Uri) (f doc) Uri -> f doc -> Set Uri
forall a s. Getting (Set a) s a -> s -> Set a
setOf ((doc -> Const (Set Uri) doc) -> f doc -> Const (Set Uri) (f doc)
forall (f :: * -> *) a. Foldable f => IndexedFold Int (f a) a
IndexedFold Int (f doc) doc
folded ((doc -> Const (Set Uri) doc) -> f doc -> Const (Set Uri) (f doc))
-> ((Uri -> Const (Set Uri) Uri) -> doc -> Const (Set Uri) doc)
-> Getting (Set Uri) (f doc) Uri
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Uri -> Const (Set Uri) Uri) -> doc -> Const (Set Uri) doc
forall s a. HasUri s a => Lens' s a
Lens' doc Uri
uri) f doc
docs
  atomically $ do
    modifyTVar' dirtyFilesV (Set.union dirtyUris)
    checkedFiles <- readTVar checkedFilesV
    -- Clear the analysis for any files which need to be re-checked.
    for_ dirtyUris \Uri
uri -> do
      case Uri -> Map Uri (TMVar FileAnalysis) -> Maybe (TMVar FileAnalysis)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Uri
uri Map Uri (TMVar FileAnalysis)
checkedFiles of
        Maybe (TMVar FileAnalysis)
Nothing -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just TMVar FileAnalysis
mvar -> STM (Maybe FileAnalysis) -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM (Maybe FileAnalysis) -> STM ())
-> STM (Maybe FileAnalysis) -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar FileAnalysis -> STM (Maybe FileAnalysis)
forall a. TMVar a -> STM (Maybe a)
tryTakeTMVar TMVar FileAnalysis
mvar

-- | Mark all files for re-checking.
--
-- We may want to do this when our names or perspective change.
markAllFilesDirty :: Lsp ()
markAllFilesDirty :: Lsp ()
markAllFilesDirty = do
  vfs <- (Env -> MVar VFS) -> Lsp (MVar VFS)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> MVar VFS
vfsVar Lsp (MVar VFS) -> (MVar VFS -> Lsp VFS) -> Lsp VFS
forall a b. Lsp a -> (a -> Lsp b) -> Lsp b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar VFS -> Lsp VFS
forall (m :: * -> *) a. MonadIO m => MVar a -> m a
readMVar
  markFilesDirty $ Map.keys (vfs ^. vfsMap)

-- | Returns the name or symbol which the provided position is contained in.
identifierAtPosition :: (Lspish m) => Uri -> Position -> MaybeT m Text
identifierAtPosition :: forall (m :: * -> *). Lspish m => Uri -> Position -> MaybeT m Text
identifierAtPosition Uri
uri Position
pos = do
  Uri -> Position -> MaybeT m (Text, Text)
forall (m :: * -> *).
Lspish m =>
Uri -> Position -> MaybeT m (Text, Text)
identifierSplitAtPosition Uri
uri Position
pos MaybeT m (Text, Text) -> ((Text, Text) -> Text) -> MaybeT m Text
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(Text
before, Text
after) -> (Text
before Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
after)

-- | Returns the prefix and suffix of the symbol which the provided position is contained in.
identifierSplitAtPosition :: (Lspish m) => Uri -> Position -> MaybeT m (Text, Text)
identifierSplitAtPosition :: forall (m :: * -> *).
Lspish m =>
Uri -> Position -> MaybeT m (Text, Text)
identifierSplitAtPosition Uri
uri Position
pos = do
  vf <- Uri -> MaybeT m VirtualFile
forall (m :: * -> *). Lspish m => Uri -> MaybeT m VirtualFile
getVirtualFile Uri
uri
  PosPrefixInfo {fullLine, cursorPos} <- MaybeT (VFS.getCompletionPrefix pos vf)
  let (before, after) = Text.splitAt (cursorPos ^. character . to fromIntegral) fullLine
  pure
    ( Text.takeWhileEnd isIdentifierChar before,
      -- names can end with '!', and it's not a force, so we include it in the identifier if it's at the end.
      Text.takeWhile (\Char
c -> Char -> Bool
isIdentifierChar Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'!') after
    )
  where
    isIdentifierChar :: Char -> Bool
isIdentifierChar Char
c =
      -- Manually exclude '!' and apostrophe, since those are usually just forces and
      -- delays, which shouldn't be replaced by auto-complete.
      (Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'!' Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'\'')
        Bool -> Bool -> Bool
&& (Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'.' Bool -> Bool -> Bool
|| Char -> Bool
Lexer.wordyIdChar Char
c Bool -> Bool -> Bool
|| Char -> Bool
Lexer.symbolyIdChar Char
c)

-- | Returns the prefix of the symbol at the provided location, and the range that prefix
-- spans.
completionPrefix :: Uri -> Position -> MaybeT Lsp (Range, Text)
completionPrefix :: Uri -> Position -> MaybeT Lsp (Range, Text)
completionPrefix Uri
uri Position
pos = do
  (before, _) <- Uri -> Position -> MaybeT Lsp (Text, Text)
forall (m :: * -> *).
Lspish m =>
Uri -> Position -> MaybeT m (Text, Text)
identifierSplitAtPosition Uri
uri Position
pos
  let posLine = Position
pos Position -> Getting UInt Position UInt -> UInt
forall s a. s -> Getting a s a -> a
^. Getting UInt Position UInt
forall s a. HasLine s a => Lens' s a
Lens' Position UInt
LSP.line
  let posChar = Position
pos Position -> Getting UInt Position UInt -> UInt
forall s a. s -> Getting a s a -> a
^. Getting UInt Position UInt
forall s a. HasCharacter s a => Lens' s a
Lens' Position UInt
LSP.character
  let range = UInt -> UInt -> UInt -> UInt -> Range
mkRange UInt
posLine (UInt
posChar UInt -> UInt -> UInt
forall a. Num a => a -> a -> a
- Int -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Text -> Int
Text.length Text
before)) UInt
posLine UInt
posChar
  pure (range, before)

--- Handlers for tracking file changes.

lspOpenFile :: Msg.TNotificationMessage 'Msg.Method_TextDocumentDidOpen -> Lsp ()
lspOpenFile :: TNotificationMessage 'Method_TextDocumentDidOpen -> Lsp ()
lspOpenFile TNotificationMessage 'Method_TextDocumentDidOpen
msg = do
  StateT VFS Lsp () -> Lsp ()
forall a. StateT VFS Lsp a -> Lsp a
usingVFS (StateT VFS Lsp () -> Lsp ())
-> (TNotificationMessage 'Method_TextDocumentDidOpen
    -> StateT VFS Lsp ())
-> TNotificationMessage 'Method_TextDocumentDidOpen
-> Lsp ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
-> TMessage 'Method_TextDocumentDidOpen -> StateT VFS Lsp ()
forall (m :: * -> *).
MonadState VFS m =>
LogAction m (WithSeverity VfsLog)
-> TMessage 'Method_TextDocumentDidOpen -> m ()
openVFS LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
vfsLogger (TNotificationMessage 'Method_TextDocumentDidOpen -> Lsp ())
-> TNotificationMessage 'Method_TextDocumentDidOpen -> Lsp ()
forall a b. (a -> b) -> a -> b
$ TNotificationMessage 'Method_TextDocumentDidOpen
msg
  [TextDocumentItem] -> Lsp ()
forall (f :: * -> *) doc.
(Foldable f, HasUri doc Uri) =>
f doc -> Lsp ()
markFilesDirty [TNotificationMessage 'Method_TextDocumentDidOpen
msg TNotificationMessage 'Method_TextDocumentDidOpen
-> Getting
     TextDocumentItem
     (TNotificationMessage 'Method_TextDocumentDidOpen)
     TextDocumentItem
-> TextDocumentItem
forall s a. s -> Getting a s a -> a
^. (DidOpenTextDocumentParams
 -> Const TextDocumentItem DidOpenTextDocumentParams)
-> TNotificationMessage 'Method_TextDocumentDidOpen
-> Const
     TextDocumentItem (TNotificationMessage 'Method_TextDocumentDidOpen)
forall s a. HasParams s a => Lens' s a
Lens'
  (TNotificationMessage 'Method_TextDocumentDidOpen)
  DidOpenTextDocumentParams
params ((DidOpenTextDocumentParams
  -> Const TextDocumentItem DidOpenTextDocumentParams)
 -> TNotificationMessage 'Method_TextDocumentDidOpen
 -> Const
      TextDocumentItem
      (TNotificationMessage 'Method_TextDocumentDidOpen))
-> ((TextDocumentItem -> Const TextDocumentItem TextDocumentItem)
    -> DidOpenTextDocumentParams
    -> Const TextDocumentItem DidOpenTextDocumentParams)
-> Getting
     TextDocumentItem
     (TNotificationMessage 'Method_TextDocumentDidOpen)
     TextDocumentItem
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TextDocumentItem -> Const TextDocumentItem TextDocumentItem)
-> DidOpenTextDocumentParams
-> Const TextDocumentItem DidOpenTextDocumentParams
forall s a. HasTextDocument s a => Lens' s a
Lens' DidOpenTextDocumentParams TextDocumentItem
textDocument]

lspCloseFile :: Msg.TNotificationMessage 'Msg.Method_TextDocumentDidClose -> Lsp ()
lspCloseFile :: TNotificationMessage 'Method_TextDocumentDidClose -> Lsp ()
lspCloseFile TNotificationMessage 'Method_TextDocumentDidClose
msg =
  StateT VFS Lsp () -> Lsp ()
forall a. StateT VFS Lsp a -> Lsp a
usingVFS (StateT VFS Lsp () -> Lsp ())
-> (TNotificationMessage 'Method_TextDocumentDidClose
    -> StateT VFS Lsp ())
-> TNotificationMessage 'Method_TextDocumentDidClose
-> Lsp ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
-> TMessage 'Method_TextDocumentDidClose -> StateT VFS Lsp ()
forall (m :: * -> *).
MonadState VFS m =>
LogAction m (WithSeverity VfsLog)
-> TMessage 'Method_TextDocumentDidClose -> m ()
closeVFS LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
vfsLogger (TNotificationMessage 'Method_TextDocumentDidClose -> Lsp ())
-> TNotificationMessage 'Method_TextDocumentDidClose -> Lsp ()
forall a b. (a -> b) -> a -> b
$ TNotificationMessage 'Method_TextDocumentDidClose
msg

lspChangeFile :: Msg.TNotificationMessage 'Msg.Method_TextDocumentDidChange -> Lsp ()
lspChangeFile :: TNotificationMessage 'Method_TextDocumentDidChange -> Lsp ()
lspChangeFile TNotificationMessage 'Method_TextDocumentDidChange
msg = do
  StateT VFS Lsp () -> Lsp ()
forall a. StateT VFS Lsp a -> Lsp a
usingVFS (StateT VFS Lsp () -> Lsp ())
-> (TNotificationMessage 'Method_TextDocumentDidChange
    -> StateT VFS Lsp ())
-> TNotificationMessage 'Method_TextDocumentDidChange
-> Lsp ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
-> TMessage 'Method_TextDocumentDidChange -> StateT VFS Lsp ()
forall (m :: * -> *).
MonadState VFS m =>
LogAction m (WithSeverity VfsLog)
-> TMessage 'Method_TextDocumentDidChange -> m ()
changeFromClientVFS LogAction (StateT VFS Lsp) (WithSeverity VfsLog)
vfsLogger (TNotificationMessage 'Method_TextDocumentDidChange -> Lsp ())
-> TNotificationMessage 'Method_TextDocumentDidChange -> Lsp ()
forall a b. (a -> b) -> a -> b
$ TNotificationMessage 'Method_TextDocumentDidChange
msg
  lastTouchedFileV <- (Env -> TVar (Maybe Uri)) -> Lsp (TVar (Maybe Uri))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TVar (Maybe Uri)
lastTouchedFileVar
  atomically $ writeTVar lastTouchedFileV (Just (msg ^. params . textDocument . uri))
  markFilesDirty [msg ^. params . textDocument]