RPC method performance improvements (#46)

* feat: experimental RPC request cancellation

* chore: ignore nix stuff

* chore: trace

* feat: memoize

* chore: adaptations for nightly-2024-02-08

* chore: bump to 2024-02-24

* fix: demo

* chore: bump to v4.7.0-rc1

---------

Co-authored-by: Scott Morrison <[email protected]>

.gitignore

@@ -3,3 +3,4 @@
 /widget/dist
 /widget/node_modules
 /widget/*.hash
+.direnv/

ProofWidgets/Cancellable.lean

@@ -0,0 +1,106 @@
+import Lean.Data.Json.FromToJson
+import Lean.Server.Rpc.RequestHandling
+import ProofWidgets.Compat
+
+/-! Experimental support for cancellable RPC requests.
+
+Note: Cancellation should eventually become a feature of the core RPC protocol,
+and the requests map should be stored in `RequestM`,
+or somewhere in the server anyway. -/
+
+namespace ProofWidgets
+open Lean Server Meta
+
+abbrev RequestId := Nat
+structure CancellableTask where
+  task : Task (Except RequestError (LazyEncodable Json))
+  /- Note: we cannot just `IO.cancel task` because it is a result of `map`.
+  See https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Should.20cancelling.20a.20purely.20mapped.20task.20cancel.20the.20original.3F -/
+  cancel : IO Unit
+
+/-- Maps the ID of each currently executing request to its task. -/
+initialize runningRequests :
+    IO.Ref (RequestId × HashMap RequestId CancellableTask) ←
+  IO.mkRef (0, HashMap.empty)
+
+/-- Transforms a request handler returning `β`
+into one that returns immediately with a `RequestId`.
+
+The ID uniquely identifies the running request:
+its results can be retrieved using `checkRequest`,
+and it can be cancelled using `cancelRequest`. -/
+def mkCancellable [RpcEncodable β] (handler : α → RequestM (RequestTask β)) :
+    α → RequestM (RequestTask RequestId) := fun a => do
+  RequestM.asTask do
+    let t ← handler a
+    let t' := t.map (·.map rpcEncode)
+    runningRequests.modifyGet fun (id, m) =>
+      (id, (id+1, m.insert id ⟨t', IO.cancel t⟩))
+
+/-- Cancel the request with ID `rid`.
+Does nothing if `rid` is invalid. -/
+@[server_rpc_method]
+def cancelRequest (rid : RequestId) : RequestM (RequestTask String) := do
+  RequestM.asTask do
+    let t? ← runningRequests.modifyGet fun (id, m) => (m.find? rid, (id, m.erase rid))
+    if let some t := t? then
+      t.cancel
+    return "ok"
+
+/-- The status of a running cancellable request. -/
+inductive CheckRequestResponse
+  | running
+  | done (result : LazyEncodable Json)
+  deriving RpcEncodable
+
+/-- Check whether a request has finished computing,
+and return the response if so.
+The request is removed from `runningRequests` the first time it is checked
+and found to have finished.
+Throws an error if the `rid` is invalid,
+or if the request itself threw an error. -/
+/- NOTE: a notification-based version would be better than this polling-based one.
+But we cannot include RPC references in notifications atm;
+another possible addition to the RPC protocol? -/
+@[server_rpc_method]
+def checkRequest (rid : RequestId) : RequestM (RequestTask CheckRequestResponse) := do
+  RequestM.asTask do
+    let (_, m) ← runningRequests.get
+    match m.find? rid with
+    | none =>
+      throw $ RequestError.invalidParams
+        s!"Request '{rid}' has already finished, or the ID is invalid."
+    | some t =>
+      if !(← IO.hasFinished t.task) then
+        return .running
+      runningRequests.modify fun (id, m) => (id, m.erase rid)
+      match t.task.get with
+      | .error e => throw e
+      | .ok v    => return .done v
+
+def cancellableSuffix : Name := `_cancellable
+
+initialize
+  registerBuiltinAttribute {
+    name := `server_rpc_method_cancellable
+    descr := "Like `server_rpc_method`, \
+    but requests for this method can be cancelled. \
+    The method should check for that using `IO.checkCanceled`. \
+    Cancellable methods are invoked differently from JavaScript: \
+    see `callCancellable` in `cancellable.ts`."
+    applicationTime := AttributeApplicationTime.afterCompilation
+    add := fun decl _ _ => Prod.fst <$> MetaM.run do
+      let name := decl ++ cancellableSuffix
+      let value ← mkAppM ``mkCancellable #[mkConst decl]
+      addAndCompile $ .defnDecl {
+        name
+        levelParams := []
+        type := ← inferType value
+        value
+        hints  := .opaque
+        safety := .safe
+      }
+      registerRpcProcedure name
+  }
+
+end ProofWidgets

ProofWidgets/Component/HtmlDisplay.lean

@@ -1,5 +1,4 @@
 import Lean.Server.Rpc.Basic
-import Std.Data.Json
 
 import ProofWidgets.Data.Html
 

ProofWidgets/Component/OfRpcMethod.lean

@@ -1,6 +1,7 @@
 import Lean.Elab.ElabRules
 import ProofWidgets.Component.Basic
 import ProofWidgets.Data.Html
+import ProofWidgets.Cancellable
 
 namespace ProofWidgets
 open Lean Server Meta Elab Term
@@ -66,10 +67,21 @@ elab "mk_rpc_widget%" fn:term : term <= expectedType => do
   let fn ← Term.elabTermEnsuringType fn arr
   let fn ← instantiateMVars fn
   if let .const nm .. := fn then
+    let cancellableNm := nm ++ cancellableSuffix
+    if (← builtinRpcProcedures.get).contains cancellableNm || userRpcProcedures.contains (← getEnv) cancellableNm then
+      -- Use the cancellable variant if possible.
+      let code : StrLit := quote $ ofRpcMethodTemplate
+        |>.replace "$RPC_METHOD" (toString cancellableNm)
+        |>.replace "window.toString()" "'true'"
+      let valStx ← `({ javascript := $code })
+      let ret ← elabTerm valStx expectedType
+      return ret
     if !(← builtinRpcProcedures.get).contains nm && !userRpcProcedures.contains (← getEnv) nm then
       throwError s!"'{nm}' is not a known RPC method. Use `@[server_rpc_method]` to register it."
     -- https://github.com/leanprover/lean4/issues/1415
-    let code : StrLit := quote $ ofRpcMethodTemplate.replace "$RPC_METHOD" (toString nm)
+    let code : StrLit := quote $ ofRpcMethodTemplate
+      |>.replace "$RPC_METHOD" (toString nm)
+      |>.replace "window.toString()" "'false'"
     let valStx ← `({ javascript := $code })
     let ret ← elabTerm valStx expectedType
     return ret

ProofWidgets/Demos/SelectInsertConv.lean

@@ -5,7 +5,6 @@ Authors: Robin Böhne, Wojciech Nawrocki, Patrick Massot
 -/
 import Lean.Meta.ExprLens
 import Std.Lean.Position
-import Std.Data.Json
 import ProofWidgets.Data.Html
 import ProofWidgets.Component.OfRpcMethod
 import ProofWidgets.Component.MakeEditLink

ProofWidgets/Demos/Venn.lean

@@ -9,16 +9,6 @@ open Lean Meta Server
 open ProofWidgets
 
 /-! # Minimal definiton of sets copied from Mathlib -/
-
-/-- Notation type class for the subset relation `⊆`. -/
-class HasSubset (α : Type u) where
-  /-- Subset relation: `a ⊆ b`  -/
-  Subset : α → α → Prop
-export HasSubset (Subset)
-
-/-- Subset relation: `a ⊆ b`  -/
-infix:50 " ⊆ " => HasSubset.Subset
-
 def Set (α : Type u) := α → Prop
 
 namespace Set
@@ -50,7 +40,7 @@ end Set
 
 /-- If `e == S ⊆ T` return `some (S, T)`, otherwise `none`. -/
 def isSubsetPred? (e : Expr) : Option (Expr × Expr) := do
-  let some (_, _, S, T) := e.app4? ``Subset | none
+  let some (_, _, S, T) := e.app4? ``HasSubset.Subset | none
   return (S, T)
 
 /-- Expressions to display as labels in a diagram. -/

lake-manifest.json

@@ -4,7 +4,7 @@
  [{"url": "https://github.com/leanprover/std4",
    "type": "git",
    "subDir": null,
-   "rev": "276953b13323ca151939eafaaec9129bf7970306",
+   "rev": "ff9850c4726f6b9fb8d8e96980c3fcb2900be8bd",
    "name": "std",
    "manifestFile": "lake-manifest.json",
    "inputRev": "main",

lean-toolchain

@@ -1 +1 @@
-leanprover/lean4:v4.6.0
+leanprover/lean4:v4.7.0-rc1

widget/package.json

@@ -12,7 +12,7 @@
   "license": "Apache-2.0",
   "type": "module",
   "dependencies": {
-    "@leanprover/infoview": "~0.4.3",
+    "@leanprover/infoview": "~0.4.5",
     "@penrose/core": "^3.1.0",
     "@react-three/drei": "^9.64.0",
     "@react-three/fiber": "^8.12.0",

widget/src/cancellable.ts

@@ -0,0 +1,41 @@
+import { RpcSessionAtPos } from '@leanprover/infoview'
+
+type RequestId = string
+type CheckRequestResponse<S> =
+  'running' | { done: { result: S } }
+
+/** An object for passing a function by reference. */
+export type Fn = {
+  fn: () => void
+}
+
+/** Returns a promise of a response to this request
+ * as well as a function you can call to cancel the request.
+ * It is request-specific what a cancelled request resolves with:
+ * it could ignore the cancellation altogether,
+ * return a null result,
+ * or throw an error. */
+export function callCancellable<T, S>(rs: RpcSessionAtPos, name: string, params: T):
+    [Promise<S>, Fn] {
+  const cancelFn = {
+    fn: () => {}
+  }
+  const promise = new Promise<S>(async (resolve, reject) => {
+    const id: RequestId = await rs.call(name, params)
+    const i = window.setInterval(async () => {
+      try {
+        const response : CheckRequestResponse<S> = await rs.call('ProofWidgets.checkRequest', id)
+        if (response === 'running') return
+        window.clearInterval(i)
+        resolve(response.done.result)
+      } catch (e) {
+        window.clearInterval(i)
+        reject(e)
+      }
+    }, 100)
+    cancelFn.fn = () => {
+      void rs.call('ProofWidgets.cancelRequest', id)
+    }
+  })
+  return [promise, cancelFn]
+}

widget/src/ofRpcMethod.tsx

@@ -1,21 +1,39 @@
 import * as React from 'react'
 import { DocumentPosition, RpcContext, mapRpcError, useAsyncPersistent } from '@leanprover/infoview'
 import { Html, default as HtmlDisplay } from './htmlDisplay'
+import { Fn, callCancellable } from './cancellable'
 
 declare type Props = any & { pos : DocumentPosition };
-// The string template is replaced in OfRpcMethod.lean
+// The string template is replaced in OfRpcMethod.lean.
 const RPC_METHOD = '$RPC_METHOD'
+// HACK: TS runs dead-code elimination
+// and eliminates the `callCancellable` branch if we use a string literal,
+// so use a runtime value.
+const RPC_CANCELLABLE = window.toString()
 
-export default function(props: Props) {
+export default React.memo((props: Props) => {
   const rs = React.useContext(RpcContext)
-  const st  = useAsyncPersistent<Html>(async () => {
-    return await rs.call(RPC_METHOD, props)
+  const cancelRef = React.useRef<Fn>({ fn: () => {} })
+  const st = useAsyncPersistent<Html>(async () => {
+    if (RPC_CANCELLABLE === 'true') {
+      cancelRef.current.fn()
+      const [res, cancel] = callCancellable<Props, Html>(rs, RPC_METHOD, props)
+      cancelRef.current = cancel
+      return res
+    } else
+      return rs.call<Props, Html>(RPC_METHOD, props)
   }, [rs, props])
 
+  // This effect runs once on startup, does nothing,
+  // and sets up a cleanup function that runs when the component is unmounted.
+  React.useEffect(() => {
+    return () => { cancelRef.current.fn() }
+  }, [])
+
   return st.state === 'rejected' ?
       <p style={{color: 'red'}}>{mapRpcError(st.error).message}</p>
     : st.state === 'loading' ?
       <>Loading..</>
     :
       <HtmlDisplay pos={props.pos} html={st.value} />
-}
+})