From ecaa004dccec38b5ab5899f2f39ee51e933275b1 Mon Sep 17 00:00:00 2001
From: larsk21 <57503246+larsk21@users.noreply.github.com>
Date: Fri, 14 Jan 2022 09:20:09 +0100
Subject: [PATCH] feat: make withWaitFindSnap consider all snaps of a document

---
 src/Lean/Server/FileWorker/Utils.lean | 7 +++++++
 src/Lean/Server/Requests.lean         | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 0a9b9363d4..a9918d6724 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -58,6 +58,13 @@ structure EditableDocument where
   cancelTk   : CancelToken
   deriving Inhabited
 
+namespace EditableDocument
+
+def allSnaps (doc : EditableDocument) : AsyncList ElabTaskError Snapshot :=
+  AsyncList.cons doc.headerSnap doc.cmdSnaps
+
+end EditableDocument
+
 structure RpcSession where
   /-- Objects that are being kept alive for the RPC client, together with their type names,
   mapped to by their RPC reference.
diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean
index fe77e4ee0e..2008ad59ea 100644
--- a/src/Lean/Server/Requests.lean
+++ b/src/Lean/Server/Requests.lean
@@ -96,7 +96,7 @@ def withWaitFindSnap (doc : EditableDocument) (p : Snapshot → Bool)
   (notFoundX : RequestM β)
   (x : Snapshot → RequestM β)
     : RequestM (RequestTask β) := do
-  let findTask ← doc.cmdSnaps.waitFind? p
+  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