From b714e087d67b6ebfde93281db224327218d021a5 Mon Sep 17 00:00:00 2001
From: "E.W.Ayers" <edward.ayers@outlook.com>
Date: Wed, 20 Jul 2022 14:54:01 +0100
Subject: [PATCH] fix: widgetSourceRegistry now stores the UserWidgetDefinition
 declaration name instead of WidgetSource

This means that the environment extension is not storing a big text object and instead the text
is retrieved from the declaration itself.
---
 src/Lean/Server/Requests.lean   | 19 +++++++++++++++++++
 src/Lean/Widget/UserWidget.lean | 32 ++++++++++++++++++--------------
 2 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean
index a3c8656e41..76f3ea7ab0 100644
--- a/src/Lean/Server/Requests.lean
+++ b/src/Lean/Server/Requests.lean
@@ -134,6 +134,25 @@ def withWaitFindSnapAtPos
     (notFoundX := throw <| RequestError.mk .invalidParams s!"no snapshot found at {lspPos}")
     f
 
+open Lean Elab Command in
+/-- Use the command state in the given snapshot to run a `CommandElabM`.-/
+def runCommand (snap : Snapshots.Snapshot) (c : CommandElabM α) : RequestM α := do
+  let r ← read
+  let ctx : Command.Context := {
+    fileName := r.doc.meta.uri,
+    fileMap := r.doc.meta.text,
+    tacticCache? := snap.tacticCache,
+  }
+  let ea ← c.run ctx |>.run snap.cmdState |> EIO.toIO'
+  match ea with
+  | Except.ok (a, _s) => return a
+  | Except.error ex => do
+    throw <| RequestError.internalError <|← ex.toMessageData.toString
+
+/-- Run a `CoreM` using the data in the given snapshot.-/
+def runCore (snap : Snapshots.Snapshot) (c : CoreM α) : RequestM α :=
+  runCommand snap <| Lean.Elab.Command.liftCoreM c
+
 end RequestM
 
 /-! # The global request handlers table -/
diff --git a/src/Lean/Widget/UserWidget.lean b/src/Lean/Widget/UserWidget.lean
index 8205d03e05..882869b55a 100644
--- a/src/Lean/Widget/UserWidget.lean
+++ b/src/Lean/Widget/UserWidget.lean
@@ -53,8 +53,8 @@ structure UserWidget where
   deriving Inhabited, ToJson, FromJson
 
 private abbrev WidgetSourceRegistry := SimplePersistentEnvExtension
-    (UInt64 × WidgetSource)
-    (Std.RBMap UInt64 WidgetSource compare)
+    (UInt64 × Name)
+    (Std.RBMap UInt64 Name compare)
 
 -- Mapping widgetSourceId to hash of sourcetext
 builtin_initialize userWidgetRegistry : MapDeclarationExtension UserWidget ← mkMapDeclarationExtension `widgetRegistry
@@ -66,21 +66,26 @@ builtin_initialize widgetSourceRegistry : WidgetSourceRegistry ←
     toArrayFn     := fun es => es.toArray
   }
 
-private unsafe def attributeImplUnsafe : AttributeImpl where
+private unsafe def getUserWidgetDefinitionUnsafe
+  (decl : Name) : CoreM UserWidgetDefinition :=
+  evalConstCheck UserWidgetDefinition ``UserWidgetDefinition decl
+
+@[implementedBy getUserWidgetDefinitionUnsafe]
+private opaque getUserWidgetDefinition
+  (decl : Name) : CoreM UserWidgetDefinition
+
+private def attributeImpl : AttributeImpl where
   name := `widget
   descr := "Mark a string as static code that can be loaded by a widget handler."
   applicationTime := AttributeApplicationTime.afterCompilation
   add decl _stx _kind := do
     let env ← getEnv
-    let defn ← evalConstCheck UserWidgetDefinition ``UserWidgetDefinition decl
+    let defn ← getUserWidgetDefinition decl
     let javascriptHash := hash defn.javascript
     let env := userWidgetRegistry.insert env decl {id := decl, name := defn.name, javascriptHash}
-    let env := widgetSourceRegistry.addEntry env (javascriptHash, {sourcetext := defn.javascript})
+    let env := widgetSourceRegistry.addEntry env (javascriptHash, decl)
     setEnv <| env
 
-@[implementedBy attributeImplUnsafe]
-opaque attributeImpl : AttributeImpl
-
 builtin_initialize registerBuiltinAttribute attributeImpl
 
 /-- Input for `getWidgetSource` RPC. -/
@@ -90,15 +95,14 @@ structure GetWidgetSourceParams where
   pos : Lean.Lsp.TextDocumentPositionParams
   deriving ToJson, FromJson
 
-open Lean.Server Lean
-
-open RequestM in
+open Lean.Server Lean RequestM in
 @[serverRpcMethod]
 def getWidgetSource (args : GetWidgetSourceParams) : RequestM (RequestTask WidgetSource) :=
   RequestM.withWaitFindSnapAtPos args.pos fun snap => do
     let env := snap.cmdState.env
-    if let some w := widgetSourceRegistry.getState env |>.find? args.hash then
-      return w
+    if let some id := widgetSourceRegistry.getState env |>.find? args.hash then
+      let d ← Lean.Server.RequestM.runCore snap <| getUserWidgetDefinition id
+      return {sourcetext := d.javascript}
     else
       throw <| RequestError.mk .invalidParams s!"No registered user-widget with hash {args.hash}"
 
@@ -133,7 +137,7 @@ structure GetWidgetsResponse where
   widgets : Array UserWidgetInstance
   deriving ToJson, FromJson
 
-open RequestM in
+open Lean Server RequestM in
 /-- Get the `UserWidget`s present at a particular position. -/
 @[serverRpcMethod]
 def getWidgets (args : Lean.Lsp.TextDocumentPositionParams) : RequestM (RequestTask (GetWidgetsResponse)) := do