[WIP] Implement a loader of Env from a ScopeTree

tools/lsp/squirrel/app/Main.hs

@@ -7,6 +7,7 @@ import           Control.Monad
 
 import           Data.Default
 import qualified Data.Text                             as Text
+import           Data.Text                               (Text)
 import           Data.String.Interpolate (i)
 
 import qualified Language.Haskell.LSP.Control          as CTRL
@@ -23,7 +24,9 @@ import qualified System.Log                            as L
 
 import           Parser
 import           Range
+import           Product
 import           AST hiding (def)
+import qualified AST.Find as Find
 import           Error
 import           Tree
 
@@ -150,11 +153,32 @@ eventLoop funs chan = do
           (J.uriToFilePath doc)
           (Just 0)
 
-      -- ReqDefinition req -> do
-
+      ReqDefinition req -> do
+        let uri = req^.J.params.J.textDocument.J.uri
+        let pos = posToRange $ req^.J.params.J.position
+        tree <- loadByURI uri
+        case Find.definitionOf pos tree of
+          Just defPos -> do
+            error "do later"
+            -- Core.sendFunc funs $ RspDefinition $ _ $ J.SingleLoc $ J.Location uri $ rangeToLoc defPos
 
       _ -> U.logs "unknown msg"
 
+posToRange :: J.Position -> Range
+posToRange (J.Position l c) = Range (l, c, 0) (l, c, 0)
+
+rangeToJRange :: Range -> J.Range
+rangeToJRange (Range (a, b, _) (c, d, _)) = J.Range (J.Position a b) (J.Position c d)
+
+rangeToLoc :: Range -> J.Range
+rangeToLoc (Range (a, b, _) (c, d, _)) = J.Range (J.Position a b) (J.Position c d)
+
+loadByURI :: J.Uri -> IO (Pascal (Product [Env, Range, [Text]]))
+loadByURI uri = do
+  case J.uriToFilePath uri of
+    Just fin -> do
+      (tree, _) <- runParser contract fin
+      return $ ascribeEnv tree
 
 collectErrors
   :: Core.LspFuncs ()

tools/lsp/squirrel/package.yaml

@@ -3,6 +3,7 @@ name: squirrel
 dependencies:
   - base
   - bytestring
+  - containers
   - data-default
   - mtl
   - pretty

tools/lsp/squirrel/src/AST/Find.hs

@@ -5,10 +5,13 @@ import Control.Monad
 
 import AST.Types
 import AST.Scope
+import AST.Parser
 
+import Parser
 import Tree
 import Range
 import Lattice
+import Pretty
 
 import Debug.Trace
 
@@ -20,9 +23,9 @@ findScopedDecl
   -> Pascal info
   -> Maybe ScopedDecl
 findScopedDecl pos tree = do
-  point <- lookupTree (\info -> pos <? getRange info) tree
+  point <- lookupTree pos tree
   let env = getEnv (infoOf point)
-  lookupEnv (void point) env
+  lookupEnv (ppToText $ void point) env
 
 definitionOf
   :: ( HasEnv info
@@ -53,3 +56,13 @@ implementationOf
   -> Maybe Range
 implementationOf pos tree =
   _sdBody =<< findScopedDecl pos tree
+
+referencesOf
+  :: ( HasEnv info
+     , HasRange info
+     )
+  => Range
+  -> Pascal info
+  -> Maybe [Range]
+referencesOf pos tree =
+  _sdRefs <$> findScopedDecl pos tree

tools/lsp/squirrel/src/AST/Parser.hs

@@ -845,8 +845,8 @@ typeTuple = do
 -- example = "../../../src/test/contracts/address.ligo"
 -- example = "../../../src/test/contracts/amount.ligo"
 -- example = "../../../src/test/contracts/annotation.ligo"
--- example = "../../../src/test/contracts/arithmetic.ligo"
-example = "../../../src/test/contracts/assign.ligo"
+example = "../../../src/test/contracts/arithmetic.ligo"
+-- example = "../../../src/test/contracts/assign.ligo"
 -- example = "../../../src/test/contracts/attributes.ligo"
 -- example = "../../../src/test/contracts/bad_timestamp.ligo"
 -- example = "../../../src/test/contracts/bad_type_operator.ligo"

tools/lsp/squirrel/src/AST/Scope.hs

@@ -3,65 +3,88 @@
 -}
 
 module AST.Scope
-  ( -- * Monad
-    ScopeM
-  , evalScopeM
-  , pinEnv
+  -- ( -- * Monad
+  --   ScopeM
+  -- , evalScopeM
+  -- , collectEnv
 
-    -- * Scope
-  , Env(..)
-  , ScopedDecl(..)
-  , Kind(..)
-  , HasEnv(..)
-  , lookupEnv
+  --   -- * Scope
+  -- , Env(..)
+  -- , ScopedDecl(..)
+  -- , Kind(..)
+  -- , HasEnv(..)
+  -- , lookupEnv
 
-    -- * Methods
-  , enter
-  , leave
-  , define
-  , defType
-  , def
-  )
+  --   -- * Methods
+  -- , enter
+  -- , leave
+  -- , define
+  -- , defType
+  -- , def
+  -- )
   where
 
 import Control.Monad.State
 
-import Data.Text (Text)
+import qualified Data.Map as Map
+import           Data.Map   (Map)
+import           Data.Text  (Text)
+import qualified Data.Text as Text
+import           Data.Maybe (fromJust)
 
 import Range
 import AST.Types
+import AST.Parser
+import Parser
 import Tree
 import Comment
 import Pretty
 import Product
 
+import Debug.Trace
+
 -- | Scope-holding monad.
-type ScopeM = State [Env]
+type ScopeM = State [Scopes]
 
 -- | Run the computation with scope starting from empty scope.
 evalScopeM :: ScopeM a -> a
-evalScopeM action = evalState action [Env []]
+evalScopeM action = evalState action []
 
--- | The environment.
-newtype Env = Env
-  { _eDecls :: [ScopedDecl]
-  }
-  deriving newtype (Semigroup, Monoid)
-  deriving Show via PP Env
+type Scopes = Tree '[ScopeTree] Range
 
-instance Pretty Env where
-  pp = vcat . map pp . _eDecls
+data ScopeTree it
+  = ScopeTree Env [it]
+  deriving stock (Functor, Foldable, Traversable)
+  deriving stock Show
+
+instance Pretty1 ScopeTree where
+  pp1 (ScopeTree e cs) =
+    pp e `indent` block cs
+
+instance {-# OVERLAPS #-} Pretty (Map Text ScopedDecl) where
+  pp = block . map aux .  Map.toList
+    where
+      aux (n, ScopedDecl o b t rs) =
+        pp o <+> (pp n <> ":") <+> pp t <+> "=" <+> pp b <+> "/" <+> (list rs <> ";")
+
+  -- | The environment.
+type Env = Map Text ScopedDecl
+  -- deriving Show via PP Env
+
+-- instance Pretty Env where
+--   pp = vcat . map pp . _eDecls
 
 -- | The type/value declaration.
 data ScopedDecl = ScopedDecl
-  { _sdName   :: (Pascal ())
-  , _sdOrigin :: Range
+  { _sdOrigin :: Range
   , _sdBody   :: Maybe Range
   , _sdType   :: Maybe (Either (Pascal ()) Kind)
+  , _sdRefs   :: [Range]
   }
+  deriving Show via PP ScopedDecl
 
 instance Pretty ScopedDecl where
-  pp (ScopedDecl n o b t) = pp o <+> "-" <+> (pp n <> ":") <+> maybe "?" (either pp pp) t <+> "=" <+> pp o
+  pp (ScopedDecl o b t refs) = pp o <+> "-" <+> maybe "?" (either pp pp) t <+> "=" <+> pp o
 
 -- | The kind.
 data Kind = Star
@@ -70,34 +93,60 @@ data Kind = Star
 instance Pretty Kind where
   pp _ = "*"
 
-lookupEnv :: Pascal () -> Env -> Maybe ScopedDecl
-lookupEnv name = go . _eDecls
-  where
-    go (sd@(ScopedDecl {_sdName}) : rest)
-      | ppToText _sdName == ppToText name = Just sd
-      | otherwise       = go rest
-    go _ = Nothing
+lookupEnv :: Text -> Env -> Maybe ScopedDecl
+lookupEnv name = Map.lookup name
 
 -- | Make a new scope out of enclosing parent one.
-enter  :: ScopeM ()
-enter = modify \(a : b) -> a : a : b
+enter :: Range -> ScopeM ()
+enter r =
+  modify \rest ->
+    mk r (ScopeTree Map.empty []) : rest
 
 -- | Leave current scope, return to parent one.
 leave  :: ScopeM ()
-leave = modify tail
+leave =
+  modify \(a : parent : rest) ->
+    fromJust do
+      (r, ScopeTree e cs) <- match parent
+      return $ mk r (ScopeTree e (a : cs)) : rest
 
 -- | Add a declaration to the current scope.
-define :: ScopedDecl -> ScopeM ()
-define d = modify \(Env a : b) -> Env (d : a) : b
+define :: Text -> ScopedDecl -> ScopeM ()
+define name d =
+  modify \(top : rest) ->
+    fromJust do
+      (r, ScopeTree a cs) <- match top
+      return $ mk r (ScopeTree (Map.insert name d a) cs) : rest
 
 -- | Add a type declaration to the current scope.
 defType :: HasRange a => Pascal a -> Kind -> Pascal a -> ScopeM ()
 defType name kind body = do
-  define $ ScopedDecl
-    (void name)
-    (getRange $ infoOf name)
-    (Just $ getRange $ infoOf body)
-    (Just (Right kind))
+  define (ppToText $ void name)
+    $ ScopedDecl
+      (getRange $ infoOf name)
+      (Just $ getRange $ infoOf body)
+      (Just (Right kind))
+      []
+
+addRef :: Text -> Range -> [Scopes] -> [Scopes]
+addRef _    _   []                  = error "addRef: empty env stack"
+addRef name pos stack@(topmost : _)
+  | Just (r, ScopeTree top cs) <- match topmost =
+    case Map.lookup name top of
+      Just ScopedDecl {_sdOrigin} -> go _sdOrigin stack
+      Nothing                     -> stack
+
+    where
+      go _     []             = []
+      go range initial@(topmost : rest)
+        | Just (r, ScopeTree e cs) <- match topmost =
+          case Map.lookup name e of
+            Just it | _sdOrigin it == range ->
+              mk r (ScopeTree (pushRef pos e) cs) : go range rest
+            _ ->
+              initial
+
+      pushRef pos' = Map.adjust (\sd -> sd { _sdRefs = pos' : _sdRefs sd }) name
 
 -- | Add a value declaration to the current scope.
 def
@@ -107,29 +156,32 @@ def
   -> Maybe (Pascal a)
   -> ScopeM ()
 def name ty body = do
-  define $ ScopedDecl
-    (void name)
-    (getRange $ infoOf name)
-    ((getRange . infoOf) <$> body)
-    ((Left . void) <$> ty)
+  define (ppToText $ void name)
+    $ ScopedDecl
+      (getRange $ infoOf name)
+      ((getRange . infoOf) <$> body)
+      ((Left . void) <$> ty)
+      []
 
-instance UpdateOver ScopeM Contract (Pascal a)
+instance UpdateOver ScopeM Contract (Pascal a) where
+  before r _ = enter r
+  after  _ _ = leave
 
 instance HasRange a => UpdateOver ScopeM Declaration (Pascal a) where
-  before = \case
+  before _ = \case
     TypeDecl ty body -> defType ty Star body
     _ -> skip
 
 instance HasRange a => UpdateOver ScopeM Binding (Pascal a) where
-  before = \case
+  before r = \case
     Function recur name _args ty body -> do
       when recur do
         def name (Just ty) (Just body)
-      enter
+      enter r
 
-    _ -> enter
+    _ -> enter r
 
-  after = \case
+  after _ = \case
     Irrefutable name    body -> do leave; def name  Nothing  (Just body)
     Var         name ty body -> do leave; def name (Just ty) (Just body)
     Const       name ty body -> do leave; def name (Just ty) (Just body)
@@ -139,7 +191,7 @@ instance HasRange a => UpdateOver ScopeM Binding (Pascal a) where
         def name (Just ty) (Just body)
 
 instance HasRange a => UpdateOver ScopeM VarDecl (Pascal a) where
-  after  (Decl _ name ty) = def name (Just ty) Nothing
+  after _ (Decl _ name ty) = def name (Just ty) Nothing
 
 instance UpdateOver ScopeM Mutable (Pascal a)
 instance UpdateOver ScopeM Type    (Pascal a)
@@ -147,21 +199,21 @@ instance UpdateOver ScopeM Variant (Pascal a)
 instance UpdateOver ScopeM TField  (Pascal a)
 
 instance HasRange a => UpdateOver ScopeM Expr (Pascal a) where
-  before = \case
-    Let    {} -> enter
-    Lambda {} -> enter
+  before r = \case
+    Let    {} -> enter r
+    Lambda {} -> enter r
     ForLoop k _ _ _ -> do
-      enter
+      enter r
       def k Nothing Nothing
 
     ForBox k mv _ _ _ -> do
-      enter
+      enter r
       def k Nothing Nothing
       maybe skip (\v -> def v Nothing Nothing) mv
 
     _ -> skip
 
-  after = \case
+  after _ = \case
     Let     {} -> leave
     Lambda  {} -> leave
     ForLoop {} -> leave
@@ -169,8 +221,8 @@ instance HasRange a => UpdateOver ScopeM Expr (Pascal a) where
     _ -> skip
 
 instance HasRange a => UpdateOver ScopeM Alt (Pascal a) where
-  before _ = enter
-  after  _ = leave
+  before r _ = enter r
+  after  _ _ = leave
 
 instance UpdateOver ScopeM LHS             (Pascal a)
 instance UpdateOver ScopeM MapBinding      (Pascal a)
@@ -179,13 +231,15 @@ instance UpdateOver ScopeM FieldAssignment (Pascal a)
 instance UpdateOver ScopeM Constant        (Pascal a)
 
 instance HasRange a => UpdateOver ScopeM Pattern (Pascal a) where
-  before = \case
+  before _ = \case
     IsVar n -> def n Nothing Nothing
     _       -> skip
 
 instance UpdateOver ScopeM QualifiedName (Pascal a)
 instance UpdateOver ScopeM Path          (Pascal a)
-instance UpdateOver ScopeM Name          (Pascal a)
+instance UpdateOver ScopeM Name          (Pascal a) where
+  before range (Name raw) = do
+    modify $ addRef raw range
 
 class HasEnv a where
   getEnv :: a -> Env
@@ -201,5 +255,43 @@ data Scope = Scope { unScope :: [Text] }
 instance HasComments Scope where
   getComments = unScope
 
-pinEnv :: Product xs -> ScopeM (Product (Env : xs))
-pinEnv xs = (`Cons` xs) <$> gets head
+-- pinEnv :: Product xs -> ScopeM (Product (Env : xs))
+-- pinEnv xs = (`Cons` xs) <$> gets head
+
+collectEnv :: Contains Range xs => Product xs -> ScopeM (Product (Scopes : xs))
+collectEnv xs = do
+  gets \case
+    st : _ -> Cons st xs
+    []     -> Cons (mk (getRange xs) $ ScopeTree Map.empty []) xs
+
+instance UpdateOver (State [Env]) ScopeTree Scopes where
+  before r (ScopeTree e _) = modify (e :)
+  after  r  _              = modify tail
+
+distributeEnv :: ScopeTree Scopes -> State [Env] (ScopeTree Scopes)
+distributeEnv (ScopeTree e' cs) = do
+  e <- gets (Map.unions . (e' :))
+  traceShowM ("distribute", e', e)
+  return $ ScopeTree e cs
+
+pinEnv :: Contains Range xs => Scopes -> Product xs -> ScopeM (Product (Env : xs))
+pinEnv scopes info = do
+  let (_, ScopeTree e _) = fromJust $ match =<< lookupTree (getElem info) scopes
+  return (Cons e info)
+
+instance HasComments Range where
+  getComments _ = []
+
+instance Pretty (Product xs) => HasComments (Product xs) where
+  getComments xs = if Text.null $ Text.strip x then [] else [x]
+    where
+      x = ppToText $ color 3 $ pp $ xs
+
+ascribeEnv :: (Contains Range xs, Pretty (Product xs)) => Pascal (Product xs) -> Pascal (Product (Env : xs))
+ascribeEnv tree =
+  let
+    Cons scopes _ = traceShowId $ infoOf $ evalScopeM $ traverseTree collectEnv tree
+    distributed   = evalState (traverseOnly distributeEnv scopes) []
+  in
+    -- distributed
+    evalScopeM $ traverseTree (pinEnv distributed) tree

tools/lsp/squirrel/src/Comment.hs

@@ -34,5 +34,5 @@ c i d =
 instance HasComments () where
   getComments () = []
 
-instance (Contains [Text] xs) => HasComments (Product xs) where
-  getComments = getElem
+-- instance (Contains [Text] xs) => HasComments (Product xs) where
+--   getComments = getElem

tools/lsp/squirrel/src/Pretty.hs

@@ -23,6 +23,7 @@ module Pretty
   , sepByDot
   , mb
   , sparseBlock
+  , color
 
     -- * Full might of pretty printing
   , module Text.PrettyPrint
@@ -69,6 +70,12 @@ instance (Pretty1 p, Functor p, Pretty a) => Pretty (p a) where
 instance Pretty1 [] where
   pp1 = list
 
+instance Pretty1 Maybe where
+  pp1 = maybe empty pp
+
+instance {-# OVERLAPS #-} (Pretty a, Pretty b) => Pretty (Either a b) where
+  pp = either pp pp
+
 -- | Common instance.
 instance Pretty Text where
   pp = text . Text.unpack
@@ -113,4 +120,12 @@ mb f = maybe empty (f . pp)
 
 -- | Pretty print as a vertical with elements separated by newline.
 sparseBlock :: Pretty a => [a] -> Doc
-sparseBlock = vcat . punctuate "\n" . map (($$ empty) . pp)
+sparseBlock = vcat . punctuate "\n" . map (($$ empty) . pp)
+
+type Color = Int
+
+color :: Color -> Doc -> Doc
+color c d = zeroWidthText begin <> d <> zeroWidthText end
+  where
+    begin = "\x1b[" ++ show (30 + c) ++ "m"
+    end   = "\x1b[0m"

tools/lsp/squirrel/src/Product.hs

@@ -1,6 +1,8 @@
 
 module Product where
 
+import qualified Data.Text as Text
+
 import Pretty
 
 data Product xs where
@@ -29,4 +31,9 @@ instance Pretty (Product '[]) where
   pp _ = "{}"
 
 instance (Pretty x, Pretty (Product xs)) => Pretty (Product (x : xs)) where
-  pp (Cons x xs) = pp x <+> "&" <+> pp xs
+  pp (Cons x xs) =
+    if Text.null $ Text.strip ppx
+    then pp xs
+    else pp ppx <+> "&" <+> pp xs
+    where
+      ppx = ppToText x

tools/lsp/squirrel/src/Range.hs

@@ -39,6 +39,9 @@ instance Pretty Range where
 class HasRange a where
   getRange :: a -> Range
 
+instance HasRange Range where
+  getRange = id
+
 instance Contains Range xs => HasRange (Product xs) where
   getRange = getElem
 
@@ -54,3 +57,7 @@ instance Lattice Range where
   Range (ll1, lc1, _) (ll2, lc2, _) <? Range (rl1, rc1, _) (rl2, rc2, _) =
     (rl1 < ll1 || rl1 == ll1 && rc1 <= lc1) &&
     (rl2 > ll2 || rl2 == ll2 && rc2 >= lc2)
+
+instance Eq Range where
+  Range (l, c, _) (r, d, _) == Range (l1, c1, _) (r1, d1, _) =
+    (l, c, r, d) == (l1, c1, r1, d1)

tools/lsp/squirrel/src/Tree.hs

@@ -10,17 +10,17 @@
 -}
 
 module Tree
-  ( -- * Tree type
-    Tree
-  , lookupTree
-  , traverseTree
-  , mk
-  , infoOf
+  -- ( -- * Tree type
+  --   Tree
+  -- , lookupTree
+  -- , traverseTree
+  -- , mk
+  -- , infoOf
 
-    -- * Callbacks on update
-  , UpdateOver (..)
-  , skip
-  )
+  --   -- * Callbacks on update
+  -- , UpdateOver (..)
+  -- , skip
+  -- )
   where
 
 import Data.Foldable
@@ -32,6 +32,7 @@ import Lattice
 import Comment
 import Pretty
 import Error
+import Range
 
 import Debug.Trace
 
@@ -97,20 +98,22 @@ lookupTree
   :: forall fs info
   .  ( Apply Foldable fs
      , Apply Functor fs
+     , HasRange info
      )
-  => (info -> Bool)
+  => Range
   -> Tree fs info
   -> Maybe (Tree fs info)
-lookupTree rightInfo = go
+lookupTree target = go
   where
     go :: Tree fs info -> Maybe (Tree fs info)
     go tree = do
-      if rightInfo (infoOf tree)
+      if target <? getRange (infoOf tree)
       then getFirst $ foldMap (First . go) (layers tree) <> First (Just tree)
       else Nothing
 
     layers :: (Apply Foldable fs) => Tree fs info -> [Tree fs info]
     layers (Tree (Right (_, ls))) = toList ls
+
 -- | Traverse the tree over some monad that exports its methods.
 --
 --   For each tree piece, will call `before` and `after` callbacks.
@@ -120,25 +123,69 @@ traverseTree
      , Apply Foldable fs
      , Apply Functor fs
      , Apply Traversable fs
+     , HasRange a
      )
   => (a -> m b) -> Tree fs a -> m (Tree fs b)
 traverseTree act = go
   where
     go (Tree (Right (a, union))) = do
       b <- act a
-      before union
+      before (getRange a) union
       union' <- traverse go union
-      after union
+      after (getRange a) union
       return (Tree (Right (b, union')))
 
     go (Tree (Left err)) = do
       err' <- traverse act err
       return (Tree (Left err'))
 
+traverseOnly
+  :: forall f fs m a
+  .  ( UpdateOver m      f   (Tree fs a)
+     , UpdateOver m (Sum fs) (Tree fs a)
+     , Element f fs
+     , Apply Foldable fs
+     , Apply Functor fs
+     , Apply Traversable fs
+     , Traversable f
+     , HasRange a
+     , Show (f (Tree fs a))
+     , Show a
+     )
+  => (f (Tree fs a) -> m (f (Tree fs a)))
+  ->     Tree fs a  -> m    (Tree fs a)
+traverseOnly act = go
+  where
+    go (match -> Just (r, fa)) = do
+      traceShowM ("traversingA", fa)
+      before (getRange r) fa
+      fb <- act fa
+      fc <- traverse go fb
+      after  (getRange r) fa
+      return $ mk r fc
+
+    go tree@(Tree (Right (r, union))) = do
+      traceShowM ("traversingB", ())
+      before (getRange r) union
+      union' <- traverse go union
+      after  (getRange r) union
+      return $ Tree $ Right (r, union')
+
+    go tree = return tree
+
 -- | Make a tree out of a layer and an info.
 mk :: (Functor f, Element f fs) => info -> f (Tree fs info) -> Tree fs info
 mk i fx = Tree $ Right (i, inject fx)
 
+match
+  :: (Functor f, Element f fs)
+  => Tree fs info
+  -> Maybe (info, f (Tree fs info))
+match (Tree (Left   _))      = Nothing
+match (Tree (Right (r, it))) = do
+  f <- project it
+  return (r, f)
+
 -- | Get info from the tree.
 infoOf :: Tree fs info -> info
 infoOf = either eInfo fst . unTree
@@ -154,11 +201,11 @@ instance Apply Foldable fs => HasErrors (Tree fs info) info where
 
 -- | Update callbacks for a @f a@ while working inside monad @m@.
 class Monad m => UpdateOver m f a where
-  before :: f a -> m ()
-  after  :: f a -> m ()
+  before :: Range -> f a -> m ()
+  after  :: Range -> f a -> m ()
 
-  before _ = skip
-  after  _ = skip
+  before _ _ = skip
+  after  _ _ = skip
 
 -- | Do nothing.
 skip :: Monad m => m ()
@@ -169,5 +216,5 @@ instance Monad m => UpdateOver m (Sum '[]) a where
   after  = error "Sum.empty"
 
 instance (UpdateOver m f a, UpdateOver m (Sum fs) a) => UpdateOver m (Sum (f : fs)) a where
-  before = either before before . decompose
-  after  = either after  after  . decompose
+  before r = either (before r) (before r) . decompose
+  after  r = either (after  r) (after  r) . decompose