Plot cell: allow async, multiple inputs, one output

src/components/cell/CellOutput.svelte

@@ -40,7 +40,7 @@
       {#if state.type === "code"}
         <FullView value={state.output} />
       {:else}
-        <PlotView value={state.output} />
+        <PlotView name={state.result.results[0]} value={state.output} />
       {/if}
     {/if}
   </div>

src/components/cell/output/FullView.svelte

@@ -1,7 +1,9 @@
 <script lang="ts">
+  import type { RelationSet } from "@/lib/types";
+
   import RelationView from "./RelationView.svelte";
 
-  export let value: Record<string, object[]>;
+  export let value: RelationSet;
 </script>
 
 <div class="flex flex-col space-y-3">

src/components/cell/output/PlotView.svelte

@@ -1,9 +1,25 @@
 <script lang="ts">
-  export let value: string;
+  import { isRenderedElement } from "@/lib/types";
+  import RelationView from "./RelationView.svelte";
+
+  export let name: string | undefined;
+  export let value: unknown;
 </script>
 
 {#if value}
-  <div class="plot">{@html value}</div>
+  {#if isRenderedElement(value)}
+    <div class="plot">{@html value.outerHTML}</div>
+  {:else if Array.isArray(value)}
+    <RelationView name={name ?? ""} values={value} />
+  {:else}
+    <div class="font-mono text-[0.95rem] text-gray-700">
+      {#if name !== undefined}
+        <span class="font-bold">{name}</span>
+        <span class="text-gray-400">:=</span>
+      {/if}
+      {JSON.stringify(value, undefined, 2)}
+    </div>
+  {/if}
 {:else}
   <span class="italic text-sm font-light">(fresh pixels for a plot...)</span>
 {/if}

src/components/cell/output/RelationView.svelte

@@ -1,8 +1,10 @@
 <script lang="ts">
+  import type { Relation } from "@/lib/types";
+
   import ValueView from "./ValueView.svelte";
 
   export let name: string;
-  export let values: object[];
+  export let values: Relation;
 
   let displaying = 0;
   $: displaying = Math.min(values.length, 5); // hide long lists

src/lib/notebook.ts

@@ -3,6 +3,7 @@ import { build } from "./runtime";
 import type { CompilerResult } from "./runtime";
 import { buildPlot } from "./plot";
 import type { PlotResult } from "./plot";
+import type { Relation, RelationSet } from "./types";
 
 export type MarkdownCell = {
   type: "markdown";
@@ -27,7 +28,7 @@ export type CellData = MarkdownCell | CodeCellData | PlotCellData;
 export type CodeCellState = CodeCellData & {
   result: CompilerResult;
   status: "stale" | "pending" | "done";
-  output?: Record<string, object[]>;
+  output?: RelationSet;
   graphErrors?: string;
   runtimeErrors?: string;
   evaluateHandle?: () => void;
@@ -36,7 +37,7 @@ export type CodeCellState = CodeCellData & {
 export type PlotCellState = PlotCellData & {
   result: PlotResult;
   status: "stale" | "pending" | "done";
-  output?: string;
+  output?: unknown;
   graphErrors?: string;
   runtimeErrors?: string;
   evaluateHandle?: () => void;
@@ -169,7 +170,7 @@ export class NotebookState {
       if (cell.graphErrors !== undefined) {
         delete cell.graphErrors;
       }
-      if (cell.result.ok && cell.type === "code") {
+      if (cell.result.ok && cell.result.results.length > 0) {
         for (const relation of cell.result.results) {
           const array = creators.get(relation) ?? [];
           array.push(id);
@@ -183,7 +184,7 @@ export class NotebookState {
       if (cellIds.length > 1) {
         for (const id of cellIds) {
           const cell = this.getCell(id);
-          if (cell.type !== "code") throw new Error("unreachable");
+          if (cell.type === "markdown") throw new Error("unreachable");
           clear(cell, "stale");
           cell.graphErrors = `Relation "${relation}" is defined in multiple cells.`;
         }
@@ -211,23 +212,29 @@ export class NotebookState {
         cell.status === "stale"
       ) {
         let depsOk = true;
-        const deps: Record<string, object[]> = {};
+        const deps: RelationSet = {};
         for (const relation of cell.result.deps) {
           const cellIds = creators.get(relation);
           if (!cellIds || cellIds.length != 1) {
             depsOk = false;
             break;
           }
           const prev = this.getCell(cellIds[0]);
-          if (prev.type !== "code") throw new Error("unreachable");
+          if (prev.type === "markdown") throw new Error("unreachable");
           if (
             prev.status === "done" &&
             prev.result.ok &&
             prev.graphErrors === undefined &&
-            prev.runtimeErrors === undefined &&
-            prev.output?.[relation]
+            prev.runtimeErrors === undefined
           ) {
-            deps[relation] = prev.output[relation];
+            if (prev.type === "code" && prev.output?.[relation]) {
+              deps[relation] = prev.output[relation];
+            } else if (prev.type === "plot" && prev.output !== undefined) {
+              deps[relation] = prev.output as Relation;
+            } else {
+              depsOk = false;
+              break;
+            }
           } else {
             depsOk = false;
             break;
@@ -254,13 +261,15 @@ export class NotebookState {
                 }
               });
           } else {
-            const promise = cell.result.evaluate(deps[cell.result.deps[0]]);
+            const depValues = cell.result.deps.map((dep) => deps[dep]);
+            const promise = cell.result.evaluate(depValues);
             cell.evaluateHandle = () => promise.cancel();
+            const results = cell.result.results; // storing for async callback
             promise
               .then((figure) => {
                 cell.output = figure;
                 cell.status = "done";
-                this.revalidate();
+                this.markUpdate(results);
               })
               .catch((err: Error) => {
                 if (err.message !== "Promise was cancelled by user") {

src/lib/plot.test.ts

@@ -0,0 +1,35 @@
+import { expect } from "chai";
+import { buildPlot } from "./plot";
+
+function check(
+  str: string,
+  deps: string[] = [],
+  result: string | undefined = undefined,
+) {
+  const context = (name: string) => `build('${str}').${name}`;
+  const plot = buildPlot(str);
+  expect(plot.ok, context("ok")).to.be.true;
+  if (plot.ok) {
+    expect(plot.deps, context("deps")).to.deep.eq(deps);
+    if (result) {
+      expect(plot.results, context("results")).to.deep.eq([result]);
+    }
+  }
+}
+
+describe("buildPlot", () => {
+  it("parses empty string", () => check(""));
+
+  it("parses a basic function", () => check("x => cool", ["x"]));
+
+  it("parses a two-arg function", () => check("(x, y) => x + y", ["x", "y"]));
+
+  it("parses a one-arg result", () =>
+    check("stuff = x => nice", ["x"], "stuff"));
+
+  it("parses a no-arg result", () =>
+    check("result = () => swanky", [], "result"));
+
+  it("parses an async function", () =>
+    check('_ = async () => import("lodash")', [], "_"));
+});

src/lib/plot.ts

@@ -1,15 +1,17 @@
 import Worker from "./plot.worker?worker&inline";
+import type { Relation } from "./types";
 
 interface CancellablePromise<T> extends Promise<T> {
   cancel: () => void;
 }
 
-type EvalPromise = CancellablePromise<string>;
+type EvalPromise = CancellablePromise<unknown>;
 
 type PlotResultOk = {
   ok: true;
-  evaluate: (data: object[]) => EvalPromise;
+  evaluate: (dependencies: Relation[]) => EvalPromise;
   deps: string[];
+  results: string[];
 };
 
 type PlotResultErr = {
@@ -30,20 +32,48 @@ export function buildPlot(src: string): PlotResult {
         return promise as EvalPromise;
       },
       deps: [],
+      results: [],
     };
   }
 
-  const result = src.match(/^\s*([a-zA-Z_$][a-zA-Z_$0-9]*)\s*=>/);
-  if (result === null) {
+  // This mess of regexps parses the first line of an arrow function declaration.
+  // See plot.test.ts for examples of valid declarations.
+  const nameFragment = `[a-zA-Z_$][a-zA-Z_$0-9]*`;
+  const resultNameFragment = `(?<resultName>${nameFragment})\\s*=\\s*`; // `resultName = `
+  const nextFunctionArg = `(?:\\s*,\\s*(${nameFragment}))?`; // `, nextDepName`
+  const either = (l: string, r: string) =>
+    "(?:" + [l, r].map((s) => `(?:${s})`).join("|") + ")";
+  const fullArgumentsList = [
+    // matches eg `(depName1, depName2)`
+    `\\(?`,
+    `(${nameFragment})`, // `depName1`
+    // If a capture group is repeated, RegExp only retains the last such capture.
+    // To support N arguments, we need N capture groups, or we can do some kind of loop w/ regex.exec(...).
+    ...new Array(10).fill(undefined).map(() => nextFunctionArg),
+    `\\)?`,
+  ].join("");
+  const emptyArgumentsList = `\\(\\s*\\)`; // `()`
+  const argumentsList = either(emptyArgumentsList, fullArgumentsList); // `()` | `dep` | (dep) | `(dep1, dep2)`
+  const regexp = new RegExp( // `resultName = async (dep1, dep2) =>`
+    `^\\s*(?:${resultNameFragment})?(?:async\\s+)?${argumentsList}\\s*=>`,
+  );
+
+  const parsed = src.match(regexp);
+  if (parsed === null) {
     return {
       ok: false,
       error: "Expected plot cell to start with `name =>` syntax",
     };
   }
 
+  const resultName = parsed.groups?.resultName;
+  const deps = Array.from(parsed.slice(2)).filter(
+    (s) => s !== undefined && s !== null,
+  );
+
   return {
     ok: true,
-    evaluate: (data: object[]) => {
+    evaluate: (data: Relation[]) => {
       const worker = new Worker();
       let rejectCb: (reason?: any) => void;
       const promise: Partial<EvalPromise> = new Promise((resolve, reject) => {
@@ -64,6 +94,7 @@ export function buildPlot(src: string): PlotResult {
       };
       return promise as EvalPromise;
     },
-    deps: [result[1]],
+    deps,
+    results: resultName ? [resultName] : [],
   };
 }

src/lib/plot.worker.ts

@@ -1,14 +1,21 @@
 import * as Plot from "@observablehq/plot";
 import domino from "domino";
+import { RenderedElement } from "./types";
 
 globalThis.document = domino.createDocument();
 
-onmessage = (event) => {
+onmessage = async (event) => {
   const { code, data } = event.data;
   const fn = new Function(
     "Plot",
-    "__percival_data",
-    `return (${code})(__percival_data);`,
+    "...__percival_data",
+    `return (${code})(...__percival_data);`,
   );
-  postMessage(fn(Plot, data).outerHTML);
+  let result = await fn(Plot, ...data);
+
+  if (result && "outerHTML" in result) {
+    result = RenderedElement(result);
+  }
+
+  postMessage(result);
 };