feat: LSP code folding support

The following constructs are foldable:
 - Sections and namespaces
 - Blocks of import/open statements
 - Multi-line commands (so mostly definitions)
 - Mutual definitions
 - Module-level doc comments
 - Top-level definition doc comments

Fixes #1012

src/Lean/Data/Lsp/Capabilities.lean

@@ -45,6 +45,7 @@ structure ServerCapabilities where
   typeDefinitionProvider : Bool := false
   referencesProvider : Bool := false
   workspaceSymbolProvider : Bool := false
+  foldingRangeProvider : Bool := false
   semanticTokensProvider? : Option SemanticTokensOptions := none
   deriving ToJson, FromJson
 

src/Lean/Data/Lsp/LanguageFeatures.lean

@@ -301,5 +301,26 @@ structure SemanticTokens where
   data : Array Nat
   deriving FromJson, ToJson
 
+structure FoldingRangeParams where
+  textDocument : TextDocumentIdentifier
+  deriving FromJson, ToJson
+
+inductive FoldingRangeKind where
+  | comment
+  | imports
+  | region
+
+instance : ToJson FoldingRangeKind where
+  toJson
+    | FoldingRangeKind.comment => "comment"
+    | FoldingRangeKind.imports => "imports"
+    | FoldingRangeKind.region => "region"
+
+structure FoldingRange where
+  startLine : Nat
+  endLine : Nat
+  kind? : Option FoldingRangeKind := none
+  deriving ToJson
+
 end Lsp
 end Lean

src/Lean/Server/FileSource.lean

@@ -72,6 +72,9 @@ instance : FileSource SemanticTokensParams :=
 instance : FileSource SemanticTokensRangeParams :=
   ⟨fun p => fileSource p.textDocument⟩
 
+instance : FileSource FoldingRangeParams :=
+  ⟨fun p => fileSource p.textDocument⟩
+
 instance : FileSource PlainGoalParams :=
   ⟨fun p => fileSource p.textDocument⟩
 

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -412,6 +412,87 @@ def handleSemanticTokensRange (p : SemanticTokensRangeParams)
   let endPos := text.lspPosToUtf8Pos p.range.end
   handleSemanticTokens beginPos endPos
 
+partial def handleFoldingRange (p : FoldingRangeParams)
+  : RequestM (RequestTask (Array FoldingRange)) := do
+  let doc ← readDoc
+  asTask do
+    let ⟨snaps, e?⟩ ← doc.allSnaps.updateFinishedPrefix
+    let mut stxs := snaps.finishedPrefix.map (·.stx)
+    match e? with
+    | some ElabTaskError.aborted =>
+      throw RequestError.fileChanged
+    | some (ElabTaskError.ioError e) =>
+      throw (e : RequestError)
+    | _ => pure ()
+
+    let lastSnap := snaps.finishedPrefix.getLastD doc.headerSnap
+    stxs := stxs ++ (← parseAhead doc.meta.text.source lastSnap).toList
+    let (_, ranges) ← StateT.run (addRanges doc.meta.text [] stxs) #[]
+    return ranges
+  where
+    isImport stx := stx.isOfKind ``Lean.Parser.Module.header || stx.isOfKind ``Lean.Parser.Command.open
+
+    addRanges (text : FileMap) sections
+    | [] => return
+    | stx::stxs => match stx with
+      | `(namespace $id)  => addRanges text (stx.getPos?::sections) stxs
+      | `(section $(id)?) => addRanges text (stx.getPos?::sections) stxs
+      | `(end $(id)?) => do
+        if let start::rest := sections then
+          addRange text FoldingRangeKind.region start stx.getTailPos?
+          addRanges text rest stxs
+        else
+          addRanges text sections stxs
+      | `(mutual $body* end) => do
+        addRangeFromSyntax text FoldingRangeKind.region stx
+        addRanges text [] body.toList
+        addRanges text sections stxs
+      | _ => do
+        if isImport stx then
+          addImportRanges text sections stx.getPos? stx.getTailPos? stxs
+        else
+          addCommandRange text stx
+          addRanges text sections stxs
+
+    addImportRanges (text : FileMap) sections startP endP stxs := do
+      if let stx::stxs := stxs then
+        if isImport stx then
+          return ← addImportRanges text sections (startP <|> stx.getPos?) (stx.getTailPos? <|> endP) stxs
+
+      addRange text FoldingRangeKind.imports startP endP
+      addRanges text sections stxs
+
+    addCommandRange text stx :=
+      match stx.getKind with
+      | `Lean.Parser.Command.moduleDoc =>
+        addRangeFromSyntax text FoldingRangeKind.comment stx
+      | ``Lean.Parser.Command.declaration => do
+        -- When visiting a declaration, attempt to fold the doc comment
+        -- separately to the main definition.
+        -- We never fold other modifiers, such as annotations.
+        if let `($dm:declModifiers $decl) := stx then
+          if let some comment := dm[0].getOptional? then
+            addRangeFromSyntax text FoldingRangeKind.comment comment
+
+          addRangeFromSyntax text FoldingRangeKind.region decl
+          return
+
+        addRangeFromSyntax text FoldingRangeKind.region stx
+      | _ =>
+        addRangeFromSyntax text FoldingRangeKind.comment stx
+
+    addRangeFromSyntax (text : FileMap) kind stx := addRange text kind stx.getPos? stx.getTailPos?
+
+    addRange (text : FileMap) kind start stop := do
+      if let (Option.some startP, Option.some endP) := (start, stop) then
+        let startP := text.utf8PosToLspPos startP
+        let endP := text.utf8PosToLspPos endP
+        if startP.line != endP.line then
+          modify fun st => st.push
+            { startLine := startP.line
+              endLine := endP.line
+              kind? := some kind }
+
 partial def handleWaitForDiagnostics (p : WaitForDiagnosticsParams)
     : RequestM (RequestTask WaitForDiagnostics) := do
   let rec waitLoop : RequestM EditableDocument := do
@@ -438,6 +519,7 @@ builtin_initialize
   registerLspRequestHandler "textDocument/documentSymbol"       DocumentSymbolParams       DocumentSymbolResult    handleDocumentSymbol
   registerLspRequestHandler "textDocument/semanticTokens/full"  SemanticTokensParams       SemanticTokens          handleSemanticTokensFull
   registerLspRequestHandler "textDocument/semanticTokens/range" SemanticTokensRangeParams  SemanticTokens          handleSemanticTokensRange
+  registerLspRequestHandler "textDocument/foldingRange"         FoldingRangeParams         (Array FoldingRange)    handleFoldingRange
   registerLspRequestHandler "$/lean/plainGoal"                  PlainGoalParams            (Option PlainGoal)      handlePlainGoal
   registerLspRequestHandler "$/lean/plainTermGoal"              PlainTermGoalParams        (Option PlainTermGoal)  handlePlainTermGoal
 

src/Lean/Server/Watchdog.lean

@@ -692,6 +692,7 @@ def mkLeanServerCapabilities : ServerCapabilities := {
   workspaceSymbolProvider := true
   documentHighlightProvider := true
   documentSymbolProvider := true
+  foldingRangeProvider := true
   semanticTokensProvider? := some {
     legend := {
       tokenTypes     := SemanticTokenType.names

tests/lean/interactive/foldingRange.lean

@@ -0,0 +1,41 @@
+--^ textDocument/foldingRange
+import Lean
+import Lean.Data
+
+open Lean
+
+namespace Foo
+
+open Std
+open Lean
+
+section Bar
+
+/-!
+  A module-level doc comment
+-/
+
+/--
+  Some documentation comment
+-/
+@[inline]
+def add (x y : Nat) :=
+  x + y
+
+inductive InductiveTy
+| a
+|
+  /--
+    Another doc comment. This one is not folded.
+  -/
+  b
+
+mutual
+  def a :=
+    1
+  def b :=
+    a
+end
+
+end Bar
+end Foo

tests/lean/interactive/foldingRange.lean.expected.out

@@ -0,0 +1,13 @@
+{"textDocument": {"uri": "file://foldingRange.lean"},
+ "position": {"line": 0, "character": 2}}
+[{"startLine": 1, "kind": "imports", "endLine": 4},
+ {"startLine": 8, "kind": "imports", "endLine": 9},
+ {"startLine": 13, "kind": "comment", "endLine": 15},
+ {"startLine": 17, "kind": "comment", "endLine": 19},
+ {"startLine": 21, "kind": "region", "endLine": 22},
+ {"startLine": 24, "kind": "region", "endLine": 30},
+ {"startLine": 32, "kind": "region", "endLine": 37},
+ {"startLine": 33, "kind": "region", "endLine": 34},
+ {"startLine": 35, "kind": "region", "endLine": 36},
+ {"startLine": 11, "kind": "region", "endLine": 39},
+ {"startLine": 6, "kind": "region", "endLine": 40}]