Rewrite scope loader

tools/lsp/squirrel/app/Main.hs

@@ -173,12 +173,12 @@ rangeToJRange (Range (a, b, _) (c, d, _)) = J.Range (J.Position a b) (J.Position
 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 :: J.Uri -> IO (Pascal (Product [[ScopedDecl], Range, [Text]]))
 loadByURI uri = do
   case J.uriToFilePath uri of
     Just fin -> do
       (tree, _) <- runParser contract fin
-      return $ ascribeEnv tree
+      return $ addLocalScopes tree
 
 collectErrors
   :: Core.LspFuncs ()

tools/lsp/squirrel/package.yaml

@@ -1,20 +1,23 @@
-name: squirrel
+name: ligo-squirrel
 
 dependencies:
   - base
   - bytestring
   - containers
   - data-default
+  - fastsum
   - mtl
   - pretty
   - template-haskell
   - text
   - tree-sitter
-  - fastsum
 
 default-extensions:
+  - AllowAmbiguousTypes
+  - ApplicativeDo
   - BangPatterns
   - BlockArguments
+  - ConstraintKinds
   - DataKinds
   - DeriveFoldable
   - DeriveFunctor
@@ -23,9 +26,11 @@ default-extensions:
   - DerivingVia
   - FlexibleContexts
   - FlexibleInstances
+  - FunctionalDependencies
   - GADTs
   - GeneralisedNewtypeDeriving
   - LambdaCase
+  - MagicHash
   - MultiParamTypeClasses
   - NamedFieldPuns
   - OverloadedStrings
@@ -33,15 +38,11 @@ default-extensions:
   - ScopedTypeVariables
   - StandaloneDeriving
   - TemplateHaskell
+  - TypeApplications
   - TypeFamilies
   - TypeOperators
   - UndecidableInstances
-  - FunctionalDependencies
   - ViewPatterns
-  - ConstraintKinds
-  - TypeApplications
-  - AllowAmbiguousTypes
-  - MagicHash
 
 ghc-options: -freverse-errors -Wall -threaded
 
@@ -58,12 +59,12 @@ library:
 executables:
   squirrel:
     dependencies:
-      - lens
-      - stm
       - haskell-lsp
-      - squirrel
       - hslogger
       - interpolate
+      - lens
+      - ligo-squirrel
+      - stm
 
     main: Main.hs
 

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

@@ -16,7 +16,7 @@ import Pretty
 import Debug.Trace
 
 findScopedDecl
-  :: ( HasEnv info
+  :: ( HasLocalScope info
      , HasRange info
      )
   => Range
@@ -24,11 +24,10 @@ findScopedDecl
   -> Maybe ScopedDecl
 findScopedDecl pos tree = do
   point <- lookupTree pos tree
-  let env = getEnv (infoOf point)
+  lookupEnv (ppToText $ void point) (getLocalScope (infoOf point))
-  lookupEnv (ppToText $ void point) env
 
 definitionOf
-  :: ( HasEnv info
+  :: ( HasLocalScope info
      , HasRange info
      )
   => Range
@@ -38,7 +37,7 @@ definitionOf pos tree =
   _sdOrigin <$> findScopedDecl pos tree
 
 typeOf
-  :: ( HasEnv info
+  :: ( HasLocalScope info
      , HasRange info
      )
   => Range
@@ -48,7 +47,7 @@ typeOf pos tree =
   _sdType =<< findScopedDecl pos tree
 
 implementationOf
-  :: ( HasEnv info
+  :: ( HasLocalScope info
      , HasRange info
      )
   => Range
@@ -58,7 +57,7 @@ implementationOf pos tree =
   _sdBody =<< findScopedDecl pos tree
 
 referencesOf
-  :: ( HasEnv info
+  :: ( HasLocalScope info
      , HasRange info
      )
   => Range

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

@@ -845,7 +845,7 @@ 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/arithmetic.ligo"
 -- example = "../../../src/test/contracts/assign.ligo"
 -- example = "../../../src/test/contracts/attributes.ligo"
 -- example = "../../../src/test/contracts/bad_timestamp.ligo"
@@ -858,10 +858,10 @@ example = "../../../src/test/contracts/arithmetic.ligo"
 -- example = "../../../src/test/contracts/bytes_arithmetic.ligo"
 -- example = "../../../src/test/contracts/bytes_unpack.ligo"
 -- example = "../../../src/test/contracts/chain_id.ligo"
--- example = "../../../src/test/contracts/coase.ligo"
+example = "../../../src/test/contracts/coase.ligo"
 -- example = "../../../src/test/contracts/failwith.ligo"
 -- example = "../../../src/test/contracts/loop.ligo"
--- example = "../../../src/test/contracts/application.ligo"
+-- example = "../../../src/test/contracts/redeclaration.ligo"
 -- example = "../../../src/test/contracts/application.ligo"
 -- example = "../../../src/test/contracts/application.ligo"
 -- example = "../../../src/test/contracts/application.ligo"

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

@@ -1,11 +1,13 @@
 
+{-# language Strict #-}
+
 {- | /The/ scope resolution system.
 -}
 
 module AST.Scope
   -- ( -- * Monad
-  --   ScopeM
+  --   CollectM
-  -- , evalScopeM
+  -- , evalCollectM
   -- , collectEnv
 
   --   -- * Scope
@@ -24,13 +26,17 @@ module AST.Scope
   -- )
   where
 
+import Control.Arrow (second)
 import Control.Monad.State
+import Control.Monad.Writer.Strict hiding (Alt, Product)
 
+import           Data.Function
 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           Data.Maybe (fromJust, listToMaybe)
+import qualified Data.List as List
 
 import Range
 import AST.Types
@@ -40,43 +46,119 @@ import Tree
 import Comment
 import Pretty
 import Product
+import Lattice
 
 import Debug.Trace
 
--- | Scope-holding monad.
+class HasLocalScope x where
-type ScopeM = State [Scopes]
+  getLocalScope :: x -> [ScopedDecl]
+
+instance Contains [ScopedDecl] xs => HasLocalScope (Product xs) where
+  getLocalScope = getElem
+
+type CollectM = State (Product [FullEnv, [Range]])
+
+type AddRefsM = State FullEnv
+
+type FullEnv = Map Range [ScopedDecl]
+
+addLocalScopes
+  :: HasRange (Product xs)
+  => Pascal (Product xs)
+  -> Pascal (Product ([ScopedDecl] : xs))
+addLocalScopes tree =
+    fmap (\xs -> Cons (envAt envWithREfs $ getRange xs) xs) tree
+  where
+    envWithREfs = flip execState env do
+      flip traverseOnly tree \r (Name t) -> do
+        modify $ addRef (getRange r) t
+        return (Name t)
+
+    env
+      = execCollectM
+      $ traverseTree pure tree
+
+envAt :: FullEnv -> Range -> [ScopedDecl]
+envAt env pos =
+    Map.elems scopes
+  where
+    ranges = List.sortBy partOrder $ filter isCovering $ Map.keys env
+    scopes = Map.unions $ (map.foldMap) toScopeMap $ map (env Map.!) ranges
+
+    isCovering = (pos <?)
+    toScopeMap sd@ScopedDecl {_sdName} = Map.singleton (ppToText _sdName) sd
+
+addRef :: Range -> Text -> FullEnv -> FullEnv
+addRef r n env = Map.union (go range) env
+  where
+    go (r' : rest) =
+      let decls = env Map.! r'
+      in
+        case updateOnly n r addRefToDecl decls of
+          (True,  decls) -> Map.singleton r' decls
+          (False, decls) -> Map.insert r' decls (go rest)
+    go [] = Map.empty
+
+    range
+      = List.sortBy partOrder
+      $ filter (r <?)
+      $ Map.keys env
+
+    -- decls' list = do
+    --   r'    <- range
+    --   decls <- Map.lookup r' env
+    --   return $ (r', updateOnly n r addRefToDecl decls)
+
+    addRefToDecl sd = sd
+      { _sdRefs = r : _sdRefs sd
+      }
+
+updateOnly
+  :: Text
+  -> Range
+  -> (ScopedDecl -> ScopedDecl)
+  -> [ScopedDecl]
+  -> (Bool, [ScopedDecl])
+updateOnly name r f = go
+  where
+    go = \case
+      d : ds
+        | ppToText (_sdName d) == name ->
+          if r == _sdOrigin d
+          then         (True,   d : ds)
+          else         (True, f d : ds)
+        | otherwise -> second (d :) (go ds)
+
+      [] -> (False, [])
+
+enter :: Range -> CollectM ()
+enter r = do
+  modify $ modElem (r :)
+
+define :: ScopedDecl -> CollectM ()
+define sd = do
+  r <- gets (head . getElem)
+  modify
+    $ modElem @FullEnv
+    $ Map.insertWith (++) r [sd]
+
+leave :: CollectM ()
+leave = modify $ modElem @[Range] tail
 
 -- | Run the computation with scope starting from empty scope.
-evalScopeM :: ScopeM a -> a
+execCollectM :: CollectM a -> FullEnv
-evalScopeM action = evalState action []
+execCollectM action = getElem $ execState action $ Cons Map.empty (Cons [] Nil)
 
-type Scopes = Tree '[ScopeTree] Range
+instance {-# OVERLAPS #-} Pretty FullEnv where
-
-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) =
+      aux (r, decls) =
-        pp o <+> (pp n <> ":") <+> pp t <+> "=" <+> pp b <+> "/" <+> (list rs <> ";")
+        pp r `indent` block decls
-
-  -- | 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
-  { _sdOrigin :: Range
+  { _sdName    :: Pascal ()
+  , _sdOrigin  :: Range
   , _sdBody    :: Maybe Range
   , _sdType    :: Maybe (Either (Pascal ()) Kind)
   , _sdRefs    :: [Range]
@@ -84,69 +166,100 @@ data ScopedDecl = ScopedDecl
   deriving Show via PP ScopedDecl
 
 instance Pretty ScopedDecl where
-  pp (ScopedDecl o b t refs) = pp o <+> "-" <+> maybe "?" (either pp pp) t <+> "=" <+> pp o
+  pp (ScopedDecl n o b t refs) = color 3 (pp n) <+> pp o <+> ":" <+> color 4 (maybe "?" (either pp pp) t) <+> "=" <+> pp refs
 
 -- | The kind.
 data Kind = Star
   deriving Show via PP Kind
 
 instance Pretty Kind where
-  pp _ = "*"
+  pp _ = "TYPE"
 
-lookupEnv :: Text -> Env -> Maybe ScopedDecl
+-- observe :: String -> CollectM a -> CollectM a
-lookupEnv name = Map.lookup name
+-- observe what act = do
+--   s <- get
+--   traceShowM (what, "BEFORE", s)
+--   a <- act
+--   s1 <- get
+--   traceShowM (what, "AFTER", s1)
+--   return a
 
--- | Make a new scope out of enclosing parent one.
+lookupEnv :: Text -> [ScopedDecl] -> Maybe ScopedDecl
-enter :: Range -> ScopeM ()
+lookupEnv name = listToMaybe . filter ((name ==) . ppToText . _sdName)
-enter r =
-  modify \rest ->
-    mk r (ScopeTree Map.empty []) : rest
 
--- | Leave current scope, return to parent one.
+-- -- | Make a new scope out of enclosing parent one.
-leave  :: ScopeM ()
+-- enter :: Range -> CollectM ()
-leave =
+-- enter r = observe "enter" do
-  modify \(a : parent : rest) ->
+--   modify \rest ->
-    fromJust do
+--     mk r (ScopeTree Map.empty []) : rest
-      (r, ScopeTree e cs) <- match parent
-      return $ mk r (ScopeTree e (a : cs)) : rest
 
--- | Add a declaration to the current scope.
+-- -- | Leave current scope, return to parent one.
-define :: Text -> ScopedDecl -> ScopeM ()
+-- leave :: CollectM ()
-define name d =
+-- leave = observe "leave" do
-  modify \(top : rest) ->
+--   modify \case
-    fromJust do
+--     (a : parent : rest) ->
-      (r, ScopeTree a cs) <- match top
+--       fromJust do
-      return $ mk r (ScopeTree (Map.insert name d a) cs) : rest
+--         -- traceShowM ("MOVE", a)
+--         -- traceShowM ("TO  ", parent)
+--         (r, ScopeTree e cs) <- match parent
+--         -- traceShowM ("==  ", mk r (ScopeTree e (a : cs)))
+--         -- traceShowM ("--")
+--         return $ mk r (ScopeTree e (a : cs)) : rest
+
+--     [x] -> error $ "trying to leave \n" ++ show x
+
+-- -- | Add a declaration to the current scope.
+-- define :: Text -> ScopedDecl -> CollectM ()
+-- define name d = observe "define" do
+--   s <- get
+--   traceShowM ("DEFINE", s)
+--   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 :: HasRange a => Pascal a -> Kind -> Pascal a -> CollectM ()
 defType name kind body = do
-  define (ppToText $ void name)
+  define
     $ ScopedDecl
+      (void name)
       (getRange $ infoOf name)
       (Just $ getRange $ infoOf body)
       (Just (Right kind))
       []
 
-addRef :: Text -> Range -> [Scopes] -> [Scopes]
+observe :: Pretty i => Pretty res => Text -> i -> res -> res
-addRef _    _   []                  = error "addRef: empty env stack"
+observe msg i res
-addRef name pos stack@(topmost : _)
+  = traceShow (pp msg, "INPUT", pp i)
-  | Just (r, ScopeTree top cs) <- match topmost =
+  $ traceShow (pp msg, "OUTPUT", pp res)
-    case Map.lookup name top of
+  $ res
-      Just ScopedDecl {_sdOrigin} -> go _sdOrigin stack
-      Nothing                     -> stack
 
-    where
+-- addRef
-      go _     []             = []
+--   :: Pascal ()
-      go range initial@(topmost : rest)
+--   -> Range
-        | Just (r, ScopeTree e cs) <- match topmost =
+--   -> FullEnv
-          case Map.lookup name e of
+--   -> FullEnv
-            Just it | _sdOrigin it == range ->
+-- addRef name pos (AppendMap envs) =
-              mk r (ScopeTree (pushRef pos e) cs) : go range rest
+--     AppendMap $ envs <> affected''
-            _ ->
+--   where
-              initial
+--     ranges = Map.keys envs
+--     (affected, other) = List.partition (pos <?) ranges
+--     affected' = foldMap (\r -> Map.singleton r (envs Map.! r)) affected
+--     affected'' = Map.map (\decls -> observe "addRef" decls $ addRefScopedDecls decls) affected'
 
-      pushRef pos' = Map.adjust (\sd -> sd { _sdRefs = pos' : _sdRefs sd }) name
+--     addRefScopedDecls :: [ScopedDecl] -> [ScopedDecl]
+--     addRefScopedDecls decls =
+--       case after of
+--         decl : after -> before ++ [addRefScopedDecl decl] ++ after
+--         [] -> before
+--       where
+--         (before, after) = break (\sd -> ppToText (_sdName sd) == ppName) decls
+
+--     addRefScopedDecl :: ScopedDecl -> ScopedDecl
+--     addRefScopedDecl decl = decl { _sdRefs = pos : _sdRefs decl }
+
+--     ppName = ppToText name
 
 -- | Add a value declaration to the current scope.
 def
@@ -154,25 +267,26 @@ def
   => Pascal a
   -> Maybe (Pascal a)
   -> Maybe (Pascal a)
-  -> ScopeM ()
+  -> CollectM ()
 def name ty body = do
-  define (ppToText $ void name)
+  define
     $ ScopedDecl
+      (void name)
       (getRange $ infoOf name)
       ((getRange . infoOf) <$> body)
       ((Left . void) <$> ty)
       []
 
-instance UpdateOver ScopeM Contract (Pascal a) where
+instance UpdateOver CollectM Contract (Pascal a) where
   before r _ = enter r
-  after  _ _ = leave
+  after  _ _ = skip
 
-instance HasRange a => UpdateOver ScopeM Declaration (Pascal a) where
+instance HasRange a => UpdateOver CollectM Declaration (Pascal a) where
   before _ = \case
     TypeDecl ty body -> defType ty Star body
     _ -> skip
 
-instance HasRange a => UpdateOver ScopeM Binding (Pascal a) where
+instance HasRange a => UpdateOver CollectM Binding (Pascal a) where
   before r = \case
     Function recur name _args ty body -> do
       when recur do
@@ -190,15 +304,15 @@ instance HasRange a => UpdateOver ScopeM Binding (Pascal a) where
       unless recur do
         def name (Just ty) (Just body)
 
-instance HasRange a => UpdateOver ScopeM VarDecl (Pascal a) where
+instance HasRange a => UpdateOver CollectM VarDecl (Pascal a) where
   after _ (Decl _ name ty) = def name (Just ty) Nothing
 
-instance UpdateOver ScopeM Mutable (Pascal a)
+instance UpdateOver CollectM Mutable (Pascal a)
-instance UpdateOver ScopeM Type    (Pascal a)
+instance UpdateOver CollectM Type    (Pascal a)
-instance UpdateOver ScopeM Variant (Pascal a)
+instance UpdateOver CollectM Variant (Pascal a)
-instance UpdateOver ScopeM TField  (Pascal a)
+instance UpdateOver CollectM TField  (Pascal a)
 
-instance HasRange a => UpdateOver ScopeM Expr (Pascal a) where
+instance HasRange a => UpdateOver CollectM Expr (Pascal a) where
   before r = \case
     Let    {} -> enter r
     Lambda {} -> enter r
@@ -220,78 +334,84 @@ instance HasRange a => UpdateOver ScopeM Expr (Pascal a) where
     ForBox  {} -> leave
     _ -> skip
 
-instance HasRange a => UpdateOver ScopeM Alt (Pascal a) where
+instance HasRange a => UpdateOver CollectM Alt (Pascal a) where
   before r _ = enter r
   after  _ _ = leave
 
-instance UpdateOver ScopeM LHS             (Pascal a)
+instance UpdateOver CollectM LHS             (Pascal a)
-instance UpdateOver ScopeM MapBinding      (Pascal a)
+instance UpdateOver CollectM MapBinding      (Pascal a)
-instance UpdateOver ScopeM Assignment      (Pascal a)
+instance UpdateOver CollectM Assignment      (Pascal a)
-instance UpdateOver ScopeM FieldAssignment (Pascal a)
+instance UpdateOver CollectM FieldAssignment (Pascal a)
-instance UpdateOver ScopeM Constant        (Pascal a)
+instance UpdateOver CollectM Constant        (Pascal a)
 
-instance HasRange a => UpdateOver ScopeM Pattern (Pascal a) where
+instance HasRange a => UpdateOver CollectM Pattern (Pascal a) where
   before _ = \case
     IsVar n -> def n Nothing Nothing
     _       -> skip
 
-instance UpdateOver ScopeM QualifiedName (Pascal a)
+instance UpdateOver CollectM QualifiedName (Pascal a)
-instance UpdateOver ScopeM Path          (Pascal a)
+instance UpdateOver CollectM Path          (Pascal a)
-instance UpdateOver ScopeM Name          (Pascal a) where
+instance UpdateOver CollectM Name          (Pascal a) where
   before range (Name raw) = do
-    modify $ addRef raw range
+    -- traceShowM ("name", raw)
+    skip
+    -- modify $ modElem $ addRef range (mk () (Name raw))
 
-class HasEnv a where
+-- class HasEnv a where
-  getEnv :: a -> Env
+--   getEnv :: a -> Env
 
-instance HasEnv Env where
+-- instance HasEnv Env where
-  getEnv = id
+--   getEnv = id
 
-instance Contains Env xs => HasEnv (Product xs) where
+-- instance Contains Env xs => HasEnv (Product xs) where
-  getEnv = getElem
+--   getEnv = getElem
 
-data Scope = Scope { unScope :: [Text] }
+-- data Scope = Scope { unScope :: [Text] }
 
-instance HasComments Scope where
+-- instance HasComments Scope where
-  getComments = unScope
+--   getComments = unScope
 
--- pinEnv :: Product xs -> ScopeM (Product (Env : xs))
+-- -- pinEnv :: Product xs -> CollectM (Product (Env : xs))
--- pinEnv xs = (`Cons` xs) <$> gets head
+-- -- pinEnv xs = (`Cons` xs) <$> gets head
 
-collectEnv :: Contains Range xs => Product xs -> ScopeM (Product (Scopes : xs))
+-- collectEnv :: Contains Range xs => Product xs -> CollectM (Product (Scopes : xs))
-collectEnv xs = do
+-- collectEnv xs = do
-  gets \case
+--   gets \case
-    st : _ -> Cons st xs
+--     st : _ -> Cons st xs
-    []     -> Cons (mk (getRange xs) $ ScopeTree Map.empty []) xs
+--     []     -> Cons (mk (getRange xs) $ ScopeTree Map.empty []) xs
 
-instance UpdateOver (State [Env]) ScopeTree Scopes where
+-- instance UpdateOver (State [Env]) ScopeTree Scopes where
-  before r (ScopeTree e _) = modify (e :)
+--   before r (ScopeTree e _) = modify (e :)
-  after  r  _              = modify tail
+--   after  r  _              = modify tail
 
-distributeEnv :: ScopeTree Scopes -> State [Env] (ScopeTree Scopes)
+-- distributeEnv :: ScopeTree Scopes -> State [Env] (ScopeTree Scopes)
-distributeEnv (ScopeTree e' cs) = do
+-- distributeEnv (ScopeTree e' cs) = do
-  e <- gets (Map.unions . (e' :))
+--   e <- gets (Map.unions . (e' :))
-  traceShowM ("distribute", e', e)
+--   return $ ScopeTree e cs
-  return $ ScopeTree e cs
 
-pinEnv :: Contains Range xs => Scopes -> Product xs -> ScopeM (Product (Env : xs))
+-- pinEnv :: Contains Range xs => Scopes -> Product xs -> CollectM (Product (Env : xs))
-pinEnv scopes info = do
+-- pinEnv scopes info = do
-  let (_, ScopeTree e _) = fromJust $ match =<< lookupTree (getElem info) scopes
+--   let (_, ScopeTree e _) = fromJust $ match =<< lookupTree (getElem info) scopes
-  return (Cons e info)
+--   return (Cons e info)
 
-instance HasComments Range where
+-- instance HasComments Range where
-  getComments _ = []
+--   getComments _ = []
 
-instance Pretty (Product xs) => HasComments (Product xs) where
+-- instance Pretty (Product xs) => HasComments (Product xs) where
-  getComments xs = if Text.null $ Text.strip x then [] else [x]
+--   getComments xs = if Text.null $ Text.strip x then [] else [x]
-    where
+--     where
-      x = ppToText $ color 3 $ pp $ xs
+--       x = ppToText $ color 3 $ pp $ xs
 
-ascribeEnv :: (Contains Range xs, Pretty (Product xs)) => Pascal (Product xs) -> Pascal (Product (Env : xs))
+-- ascribeEnv :: (Contains Range xs, Pretty (Product xs)) => Pascal (Product xs) -> Scopes -- Pascal (Product (Env : xs))
-ascribeEnv tree =
+-- ascribeEnv tree =
-  let
+--   let
-    Cons scopes _ = traceShowId $ infoOf $ evalScopeM $ traverseTree collectEnv tree
+--     scopes =
-    distributed   = evalState (traverseOnly distributeEnv scopes) []
+--       evalCollectM do
-  in
+--         traverseTree collectEnv tree
-    -- distributed
+--         gets head
-    evalScopeM $ traverseTree (pinEnv distributed) tree
+
+--     -- distributed   = evalState (traverseOnly distributeEnv scopes) []
+--   in
+--     scopes
+--     -- distributed
+--     -- evalCollectM $ 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
+instance (Contains [Text] xs) => HasComments (Product xs) where
---   getComments = getElem
+  getComments = getElem

tools/lsp/squirrel/src/Lattice.hs

@@ -4,6 +4,7 @@
 
 module Lattice
   ( Lattice(..)
+  , partOrder
   )
   where
 
@@ -16,3 +17,10 @@ class Lattice l where
   (<?) = flip (?>)
 
   {-# minimal (?>) | (<?) #-}
+
+partOrder :: Lattice l => l -> l -> Ordering
+partOrder a b | a <? b && b <? a = EQ
+partOrder a b | a <? b           = LT
+partOrder a b |           b <? a = GT
+partOrder a b                    = error "partOrder: Non-orderable"
+

tools/lsp/squirrel/src/Pretty.hs

@@ -76,6 +76,9 @@ instance Pretty1 Maybe where
 instance {-# OVERLAPS #-} (Pretty a, Pretty b) => Pretty (Either a b) where
   pp = either pp pp
 
+instance Pretty Int where
+  pp = int
+
 -- | Common instance.
 instance Pretty Text where
   pp = text . Text.unpack
@@ -108,7 +111,7 @@ train sep' = fsep . punctuate sep' . map pp
 
 -- | Pretty print as a vertical block.
 block :: Pretty p => [p] -> Doc
-block = vcat . map pp
+block = foldr ($+$) empty . map pp
 
 -- | For pretty-printing qualified names.
 sepByDot :: Pretty p => [p] -> Doc

tools/lsp/squirrel/src/Product.hs

@@ -24,8 +24,8 @@ instance Contains x xs => Contains x (y : xs) where
   getElem   (Cons _ xs) = getElem xs
   putElem x (Cons y xs) = Cons y (putElem x xs)
 
-modifyElem :: Contains x xs => (x -> x) -> Product xs -> Product xs
+modElem :: Contains x xs => (x -> x) -> Product xs -> Product xs
-modifyElem f xs = putElem (f $ getElem xs) xs
+modElem f xs = putElem (f $ getElem xs) xs
 
 instance Pretty (Product '[]) where
   pp _ = "{}"

tools/lsp/squirrel/src/Range.hs

@@ -25,6 +25,7 @@ data Range = Range
   , rFinish :: (Int, Int, Int)  -- ^ ... End: line, col, byte-offset).
   }
   deriving (Show) via PP Range
+  deriving stock (Ord)
 
 -- | TODO: Ugh. Purge it.
 diffRange :: Range -> Range -> Range
@@ -32,6 +33,7 @@ diffRange (Range ws wf) (Range ps _) = Range (max ws ps) wf
 
 instance Pretty Range where
   pp (Range (ll, lc, _) (rl, rc, _)) =
+    color 2 do
       brackets do
         int ll <> ":" <> int lc <> "-" <> int rl <> ":" <> int rc
 

tools/lsp/squirrel/src/Tree.hs

@@ -68,6 +68,22 @@ instance Apply Foldable layers => Foldable (Tree layers) where
       go (Tree (Left   err))      = foldMap f err
       go (Tree (Right (a, rest))) = f a <> foldMap go rest
 
+instance
+    ( Apply Traversable layers
+    , Apply Foldable    layers
+    , Apply Functor     layers
+    )
+  =>
+    Traversable (Tree layers)
+  where
+    traverse f = go
+      where
+        go (Tree (Left   err))      = (Tree . Left) <$> traverse f err
+        go (Tree (Right (a, rest))) = do
+          a'    <- f a
+          rest' <- (traverse.traverse) f rest
+          return $ Tree $ Right (a', rest')
+
 instance
     ( Apply Functor layers
     , HasComments info
@@ -140,38 +156,63 @@ traverseTree act = go
       return (Tree (Left err'))
 
 traverseOnly
-  :: forall f fs m a
+  :: forall f a b fs m
-  .  ( UpdateOver m      f   (Tree fs a)
+  .  ( Monad m
-     , UpdateOver m (Sum fs) (Tree fs a)
+     , Monad m
      , 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)))
+  => (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)
+      fb <- act r fa
-      before (getRange r) fa
-      fb <- act fa
       fc <- traverse go fb
-      after  (getRange r) fa
+      pure $ mk r fc
-      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
+      pure $ Tree $ Right (r, union')
-      return $ Tree $ Right (r, union')
 
-    go tree = return tree
+    go tree = pure tree
+
+data Visit fs a m where
+  Visit
+    :: (Element f fs, Traversable f)
+    => (a -> f (Tree fs a) -> m (f (Tree fs a)))
+    -> Visit fs a m
+
+traverseMany
+  :: ( Apply Functor fs
+     , Apply Foldable fs
+     , Apply Traversable fs
+     , Monad m
+     )
+  => [Visit fs a m]
+  -> Tree fs a
+  -> m (Tree fs a)
+traverseMany visitors = go
+  where
+    go tree = aux visitors
+      where
+        aux (Visit visitor : rest) = do
+          case match tree of
+            Just (r, fa) -> do
+              fa'  <- visitor r fa
+              fa'' <- traverse go fa'
+              return $ mk r fa''
+            Nothing -> do
+              aux rest
+        aux [] = do
+          case tree of
+            Tree (Right (r, union)) -> do
+              union' <- traverse go union
+              return $ Tree (Right (r, union'))
 
 -- | Make a tree out of a layer and an info.
 mk :: (Functor f, Element f fs) => info -> f (Tree fs info) -> Tree fs info