From 11cce61e4df3b8cd35de0d4d9197a36d602b885e Mon Sep 17 00:00:00 2001
From: Jonathan Coates <git@squiddev.cc>
Date: Sat, 12 Feb 2022 13:09:15 +0000
Subject: [PATCH] chore: Clean up LSP folding a little

 - Wait for all terms to be elaborated before showing folding regions.
   May want to change this to support partial results.
 - Use .span to extract import statements, rather than mutually
   recursive functions.
 - Some tiny other bits of cleanup
---
 .../Server/FileWorker/RequestHandling.lean    | 39 ++++++-------------
 tests/lean/interactive/foldingRange.lean      |  6 +++
 .../foldingRange.lean.expected.out            |  3 +-
 3 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean
index bc9d4e3936..b77be05bd6 100644
--- a/src/Lean/Server/FileWorker/RequestHandling.lean
+++ b/src/Lean/Server/FileWorker/RequestHandling.lean
@@ -415,18 +415,9 @@ def handleSemanticTokensRange (p : SemanticTokensRangeParams)
 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 t ← doc.allSnaps.waitAll
+  mapTask t fun (snaps, _) => do
+    let stxs := snaps.map (·.stx)
     let (_, ranges) ← StateT.run (addRanges doc.meta.text [] stxs) #[]
     return ranges
   where
@@ -449,19 +440,14 @@ partial def handleFoldingRange (p : FoldingRangeParams)
         addRanges text sections stxs
       | _ => do
         if isImport stx then
-          addImportRanges text sections stx.getPos? stx.getTailPos? stxs
+          let (imports, stxs) := stxs.span isImport
+          let last := imports.getLastD stx
+          addRange text FoldingRangeKind.imports stx.getPos? last.getTailPos?
+          addRanges text sections 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 =>
@@ -475,16 +461,15 @@ partial def handleFoldingRange (p : FoldingRangeParams)
             addRangeFromSyntax text FoldingRangeKind.comment comment
 
           addRangeFromSyntax text FoldingRangeKind.region decl
-          return
-
-        addRangeFromSyntax text FoldingRangeKind.region stx
+        else
+          addRangeFromSyntax text FoldingRangeKind.region stx
       | _ =>
-        addRangeFromSyntax text FoldingRangeKind.comment stx
+        addRangeFromSyntax text FoldingRangeKind.region 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
+    addRange (text : FileMap) kind start? stop? := do
+      if let (some startP, some endP) := (start?, stop?) then
         let startP := text.utf8PosToLspPos startP
         let endP := text.utf8PosToLspPos endP
         if startP.line != endP.line then
diff --git a/tests/lean/interactive/foldingRange.lean b/tests/lean/interactive/foldingRange.lean
index 0b58d7a175..8ace1baca7 100644
--- a/tests/lean/interactive/foldingRange.lean
+++ b/tests/lean/interactive/foldingRange.lean
@@ -39,3 +39,9 @@ end
 
 end Bar
 end Foo
+
+#check #[
+  1,
+  2,
+  3
+]
diff --git a/tests/lean/interactive/foldingRange.lean.expected.out b/tests/lean/interactive/foldingRange.lean.expected.out
index 7ddd42cbe9..e33ef5c371 100644
--- a/tests/lean/interactive/foldingRange.lean.expected.out
+++ b/tests/lean/interactive/foldingRange.lean.expected.out
@@ -10,4 +10,5 @@
  {"startLine": 33, "kind": "region", "endLine": 34},
  {"startLine": 35, "kind": "region", "endLine": 36},
  {"startLine": 11, "kind": "region", "endLine": 39},
- {"startLine": 6, "kind": "region", "endLine": 40}]
+ {"startLine": 6, "kind": "region", "endLine": 40},
+ {"startLine": 42, "kind": "region", "endLine": 46}]