Wingman: streaming tactic solutions (#2102)

* WIP abstract LSP, take the pain out of writing LSP stuff

* Finish making commands

* Separate code lenses and actions

* Pull out types

* Finalize the abstract API

* Bug fix in JSON; first connected abstract handler

* Add ContinuationResult for better control over how edits work

* Remove IO from TacticProviders; use LspEnv instead

* installInteractions

* Pull TacticCodeActions into their own file

* Misc cleanup

* Haddock

* Fix bug in codelens

* Port EmptyCase to Interaction

* Rename makeTacticCodeAction -> makeTacticInteraction

* Support for partial timeouts in upcoming refinery v5

* asum instead of choice for assumption

* Don't count it as using a term if you only destruct it

* Let interactions return multiple results --- aka also info messages

* Update refinery lower bounds

* Revert "Update refinery lower bounds"

This reverts commit 53199b3.

* Pull refinery from the future

* Fix tests

* Add -XNumDecimals

* Fix AutoTypeLevel test

* Continue to emit errors

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

plugins/hls-tactics-plugin/hls-tactics-plugin.cabal

@@ -27,6 +27,7 @@ library
   hs-source-dirs:     src
   exposed-modules:
     Ide.Plugin.Tactic
+    Refinery.Future
     Wingman.AbstractLSP
     Wingman.AbstractLSP.TacticActions
     Wingman.AbstractLSP.Types
@@ -93,6 +94,7 @@ library
     , refinery              ^>=0.4
     , retrie                >=0.1.1.0
     , syb
+    , unagi-chan
     , text
     , transformers
     , unordered-containers

plugins/hls-tactics-plugin/src/Refinery/Future.hs

@@ -0,0 +1,140 @@
+{-# LANGUAGE RankNTypes #-}
+
+------------------------------------------------------------------------------
+-- | Things that belong in the future release of refinery v5.
+module Refinery.Future
+  ( runStreamingTacticT
+  , hoistListT
+  , consume
+  ) where
+
+import Control.Applicative
+import Control.Monad (ap, (>=>))
+import Control.Monad.State.Lazy (runStateT)
+import Control.Monad.Trans
+import Data.Either (isRight)
+import Data.Functor ((<&>))
+import Data.Tuple (swap)
+import Refinery.ProofState
+import Refinery.Tactic.Internal
+
+
+
+hoistElem :: Functor m => (forall x. m x -> n x) -> Elem m a -> Elem n a
+hoistElem _ Done = Done
+hoistElem f (Next a lt) = Next a $ hoistListT f lt
+
+
+hoistListT :: Functor m => (forall x. m x -> n x) -> ListT m a -> ListT n a
+hoistListT f t = ListT $ f $ fmap (hoistElem f) $ unListT t
+
+
+consume :: Monad m => ListT m a -> (a -> m ()) -> m ()
+consume lt f = unListT lt >>= \case
+  Done -> pure ()
+  Next a lt' -> f a >> consume lt' f
+
+
+newHole :: MonadExtract meta ext err s m => s -> m (s, (meta, ext))
+newHole = fmap swap . runStateT hole
+
+runStreamingTacticT :: (MonadExtract meta ext err s m) => TacticT jdg ext err s m () -> jdg -> s -> ListT m (Either err (Proof s meta jdg ext))
+runStreamingTacticT t j s = streamProofs s $ fmap snd $ proofState t j
+
+data Elem m a
+    = Done
+    | Next a (ListT m a)
+    deriving stock Functor
+
+
+point :: Applicative m => a -> Elem m a
+point a = Next a $ ListT $ pure Done
+
+newtype ListT m a = ListT { unListT :: m (Elem m a) }
+
+cons :: (Applicative m) => a -> ListT m a -> ListT m a
+cons x xs = ListT $ pure $ Next x xs
+
+instance Functor m => Functor (ListT m) where
+    fmap f (ListT xs) = ListT $ xs <&> \case
+        Done -> Done
+        Next a xs -> Next (f a) (fmap f xs)
+
+instance (Monad m) => Applicative (ListT m) where
+    pure = return
+    (<*>) = ap
+
+instance (Monad m) => Alternative (ListT m) where
+    empty = ListT $ pure Done
+    (ListT xs) <|> (ListT ys) =
+        ListT $ xs >>= \case
+          Done -> ys
+          Next x xs -> pure (Next x (xs <|> ListT ys))
+
+instance (Monad m) => Monad (ListT m) where
+    return a = cons a empty
+    (ListT xs) >>= k =
+        ListT $ xs >>= \case
+          Done -> pure Done
+          Next x xs -> unListT $ k x <|> (xs >>= k)
+
+
+instance MonadTrans ListT where
+    lift m = ListT $ fmap (\x -> Next x empty) m
+
+
+interleaveT :: (Monad m) => Elem m a -> Elem m a -> Elem m a
+interleaveT xs ys =
+    case xs of
+      Done -> ys
+      Next x xs -> Next x $ ListT $ fmap (interleaveT ys) $ unListT xs
+
+--         ys <&> \case
+--           Done -> Next x xs
+--           Next y ys -> Next x (cons y (interleaveT xs ys))
+
+force :: (Monad m) => Elem m a -> m [a]
+force = \case
+    Done -> pure []
+    Next x xs' -> (x:) <$> (unListT xs' >>= force)
+
+ofList :: Monad m => [a] -> Elem m a
+ofList [] = Done
+ofList (x:xs) = Next x $ ListT $ pure $ ofList xs
+
+streamProofs :: forall ext err s m goal meta. (MonadExtract meta ext err s m) => s -> ProofStateT ext ext err s m goal -> ListT m (Either err (Proof s meta goal ext))
+streamProofs s p = ListT $ go s [] pure p
+    where
+      go :: s -> [(meta, goal)] -> (err -> m err) -> ProofStateT ext ext err s m goal -> m (Elem m (Either err (Proof s meta goal ext)))
+      go s goals _ (Subgoal goal k) = do
+         (s', (meta, h)) <- newHole s
+         -- Note [Handler Reset]:
+         -- We reset the handler stack to avoid the handlers leaking across subgoals.
+         -- This would happen when we had a handler that wasn't followed by an error call.
+         --     pair >> goal >>= \g -> (handler_ $ \_ -> traceM $ "Handling " <> show g) <|> failure "Error"
+         -- We would see the "Handling a" message when solving for b.
+         (go s' (goals ++ [(meta, goal)]) pure $ k h)
+      go s goals handlers (Effect m) = m >>= go s goals handlers
+      go s goals handlers (Stateful f) =
+          let (s', p) = f s
+          in go s' goals handlers p
+      go s goals handlers (Alt p1 p2) =
+          unListT $ ListT (go s goals handlers p1) <|> ListT (go s goals handlers p2)
+      go s goals handlers (Interleave p1 p2) =
+          interleaveT <$> (go s goals handlers p1) <*> (go s goals handlers p2)
+      go s goals handlers (Commit p1 p2) = do
+          solns <- force =<< go s goals handlers p1
+          if (any isRight solns) then pure $ ofList solns else go s goals handlers p2
+      go _ _ _ Empty = pure Done
+      go _ _ handlers (Failure err _) = do
+          annErr <- handlers err
+          pure $ point $ Left annErr
+      go s goals handlers (Handle p h) =
+          -- Note [Handler ordering]:
+          -- If we have multiple handlers in scope, then we want the handlers closer to the error site to
+          -- run /first/. This allows the handlers up the stack to add their annotations on top of the
+          -- ones lower down, which is the behavior that we desire.
+          -- IE: for @handler f >> handler g >> failure err@, @g@ ought to be run before @f@.
+          go s goals (h >=> handlers) p
+      go s goals _ (Axiom ext) = pure $ point $ Right (Proof ext s goals)
+

plugins/hls-tactics-plugin/src/Wingman/AbstractLSP.hs

@@ -13,8 +13,11 @@ import           Control.Monad.IO.Class
 import           Control.Monad.Trans (lift)
 import           Control.Monad.Trans.Maybe (MaybeT, mapMaybeT)
 import qualified Data.Aeson as A
+import           Data.Coerce
 import           Data.Foldable (traverse_)
+import           Data.Monoid (Last (..))
 import qualified Data.Text as T
+import           Data.Traversable (for)
 import           Data.Tuple.Extra (uncurry3)
 import           Development.IDE (IdeState)
 import           Development.IDE.Core.UseStale
@@ -93,26 +96,30 @@ runContinuation plId cont state (fc, b) = do
       env@LspEnv{..} <- buildEnv state plId fc
       let stale a = runStaleIde "runContinuation" state (fc_nfp le_fileContext) a
       args <- fetchTargetArgs @a env
-      c_runCommand cont env args fc b >>= \case
-        ErrorMessages errs -> do
-          traverse_ showUserFacingMessage errs
-          pure $ Right A.Null
-        RawEdit edits -> do
-          sendEdits edits
-          pure $ Right A.Null
-        GraftEdit gr -> do
-          ccs <- lift getClientCapabilities
-          TrackedStale pm _ <- mapMaybeT liftIO $ stale GetAnnotatedParsedSource
-          case mkWorkspaceEdits le_dflags ccs (fc_uri le_fileContext) (unTrack pm) gr of
-            Left errs ->
-              pure $ Left $ ResponseError
-                { _code    = InternalError
-                , _message = T.pack $ show errs
-                , _xdata   = Nothing
-                }
-            Right edits -> do
-              sendEdits edits
-              pure $ Right A.Null
+      res <- c_runCommand cont env args fc b
+
+      -- This block returns a maybe error.
+      fmap (maybe (Right $ A.Null) Left . coerce . foldMap Last) $
+        for res $ \case
+          ErrorMessages errs -> do
+            traverse_ showUserFacingMessage errs
+            pure Nothing
+          RawEdit edits -> do
+            sendEdits edits
+            pure Nothing
+          GraftEdit gr -> do
+            ccs <- lift getClientCapabilities
+            TrackedStale pm _ <- mapMaybeT liftIO $ stale GetAnnotatedParsedSource
+            case mkWorkspaceEdits le_dflags ccs (fc_uri le_fileContext) (unTrack pm) gr of
+              Left errs ->
+                pure $ Just $ ResponseError
+                  { _code    = InternalError
+                  , _message = T.pack $ show errs
+                  , _xdata   = Nothing
+                  }
+              Right edits -> do
+                sendEdits edits
+                pure $ Nothing
 
 
 ------------------------------------------------------------------------------

plugins/hls-tactics-plugin/src/Wingman/AbstractLSP/TacticActions.hs

@@ -52,23 +52,38 @@ makeTacticInteraction cmd =
         pm_span <- liftMaybe $ mapAgeFrom pmmap span
         let t = commandTactic cmd var_name
 
-        res <- liftIO $ timeout (cfg_timeout_seconds le_config * seconds) $ do
-          runTactic hj_ctx hj_jdg t >>= \case
-            Left err -> pure $ ErrorMessages $ pure $ mkUserFacingMessage err
-            Right rtr ->
-              case rtr_extract rtr of
-                L _ (HsVar _ (L _ rdr)) | isHole (occName rdr) ->
-                  pure $ ErrorMessages [NothingToDo]
-                _ -> do
-                  for_ (rtr_other_solns rtr) $ \soln -> do
-                    traceMX "other solution" $ syn_val soln
-                    traceMX "with score" $ scoreSolution soln (rtr_jdg rtr) []
-                  traceMX "solution" $ rtr_extract rtr
-                  pure $ GraftEdit $ graftHole (RealSrcSpan $ unTrack pm_span) rtr
-
-        pure $ case res of
-          Nothing -> ErrorMessages $ pure TimedOut
-          Just c  -> c
+        liftIO $ runTactic (cfg_timeout_seconds le_config * seconds) hj_ctx hj_jdg t >>= \case
+          Left err ->
+            pure
+              $ pure
+              $ ErrorMessages
+              $ pure
+              $ mkUserFacingMessage err
+          Right rtr ->
+            case rtr_extract rtr of
+              L _ (HsVar _ (L _ rdr)) | isHole (occName rdr) ->
+                pure
+                  $ addTimeoutMessage rtr
+                  $ pure
+                  $ ErrorMessages
+                  $ pure NothingToDo
+              _ -> do
+                for_ (rtr_other_solns rtr) $ \soln -> do
+                  traceMX "other solution" $ syn_val soln
+                  traceMX "with score" $ scoreSolution soln (rtr_jdg rtr) []
+                traceMX "solution" $ rtr_extract rtr
+                pure
+                  $ addTimeoutMessage rtr
+                  $ pure
+                  $ GraftEdit
+                  $ graftHole (RealSrcSpan $ unTrack pm_span) rtr
+
+
+addTimeoutMessage :: RunTacticResults -> [ContinuationResult] -> [ContinuationResult]
+addTimeoutMessage rtr = mappend
+  [ ErrorMessages $ pure TimedOut
+  | rtr_timed_out rtr
+  ]
 
 
 ------------------------------------------------------------------------------
@@ -82,6 +97,7 @@ seconds = 1e6
 mkUserFacingMessage :: [TacticError] -> UserFacingMessage
 mkUserFacingMessage errs
   | elem OutOfGas errs = NotEnoughGas
+mkUserFacingMessage [] = NothingToDo
 mkUserFacingMessage _ = TacticErrors
 
 

plugins/hls-tactics-plugin/src/Wingman/AbstractLSP/Types.hs

@@ -113,7 +113,7 @@ data Continuation sort target payload = Continuation
         -> TargetArgs target
         -> FileContext
         -> payload
-        -> MaybeT (LspM Plugin.Config) ContinuationResult
+        -> MaybeT (LspM Plugin.Config) [ContinuationResult]
   }
 
 

plugins/hls-tactics-plugin/src/Wingman/EmptyCase.hs

@@ -84,7 +84,7 @@ emptyCaseInteraction = Interaction $
           , edits
           )
     )
-  $ (\ _ _ _ we -> pure $ RawEdit we)
+  $ (\ _ _ _ we -> pure $ pure $ RawEdit we)
 
 
 scrutinzedType :: EmptyCaseSort Type -> Maybe Type