From 96770b4d83d00736c8442c21dbdf4ae80f574bdb Mon Sep 17 00:00:00 2001
From: Wojciech Nawrocki <wjnawrocki@protonmail.com>
Date: Fri, 25 Mar 2022 17:10:41 -0400
Subject: [PATCH] refactor: remove some code duplication

---
 src/Lean/Server/Requests.lean            | 31 ++++++++++++++++--------
 src/Lean/Server/Rpc/RequestHandling.lean | 15 ++++--------
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean
index bf3809c55a..e4588cf6c5 100644
--- a/src/Lean/Server/Requests.lean
+++ b/src/Lean/Server/Requests.lean
@@ -89,6 +89,18 @@ def mapTask (t : Task α) (f : α → RequestM β) : RequestM (RequestTask β) :
 def bindTask (t : Task α) (f : α → RequestM (RequestTask β)) : RequestM (RequestTask β) := fun rc => do
   EIO.bindTask t (f · rc)
 
+def waitFindSnapAux (notFoundX : RequestM α) (x : Snapshot → RequestM α)
+    : Except ElabTaskError (Option Snapshot) → RequestM α
+  /- The elaboration task that we're waiting for may be aborted if the file contents change.
+  In that case, we reply with the `fileChanged` error. Thanks to this, the server doesn't
+  get bogged down in requests for an old state of the document. -/
+  | Except.error FileWorker.ElabTaskError.aborted =>
+    throwThe RequestError RequestError.fileChanged
+  | Except.error (FileWorker.ElabTaskError.ioError e) =>
+    throw (e : RequestError)
+  | Except.ok none => notFoundX
+  | Except.ok (some snap) => x snap
+
 /-- Create a task which waits for the first snapshot matching `p`, handles various errors,
 and if a matching snapshot was found executes `x` with it. If not found, the task executes
 `notFoundX`. -/
@@ -97,16 +109,15 @@ def withWaitFindSnap (doc : EditableDocument) (p : Snapshot → Bool)
   (x : Snapshot → RequestM β)
     : RequestM (RequestTask β) := do
   let findTask ← doc.allSnaps.waitFind? p
-  mapTask findTask fun
-    /- The elaboration task that we're waiting for may be aborted if the file contents change.
-    In that case, we reply with the `fileChanged` error. Thanks to this, the server doesn't
-    get bogged down in requests for an old state of the document. -/
-    | Except.error FileWorker.ElabTaskError.aborted =>
-      throwThe RequestError RequestError.fileChanged
-    | Except.error (FileWorker.ElabTaskError.ioError e) =>
-      throw (e : RequestError)
-    | Except.ok none => notFoundX
-    | Except.ok (some snap) => x snap
+  mapTask findTask <| waitFindSnapAux notFoundX x
+
+/-- See `withWaitFindSnap`. -/
+def bindWaitFindSnap (doc : EditableDocument) (p : Snapshot → Bool)
+  (notFoundX : RequestM (RequestTask β))
+  (x : Snapshot → RequestM (RequestTask β))
+    : RequestM (RequestTask β) := do
+  let findTask ← doc.allSnaps.waitFind? p
+  bindTask findTask <| waitFindSnapAux notFoundX x
 
 end RequestM
 
diff --git a/src/Lean/Server/Rpc/RequestHandling.lean b/src/Lean/Server/Rpc/RequestHandling.lean
index d604b95105..90832f7085 100644
--- a/src/Lean/Server/Rpc/RequestHandling.lean
+++ b/src/Lean/Server/Rpc/RequestHandling.lean
@@ -28,16 +28,11 @@ private def handleRpcCall (p : Lsp.RpcCallParams) : RequestM (RequestTask Json)
   let doc ← readDoc
   let text := doc.meta.text
   let callPos := text.lspPosToUtf8Pos p.position
-  let findTask ← doc.allSnaps.waitFind? (fun s => s.endPos >= callPos)
-  bindTask findTask fun
-    | Except.error FileWorker.ElabTaskError.aborted =>
-      throwThe RequestError RequestError.fileChanged
-    | Except.error (FileWorker.ElabTaskError.ioError e) =>
-      throw (e : RequestError)
-    | Except.ok none =>
-      throwThe RequestError { code := JsonRpc.ErrorCode.invalidParams
-                              message := s!"Incorrect position '{p.toTextDocumentPositionParams}' in RPC call" }
-    | Except.ok (some snap) => do
+  bindWaitFindSnap doc (fun s => s.endPos >= callPos)
+    (notFoundX := throwThe RequestError
+      { code := JsonRpc.ErrorCode.invalidParams
+        message := s!"Incorrect position '{p.toTextDocumentPositionParams}' in RPC call" })
+    fun snap => do
       if let some proc := (← builtinRpcProcedures.get).find? p.method then
         proc.wrapper p.sessionId p.params
       else if let some proc := userRpcProcedures.getState snap.env |>.find? p.method then