Map code action edits back to the notebook

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased changes
 
 * Support textDocument/inlayHint and textDocument/onTypeFormatting
+* Map code action edits back to the notebook, instead of returning the shadow file's uri and lines
 
 ## 0.2.2.0
 

src/Transform/Common.hs

@@ -4,11 +4,13 @@
 module Transform.Common where
 
 import Control.Lens hiding (List)
+import qualified Data.Map as M
 import Data.Maybe
 import Language.LSP.Notebook
 import Language.LSP.Protocol.Lens as Lens
 import Language.LSP.Protocol.Types
 import Language.LSP.Transformer
+import Transform.Util
 
 
 -- * Transform
@@ -45,6 +47,43 @@ untransformInlayHint tx hint =
     & traverseOf position (untransformPosition (getParams tx) tx)
     & fmap (over (textEdits . _Just) (mapMaybe (untransformRanged tx)))
 
+untransformWorkspaceEdit :: DocumentState -> WorkspaceEdit -> Maybe WorkspaceEdit
+untransformWorkspaceEdit (DocumentState {transformer=tx, origUri=notebookUri, newUri=shadowUri}) wsEdit = do
+  changes' <- traverse untransformChanges (wsEdit ^. changes)
+  documentChanges' <- traverse (traverse untransformDocumentChange) (wsEdit ^. documentChanges)
+  pure $ wsEdit
+    & set changes changes'
+    & set documentChanges documentChanges'
+  where
+    untransformChanges m = M.fromList <$> traverse untransformEntry (M.toList m)
+
+    untransformEntry (u, es)
+      | u == shadowUri = (,) notebookUri <$> traverse (untransformRanged tx) es
+      | otherwise = Nothing
+
+    untransformDocumentChange (InL textDocumentEdit) = InL <$> untransformTextDocumentEdit textDocumentEdit
+    untransformDocumentChange (InR _createRenameOrDelete) = Nothing
+
+    untransformTextDocumentEdit tde
+      | tde ^. textDocument . uri /= shadowUri = Nothing
+      | otherwise = tde
+          & set (textDocument . uri) notebookUri
+          & set (textDocument . version) (InR Null)
+          & traverseOf (edits . traverse) untransformEdit
+
+    untransformEdit (InL textEdit) = InL <$> untransformRanged tx textEdit
+    untransformEdit (InR annotatedTextEdit) = InR <$> untransformRanged tx annotatedTextEdit
+
+untransformCodeAction :: DocumentState -> CodeAction -> Maybe CodeAction
+untransformCodeAction ds codeAction = case codeAction ^. edit of
+  Nothing -> Just codeAction
+  Just wsEdit -> (\wsEdit' -> codeAction & set edit (Just wsEdit')) <$> untransformWorkspaceEdit ds wsEdit
+
+workspaceEditUris :: WorkspaceEdit -> [Uri]
+workspaceEditUris wsEdit =
+  M.keys (fromMaybe mempty (wsEdit ^. changes))
+  <> [tde ^. textDocument . uri | InL tde <- fromMaybe [] (wsEdit ^. documentChanges)]
+
 -- * Orphan (wish this was in lsp-types)
 
 instance HasRange Range Range where

src/Transform/ServerRsp.hs

@@ -105,11 +105,25 @@ transformServerRsp' SMethod_TextDocumentDocumentSymbol initialParams result =
 
 transformServerRsp' SMethod_TextDocumentCodeAction _initialParams (InR null) = return $ InR null
 transformServerRsp' SMethod_TextDocumentCodeAction initialParams (InL result) = (InL <$>) $
-  whenAnythingByInitialParams initialParams result $ withTransformer result $ \(DocumentState {}) -> do
-    filterM (fmap not . isInternalReferringCodeAction) result
+  whenAnythingByInitialParams initialParams result $ withTransformer result $ \ds ->
+    return $ mapMaybe (untransformCommandOrCodeAction ds) result
   where
-    isInternalReferringCodeAction (InL _command) = pure False
-    isInternalReferringCodeAction (InR _codeAction) = pure False
+    untransformCommandOrCodeAction _ds c@(InL _command) = Just c
+    untransformCommandOrCodeAction ds (InR codeAction) = InR <$> untransformCodeAction ds codeAction
+
+transformServerRsp' SMethod_CodeActionResolve _initialParams codeAction = case codeAction ^. edit of
+  Nothing -> return codeAction
+  Just wsEdit -> case workspaceEditUris wsEdit of
+    [] -> return codeAction
+    (shadowUri:_) -> lookupTransformerByNewUri shadowUri >>= \case
+      Nothing -> do
+        logWarnN [i|Resolved code action named #{shadowUri}, which is no document of ours; dropping its edit|]
+        return $ codeAction & set edit Nothing
+      Just ds -> case untransformCodeAction ds codeAction of
+        Just codeAction' -> return codeAction'
+        Nothing -> do
+          logWarnN [i|Resolved code action '#{codeAction ^. title}' doesn't map back into the notebook; dropping its edit|]
+          return $ codeAction & set edit Nothing
 
 transformServerRsp' SMethod_TextDocumentCodeLens _initialParams (InR null) = return $ InR null
 transformServerRsp' SMethod_TextDocumentCodeLens initialParams (InL result) = (InL <$>) $

src/Transform/Util.hs

@@ -128,6 +128,11 @@ lookupTransformer uri = do
   TransformerState {..} <- ask
   M.lookup (getUri uri) <$> readMVar transformerDocuments
 
+lookupTransformerByNewUri :: TransformerMonad m => Uri -> m (Maybe DocumentState)
+lookupTransformerByNewUri uri = do
+  TransformerState {..} <- ask
+  L.find (\DocumentState {newUri=u} -> u == uri) . M.elems <$> readMVar transformerDocuments
+
 withTransformer :: (TransformerMonad n) => a -> (DocumentState -> n a) -> Uri -> n a
 withTransformer def cb uri = do
   lookupTransformer uri >>= \case