Fix Folding Ranges

unison-cli/src/Unison/Codebase/Editor/HandleInput.hs

@@ -67,6 +67,7 @@ import Unison.Codebase.Editor.HandleInput.Branch (handleBranch)
 import Unison.Codebase.Editor.HandleInput.BranchRename (handleBranchRename)
 import Unison.Codebase.Editor.HandleInput.Branches (handleBranches)
 import Unison.Codebase.Editor.HandleInput.DebugDefinition qualified as DebugDefinition
+import Unison.Codebase.Editor.HandleInput.DebugFoldRanges qualified as DebugFoldRanges
 import Unison.Codebase.Editor.HandleInput.DeleteBranch (handleDeleteBranch)
 import Unison.Codebase.Editor.HandleInput.DeleteProject (handleDeleteProject)
 import Unison.Codebase.Editor.HandleInput.EditNamespace (handleEditNamespace)
@@ -1143,6 +1144,8 @@ loop e = do
               for_ (Relation.toList . Branch.deepTerms $ rootBranch0) \(r, name) ->
                 traceM $ show name ++ ",Term," ++ Text.unpack (Referent.toText r)
             DebugTermI isVerbose hqName -> DebugDefinition.debugTerm isVerbose hqName
+            DebugLSPFoldRangesI -> do
+              DebugFoldRanges.debugFoldRanges
             DebugTypeI hqName -> DebugDefinition.debugDecl hqName
             DebugClearWatchI {} ->
               Cli.runTransaction Codebase.clearWatches
@@ -1363,6 +1366,7 @@ inputDescription input =
         then pure ("debug.term.verbose " <> HQ.toText hqName)
         else pure ("debug.term " <> HQ.toText hqName)
     DebugTypeI hqName -> pure ("debug.type " <> HQ.toText hqName)
+    DebugLSPFoldRangesI -> pure "debug.lsp.fold-ranges"
     DebugNameDiffI {} -> wat
     DebugNumberedArgsI {} -> wat
     DebugTabCompletionI _input -> wat

unison-cli/src/Unison/Codebase/Editor/HandleInput/DebugFoldRanges.hs

@@ -0,0 +1,61 @@
+module Unison.Codebase.Editor.HandleInput.DebugFoldRanges (debugFoldRanges) where
+
+import Control.Lens
+import Control.Monad.Reader
+import Data.Text qualified as Text
+import Language.LSP.Protocol.Lens
+import Language.LSP.Protocol.Types qualified as LSP
+import Unison.Cli.Monad (Cli)
+import Unison.Cli.Monad qualified as Cli
+import Unison.Cli.MonadUtils qualified as Cli
+import Unison.Codebase.Editor.HandleInput.FormatFile (TextReplacement (..))
+import Unison.Codebase.Editor.HandleInput.FormatFile qualified as FormatFile
+import Unison.Codebase.Editor.Output
+import Unison.LSP.Conversions qualified as CV
+import Unison.LSP.FoldingRange (foldingRangesForFile)
+import Unison.Prelude
+import Unison.Util.Range qualified as U
+
+debugFoldRanges :: Cli ()
+debugFoldRanges = do
+  Cli.Env {loadSource} <- ask
+  (filePath, _) <- Cli.expectLatestFile
+  parsedFile <- Cli.expectLatestParsedFile
+  let foldingRanges =
+        foldingRangesForFile parsedFile
+          & fmap
+            ( \fr ->
+                LSP.Range
+                  (LSP.Position (fr ^. startLine) (fromMaybe 0 $ fr ^. startCharacter))
+                  ( case (fr ^. endCharacter) of
+                      Just c -> LSP.Position (fr ^. endLine) c
+                      -- If there's no end char specified, go all the way to the beginning of the next line
+                      Nothing -> LSP.Position ((fr ^. endLine) + 1) 0
+                  )
+            )
+  sourceTxt <-
+    liftIO (loadSource (Text.pack filePath)) >>= \case
+      Cli.InvalidSourceNameError -> Cli.returnEarly $ InvalidSourceName filePath
+      Cli.LoadError -> Cli.returnEarly $ SourceLoadFailed filePath
+      Cli.LoadSuccess contents -> pure contents
+  Cli.respond $ AnnotatedFoldRanges $ annotateRanges sourceTxt foldingRanges
+
+-- | Annotate the bounds of a range within text using 《 and 》.
+--
+-- Useful for checking that computed ranges make sense against the source text.
+--
+-- >>> annotateRanges "one\ntwo\nthree\nfour" [ LSP.Range (LSP.Position 1 0) (LSP.Position 2 3) ]
+-- "one\n\12298two\nthr\12299ee\nfour"
+annotateRanges :: Text -> [LSP.Range] -> Text
+annotateRanges txt ranges =
+  let replacements =
+        ranges
+          & foldMap
+            ( \(LSP.Range start end) ->
+                let startPos = CV.lspToUPos start
+                    endPos = CV.lspToUPos end
+                 in [ TextReplacement "《" (U.Range startPos startPos),
+                      TextReplacement "》" (U.Range endPos endPos)
+                    ]
+            )
+   in FormatFile.applyTextReplacements replacements txt

unison-cli/src/Unison/Codebase/Editor/Input.hs

@@ -222,6 +222,7 @@ data Input
   | DebugDumpNamespaceSimpleI
   | DebugTermI (Bool {- Verbose mode -}) (HQ.HashQualified Name)
   | DebugTypeI (HQ.HashQualified Name)
+  | DebugLSPFoldRangesI
   | DebugClearWatchI
   | DebugDoctorI
   | DebugNameDiffI ShortCausalHash ShortCausalHash

unison-cli/src/Unison/Codebase/Editor/Output.hs

@@ -327,6 +327,7 @@ data Output
   | DebugFuzzyOptionsNoResolver
   | DebugTerm (Bool {- verbose mode -}) (Either (Text {- A builtin hash -}) (Term Symbol Ann))
   | DebugDecl (Either (Text {- A builtin hash -}) (DD.Decl Symbol Ann)) (Maybe ConstructorId {- If 'Just' we're debugging a constructor of the given decl -})
+  | AnnotatedFoldRanges Text
   | ClearScreen
   | PulledEmptyBranch (ReadRemoteNamespace Share.RemoteProjectBranch)
   | CreatedProject Bool {- randomly-generated name? -} ProjectName
@@ -574,6 +575,7 @@ isFailure o = case o of
   DebugFuzzyOptionsNoResolver {} -> True
   DebugTerm {} -> False
   DebugDecl {} -> False
+  AnnotatedFoldRanges {} -> False
   DisplayDebugNameDiff {} -> False
   ClearScreen -> False
   PulledEmptyBranch {} -> False

unison-cli/src/Unison/CommandLine/InputPatterns.hs

@@ -2314,6 +2314,16 @@ debugType =
         _ -> Left (I.help debugType)
     )
 
+debugLSPFoldRanges :: InputPattern
+debugLSPFoldRanges =
+  InputPattern
+    "debug.lsp.fold-ranges"
+    []
+    I.Hidden
+    []
+    "Output the source from the most recently parsed file, but annotated with the computed fold ranges."
+    (const $ Right Input.DebugLSPFoldRangesI)
+
 debugClearWatchCache :: InputPattern
 debugClearWatchCache =
   InputPattern
@@ -3043,6 +3053,7 @@ validInputs =
       debugTerm,
       debugTermVerbose,
       debugType,
+      debugLSPFoldRanges,
       debugFileHashes,
       debugNameDiff,
       debugNumberedArgs,

unison-cli/src/Unison/CommandLine/OutputMessages.hs

@@ -1790,6 +1790,7 @@ notifyUser dir = \case
         <> case typ of
           Left builtinTxt -> "Builtin type: ##" <> P.text builtinTxt
           Right decl -> either (P.text . TL.toStrict . pShowNoColor) (P.text . TL.toStrict . pShowNoColor) decl
+  AnnotatedFoldRanges txt -> pure $ P.text txt
   DisplayDebugNameDiff NameChanges {termNameAdds, termNameRemovals, typeNameAdds, typeNameRemovals} -> do
     let referentText =
           -- We don't use the constructor type in the actual output here, so there's no
@@ -1953,15 +1954,15 @@ notifyUser dir = \case
       prettyProjectAndBranchName projectAndBranch <> "does not exist on" <> prettyURI host
   RemoteProjectBranchDoesntExist'Push host projectAndBranch ->
     let push = P.group . P.backticked . IP.patternName $ IP.push
-    in pure . P.wrap $
-      "The previous push target named"
-        <> prettyProjectAndBranchName projectAndBranch
-        <> "has been deleted from"
-        <> P.group (prettyURI host <> ".")
-        <> "I've deleted the invalid push target."
-        <> "Run the"
-        <> push
-        <> "command again to push to a new target."
+     in pure . P.wrap $
+          "The previous push target named"
+            <> prettyProjectAndBranchName projectAndBranch
+            <> "has been deleted from"
+            <> P.group (prettyURI host <> ".")
+            <> "I've deleted the invalid push target."
+            <> "Run the"
+            <> push
+            <> "command again to push to a new target."
   RemoteProjectBranchHeadMismatch host projectAndBranch ->
     pure . P.wrap $
       prettyProjectAndBranchName projectAndBranch

unison-cli/src/Unison/LSP/FoldingRange.hs

@@ -1,50 +1,71 @@
 {-# LANGUAGE DataKinds #-}
 
-module Unison.LSP.FoldingRange where
+module Unison.LSP.FoldingRange
+  ( foldingRangeRequest,
+    foldingRangesForFile,
+  )
+where
 
 import Control.Lens hiding (List)
+import Data.Map qualified as Map
+import Data.Text qualified as Text
 import Language.LSP.Protocol.Lens hiding (id, to)
 import Language.LSP.Protocol.Message qualified as Msg
 import Language.LSP.Protocol.Types
-import Unison.ABT qualified as ABT
 import Unison.DataDeclaration qualified as DD
-import Unison.Debug qualified as Debug
 import Unison.LSP.Conversions (annToRange)
 import Unison.LSP.FileAnalysis (getFileAnalysis)
 import Unison.LSP.Types
+import Unison.Parser.Ann (Ann)
 import Unison.Prelude
+import Unison.Symbol (Symbol)
 import Unison.UnisonFile (UnisonFile (..))
+import Unison.UnisonFile qualified as UF
+import Unison.Var qualified as Var
 
 foldingRangeRequest :: Msg.TRequestMessage 'Msg.Method_TextDocumentFoldingRange -> (Either Msg.ResponseError (Msg.MessageResult 'Msg.Method_TextDocumentFoldingRange) -> Lsp ()) -> Lsp ()
 foldingRangeRequest m respond = do
-  foldRanges <- foldingRangesForFile (m ^. params . textDocument . uri)
-  Debug.debugM Debug.LSP "Folding Ranges" foldRanges
+  let fileUri = m ^. params . textDocument . uri
+  foldRanges <-
+    fromMaybe [] <$> runMaybeT do
+      FileAnalysis {parsedFile = mayParsedFile} <- getFileAnalysis fileUri
+      parsedFile <- hoistMaybe mayParsedFile
+      pure $ foldingRangesForFile parsedFile
   respond . Right . InL $ foldRanges
 
 -- | Return a folding range for each top-level definition
-foldingRangesForFile :: Uri -> Lsp [FoldingRange]
-foldingRangesForFile fileUri =
-  fromMaybe []
-    <$> runMaybeT do
-      FileAnalysis {parsedFile} <- getFileAnalysis fileUri
-      UnisonFileId {dataDeclarationsId, effectDeclarationsId, terms} <- MaybeT $ pure parsedFile
-      let dataFolds = dataDeclarationsId ^.. folded . _2 . to dataDeclSpan
-      let abilityFolds = effectDeclarationsId ^.. folded . _2 . to DD.toDataDecl . to dataDeclSpan
-      let termFolds = terms ^.. folded . _3 . to ABT.annotation
-      let folds = dataFolds <> abilityFolds <> termFolds
-      let ranges = mapMaybe annToRange folds
-      pure $
-        ranges <&> \r ->
-          FoldingRange
-            { _startLine = r ^. start . line,
-              _startCharacter = Just (r ^. start . character),
-              _endLine = r ^. end . line,
-              _endCharacter = Just (r ^. end . character),
-              _kind = Just FoldingRangeKind_Region,
-              _collapsedText = Nothing
-            }
-  where
-    dataDeclSpan dd =
-      -- We don't have a proper Annotation for data decls so we take the span of all the
-      -- constructors using their monoid instance.
-      DD.annotation dd <> DD.constructors' dd ^. folded . to (\(a, _v, typ) -> a <> ABT.annotation typ)
+foldingRangesForFile :: UF.UnisonFile Symbol Ann -> [FoldingRange]
+foldingRangesForFile UnisonFileId {dataDeclarationsId, effectDeclarationsId, terms, watches} =
+  let dataFolds =
+        dataDeclarationsId
+          & Map.toList
+          & map \(sym, (_typ, decl)) -> (Just sym, DD.annotation decl)
+      abilityFolds =
+        effectDeclarationsId
+          & Map.toList
+          & map \(sym, (_typ, decl)) -> (Just sym, DD.annotation . DD.toDataDecl $ decl)
+      termFolds = terms & fmap \(sym, ann, _trm) -> (Just sym, ann)
+      watchFolds =
+        watches
+          & fold
+          & fmap
+            ( \(_sym, ann, _trm) ->
+                -- We don't use the symbol here because watch symbols are often auto-generated
+                -- and ugly.
+                (Nothing, ann)
+            )
+      folds =
+        dataFolds <> abilityFolds <> termFolds <> watchFolds
+      ranges =
+        folds
+          & mapMaybe \(sym, range) ->
+            (Text.pack . Var.nameStr <$> sym,) <$> annToRange range
+   in ranges <&> \(maySym, r) ->
+        FoldingRange
+          { _startLine = r ^. start . line,
+            _startCharacter = Just (r ^. start . character),
+            _endLine = r ^. end . line,
+            _endCharacter = Just (r ^. end . character),
+            _kind = Just FoldingRangeKind_Region,
+            _collapsedText = maySym
+          }

unison-cli/unison-cli.cabal

@@ -51,6 +51,7 @@ library
       Unison.Codebase.Editor.HandleInput.Branches
       Unison.Codebase.Editor.HandleInput.BranchRename
       Unison.Codebase.Editor.HandleInput.DebugDefinition
+      Unison.Codebase.Editor.HandleInput.DebugFoldRanges
       Unison.Codebase.Editor.HandleInput.DeleteBranch
       Unison.Codebase.Editor.HandleInput.DeleteProject
       Unison.Codebase.Editor.HandleInput.EditNamespace

unison-src/transcripts/lsp-fold-ranges.md

@@ -0,0 +1,33 @@
+```ucm:hide
+.> builtins.mergeio
+```
+
+```unison:hide
+
+{{ Type doc }}
+structural type Optional a =
+  None
+  | Some a
+
+{{
+  Multi line
+
+  Term doc
+}}
+List.map : 
+  (a -> b) 
+  -> [a] 
+  -> [b]
+List.map f = cases
+  (x +: xs) -> f x +: List.map f xs
+  [] -> []
+
+test> z = let
+  x = "hello"
+  y = "world"
+  [Ok (x ++ y)]
+```
+
+```ucm
+.> debug.lsp.fold-ranges
+```

unison-src/transcripts/lsp-fold-ranges.output.md

@@ -0,0 +1,52 @@
+```unison
+{{ Type doc }}
+structural type Optional a =
+  None
+  | Some a
+
+{{
+  Multi line
+
+  Term doc
+}}
+List.map : 
+  (a -> b) 
+  -> [a] 
+  -> [b]
+List.map f = cases
+  (x +: xs) -> f x +: List.map f xs
+  [] -> []
+
+test> z = let
+  x = "hello"
+  y = "world"
+  [Ok (x ++ y)]
+```
+
+```ucm
+.> debug.lsp.fold-ranges
+
+  《{{ Type doc }}》
+  《structural type Optional a =
+    None
+    | Some a》
+
+  《{{
+    Multi line
+
+    Term doc
+  }}》
+  《List.map : 
+    (a -> b) 
+    -> [a] 
+    -> [b]
+  List.map f = cases
+    (x +: xs) -> f x +: List.map f xs
+    [] -> []》
+
+  《test> z = let
+    x = "hello"
+    y = "world"
+    [Ok (x ++ y)]》
+
+```