feat: reusable rpc refs

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -346,7 +346,7 @@ def getInteractiveTermGoal (p : Lsp.PlainTermGoalParams)
     let goal ← ci.runMetaM lctx' do
       Widget.goalToInteractive (← Meta.mkFreshExprMVar ty).mvarId!
     let range := if let some r := i.range? then r.toLspRange text else ⟨p.position, p.position⟩
-    return some { goal with range, term := ⟨ti⟩ }
+    return some { goal with range, term := ← WithRpcRef.mk ti }
 
 def handlePlainTermGoal (p : PlainTermGoalParams)
     : RequestM (RequestTask (Option PlainTermGoal)) := do

src/Lean/Server/FileWorker/WidgetRequests.lean

@@ -28,7 +28,7 @@ builtin_initialize
     `Lean.Widget.InteractiveDiagnostics.msgToInteractive
     MsgToInteractive
     (TaggedText MsgEmbed)
-    fun ⟨⟨m⟩, i⟩ => RequestM.pureTask do msgToInteractive m i (hasWidgets := true)
+    fun ⟨m, i⟩ => RequestM.pureTask do msgToInteractive m.val i (hasWidgets := true)
 
 /-- The information that the infoview uses to render a popup
 for when the user hovers over an expression.
@@ -48,7 +48,8 @@ The intended usage of this is for the infoview to pass the `InfoWithCtx` which
 was stored for a particular `SubexprInfo` tag in a `TaggedText` generated with `ppExprTagged`.
  -/
 def makePopup : WithRpcRef InfoWithCtx → RequestM (RequestTask InfoPopup)
-  | ⟨i⟩ => RequestM.pureTask do
+  | i => RequestM.pureTask do
+    let i := i.val
     i.ctx.runMetaM i.info.lctx do
       let type? ← match (← i.info.type?) with
         | some type => some <$> ppExprTagged type
@@ -124,7 +125,8 @@ builtin_initialize
     `Lean.Widget.getGoToLocation
     GetGoToLocationParams
     (Array Lsp.LocationLink)
-    fun ⟨kind, ⟨i⟩⟩ => RequestM.pureTask do
+    fun ⟨kind, i⟩ => RequestM.pureTask do
+      let i := i.val
       let rc ← read
       let ls ← FileWorker.locationLinksOfInfo kind i
       if !ls.isEmpty then return ls
@@ -138,9 +140,9 @@ builtin_initialize
 def lazyTraceChildrenToInteractive (children : WithRpcRef LazyTraceChildren) :
     RequestM (RequestTask (Array (TaggedText MsgEmbed))) :=
   RequestM.pureTask do
-    let ⟨indent, children⟩ := children
-    children.mapM fun ⟨child⟩ =>
-      msgToInteractive child (hasWidgets := true) (indent := indent)
+    let ⟨indent, children⟩ := children.val
+    children.mapM fun child =>
+      msgToInteractive child.val (hasWidgets := true) (indent := indent)
 
 builtin_initialize registerBuiltinRpcProcedure ``lazyTraceChildrenToInteractive _ _ lazyTraceChildrenToInteractive
 

src/Lean/Server/Rpc/Basic.lean

@@ -19,12 +19,19 @@ first connect to the session using `$/lean/rpc/connect`. -/
 
 namespace Lean.Lsp
 
-/-- An object which RPC clients can refer to without marshalling. -/
+/--
+An object which RPC clients can refer to without marshalling.
+
+The language server may serve the same `RpcRef` multiple times and maintains a reference count
+to track how many times it has served the reference.
+If clients want to release the object associated with an `RpcRef`,
+they must release the reference as many times as they have received it from the server.
+-/
 structure RpcRef where
   /- NOTE(WN): It is important for this to be a single-field structure
   in order to deserialize as an `Object` on the JS side. -/
   p : USize
-  deriving BEq, Hashable, FromJson, ToJson
+  deriving Inhabited, BEq, Hashable, FromJson, ToJson
 
 instance : ToString RpcRef where
   toString r := toString r.p
@@ -33,38 +40,100 @@ end Lean.Lsp
 
 namespace Lean.Server
 
-structure RpcObjectStore : Type where
-  /-- Objects that are being kept alive for the RPC client, together with their type names,
-  mapped to by their RPC reference.
+/--
+Marks values to be encoded as opaque references in RPC packets.
+Two `WithRpcRef`s with the same `id` will yield the same client-side reference.
+
+See also the docstring for `RpcEncodable`.
+-/
+structure WithRpcRef (α : Type u) where
+  private mk' ::
+    val : α
+    private id : USize
+  deriving Inhabited
 
-  Note that we may currently have multiple references to the same object. It is only disposed
-  of once all of those are gone. This simplifies the client a bit as it can drop every reference
-  received separately. -/
-  aliveRefs : PersistentHashMap Lsp.RpcRef Dynamic := {}
-  /-- Value to use for the next `RpcRef`. It is monotonically increasing to avoid any possible
-  bugs resulting from its reuse. -/
-  nextRef   : USize := 0
+builtin_initialize freshWithRpcRefId : IO.Ref USize ← IO.mkRef 1
 
-def rpcStoreRef (any : Dynamic) : StateM RpcObjectStore Lsp.RpcRef := do
-  let st ← get
-  set { st with
-    aliveRefs := st.aliveRefs.insert ⟨st.nextRef⟩ any
-    nextRef := st.nextRef + 1
-  }
-  return ⟨st.nextRef⟩
+/--
+Creates an `WithRpcRef` instance with a unique `id`.
+As long as the client holds at least one reference to this `WithRpcRef`,
+serving it again will yield the same client-side reference.
+Thus, when used as React deps,
+client-side references can help preserve UI state across RPC requests.
+-/
+def WithRpcRef.mk (val : α) : BaseIO (WithRpcRef α) := do
+  let id ← freshWithRpcRefId.modifyGet fun id => (id, id + 1)
+  return { val, id }
 
-def rpcGetRef (r : Lsp.RpcRef) : ReaderT RpcObjectStore Id (Option Dynamic) :=
-  return (← read).aliveRefs.find? r
+structure ReferencedObject where
+  obj : Dynamic
+  id  : USize
+  rc  : Nat
+
+structure RpcObjectStore : Type where
+  /--
+  Objects that are being kept alive for the RPC client, together with their type names,
+  mapped to by their RPC reference.
+  -/
+  aliveRefs : PersistentHashMap Lsp.RpcRef ReferencedObject := {}
+  /--
+  Unique `RpcRef` for the ID of an object that is being referenced through RPC.
+  We store this mapping so that we can reuse `RpcRef`s for the same object.
+  Reusing `RpcRef`s is helpful because it enables clients to reuse their UI state.
+  -/
+  refsById : PersistentHashMap USize Lsp.RpcRef := {}
+  /--
+  Value to use for the next fresh `RpcRef`, monotonically increasing.
+  -/
+  nextRef : USize := 0
+
+def rpcStoreRef [TypeName α] (obj : WithRpcRef α) : StateM RpcObjectStore Lsp.RpcRef := do
+  let st ← get
+  let reusableRef? : Option Lsp.RpcRef := st.refsById.find? obj.id
+  match reusableRef? with
+  | some ref =>
+    -- Reuse `RpcRef` for this `obj` so that clients can reuse their UI state for it.
+    -- We maintain a reference count so that we only free `obj` when the client has released
+    -- all of its instances of the `RpcRef` for `obj`.
+    let some referencedObj := st.aliveRefs.find? ref
+      | return panic! "Found object ID in `refsById` but not in `aliveRefs`."
+    let referencedObj := { referencedObj with rc := referencedObj.rc + 1 }
+    set { st with aliveRefs := st.aliveRefs.insert ref referencedObj }
+    return ref
+  | none =>
+    let ref : Lsp.RpcRef := ⟨st.nextRef⟩
+    set { st with
+      aliveRefs := st.aliveRefs.insert ref ⟨.mk obj.val, obj.id, 1⟩
+      refsById := st.refsById.insert obj.id ref
+      nextRef := st.nextRef + 1
+    }
+    return ref
+
+def rpcGetRef (α) [TypeName α] (r : Lsp.RpcRef)
+    : ReaderT RpcObjectStore (ExceptT String Id) (WithRpcRef α) := do
+  let some referencedObj := (← read).aliveRefs.find? r
+    | throw s!"RPC reference '{r}' is not valid"
+  let some val := referencedObj.obj.get? α
+    | throw <| s!"RPC call type mismatch in reference '{r}'\nexpected '{TypeName.typeName α}', " ++
+        s!"got '{referencedObj.obj.typeName}'"
+  return { val, id := referencedObj.id }
 
 def rpcReleaseRef (r : Lsp.RpcRef) : StateM RpcObjectStore Bool := do
   let st ← get
-  if st.aliveRefs.contains r then
-    set { st with aliveRefs := st.aliveRefs.erase r }
-    return true
+  let some referencedObj := st.aliveRefs.find? r
+    | return false
+  let referencedObj := { referencedObj with rc := referencedObj.rc - 1 }
+  if referencedObj.rc == 0 then
+    set { st with
+      aliveRefs := st.aliveRefs.erase r
+      refsById := st.refsById.erase referencedObj.id
+    }
   else
-    return false
+    set { st with aliveRefs := st.aliveRefs.insert r referencedObj }
+  return true
 
-/-- `RpcEncodable α` means that `α` can be deserialized from and serialized into JSON
+/--
+`RpcEncodable α` means that `α` can be deserialized from and serialized into JSON
 for the purpose of receiving arguments to and sending return values from
 Remote Procedure Calls (RPCs).
 
@@ -79,13 +148,15 @@ For such data, we use the `WithRpcRef` marker.
 Note that for `WithRpcRef α` to be `RpcEncodable`,
 `α` must have a `TypeName` instance
 
-On the server side, `WithRpcRef α` is just a structure
-containing a value of type `α`.
+On the server side, `WithRpcRef α` is a structure containing a value of type `α` and an associated
+`id`.
 On the client side, it is an opaque reference of (structural) type `Lsp.RpcRef`.
 Thus, `WithRpcRef α` is cheap to transmit over the network
 but may only be accessed on the server side.
 In practice, it is used by the client to pass data
-between various RPC methods provided by the server. -/
+between various RPC methods provided by the server.
+Two `WithRpcRef`s with the same `id` will yield the same client-side reference.
+-/
 -- TODO(WN): for Lean.js, compile `WithRpcRef` to "opaque reference" on the client
 class RpcEncodable (α : Type) where
   rpcEncode : α → StateM RpcObjectStore Json
@@ -121,26 +192,11 @@ instance [RpcEncodable α] : RpcEncodable (StateM RpcObjectStore α) where
     let a : α ← rpcDecode j
     return return a
 
-/-- Marks values to be encoded as opaque references in RPC packets.
-
-See the docstring for `RpcEncodable`. -/
-structure WithRpcRef (α : Type u) where
-  val : α
-  deriving Inhabited
-
 instance [TypeName α] : RpcEncodable (WithRpcRef α) :=
   { rpcEncode, rpcDecode }
 where
   -- separate definitions to prevent inlining
-  rpcEncode r := toJson <$> rpcStoreRef (.mk r.val)
-  rpcDecode j := do
-    let r ← fromJson? j
-    match (← rpcGetRef r) with
-      | none => throw s!"RPC reference '{r}' is not valid"
-      | some any =>
-        if let some obj := any.get? α then
-          return ⟨obj⟩
-        else
-          throw s!"RPC call type mismatch in reference '{r}'\nexpected '{TypeName.typeName α}', got '{any.typeName}'"
+  rpcEncode r := toJson <$> rpcStoreRef r
+  rpcDecode j := do rpcGetRef α (← fromJson? j)
 
 end Lean.Server

src/Lean/Widget/InteractiveCode.lean

@@ -59,19 +59,19 @@ def SubexprInfo.withDiffTag (tag : DiffTag) (c : SubexprInfo) : SubexprInfo :=
 
 /-- Tags pretty-printed code with infos from the delaborator. -/
 partial def tagCodeInfos (ctx : Elab.ContextInfo) (infos : SubExpr.PosMap Elab.Info) (tt : TaggedText (Nat × Nat))
-    : CodeWithInfos :=
+    : BaseIO CodeWithInfos :=
   go tt
 where
-  go (tt : TaggedText (Nat × Nat)) :=
-    tt.rewrite fun (n, _) subTt =>
+  go (tt : TaggedText (Nat × Nat)) : BaseIO (TaggedText SubexprInfo) :=
+    tt.rewriteM fun (n, _) subTt => do
       match infos.find? n with
       | none   => go subTt
       | some i =>
         let t : SubexprInfo := {
-          info := WithRpcRef.mk { ctx, info := i, children := .empty }
+          info := ← WithRpcRef.mk { ctx, info := i, children := .empty }
           subexprPos := n
         }
-        TaggedText.tag t (go subTt)
+        return TaggedText.tag t (← go subTt)
 
 open PrettyPrinter Delaborator in
 /--
@@ -93,6 +93,6 @@ def ppExprTagged (e : Expr) (delab : Delab := Delaborator.delab) : MetaM CodeWit
     fileMap       := default
     ngen          := (← getNGen)
   }
-  return tagCodeInfos ctx infos tt
+  tagCodeInfos ctx infos tt
 
 end Lean.Widget

src/Lean/Widget/InteractiveDiagnostic.lean

@@ -192,8 +192,8 @@ partial def msgToInteractive (msgData : MessageData) (hasWidgets : Bool) (indent
   let rec fmtToTT (fmt : Format) (indent : Nat) : IO (TaggedText MsgEmbed) :=
     (TaggedText.prettyTagged fmt indent).rewriteM fun (n, col) tt =>
       match embeds[n]! with
-        | .code ctx infos =>
-          return .tag (.expr (tagCodeInfos ctx infos tt)) default
+        | .code ctx infos => do
+          return .tag (.expr (← tagCodeInfos ctx infos tt)) default
         | .goal ctx lctx g =>
           ctx.runMetaM lctx do
             return .tag (.goal (← goalToInteractive g)) default
@@ -205,7 +205,10 @@ partial def msgToInteractive (msgData : MessageData) (hasWidgets : Bool) (indent
           let col := indent + col
           let children ←
             match children with
-              | .lazy children => pure <| .lazy ⟨{indent := col+2, children := children.map .mk}⟩
+              | .lazy children => pure <| .lazy <| ← WithRpcRef.mk {
+                  indent := col+2
+                  children := ← children.mapM (WithRpcRef.mk ·)
+                }
               | .strict children => pure <| .strict (← children.mapM (fmtToTT · (col+2)))
           return .tag (.trace indent cls (← fmtToTT msg col) collapsed children) default
         | .ignoreTags => return .text tt.stripTags

src/Lean/Widget/InteractiveGoal.lean

@@ -192,7 +192,7 @@ def goalToInteractive (mvarId : MVarId) : MetaM InteractiveGoal := do
     return {
       hyps
       type := goalFmt
-      ctx := ⟨{← Elab.CommandContextInfo.save with }⟩
+      ctx := ← WithRpcRef.mk {← Elab.CommandContextInfo.save with }
       userName?
       goalPrefix := getGoalPrefix mvarDecl
       mvarId

tests/lean/derivingRpcEncoding.lean

@@ -16,7 +16,7 @@ structure FooRef where
   a : Array Nat
   deriving Inhabited, TypeName
 
-#eval test (WithRpcRef FooRef) default
+#eval do return test (WithRpcRef FooRef) (← WithRpcRef.mk default)
 
 structure FooJson where
   s : String
@@ -27,13 +27,18 @@ structure Bar where
   fooJson : FooJson
   deriving RpcEncodable, Inhabited
 
-#eval test Bar default
+def Bar.mkDefault : BaseIO Bar := do return {
+  fooRef  := ← WithRpcRef.mk default
+  fooJson := default
+}
+
+#eval do return test Bar (← Bar.mkDefault)
 
 structure BarTrans where
   bar : Bar
   deriving RpcEncodable, Inhabited
 
-#eval test BarTrans default
+#eval do return test BarTrans { bar := ← Bar.mkDefault }
 
 structure Baz where
   arr : Array String -- non-constant field
@@ -53,15 +58,15 @@ inductive BazInductive
   | baz (arr : Array Bar)
   deriving RpcEncodable, Inhabited
 
-#eval test BazInductive ⟨#[default, default]⟩
+#eval do return test BazInductive ⟨#[← Bar.mkDefault, ← Bar.mkDefault]⟩
 
 inductive FooInductive (α : Type) where
   | a : α → WithRpcRef FooRef → FooInductive α
   | b : (n : Nat) → (a : α) → (m : Nat) → FooInductive α
   deriving RpcEncodable, Inhabited
 
-#eval test (FooInductive BazInductive) (.a default default)
-#eval test (FooInductive BazInductive) (.b 42 default default)
+#eval do return test (FooInductive BazInductive) (.a default (← WithRpcRef.mk default))
+#eval do return test (FooInductive BazInductive) (.b 42 default default)
 
 inductive FooNested (α : Type) where
   | a : α → Array (FooNested α) → FooNested α