feat: reusable rpc refs (#8105)

This PR adds support for server-sided `RpcRef` reuse and fixes a bug
where trace nodes in the InfoView would close while the file was still
being processed.

The core of the trace node issue is that the server always serves new
RPC references in every single response to the client, which means that
the client is forced to reset its UI state.

In a previous attempt at fixing this (#8056), the server would memorize
the RPC-encoded JSON of interactive diagnostics (which includes RPC
references) and serve it for as long as it could reuse the snapshot
containing the diagnostics, so that RPC references are reused. The
problem with this was that the client then had multiple finalizers
registered for the same RPC reference (one for every reused RPC
reference that was served), and once the first reference was
garbage-collected, all other reused references would point into the
void.

This PR takes a different approach to resolve the issue: The meaning of
`$/lean/rpc/release` is relaxed from "Free the object pointed to by this
RPC reference" to "Decrement the RPC reference count of the object
pointed to by this RPC reference", and the server now maintains a
reference count to track how often a given `RpcRef` was served. Only
when every single served instance of the `RpcRef` has been released, the
object is freed. Additionally, the reuse mechanism is generalized from
being only supported for interactive diagnostics, to being supported for
any object using `WithRpcRef`. In order to make use of reusable RPC
references, downstream users still need to memorize the `WithRpcRef`
instances accordingly.

Closes #8053.

### Breaking changes

Since `WithRpcRef` is now capable of tracking its identity to decide
which `WithRpcRef` usage constitutes a reuse, the constructor of
`WithRpcRef` has been made `private` to discourage downstream users from
creating `WithRpcRef` instances with manually-set `id`s. Instead,
`WithRpcRef.mk` (which lives in `BaseIO`) is now the preferred way to
create `WithRpcRef` instances.

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
 
+/--
+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
+
+builtin_initialize freshWithRpcRefId : IO.Ref USize ← IO.mkRef 1
+
+/--
+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 }
+
+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,
+  /--
+  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
 
-  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
-
-def rpcStoreRef (any : Dynamic) : StateM RpcObjectStore Lsp.RpcRef := do
+def rpcStoreRef [TypeName α] (obj : WithRpcRef α) : 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⟩
+  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 (r : Lsp.RpcRef) : ReaderT RpcObjectStore Id (Option Dynamic) :=
-  return (← read).aliveRefs.find? r
+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

@@ -205,7 +205,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 α