[AD-458] Use single StateT with lenses

knit/src/Knit/Autocompletion.hs

@@ -28,92 +28,17 @@ import Knit.Procedure
 import Knit.Syntax
 import Knit.Tokenizer
 
-suggestions
-  :: forall components proxy.
-     ( KnownSpine components
-     , AllConstrained (ComponentTokenizer components) components
-     , AllConstrained (ComponentTokenToLit components) components
-     , AllConstrained (ComponentCommandProcs components) components
-     , AllConstrained ComponentPrinter components
-     )
-  => proxy components
-  -> Cursor
-  -> T.Text
-  -> [(Cursor, T.Text)]
-suggestions _ cursor cmd =
-  let
-    (spaceBefore, tokens) = tokenize cmd
-
-    tokenBalance = \case
-      TokenParenthesis bs -> withBracketSide 1 (-1) bs
-      _ -> 0
-    parensBalance = sum $ map (tokenBalance . _lItem . _twsToken) tokens
-    cmdLines = T.splitOn "\n" cmd
-    endLoc = (length cmdLines, T.length (last cmdLines) + 1)
-
-    loc' a b = loc (fromIntegral a) (fromIntegral b)
-    closingBracket i t = TokenWithSpace
-      { _twsToken = Located
-        { _lSpan = spanFromTo
-            (loc' (fst endLoc) (snd endLoc + i))
-            (loc' (fst endLoc) (snd endLoc + i + 1))
-        , _lItem = t
-        }
-      , _twsSpaceAfter = def
-      }
-
-    tokens' = tokens ++ zipWith closingBracket [0..]
-      (replicate parensBalance (TokenParenthesis BracketSideClosing))
-  in
-    case fullParses (pExpr @components) tokens' of
-      ([], _) -> [(cursor, cmd)]
-      (tree:_, _) ->
-        let
-          (formattedExpr, spaceAfter) = parseTreeToFormattedExpr (Selection $ Just cursor) tree
-          makeSws = maybe def $ \space ->
-            SpaceWithSelection
-              (Space $ toList $ getSkipped $ _lItem space)
-              (selectionInSpan (getCursor cursor) (_lSpan space))
-          makeRightSws x =
-            if cmd == ""
-              then SpaceWithSelection def (Selection $ Just def)
-              else makeSws x
-          processSuggestion ((suggestion, LeftSpace swsBefore), RightSpace swsAfter) =
-            let
-              (suggestionStr, selection) = ppFormattedExpr swsBefore suggestion swsAfter
-            in
-              ( maybe
-                  (error "Core invariant violated: no cursor after autocompletion")
-                  id
-                $ getSelection selection
-              , T.dropEnd parensBalance suggestionStr
-              )
-        in
-          map processSuggestion $
-            runStateT
-              (runStateT
-                (suggestionExprs commandProcs formattedExpr)
-                (LeftSpace $ makeSws spaceBefore))
-              (RightSpace $ makeRightSws spaceAfter)
-
-newtype LeftSpace = LeftSpace { getLeftSpace :: SpaceWithSelection }
-  deriving (Default)
-
-newtype RightSpace = RightSpace { getRightSpace :: SpaceWithSelection }
-  deriving (Default)
-
-type SuggestionMonad = StateT LeftSpace (StateT RightSpace [])
-
-liftLeft :: SuggestionMonad a -> SuggestionMonad a
-liftLeft = id
+data SuggestionCtx = SuggestionCtx
+  { _leftSpace :: SpaceWithSelection
+  , _rightSpace :: SpaceWithSelection
+  }
+makeLenses ''SuggestionCtx
 
-liftRight :: StateT RightSpace [] a -> SuggestionMonad a
-liftRight = lift
+instance Default SuggestionCtx where
+  def = SuggestionCtx def def
 
-liftList :: [a] -> SuggestionMonad a
-liftList = lift . lift
+type SuggestionMonad = StateT SuggestionCtx []
 
-{-# ANN module ("HLint: ignore Reduce duplication" :: T.Text) #-}
 -- | Returns completion suggestions assuming that the expression was completed
 -- with missing parentheses before parsing.
 suggestionExprs
@@ -130,17 +55,14 @@ suggestionExprs procs = goExpr
       lit@(ExprLit _ _) -> pure lit
       ExprProcCall NoExt p -> ExprProcCall NoExt <$> goPc p
       XExpr (ExprInBrackets l e r) -> do
-        oldLeft <- liftLeft get
-        oldRight <- liftRight get
-        liftLeft $ put $ LeftSpace $ fbSpace l
-        liftRight $ put $ RightSpace $ fbSpace r
+        oldCtx <- get
+        leftSpace .= fbSpace l
+        rightSpace .= fbSpace r
           <> SpaceWithSelection def (selectionAtTheStart $ fbBracketSelection r)
         e' <- goExpr e
-        l' <- liftLeft $ gets $ \(LeftSpace lSpace) -> l { fbSpace = lSpace }
-        r' <- liftRight $ gets $ \(RightSpace rSpace) ->
-          FormattedBracket (selectionAtTheEnd rSpace) rSpace
-        liftLeft $ put oldLeft
-        liftRight $ put oldRight
+        l' <- uses leftSpace $ \lSpace -> l { fbSpace = lSpace }
+        r' <- uses rightSpace $ \rSpace -> FormattedBracket (selectionAtTheEnd rSpace) rSpace
+        put oldCtx
         pure $ XExpr $ ExprInBrackets l' e' r'
 
     goPc
@@ -149,38 +71,36 @@ suggestionExprs procs = goExpr
     goPc pc@(ProcCall pcSelection cmd args) =
       case cmd of
         CommandIdOperator OpUnit -> do
-          leftSpace <- liftLeft $ gets getLeftSpace
-          rightSpace <- liftRight $ gets getRightSpace
-          let space = leftSpace <> rightSpace
+          lSpace <- use leftSpace
+          rSpace <- use rightSpace
+          let space = lSpace <> rSpace
           case getSelection $ swsSelection space of
             Nothing -> pure pc
             Just c -> do
-              let (leftSpace', rightSpace') = splitAtCursor c $ getSpace $ swsSpace space
-              liftLeft $ put $ LeftSpace $ SpaceWithSelection (Space leftSpace') def
-              liftRight $ put $ RightSpace
-                $ SpaceWithSelection (Space rightSpace') (Selection $ Just def)
-              liftList $ pc : map (toProcCall . fst) suggestableProcs
+              let (lSpace', rSpace') = splitAtCursor c $ getSpace $ swsSpace space
+              leftSpace .= SpaceWithSelection (Space lSpace') def
+              rightSpace .= SpaceWithSelection (Space rSpace') (Selection $ Just def)
+              lift $ pc : map (toProcCall . fst) suggestableProcs
         CommandIdOperator OpAndThen ->
           case args of
             [ArgPos (ArgPosSpace lRightSpace) lhs, ArgPos (ArgPosSpace rLeftSpace) rhs] -> do
               -- lLeftSpace lhs lRightSpace ; rLeftSpace rhs rRightSpace
 
-              rRightSpace <- liftRight get
-              liftRight $ put $ RightSpace $ lRightSpace
-                <> SpaceWithSelection def (selectionAtTheStart pcSelection)
+              rRightSpace <- use rightSpace
+              rightSpace .= lRightSpace <> SpaceWithSelection def (selectionAtTheStart pcSelection)
               lhs' <- goExpr lhs
-              lRightSpace' <- liftRight get
-              liftRight $ put rRightSpace
+              lRightSpace' <- use rightSpace
+              rightSpace .= rRightSpace
 
-              lLeftSpace <- liftLeft get
-              liftLeft $ put $ LeftSpace rLeftSpace
+              lLeftSpace <- use leftSpace
+              leftSpace .= rLeftSpace
               rhs' <- goExpr rhs
-              rLeftSpace' <- liftLeft get
-              liftLeft $ put lLeftSpace
+              rLeftSpace' <- use leftSpace
+              leftSpace .= lLeftSpace
 
-              pure $ ProcCall (selectionAtTheEnd $ getRightSpace lRightSpace') cmd
-                [ ArgPos (ArgPosSpace $ getRightSpace lRightSpace') lhs'
-                , ArgPos (ArgPosSpace $ getLeftSpace rLeftSpace') rhs'
+              pure $ ProcCall (selectionAtTheEnd lRightSpace') cmd
+                [ ArgPos (ArgPosSpace lRightSpace') lhs'
+                , ArgPos (ArgPosSpace rLeftSpace') rhs'
                 ]
             _ -> invalidOperatorApplication
         CommandIdName name -> do
@@ -202,52 +122,49 @@ suggestionExprs procs = goExpr
 
 
       ProcCall pcSelection pcName (ArgPos space e : argsRest) -> do
-        oldRight <- liftRight get
-        liftRight $ put $ RightSpace $ getArgPosSpace space
+        oldCtx <- get
+
+        rightSpace .= getArgPosSpace space
         (ProcCall pcSelection' pcName' argsRest') <-
           goPcName False pc (ProcCall pcSelection pcName argsRest)
-        space' <- liftRight $ gets getRightSpace
+        space' <- use rightSpace
 
         let
           suggestKeyword
             :: Expr FormattedExprExt CommandId components
             -> SuggestionMonad (Arg' FormattedExprExt CommandId components)
           suggestKeyword (ExprProcCall NoExt (ProcCall aPcSelection (CommandIdName aPcName) []))
             | Just c <- getSelection aPcSelection =
-              liftList $ ArgPos (ArgPosSpace space') e : do
+              lift $ ArgPos (ArgPosSpace space') e : do
                 let (toComplete, _) = splitAtCursor c $ nameStr aPcName
                 pcParam <- procParams pcName
                 guard $ isPrefixOf toComplete $ nameStr pcParam
                 let selection = Selection $ Just $ cursorAfter $ nameStr pcParam ++ ":"
                 pure $ ArgKw (ArgKwSpace space' selection def) pcParam $
                   ExprProcCall NoExt (ProcCall def (CommandIdOperator OpUnit) [])
-          suggestKeyword _ = liftList []
+          suggestKeyword _ = lift []
 
-        oldLeft <- liftLeft get
-        liftLeft $ put def
-        liftRight $ put def
+        put def
         arg' <- suggestKeyword e <|> ArgPos (ArgPosSpace space') <$> goExpr e
 
-        liftLeft $ put oldLeft
-        liftRight $ put oldRight
+        put oldCtx
 
         pure $ ProcCall pcSelection' pcName' (arg' : argsRest')
 
 
       ProcCall pcSelection pcName (ArgKw space keyword e : argsRest) -> do
-        oldRight <- liftRight get
-        liftRight $ put $ RightSpace $ aksPrefix space
+        oldCtx <- get
+
+        rightSpace .= aksPrefix space
         (ProcCall pcSelection' pcName' argsRest') <-
           goPcName False pc (ProcCall pcSelection pcName argsRest)
-        space' <- liftRight $ gets getRightSpace
+        space' <- use rightSpace
 
-        oldLeft <- liftLeft get
-        liftLeft $ put def
-        liftRight $ put def
+        leftSpace .= def
+        rightSpace .= def
         e' <- goExpr e
 
-        liftLeft $ put oldLeft
-        liftRight $ put oldRight
+        put oldCtx
 
         pure $ ProcCall pcSelection' pcName'
           (ArgKw space { aksPrefix = space' } keyword e' : argsRest')
@@ -261,7 +178,7 @@ suggestionExprs procs = goExpr
               pcNameStr = nameStr pcName
               (toComplete, _) = splitAtCursor c pcNameStr
             in
-              liftList $ (pure $ ProcCall pcSelection (CommandIdName pcName) []) <|> do
+              lift $ (pure $ ProcCall pcSelection (CommandIdName pcName) []) <|> do
                 (pcName', _) <- suggestableProcs
                 guard $ isPrefixOf toComplete $ nameStr pcName'
                 guard $ pcName' /= pcName || c /= cursorAfter pcNameStr
@@ -273,20 +190,19 @@ suggestionExprs procs = goExpr
       -> ProcCall' FormattedExprExt CommandId components
       -> SuggestionMonad (ProcCall' FormattedExprExt CommandId components)
     goPcRightSpace firstCall (ProcCall _ pcName _) pc@(ProcCall pcSelection pcCmd pcArgs) = do
-      RightSpace right <- liftRight get
+      right <- use rightSpace
       case getSelection $ swsSelection right of
         Nothing -> pure pc
         Just c ->
           if c == def || not firstCall && c == cursorAfterSpace (swsSpace right)
             then pure pc
             else pure pc <|> do
-              let (leftSpace', rightSpace') = splitAtCursor c $ getSpace $ swsSpace right
-              liftRight $ put $ RightSpace
-                $ SpaceWithSelection (Space rightSpace') (Selection $ Just def)
-              pcParam <- liftList $ procParams pcName
+              let (lSpace', rSpace') = splitAtCursor c $ getSpace $ swsSpace right
+              rightSpace .= SpaceWithSelection (Space rSpace') (Selection $ Just def)
+              pcParam <- lift $ procParams pcName
               let
                 arg' =
-                  ArgKw (ArgKwSpace (SpaceWithSelection (Space leftSpace') def) def def) pcParam $
+                  ArgKw (ArgKwSpace (SpaceWithSelection (Space lSpace') def) def def) pcParam $
                     ExprProcCall NoExt (ProcCall def (CommandIdOperator OpUnit) [])
               pure $ ProcCall pcSelection pcCmd $ arg' : pcArgs
 
@@ -324,3 +240,69 @@ suggestionExprs procs = goExpr
       []
 
     nameStr = unpack . toLazyText . build
+
+suggestions
+  :: forall components proxy.
+     ( KnownSpine components
+     , AllConstrained (ComponentTokenizer components) components
+     , AllConstrained (ComponentTokenToLit components) components
+     , AllConstrained (ComponentCommandProcs components) components
+     , AllConstrained ComponentPrinter components
+     )
+  => proxy components
+  -> Cursor
+  -> T.Text
+  -> [(Cursor, T.Text)]
+suggestions _ cursor cmd =
+  let
+    (spaceBefore, tokens) = tokenize cmd
+
+    tokenBalance = \case
+      TokenParenthesis bs -> withBracketSide 1 (-1) bs
+      _ -> 0
+    parensBalance = sum $ map (tokenBalance . _lItem . _twsToken) tokens
+    cmdLines = T.splitOn "\n" cmd
+    endLoc = (length cmdLines, T.length (last cmdLines) + 1)
+
+    loc' a b = loc (fromIntegral a) (fromIntegral b)
+    closingBracket i t = TokenWithSpace
+      { _twsToken = Located
+        { _lSpan = spanFromTo
+            (loc' (fst endLoc) (snd endLoc + i))
+            (loc' (fst endLoc) (snd endLoc + i + 1))
+        , _lItem = t
+        }
+      , _twsSpaceAfter = def
+      }
+
+    tokens' = tokens ++ zipWith closingBracket [0..]
+      (replicate parensBalance (TokenParenthesis BracketSideClosing))
+  in
+    case fullParses (pExpr @components) tokens' of
+      ([], _) -> [(cursor, cmd)]
+      (tree:_, _) ->
+        let
+          (formattedExpr, spaceAfter) = parseTreeToFormattedExpr (Selection $ Just cursor) tree
+          makeSws = maybe def $ \space ->
+            SpaceWithSelection
+              (Space $ toList $ getSkipped $ _lItem space)
+              (selectionInSpan (getCursor cursor) (_lSpan space))
+          makeRightSws x =
+            if cmd == ""
+              then SpaceWithSelection def (Selection $ Just def)
+              else makeSws x
+          processSuggestion (suggestion, SuggestionCtx swsBefore swsAfter) =
+            let
+              (suggestionStr, selection) = ppFormattedExpr swsBefore suggestion swsAfter
+            in
+              ( maybe
+                  (error "Core invariant violated: no cursor after autocompletion")
+                  id
+                $ getSelection selection
+              , T.dropEnd parensBalance suggestionStr
+              )
+        in
+          map processSuggestion $
+            runStateT
+              (suggestionExprs commandProcs formattedExpr)
+              (SuggestionCtx (makeSws spaceBefore) (makeRightSws spaceAfter))