Implement autocompletion in waspls (#681)

* Refactor diagnostics

Now stores diagnostics info in the state instead of computing when diagnostics
are being published.

* Refactor Wasp.LSP.State into multiple files

* Move all LSP type imports to qualified

More consistent and more clear where things come from

* Use new concrete parser in diagnostics

Before running full analyzer, runs CST parser to check if it parses correctly
and to collect many parse errors. The CST is always put into the state. Then, if
successful, the full analyzer gets run as before.

* Add basic autocompletion

Always suggests every declaration name in the file, with no context-sensitivity.

* Refactor some names

* Syntax tree traversal code

* Attempt at context-sensitive autocomplete

Not done, but doing some more refactoring before going back to this

* Refactor Control.Syntax.Traverse to use total functions

* Make autocompletion context sensitive

It actually works now

* Miscellaneous code clean up

* Fix incorrect type name

* Fix syntax traversal bug with back

* Remove unreadable symbols from syntax traversal

* Refactor completion code

* Refactor syntax traversal code

* Apply suggestions from code review

Co-authored-by: Martin Šošić <Martinsos@users.noreply.github.com>

* Feedback from code review

* Add tests for Control.Monad.Loops.untilM

Co-authored-by: Martin Šošić <Martinsos@users.noreply.github.com>

waspls/src/Control/Monad/Loops.hs

@@ -0,0 +1,11 @@
+module Control.Monad.Loops
+  ( untilM,
+  )
+where
+
+-- | Analogue of 'until'. @'untilM' p f b@ yields the result of applying @f@
+-- until @p@ is true.
+untilM :: Monad m => (a -> Bool) -> (a -> m a) -> a -> m a
+untilM predicate f base
+  | predicate base = return base
+  | otherwise = f base >>= untilM predicate f

waspls/src/Control/Syntax/Traverse.hs

@@ -0,0 +1,341 @@
+module Control.Syntax.Traverse
+  ( -- * Syntax tree traversal
+
+    -- | Library for traversing around a concrete syntax trees. The main
+    -- benefits are:
+    --
+    -- - Easier to find nodes relative to a particular node
+    -- - Keeps track of absolute source offset
+    Traversal,
+
+    -- * Constructors
+    fromSyntax,
+    fromSyntaxForest,
+
+    -- * Traversal operations
+
+    -- | See the section on composition functions on how to compose these.
+    bottom,
+    down,
+    up,
+    left,
+    right,
+    next,
+    previous,
+
+    -- * Composition functions
+
+    -- | These functions can be used to combine many traversal operators
+    -- together in a more readable form.
+    --
+    -- Using @&@ is recommended so that expressions start with the traversal and
+    -- then have the operations. @>=>@ can be used for left-to-right composition
+    -- of two operations.
+    --
+    -- For example, you could write @traversal & pipe (replicate 3 next)@ to move
+    -- 3 times to the next position in the traversal.
+    --
+    -- @&?@ is also exported for the same reason as @&@, but for use with
+    -- @Maybe Traversal@.
+    (&),
+    (&?),
+    (>=>),
+    pipe,
+
+    -- * Collectors
+    kindAt,
+    widthAt,
+    offsetAt,
+    offsetAfter,
+    parentKind,
+    nodeAt,
+    parentNode,
+    ancestors,
+    siblings,
+    leftSiblings,
+    rightSiblings,
+    children,
+
+    -- * Predicates
+    hasChildren,
+    hasLeftSiblings,
+    hasRightSiblings,
+    hasNext,
+    hasPrevious,
+    hasAncestors,
+  )
+where
+
+import Control.Monad ((>=>))
+import Control.Monad.Loops (untilM)
+import Data.Foldable (Foldable (foldl'))
+import Data.Function ((&))
+import Data.List (unfoldr)
+import Data.List.NonEmpty (NonEmpty ((:|)))
+import Data.Maybe (isJust)
+import Wasp.Backend.ConcreteSyntax (SyntaxKind, SyntaxNode (snodeChildren, snodeKind, snodeWidth))
+
+-- | An in-progress traversal through some tree @f@.
+data Traversal = Traversal
+  { ancestorLevels :: [TraversalLevel],
+    currentLevel :: TraversalLevel
+  }
+  deriving (Eq, Ord, Show)
+
+data TraversalLevel = TraversalLevel
+  { tlCurrentNode :: !SyntaxNode,
+    tlCurrentOffset :: !Int,
+    tlLeftSiblings :: [SyntaxNode],
+    tlRightSiblings :: [SyntaxNode]
+  }
+  deriving (Eq, Show, Ord)
+
+tLeftSiblings :: Traversal -> [SyntaxNode]
+tLeftSiblings t = tlLeftSiblings $ currentLevel t
+
+tRightSiblings :: Traversal -> [SyntaxNode]
+tRightSiblings t = tlRightSiblings $ currentLevel t
+
+tChildren :: Traversal -> [SyntaxNode]
+tChildren t = snodeChildren $ nodeAt t
+
+-- | Create a new "Traversal" from a "SyntaxNode", starting at the root.
+fromSyntax :: SyntaxNode -> Traversal
+fromSyntax t = fromSyntaxForest [t]
+
+-- | Create a new "Traversal" from a forest of "TraversableTree"s, starting
+-- at the first tree in the list.
+--
+-- This function is not total. Invariant: @not (null forest)@
+fromSyntaxForest :: [SyntaxNode] -> Traversal
+fromSyntaxForest [] = error "Control.Tree.Traversal.fromTraversableForest on empty list"
+fromSyntaxForest (t : ts) =
+  Traversal
+    { ancestorLevels = [],
+      currentLevel = levelFromTraversableTree 0 (t :| ts)
+    }
+
+-- | Create a new "TraversalLevel" from a non-empty list of nodes.
+levelFromTraversableTree :: Int -> NonEmpty SyntaxNode -> TraversalLevel
+levelFromTraversableTree offset (node :| rSiblings) =
+  TraversalLevel
+    { tlCurrentNode = node,
+      tlCurrentOffset = offset,
+      tlLeftSiblings = [],
+      tlRightSiblings = rSiblings
+    }
+
+-- | Left-to-right composition of several traversal operations.
+pipe :: [Traversal -> Maybe Traversal] -> (Traversal -> Maybe Traversal)
+pipe ops = foldl' (>=>) Just ops
+
+-- | Synonym for @>>=@. Meant to be more visually similar to @&@, since they are
+-- used for essentially the same purpose in this library.
+(&?) :: Maybe Traversal -> (Traversal -> Maybe Traversal) -> Maybe Traversal
+t &? op = t >>= op
+
+-- | Move down the tree to the deepest left-most leaf
+bottom :: Traversal -> Traversal
+bottom t = maybe t bottom $ t & down
+
+-- | Move down a level in the tree, to the first child of the current position.
+down :: Traversal -> Maybe Traversal
+down t = case tChildren t of
+  [] -> Nothing
+  (c : cs) ->
+    Just $
+      Traversal
+        { ancestorLevels = currentLevel t : ancestorLevels t,
+          currentLevel = levelFromTraversableTree (offsetAt t) (c :| cs)
+        }
+
+-- | Move up a level in the tree, to the parent of the current position.
+up :: Traversal -> Maybe Traversal
+up t = case ancestorLevels t of
+  [] -> Nothing
+  (a : as) ->
+    Just $
+      Traversal
+        { ancestorLevels = as,
+          currentLevel = a
+        }
+
+-- | Move to the sibling left of the current position.
+left :: Traversal -> Maybe Traversal
+left t = case tLeftSiblings t of
+  [] -> Nothing
+  (l : ls) ->
+    Just $
+      t
+        { currentLevel =
+            TraversalLevel
+              { tlCurrentNode = l,
+                tlCurrentOffset = offsetAt t - snodeWidth l,
+                tlLeftSiblings = ls,
+                tlRightSiblings = nodeAt t : tRightSiblings t
+              }
+        }
+
+-- | Move to the sibling right of the current position.
+right :: Traversal -> Maybe Traversal
+right t = case tRightSiblings t of
+  [] -> Nothing
+  (r : rs) ->
+    Just $
+      t
+        { currentLevel =
+            TraversalLevel
+              { tlCurrentNode = r,
+                tlCurrentOffset = offsetAt t + widthAt t,
+                tlLeftSiblings = nodeAt t : tLeftSiblings t,
+                tlRightSiblings = rs
+              }
+        }
+
+-- | Move to the next node in the tree.
+--
+-- The next node is the first childless node encountered after the current
+-- position in a left-to-right depth-first-search of the tree.
+--
+-- __Examples:__
+--
+-- Moving to the child of the current node. This looks slightly surprising (it
+-- is moving to the left), but this is just an artifact of the diagram. All
+-- children are considered to occur after their parent.
+--
+-- @
+--     ┌───B───┐
+--     │       │
+--   ┌─C─┐   ┌─F─┐
+--   │ ▲ │   │   │
+-- ┌►D │ E   G   H
+-- │   │
+-- │ start
+-- │
+-- └─start & next
+-- @
+--
+-- Even though @J@ is a level higher than @H@, it is the next node.
+--
+-- @
+--       ┌───────A────┐
+--       │            │
+--   ┌───B───┐     ┌──I──┐
+--   │       │     │     │
+-- ┌─C─┐   ┌─F─┐   J   ┌─K─┐
+-- │   │   │   │   ▲   │   │
+-- D   E   G   H   │   L   M
+--             ▲   │
+--           start │
+--                 │
+--            start & next
+-- @
+next :: Traversal -> Maybe Traversal
+next t
+  | hasChildren t = untilM (not . hasChildren) down t
+  | hasAncestors t = case untilM hasRightSiblings up t of
+    Nothing -> Nothing
+    Just t' -> t' & pipe [right, untilM (not . hasChildren) down]
+  | otherwise = Nothing
+
+-- | Move to the previous node in a tree. This is 'next', but moves left instead
+-- of right.
+previous :: Traversal -> Maybe Traversal
+previous t
+  | hasChildren t = untilM (not . hasChildren) down t
+  | hasAncestors t = case untilM hasLeftSiblings up t of
+    Nothing -> Nothing
+    Just t' -> t' & pipe [left, untilM (not . hasChildren) $ down >=> rightMostSibling]
+  | otherwise = Nothing
+  where
+    rightMostSibling = untilM (not . hasRightSiblings) right
+
+-- | Get the "SyntaxKind" at the current position.
+kindAt :: Traversal -> SyntaxKind
+kindAt t = snodeKind $ nodeAt t
+
+-- | Get the width of the current node.
+widthAt :: Traversal -> Int
+widthAt t = snodeWidth $ nodeAt t
+
+-- | Get the offset of the start of the current node in the source text.
+offsetAt :: Traversal -> Int
+offsetAt t = tlCurrentOffset (currentLevel t)
+
+-- | Get the offset of the end of the current node in the source text.
+offsetAfter :: Traversal -> Int
+offsetAfter t = offsetAt t + widthAt t
+
+-- | Get the "SyntaxKind" of the parent of the current position.
+--
+-- [Property] @'parentKind' t == 'contentAt' (t & 'up')@
+parentKind :: Traversal -> Maybe SyntaxKind
+parentKind t = kindAt <$> up t
+
+-- | Get the node at the current position.
+nodeAt :: Traversal -> SyntaxNode
+nodeAt t = tlCurrentNode (currentLevel t)
+
+-- | Get the parent node of the current position.
+--
+-- [Property] @'parentNode' t == 'nodeAt' (t & 'up')@
+parentNode :: Traversal -> Maybe SyntaxNode
+parentNode t = nodeAt <$> up t
+
+-- | Get the ancestors of the current position.
+ancestors :: Traversal -> [Traversal]
+ancestors t = unfoldr step (t & up)
+  where
+    step Nothing = Nothing
+    step (Just t') = Just (t', t' & up)
+
+-- | Get the siblings of the current position (not including the current node).
+--
+-- [Property] @'siblings' t == 'leftSiblings' t ++ 'rightSiblings' t@
+siblings :: Traversal -> [Traversal]
+siblings t = leftSiblings t ++ rightSiblings t
+
+-- | Get siblings left of the current position.
+leftSiblings :: Traversal -> [Traversal]
+leftSiblings t = reverse $ unfoldr step t
+  where
+    step = left >=> (\x -> return (x, x))
+
+-- | Get siblings right of the current position.
+rightSiblings :: Traversal -> [Traversal]
+rightSiblings t = unfoldr step t
+  where
+    step = right >=> (\x -> return (x, x))
+
+-- | Get the children of the current position.
+children :: Traversal -> [Traversal]
+children t = unfoldr step (t & down)
+  where
+    step Nothing = Nothing
+    step (Just t') = Just (t', t' & right)
+
+-- | Check if the current position has children.
+hasChildren :: Traversal -> Bool
+hasChildren t = not $ null $ children t
+
+-- | Check if the current position has siblings to the left.
+hasLeftSiblings :: Traversal -> Bool
+hasLeftSiblings t = not $ null $ leftSiblings t
+
+-- | Check if the current position has siblings to the right.
+hasRightSiblings :: Traversal -> Bool
+hasRightSiblings t = not $ null $ rightSiblings t
+
+-- | Check if the current position has a next position. See the documentation for
+-- 'next' for a definition of what this means.
+hasNext :: Traversal -> Bool
+hasNext t = isJust $ next t
+
+-- | Check if the current position has a previous position. Analogue for 'previous'
+-- of 'hasNext'.
+hasPrevious :: Traversal -> Bool
+hasPrevious t = isJust $ previous t
+
+-- | Check if the current position has at least one parent.
+hasAncestors :: Traversal -> Bool
+hasAncestors t = not $ null $ ancestorLevels t

waspls/src/Wasp/LSP/Completion.hs

@@ -0,0 +1,97 @@
+module Wasp.LSP.Completion
+  ( getCompletionsAtPosition,
+  )
+where
+
+import Control.Lens ((?~), (^.))
+import Control.Syntax.Traverse
+import Data.Maybe (maybeToList)
+import qualified Data.Text as Text
+import qualified Language.LSP.Types as LSP
+import qualified Language.LSP.Types.Lens as LSP
+import Wasp.Backend.ConcreteSyntax (SyntaxNode)
+import qualified Wasp.Backend.ConcreteSyntax as S
+import Wasp.LSP.ServerM
+import Wasp.LSP.ServerState
+import Wasp.LSP.Syntax (findChild, isAtExprPlace, lexemeAt, lspPositionToOffset, showNeighborhood, toOffset)
+
+-- | Get the list of completions at a (line, column) position in the source.
+getCompletionsAtPosition :: LSP.Position -> ServerM [LSP.CompletionItem]
+getCompletionsAtPosition position = do
+  src <- gets (^. currentWaspSource)
+  maybeSyntax <- gets (^. cst)
+  case maybeSyntax of
+    -- If there is no syntax tree, make no completions
+    Nothing -> return []
+    Just syntax -> do
+      let offset = lspPositionToOffset src position
+      -- 'location' is a traversal through the syntax tree that points to 'position'
+      let location = toOffset offset (fromSyntaxForest syntax)
+      logM $ "[getCompletionsAtPosition] neighborhood=\n" ++ showNeighborhood location
+      exprCompletions <-
+        if isAtExprPlace location
+          then do
+            logM $ "[getCompletionsAtPosition] position=" ++ show position ++ " atExpr=True"
+            getExprCompletions src syntax
+          else do
+            logM $ "[getCompletionsAtPosition] position=" ++ show position ++ " atExpr=False"
+            return []
+      let completions = exprCompletions
+      return completions
+
+-- | If the location is at an expression, find declaration names in the file
+-- and return them as autocomplete suggestions
+--
+-- TODO: include completions for enum variants (use standard type defs from waspc)
+getExprCompletions :: String -> [SyntaxNode] -> ServerM [LSP.CompletionItem]
+getExprCompletions src syntax = do
+  let declNames = findDeclNames src syntax
+  logM $ "[getExprCompletions] declnames=" ++ show declNames
+  return $
+    map
+      ( \(name, typ) ->
+          makeBasicCompletionItem (Text.pack name)
+            & (LSP.kind ?~ LSP.CiVariable)
+            & (LSP.detail ?~ Text.pack (":: " ++ typ ++ " (declaration type)"))
+      )
+      declNames
+
+-- | Search through the CST and collect all @(declName, declType)@ pairs.
+findDeclNames :: String -> [SyntaxNode] -> [(String, String)]
+findDeclNames src syntax = traverseForDeclNames $ fromSyntaxForest syntax
+  where
+    traverseForDeclNames :: Traversal -> [(String, String)]
+    traverseForDeclNames t = case kindAt t of
+      S.Program -> maybe [] traverseForDeclNames $ down t
+      S.Decl ->
+        let declNameAndType = maybeToList $ getDeclNameAndType t
+         in declNameAndType ++ maybe [] traverseForDeclNames (right t)
+      _ -> maybe [] traverseForDeclNames $ right t
+    getDeclNameAndType :: Traversal -> Maybe (String, String)
+    getDeclNameAndType t = do
+      nameT <- findChild S.DeclName t
+      typeT <- findChild S.DeclType t
+      return (lexemeAt src nameT, lexemeAt src typeT)
+
+-- | Create a completion item containing only a label.
+makeBasicCompletionItem :: Text.Text -> LSP.CompletionItem
+makeBasicCompletionItem name =
+  LSP.CompletionItem
+    { _label = name,
+      _kind = Nothing,
+      _tags = Nothing,
+      _detail = Nothing,
+      _documentation = Nothing,
+      _deprecated = Nothing,
+      _preselect = Nothing,
+      _sortText = Nothing,
+      _filterText = Nothing,
+      _insertText = Nothing,
+      _insertTextFormat = Nothing,
+      _insertTextMode = Nothing,
+      _textEdit = Nothing,
+      _additionalTextEdits = Nothing,
+      _commitCharacters = Nothing,
+      _command = Nothing,
+      _xdata = Nothing
+    }

waspls/src/Wasp/LSP/Core.hs

@@ -1,52 +0,0 @@
-module Wasp.LSP.Core
-  ( ServerM,
-    ServerError (..),
-    Severity (..),
-    ServerState,
-    ServerConfig,
-  )
-where
-
-import Control.Monad.Trans.Except (ExceptT)
-import Control.Monad.Trans.State.Strict (StateT)
-import Data.Aeson
-import Data.Aeson.Types (prependFailure, typeMismatch)
-import Data.Default (Default (def))
-import Data.Text (Text)
-import Language.LSP.Server (LspT)
-
-type ServerM =
-  ExceptT ServerError (StateT ServerState (LspT ServerConfig IO))
-
--- | The type for a language server error. These are separate from diagnostics
--- and should be reported when the server fails to process a request/notification
--- for some reason.
-data ServerError = ServerError Severity Text
-
--- | Error severity levels
-data Severity
-  = -- | Displayed to user as an error
-    Error
-  | -- | Displayed to user as a warning
-    Warning
-  | -- | Displayed to user
-    Info
-  | -- | Not displayed to the user
-    Log
-
-data ServerConfig = ServerConfig {}
-
-instance Default ServerConfig where
-  def = ServerConfig {}
-
-instance FromJSON ServerConfig where
-  parseJSON (Object _) = pure ServerConfig
-  parseJSON invalid =
-    prependFailure
-      "parsing ServerConfig failed, "
-      (typeMismatch "Object" invalid)
-
-data ServerState = ServerState {}
-
-instance Default ServerState where
-  def = ServerState {}

waspls/src/Wasp/LSP/Diagnostic.hs

@@ -0,0 +1,72 @@
+module Wasp.LSP.Diagnostic
+  ( waspErrorToDiagnostic,
+    concreteParseErrorToDiagnostic,
+  )
+where
+
+import Data.Text (Text)
+import qualified Data.Text as Text
+import qualified Language.LSP.Types as LSP
+import qualified Wasp.Analyzer.AnalyzeError as W
+import qualified Wasp.Analyzer.Parser as W
+import qualified Wasp.Backend.ParseError as C
+import Wasp.LSP.ServerM (ServerM, logM)
+import Wasp.LSP.Util (waspSourceRegionToLspRange)
+
+concreteParseErrorToDiagnostic :: String -> C.ParseError -> ServerM LSP.Diagnostic
+concreteParseErrorToDiagnostic src err =
+  let message = Text.pack $ C.showError src err
+      source = "parse"
+      range = concreteErrorRange err
+   in logM ("[concreteParseErroToDiagnostic] _range=" ++ show range)
+        >> return
+          ( LSP.Diagnostic
+              { _range = range,
+                _severity = Nothing,
+                _code = Nothing,
+                _source = Just source,
+                _message = message,
+                _tags = Nothing,
+                _relatedInformation = Nothing
+              }
+          )
+  where
+    concreteErrorRange e = case C.errorRegion e of
+      C.Region start end ->
+        let startPos = C.offsetToSourcePos src start
+            endPos = C.offsetToSourcePos src end
+         in LSP.Range (concretePosToLSPPos startPos) (concretePosToLSPPos endPos)
+    concretePosToLSPPos (C.SourcePos l c) =
+      LSP.Position (fromIntegral l - 1) (fromIntegral c - 1)
+
+waspErrorToDiagnostic :: W.AnalyzeError -> LSP.Diagnostic
+waspErrorToDiagnostic err =
+  let message = waspErrorAsPrettyEditorMessage err
+      source = waspErrorSource err
+      range = waspErrorRange err
+   in LSP.Diagnostic
+        { _range = range,
+          _severity = Nothing,
+          _code = Nothing,
+          _source = Just source,
+          _message = message,
+          _tags = Nothing,
+          _relatedInformation = Nothing
+        }
+
+-- | Convert a wasp error to a message to display to the developer.
+--
+-- TODO: Write a new conversion from error to text here that is better suited
+-- for in-editor display
+waspErrorAsPrettyEditorMessage :: W.AnalyzeError -> Text
+waspErrorAsPrettyEditorMessage = Text.pack . fst . W.getErrorMessageAndCtx
+
+waspErrorSource :: W.AnalyzeError -> Text
+waspErrorSource (W.ParseError _) = "parse"
+waspErrorSource (W.TypeError _) = "typecheck"
+waspErrorSource (W.EvaluationError _) = "evaluate"
+
+waspErrorRange :: W.AnalyzeError -> LSP.Range
+waspErrorRange err =
+  let (_, W.Ctx rgn) = W.getErrorMessageAndCtx err
+   in waspSourceRegionToLspRange rgn

waspls/src/Wasp/LSP/Handlers.hs

@@ -3,13 +3,11 @@ module Wasp.LSP.Handlers
     didOpenHandler,
     didChangeHandler,
     didSaveHandler,
+    completionHandler,
   )
 where
 
-import Control.Lens ((+~), (^.))
-import Control.Monad.Trans (lift)
-import Control.Monad.Trans.Except (throwE)
-import Data.Function ((&))
+import Control.Lens ((.~), (?~), (^.))
 import Data.Text (Text)
 import qualified Data.Text as T
 import Language.LSP.Server (Handlers, LspT)
@@ -17,11 +15,14 @@ import qualified Language.LSP.Server as LSP
 import qualified Language.LSP.Types as LSP
 import qualified Language.LSP.Types.Lens as LSP
 import Language.LSP.VFS (virtualFileText)
-import qualified Wasp.Analyzer
-import qualified Wasp.Analyzer.AnalyzeError as WE
-import Wasp.Analyzer.Parser (Ctx (Ctx))
-import Wasp.Analyzer.Parser.SourceRegion (getRgnEnd, getRgnStart)
-import Wasp.LSP.Core (ServerConfig, ServerError (ServerError), ServerM, Severity (..))
+import Wasp.Analyzer (analyze)
+import Wasp.Backend.ConcreteParser (parseCST)
+import qualified Wasp.Backend.Lexer as L
+import Wasp.LSP.Completion (getCompletionsAtPosition)
+import Wasp.LSP.Diagnostic (concreteParseErrorToDiagnostic, waspErrorToDiagnostic)
+import Wasp.LSP.ServerConfig (ServerConfig)
+import Wasp.LSP.ServerM (ServerError (..), ServerM, Severity (..), gets, lift, modify, throwError)
+import Wasp.LSP.ServerState (cst, currentWaspSource, latestDiagnostics)
 
 -- LSP notification and request handlers
 
@@ -59,6 +60,12 @@ didSaveHandler :: Handlers ServerM
 didSaveHandler =
   LSP.notificationHandler LSP.STextDocumentDidSave $ diagnoseWaspFile . extractUri
 
+completionHandler :: Handlers ServerM
+completionHandler =
+  LSP.requestHandler LSP.STextDocumentCompletion $ \request respond -> do
+    completions <- getCompletionsAtPosition $ request ^. LSP.params . LSP.position
+    respond $ Right $ LSP.InL $ LSP.List completions
+
 -- | Does not directly handle a notification or event, but should be run when
 -- text document content changes.
 --
@@ -67,50 +74,41 @@ didSaveHandler =
 -- file in "Wasp.LSP.State.State".
 diagnoseWaspFile :: LSP.Uri -> ServerM ()
 diagnoseWaspFile uri = do
-  src <- readVFSFile uri
-  let appSpecOrError = Wasp.Analyzer.analyze $ T.unpack src
-  diagnostics <- case appSpecOrError of
-    -- Valid wasp file, send no diagnostics
-    Right _ -> return $ LSP.List []
-    -- Report the error (for now, just one error per analyze is possible)
-    Left err ->
-      return $
-        LSP.List
-          [ waspErrorToLspDiagnostic err
-          ]
+  analyzeWaspFile uri
+  currentDiagnostics <- gets (^. latestDiagnostics)
   liftLSP $
     LSP.sendNotification LSP.STextDocumentPublishDiagnostics $
-      LSP.PublishDiagnosticsParams uri Nothing diagnostics
+      LSP.PublishDiagnosticsParams uri Nothing (LSP.List currentDiagnostics)
+
+analyzeWaspFile :: LSP.Uri -> ServerM ()
+analyzeWaspFile uri = do
+  srcString <- readAndStoreSourceString
+  let (concreteErrorMessages, concreteSyntax) = parseCST $ L.lex srcString
+  modify (cst ?~ concreteSyntax)
+  if not $ null concreteErrorMessages
+    then storeCSTErrors concreteErrorMessages
+    else runWaspAnalyzer srcString
   where
-    waspErrorToLspDiagnostic :: WE.AnalyzeError -> LSP.Diagnostic
-    waspErrorToLspDiagnostic err =
-      let errSrc = case err of
-            WE.ParseError _ -> "parse"
-            WE.TypeError _ -> "typecheck"
-            WE.EvaluationError _ -> "evaluate"
-          (errMsg, errCtx) = WE.getErrorMessageAndCtx err
-       in LSP.Diagnostic
-            { _range = waspCtxToLspRange errCtx,
-              _severity = Nothing,
-              _code = Nothing,
-              _source = Just errSrc,
-              _message = T.pack errMsg,
-              _tags = Nothing,
-              _relatedInformation = Nothing
-            }
+    readAndStoreSourceString = do
+      srcString <- T.unpack <$> readVFSFile uri
+      modify (currentWaspSource .~ srcString)
+      return srcString
 
-    waspCtxToLspRange :: Ctx -> LSP.Range
-    waspCtxToLspRange (Ctx region) =
-      LSP.Range
-        { _start = waspSourcePositionToLspPosition (getRgnStart region),
-          -- Increment end character by 1: Wasp uses an inclusive convention for
-          -- the end position, but LSP considers end position to not be part of
-          -- the range.
-          _end = waspSourcePositionToLspPosition (getRgnEnd region) & LSP.character +~ (1 :: LSP.UInt)
-        }
+    storeCSTErrors concreteErrorMessages = do
+      srcString <- gets (^. currentWaspSource)
+      newDiagnostics <- mapM (concreteParseErrorToDiagnostic srcString) concreteErrorMessages
+      modify (latestDiagnostics .~ newDiagnostics)
 
-    waspSourcePositionToLspPosition (WE.SourcePosition l c) =
-      LSP.Position (fromIntegral $ l - 1) (fromIntegral $ c - 1)
+    runWaspAnalyzer srcString = do
+      let analyzeResult = analyze srcString
+      case analyzeResult of
+        Right _ -> do
+          modify (latestDiagnostics .~ [])
+        Left err -> do
+          let newDiagnostics =
+                [ waspErrorToDiagnostic err
+                ]
+          modify (latestDiagnostics .~ newDiagnostics)
 
 -- | Run a LSP function in the "ServerM" monad.
 liftLSP :: LspT ServerConfig IO a -> ServerM a
@@ -123,7 +121,7 @@ readVFSFile uri = do
   mVirtualFile <- liftLSP $ LSP.getVirtualFile $ LSP.toNormalizedUri uri
   case mVirtualFile of
     Just virtualFile -> return $ virtualFileText virtualFile
-    Nothing -> throwE $ ServerError Error $ "Could not find " <> T.pack (show uri) <> " in VFS."
+    Nothing -> throwError $ ServerError Error $ "Could not find " <> T.pack (show uri) <> " in VFS."
 
 -- | Get the "Uri" from an object that has a "TextDocument".
 extractUri :: (LSP.HasParams a b, LSP.HasTextDocument b c, LSP.HasUri c LSP.Uri) => a -> LSP.Uri

waspls/src/Wasp/LSP/Server.hs

@@ -18,8 +18,20 @@ import qualified Language.LSP.Server as LSP
 import qualified Language.LSP.Types as LSP
 import System.Exit (ExitCode (ExitFailure), exitWith)
 import qualified System.Log.Logger
-import Wasp.LSP.Core (ServerConfig, ServerError (ServerError), ServerM, ServerState, Severity (..))
 import Wasp.LSP.Handlers
+import Wasp.LSP.ServerConfig (ServerConfig)
+import Wasp.LSP.ServerM (ServerError (..), ServerM, Severity (..))
+import Wasp.LSP.ServerState (ServerState)
+
+lspServerHandlers :: LSP.Handlers ServerM
+lspServerHandlers =
+  mconcat
+    [ initializedHandler,
+      didOpenHandler,
+      didSaveHandler,
+      didChangeHandler,
+      completionHandler
+    ]
 
 serve :: Maybe FilePath -> IO ()
 serve maybeLogFile = do
@@ -94,15 +106,6 @@ lspServerOptions =
       LSP.completionTriggerCharacters = Just [':']
     }
 
-lspServerHandlers :: LSP.Handlers ServerM
-lspServerHandlers =
-  mconcat
-    [ initializedHandler,
-      didOpenHandler,
-      didSaveHandler,
-      didChangeHandler
-    ]
-
 -- | Options to tell the client how to update the server about the state of text
 -- documents in the workspace.
 syncOptions :: LSP.TextDocumentSyncOptions

waspls/src/Wasp/LSP/ServerConfig.hs

@@ -0,0 +1,20 @@
+module Wasp.LSP.ServerConfig
+  ( ServerConfig (..),
+  )
+where
+
+import Data.Aeson (FromJSON (parseJSON), Value (Object))
+import Data.Aeson.Types (prependFailure, typeMismatch)
+import Data.Default (Default (def))
+
+data ServerConfig = ServerConfig {}
+
+instance Default ServerConfig where
+  def = ServerConfig {}
+
+instance FromJSON ServerConfig where
+  parseJSON (Object _) = pure ServerConfig
+  parseJSON invalid =
+    prependFailure
+      "parsing ServerConfig failed, "
+      (typeMismatch "Object" invalid)

waspls/src/Wasp/LSP/ServerM.hs

@@ -0,0 +1,58 @@
+module Wasp.LSP.ServerM
+  ( ServerM,
+    ServerError (..),
+    Severity (..),
+    -- | You should usually use lenses for accessing the state.
+    --
+    -- __Examples:__
+    --
+    -- > import Control.Lens ((^.))
+    -- > gets (^. diagnostics) -- Gets the list of diagnostics
+    --
+    -- > import Control.Lens ((.~))
+    -- > modify (diagnostics .~ []) -- Clears diagnostics in the state
+    StateT.gets,
+    StateT.modify,
+    logM,
+    lift,
+    catchError,
+    throwError,
+  )
+where
+
+import Control.Monad.Except (ExceptT, catchError, throwError)
+import Control.Monad.State.Strict (StateT)
+import qualified Control.Monad.State.Strict as StateT
+import Control.Monad.Trans (MonadIO (liftIO), lift)
+import Data.Text (Text)
+import Language.LSP.Server (LspT)
+import qualified System.Log.Logger as L
+import Wasp.LSP.ServerConfig (ServerConfig)
+import Wasp.LSP.ServerState (ServerState)
+
+type ServerM = ExceptT ServerError (StateT ServerState (LspT ServerConfig IO))
+
+-- | Log a string.
+--
+-- Behavior depends on the "--log" command line flag. If set to "[OUTPUT]",
+-- logged messages will be displayed in the LSP client (e.g. for VSCode, in the
+-- "Wasp Language Extension" output panel). Otherwise, it may be sent to a file
+-- or not recorded at all.
+logM :: String -> ServerM ()
+logM = liftIO . L.logM "haskell-lsp" L.DEBUG
+
+-- | The type for a language server error. These are separate from diagnostics
+-- and should be reported when the server fails to process a request/notification
+-- for some reason.
+data ServerError = ServerError Severity Text
+
+-- | Error severity levels
+data Severity
+  = -- | Displayed to user as an error
+    Error
+  | -- | Displayed to user as a warning
+    Warning
+  | -- | Displayed to user
+    Info
+  | -- | Not displayed to the user
+    Log

waspls/src/Wasp/LSP/ServerState.hs

@@ -0,0 +1,39 @@
+{-# LANGUAGE TemplateHaskell #-}
+
+module Wasp.LSP.ServerState
+  ( ServerState (..),
+    currentWaspSource,
+    latestDiagnostics,
+    cst,
+  )
+where
+
+import Control.Lens (makeClassy)
+import Data.Default (Default (def))
+import qualified Language.LSP.Types as LSP
+import Wasp.Backend.ConcreteSyntax (SyntaxNode)
+
+-- | LSP State preserved between handlers.
+--
+-- The server assumes the project has only wasp file. This state will not
+-- work correctly if there are multiple wasp files.
+--
+-- Recommended to use the lenses for accessing the fields.
+data ServerState = ServerState
+  { -- | Source text for wasp file.
+    _currentWaspSource :: String,
+    -- | List of diagnostics generated by waspc after the last file change.
+    _latestDiagnostics :: [LSP.Diagnostic],
+    -- | Concrete syntax tree representing '_currentWaspSource'.
+    _cst :: Maybe [SyntaxNode]
+  }
+
+makeClassy 'ServerState
+
+instance Default ServerState where
+  def =
+    ServerState
+      { _currentWaspSource = "",
+        _latestDiagnostics = [],
+        _cst = Nothing
+      }

waspls/src/Wasp/LSP/Syntax.hs

@@ -0,0 +1,89 @@
+module Wasp.LSP.Syntax
+  ( -- * Syntax
+
+    -- | Module with utilities for working with/looking for patterns in CSTs
+    lspPositionToOffset,
+    toOffset,
+    isAtExprPlace,
+    lexemeAt,
+    findChild,
+    -- | Printing
+    showNeighborhood,
+  )
+where
+
+import Control.Syntax.Traverse
+import Data.List (find, intercalate)
+import qualified Language.LSP.Types as J
+import qualified Wasp.Backend.ConcreteSyntax as S
+
+-- | @lspPositionToOffset srcString position@ returns 0-based offset from the
+-- start of @srcString@ to the specified line and column.
+lspPositionToOffset :: String -> J.Position -> Int
+lspPositionToOffset srcString (J.Position l c) =
+  let linesBefore = take (fromIntegral l) (lines srcString)
+   in -- We add 1 to the length of each line to make sure to count the newline
+      sum (map ((+ 1) . length) linesBefore) + fromIntegral c
+
+-- | Move to the node containing the offset.
+--
+-- This tries to prefer non-trivia tokens where possible. If the offset falls
+-- exactly between two tokens, it choses the left-most non-trivia token.
+toOffset :: Int -> Traversal -> Traversal
+toOffset targetOffset start = go $ bottom start
+  where
+    go :: Traversal -> Traversal
+    go at
+      | offsetAt at == targetOffset = at
+      | offsetAfter at > targetOffset = at
+      | offsetAfter at == targetOffset && not (S.syntaxKindIsTrivia (kindAt at)) =
+        at
+      -- If @at & next@ fails, the input doesn't contain the offset, so just
+      -- return the last node instead.
+      | otherwise = maybe at go $ at & next
+
+-- | Check whether a position in a CST is somewhere an expression belongs. These
+-- locations (as of now) are:
+--
+-- - Parent is DictEntry, has a DictKey left siblings
+-- - Parent is Decl, has DeclType and DeclName left siblings
+-- - Parent is a List
+-- - Parent is a Tuple
+isAtExprPlace :: Traversal -> Bool
+isAtExprPlace t =
+  (parentIs S.DictEntry && hasLeft S.DictKey)
+    || parentIs S.List
+    || (parentIs S.Decl && hasLeft S.DeclType && hasLeft S.DeclName)
+    || parentIs S.Tuple
+  where
+    parentIs k = Just k == parentKind t
+    hasLeft k = k `elem` map kindAt (leftSiblings t)
+
+-- | Show the nodes around the current position
+--
+-- Used for debug purposes
+showNeighborhood :: Traversal -> String
+showNeighborhood t =
+  let parentStr = case t & up of
+        Nothing -> "<ROOT>"
+        Just parent -> showNode "" parent
+      leftSiblingLines = map (showNode "  ") $ leftSiblings t
+      currentStr = showNode "  " t ++ " <--"
+      rightSiblingLines = map (showNode "  ") $ rightSiblings t
+   in intercalate "\n" $ parentStr : leftSiblingLines ++ [currentStr] ++ rightSiblingLines
+  where
+    showNode indent node =
+      indent
+        ++ show (kindAt node)
+        ++ "@"
+        ++ show (offsetAt node)
+        ++ ".."
+        ++ show (offsetAfter node)
+
+-- | Search for a child node with the matching "SyntaxKind".
+findChild :: S.SyntaxKind -> Traversal -> Maybe Traversal
+findChild skind t = find ((== skind) . kindAt) $ children t
+
+-- | @lexeme src traversal@
+lexemeAt :: String -> Traversal -> String
+lexemeAt src t = take (widthAt t) $ drop (offsetAt t) src

waspls/src/Wasp/LSP/Util.hs

@@ -0,0 +1,22 @@
+module Wasp.LSP.Util (waspSourceRegionToLspRange, waspPositionToLspPosition) where
+
+import Control.Lens ((+~))
+import Data.Function ((&))
+import qualified Language.LSP.Types as LSP
+import qualified Language.LSP.Types.Lens as LSP
+import qualified Wasp.Analyzer.Parser as W
+import qualified Wasp.Analyzer.Parser.SourceRegion as W
+
+waspSourceRegionToLspRange :: W.SourceRegion -> LSP.Range
+waspSourceRegionToLspRange rgn =
+  LSP.Range
+    { _start = waspPositionToLspPosition (W.getRgnStart rgn),
+      _end = waspPositionToLspPosition (W.getRgnEnd rgn) & LSP.character +~ 1
+    }
+
+waspPositionToLspPosition :: W.SourcePosition -> LSP.Position
+waspPositionToLspPosition (W.SourcePosition ln col) =
+  LSP.Position
+    { _line = fromIntegral ln - 1,
+      _character = fromIntegral col - 1
+    }

waspls/test/Control/Monad/LoopsTest.hs

@@ -0,0 +1,26 @@
+{-# LANGUAGE DeriveFunctor #-}
+
+module Control.Monad.LoopsTest where
+
+import Control.Monad.Loops
+import Test.Tasty.Hspec
+
+data Identity a = Identity {runIdentity :: a} deriving (Functor)
+
+instance Applicative Identity where
+  pure = Identity
+  Identity f <*> Identity x = Identity (f x)
+
+instance Monad Identity where
+  return = Identity
+  Identity x >>= f = f x
+
+spec_untilM :: Spec
+spec_untilM = do
+  it "Is identical to until with the Identity monad" $ do
+    runIdentity (untilM (== 0) (return . pred) 10) `shouldBe` until (== 0) pred (10 :: Int)
+
+  it "Works with Maybe monad" $ do
+    let test n = untilM (== 0) (\x -> if x < 0 then Nothing else return (pred x)) n
+    test (10 :: Int) `shouldBe` Just 0
+    test (-2 :: Int) `shouldBe` Nothing

waspls/test/Control/Syntax/TraverseTest.hs

@@ -0,0 +1,55 @@
+{-# LANGUAGE TypeFamilies #-}
+
+module Control.Syntax.TraverseTest where
+
+import Control.Syntax.Traverse
+import Test.Tasty.Hspec (Spec, it, shouldBe)
+import Wasp.Backend.ConcreteParser (parseCST)
+import Wasp.Backend.ConcreteSyntax (SyntaxKind (..), SyntaxNode)
+import qualified Wasp.Backend.Lexer as L
+import qualified Wasp.Backend.Token as T
+
+example1 :: [SyntaxNode]
+example1 =
+  snd $
+    parseCST $
+      L.lex $
+        unlines
+          [ "app Main {",
+            "  test: 5",
+            "}"
+          ]
+
+spec_Traverse :: Spec
+spec_Traverse = do
+  it "Can traverse around a tree" $ do
+    let root = fromSyntaxForest example1
+    kindAt root `shouldBe` Program
+    let declName = root & pipe [down, down, right, right]
+    (kindAt <$> declName) `shouldBe` Just DeclName
+    let dictKey = declName &? pipe [right, right, down, right, right, right, down]
+    (kindAt <$> dictKey) `shouldBe` Just DictKey
+
+  it "next goes as deep as possible" $ do
+    let start = fromSyntaxForest example1 & next
+    (kindAt <$> start) `shouldBe` Just DeclType
+
+  it "Can traverse across a tree" $ do
+    let start = fromSyntaxForest example1 & pipe [down, down]
+    (kindAt <$> start) `shouldBe` Just DeclType
+    let end = start &? pipe (replicate 4 next)
+    (kindAt <$> end) `shouldBe` Just (Token T.LCurly)
+    let start' = end &? pipe (replicate 4 previous)
+    (kindAt <$> start') `shouldBe` (kindAt <$> start)
+
+  it "Keeps track of offsets correctly" $ do
+    let start = fromSyntaxForest example1 & next -- Get to first leaf token
+    (offsetAt <$> start) `shouldBe` Just 0
+    -- Check offset at {
+    (kindAt <$> (start &? pipe (replicate 4 next))) `shouldBe` Just (Token T.LCurly)
+    (offsetAt <$> (start &? pipe (replicate 4 next))) `shouldBe` Just 9
+    -- Check offset after moving around
+    (offsetAt <$> (start &? pipe [right, left])) `shouldBe` Just 0
+    -- Check offset after moving up and down
+    let dictStart = start &? pipe (replicate 4 right)
+    (offsetAt <$> (dictStart &? pipe [down, up])) `shouldBe` (offsetAt <$> dictStart)

waspls/waspls.cabal

@@ -42,6 +42,8 @@ library
   import:           common-all
   exposed-modules:
     Wasp.LSP.Server
+    Control.Monad.Loops
+    Control.Syntax.Traverse
     Wasp.Backend.Token
     Wasp.Backend.Lexer
     Wasp.Backend.TokenSet
@@ -51,8 +53,14 @@ library
     Wasp.Backend.ConcreteParser.Internal
   other-modules:
     Paths_waspls
-    Wasp.LSP.Core
+    Wasp.LSP.ServerState
+    Wasp.LSP.ServerConfig
+    Wasp.LSP.ServerM
     Wasp.LSP.Handlers
+    Wasp.LSP.Diagnostic
+    Wasp.LSP.Completion
+    Wasp.LSP.Util
+    Wasp.LSP.Syntax
     Wasp.Backend.Lexer.Internal
     Wasp.Backend.Lexer.Lexer
   hs-source-dirs:
@@ -98,6 +106,8 @@ test-suite waspls-test
     , tasty-quickcheck      ^>= 0.10
   other-modules:
     TestUtil
+    Control.Syntax.TraverseTest
+    Control.Monad.LoopsTest
     Wasp.Backend.LexerTest
     Wasp.Backend.ConcreteParserTest
 

waspls/watch

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+# Rebuilds project when a source file changes
+
+inotifywait -m -e close_write -r src/ exe/ waspls.cabal |
+  while read directory action file; do
+    printf "\033[33mFile changed. Starting recompile...\033[0m\n"
+    cabal build
+    printf "\033[1;32mFinished\033[0m\n"
+  done