From 29bc1721d1ea084d502d23eab5cf6c6da96d64f8 Mon Sep 17 00:00:00 2001
From: Kirill Andreev <hindmost.one@gmail.com>
Date: Tue, 7 Jul 2020 17:41:14 +0400
Subject: [PATCH] Load from VFS

---
 tools/lsp/squirrel/app/Main.hs      | 20 ++++++++++++++++----
 tools/lsp/squirrel/src/ParseTree.hs | 19 ++++++++++++++++---
 tools/lsp/squirrel/src/Parser.hs    |  8 ++++----
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/tools/lsp/squirrel/app/Main.hs b/tools/lsp/squirrel/app/Main.hs
index 81dd78c72..6358849bc 100644
--- a/tools/lsp/squirrel/app/Main.hs
+++ b/tools/lsp/squirrel/app/Main.hs
@@ -18,12 +18,13 @@ import           Language.Haskell.LSP.Messages         as Msg
 import qualified Language.Haskell.LSP.Types            as J
 import qualified Language.Haskell.LSP.Types.Lens       as J
 import qualified Language.Haskell.LSP.Utility          as U
--- import           Language.Haskell.LSP.VFS
+import           Language.Haskell.LSP.VFS
 
 import           System.Exit
 import qualified System.Log                            as L
 
 import           Parser
+import           ParseTree
 import           Range
 import           Product
 import           AST hiding (def)
@@ -168,7 +169,7 @@ eventLoop funs chan = do
       ReqFindReferences req -> do
         let uri = req^.J.params.J.textDocument.J.uri
         let pos = posToRange $ req^.J.params.J.position
-        tree <- loadByURI uri
+        tree <- loadFromVFS funs uri
         case Find.referencesOf pos tree of
           Just refs -> do
             let locations = J.Location uri . rangeToLoc <$> refs
@@ -194,11 +195,22 @@ rangeToLoc (Range (a, b, _) (c, d, _) _) =
     (J.Position (a - 1) (b - 1))
     (J.Position (c - 1) (d - 1))
 
+loadFromVFS
+  :: Core.LspFuncs ()
+  -> J.Uri
+  -> IO (Pascal (Product [[ScopedDecl], Range, [Text]]))
+loadFromVFS funs uri = do
+  Just vf <- Core.getVirtualFileFunc funs $ J.toNormalizedUri uri
+  let txt = virtualFileText vf
+  let Just fin = J.uriToFilePath uri
+  (tree, _) <- runParser contract (Text fin txt)
+  return $ addLocalScopes tree
+
 loadByURI :: J.Uri -> IO (Pascal (Product [[ScopedDecl], Range, [Text]]))
 loadByURI uri = do
   case J.uriToFilePath uri of
     Just fin -> do
-      (tree, _) <- runParser contract fin
+      (tree, _) <- runParser contract (Path fin)
       return $ addLocalScopes tree
 
 collectErrors
@@ -210,7 +222,7 @@ collectErrors
 collectErrors funs uri path version = do
   case path of
     Just fin -> do
-      (tree, errs) <- runParser contract fin
+      (tree, errs) <- runParser contract (Path fin)
       Core.publishDiagnosticsFunc funs 100 uri version
         $ partitionBySource
         $ map errorToDiag (errs <> errors tree)
diff --git a/tools/lsp/squirrel/src/ParseTree.hs b/tools/lsp/squirrel/src/ParseTree.hs
index 9869a265f..98c00dd7f 100644
--- a/tools/lsp/squirrel/src/ParseTree.hs
+++ b/tools/lsp/squirrel/src/ParseTree.hs
@@ -11,6 +11,7 @@ module ParseTree
   ( -- * Tree/Forest
     ParseTree(..)
   , ParseForest(..)
+  , Source(..)
 
     -- * Invoke the TreeSitter and get the tree it outputs
   , toParseTree
@@ -21,6 +22,7 @@ import Data.ByteString (ByteString)
 import Data.IORef
 import qualified Data.ByteString as BS
 import qualified Data.Text as Text
+import qualified Data.Text.Encoding as Text
 import Data.Text (Text)
 import Data.Traversable (for)
 
@@ -49,6 +51,17 @@ import Pretty
 
 foreign import ccall unsafe tree_sitter_PascaLigo :: Ptr Language
 
+data Source
+  = Path       { srcPath :: FilePath }
+  | Text       { srcPath :: FilePath, srcText :: Text }
+  | ByteString { srcPath :: FilePath, srcBS   :: ByteString }
+
+srcToBytestring :: Source -> IO ByteString
+srcToBytestring = \case
+  Path       p   -> BS.readFile p
+  Text       _ t -> return $ Text.encodeUtf8 t
+  ByteString _ s -> return s
+
 -- | The tree tree-sitter produces.
 data ParseTree = ParseTree
   { ptID       :: Int          -- ^ Unique number, for fast comparison.
@@ -85,12 +98,12 @@ instance Pretty ParseForest where
         else hang (text (Text.unpack field) <> ": ") 2 (pp tree)
 
 -- | Feed file contents into PascaLIGO grammar recogniser.
-toParseTree :: FilePath -> IO ParseForest
+toParseTree :: Source -> IO ParseForest
 toParseTree fin = do
   parser <- ts_parser_new
   True   <- ts_parser_set_language parser tree_sitter_PascaLigo
 
-  src <- BS.readFile fin
+  src <- srcToBytestring fin
 
   idCounter <- newIORef 0
 
@@ -146,7 +159,7 @@ toParseTree fin = do
                   , i $ pointColumn finish2D + 1
                   , i $ nodeEndByte node
                   )
-              , rFile = takeFileName fin
+              , rFile = takeFileName $ srcPath fin
               }
 
           return $ ParseTree
diff --git a/tools/lsp/squirrel/src/Parser.hs b/tools/lsp/squirrel/src/Parser.hs
index 02ce726c9..7d9dcb677 100644
--- a/tools/lsp/squirrel/src/Parser.hs
+++ b/tools/lsp/squirrel/src/Parser.hs
@@ -103,12 +103,12 @@ type ASTInfo = Product [Range, [Text]]
 
 runParser
   :: Parser a
-  -> FilePath
+  -> Source
   -> IO (a, [Error ASTInfo])
 runParser parser fin = do
   pforest <- toParseTree fin
 
-  let dir = takeDirectory fin
+  let dir = takeDirectory $ srcPath fin
 
   runWriterT parser `evalStateT`
      Cons pforest
@@ -128,7 +128,7 @@ restart p fin = do
     fallback "recusive imports"
   else do
     (a, errs) <- liftIO do
-      flip runParser full do
+      flip runParser (Path full) do
         put' (Set.insert full set)
         p
     tell errs
@@ -319,7 +319,7 @@ some p = some'
 
 -- | Run parser on given file and pretty-print stuff.
 --
-debugParser :: Show a => Parser a -> FilePath -> IO ()
+debugParser :: Show a => Parser a -> Source -> IO ()
 debugParser parser fin = do
   (res, errs) <- runParser parser fin
   putStrLn "Result:"