Merge pull request #2988 from quantified-uncertainty/decorated-statem…

…ents-fixes

Find "Find in editor" and value tooltips for decorated statements

.changeset/angry-bears-dress.md

@@ -0,0 +1,6 @@
+---
+"@quri/squiggle-lang": patch
+"@quri/squiggle-components": patch
+---
+
+Fixes error with "Find in editor" and tooltips not working for decorated values

packages/components/src/components/CodeEditor/languageSupport/autocomplete.tsx

@@ -5,11 +5,11 @@ import {
 } from "@codemirror/autocomplete";
 import { syntaxTree } from "@codemirror/language";
 import { SyntaxNode, Tree } from "@lezer/common";
-import { createRoot } from "react-dom/client";
 
 import { SqProject } from "@quri/squiggle-lang";
 
 import { FnDocumentationFromName } from "../../ui/FnDocumentation.js";
+import { reactAsDom } from "../utils.js";
 
 type NameNode = {
   node: SyntaxNode;
@@ -78,12 +78,7 @@ export function makeCompletionSource(project: SqProject) {
   const decoratorCompletions: Completion[] = [];
 
   const getInfoFunction = (name: string): Completion["info"] => {
-    return () => {
-      const dom = document.createElement("div");
-      const root = createRoot(dom);
-      root.render(<FnDocumentationFromName functionName={name} />);
-      return { dom };
-    };
+    return () => reactAsDom(<FnDocumentationFromName functionName={name} />);
   };
 
   for (const [name, value] of project.getStdLib().entrySeq()) {

packages/components/src/components/CodeEditor/useTooltipsExtension.tsx

@@ -2,7 +2,6 @@ import { syntaxTree } from "@codemirror/language";
 import { EditorView, hoverTooltip, repositionTooltips } from "@codemirror/view";
 import { SyntaxNode } from "@lezer/common";
 import { FC, PropsWithChildren, useEffect } from "react";
-import { createRoot } from "react-dom/client";
 
 import {
   getFunctionDocumentation,
@@ -18,6 +17,7 @@ import {
 } from "../SquiggleViewer/ViewerProvider.js";
 import { FnDocumentation } from "../ui/FnDocumentation.js";
 import { useReactiveExtension } from "./codemirrorHooks.js";
+import { reactAsDom } from "./utils.js";
 
 type Hover = NonNullable<ReturnType<typeof getFunctionDocumentation>>;
 
@@ -100,12 +100,7 @@ function buildWordHoverExtension({
         pos: node.from,
         end: node.to,
         above: true,
-        create() {
-          const dom = document.createElement("div");
-          const root = createRoot(dom);
-          root.render(<HoverTooltip hover={hover} view={view} />);
-          return { dom };
-        },
+        create: () => reactAsDom(<HoverTooltip hover={hover} view={view} />),
       };
     };
 
@@ -117,12 +112,7 @@ function buildWordHoverExtension({
         pos: node.from,
         end: node.to,
         above: true,
-        create() {
-          const dom = document.createElement("div");
-          const root = createRoot(dom);
-          root.render(<ValueTooltip value={value} view={view} />);
-          return { dom };
-        },
+        create: () => reactAsDom(<ValueTooltip value={value} view={view} />),
       };
     };
 
@@ -163,11 +153,15 @@ function buildWordHoverExtension({
         const value = bindings.value.get(name);
         if (!value) return null;
 
-        // Should be LetStatement or DefunStatement
+        // Should be a statement
         const valueAst = value.context?.valueAst;
 
+        if (!valueAst) {
+          return null;
+        }
+
         if (
-          valueAst &&
+          // Note that `valueAst` can't be "DecoratedStatement", we skip those in `SqValueContext` and AST symbols
           (valueAst.type === "LetStatement" ||
             valueAst.type === "DefunStatement") &&
           // If these don't match then variable was probably shadowed by a later statement and we can't show its value.

packages/components/src/components/CodeEditor/utils.tsx

@@ -0,0 +1,10 @@
+import { ReactNode } from "react";
+import { createRoot } from "react-dom/client";
+
+export function reactAsDom(node: ReactNode): { dom: HTMLDivElement } {
+  const dom = document.createElement("div");
+  const root = createRoot(dom);
+  root.render(node);
+  // This is compatible with `CompletionInfo` and `TooltipView` CodeMirror types
+  return { dom };
+}

packages/squiggle-lang/__tests__/Jstat_test.ts

@@ -11,8 +11,8 @@ describe("cumulative density function of a normal distribution", () => {
         async (mean, stdev) => {
           const threeStdevsAboveMean = mean + 3 * stdev;
           const squiggleString = `cdf(Sym.normal(${mean}, ${stdev}), ${threeStdevsAboveMean})`;
-          const squiggleResult = await testRun(squiggleString);
-          expect(squiggleResult.value).toBeCloseTo(1);
+          const { result } = await testRun(squiggleString);
+          expect(result.value).toBeCloseTo(1);
         }
       )
     );
@@ -26,8 +26,8 @@ describe("cumulative density function of a normal distribution", () => {
         async (mean, stdev) => {
           const threeStdevsBelowMean = mean - 3 * stdev;
           const squiggleString = `cdf(Sym.normal(${mean}, ${stdev}), ${threeStdevsBelowMean})`;
-          const squiggleResult = await testRun(squiggleString);
-          expect(squiggleResult.value).toBeCloseTo(0);
+          const { result } = await testRun(squiggleString);
+          expect(result.value).toBeCloseTo(0);
         }
       )
     );

packages/squiggle-lang/__tests__/SqValue/asJS_test.ts

@@ -0,0 +1,25 @@
+import { testRun } from "../helpers/helpers.js";
+
+describe("SqValue.asJS", () => {
+  test("SqDict -> Map", async () => {
+    const value = (
+      await testRun('{ x: 5, y: [3, "foo", { dist: normal(5,2) } ] }')
+    ).result.asJS();
+
+    expect(value).toBeInstanceOf(Object);
+  });
+
+  test("Dict fields", async () => {
+    const value = (await testRun("{ x: 5 }")).result.asJS();
+
+    expect((value as any).value.x).toBe(5);
+  });
+
+  test("Deeply nested dist", async () => {
+    const value = (
+      await testRun('{ x: 5, y: [3, "foo", { dist: normal(5,2) } ] }')
+    ).result.asJS();
+
+    expect((value as any).value.y[2].value.dist).toBeInstanceOf(Array);
+  });
+});

packages/squiggle-lang/__tests__/SqValue/context_test.ts

@@ -0,0 +1,45 @@
+import { nodeToString } from "../../src/ast/parse.js";
+import { assertTag, testRun } from "../helpers/helpers.js";
+
+describe("SqValueContext", () => {
+  test("valueAst for nested nodes", async () => {
+    const { bindings } = await testRun(`
+x = { foo: 5 }
+
+y = 6
+`);
+
+    const x = bindings.get("x");
+    assertTag(x, "Dict");
+    expect(nodeToString(x.context!.valueAst)).toBe(
+      "(LetStatement :x (Block (Dict (KeyValue 'foo' 5))))"
+    );
+    expect(x.context?.valueAstIsPrecise).toBe(true);
+
+    const foo = x.value.get("foo");
+    expect(nodeToString(foo!.context!.valueAst)).toBe("(KeyValue 'foo' 5)");
+    expect(foo!.context!.valueAstIsPrecise).toBe(true);
+
+    const y = bindings.get("y");
+    assertTag(y, "Number");
+
+    expect(nodeToString(y.context!.valueAst)).toBe(
+      "(LetStatement :y (Block 6))"
+    );
+    expect(y!.context!.valueAstIsPrecise).toBe(true);
+  });
+
+  test("valueAst for decorated statements", async () => {
+    const { bindings } = await testRun(`
+@name("Z")
+z = 5
+`);
+
+    const z = bindings.get("z");
+    assertTag(z, "Number");
+
+    expect(nodeToString(z.context!.valueAst)).toBe(
+      "(LetStatement :z (Block 5))"
+    );
+  });
+});

packages/squiggle-lang/__tests__/public/SqValue_test.ts → packages/squiggle-lang/__tests__/SqValue/docstring_test.ts

@@ -1,29 +1,4 @@
 import { run, sq } from "../../src/index.js";
-import { testRun } from "../helpers/helpers.js";
-
-describe("SqValue.asJS", () => {
-  test("SqDict -> Map", async () => {
-    const value = (
-      await testRun('{ x: 5, y: [3, "foo", { dist: normal(5,2) } ] }')
-    ).asJS();
-
-    expect(value).toBeInstanceOf(Object);
-  });
-
-  test("Dict fields", async () => {
-    const value = (await testRun("{ x: 5 }")).asJS();
-
-    expect((value as any).value.x).toBe(5);
-  });
-
-  test("Deeply nested dist", async () => {
-    const value = (
-      await testRun('{ x: 5, y: [3, "foo", { dist: normal(5,2) } ] }')
-    ).asJS();
-
-    expect((value as any).value.y[2].value.dist).toBeInstanceOf(Array);
-  });
-});
 
 describe("docstrings", () => {
   const runToResult = async (code: string) => {

packages/squiggle-lang/__tests__/helpers/helpers.ts

@@ -1,4 +1,4 @@
-import { run } from "../../src/index.js";
+import { run, SqValue } from "../../src/index.js";
 
 export async function testRun(x: string) {
   const outputR = await run(x, {
@@ -9,7 +9,7 @@ export async function testRun(x: string) {
   });
 
   if (outputR.ok) {
-    return outputR.value.result;
+    return outputR.value;
   } else {
     throw new Error(
       `Expected squiggle expression to evaluate but got error: ${outputR.value}`
@@ -39,3 +39,15 @@ export function expectErrorToBeBounded(
   const error = distance / normalizingDenom;
   expect(error).toBeLessThanOrEqual(epsilon);
 }
+
+export function assertTag<T extends SqValue["tag"]>(
+  value: SqValue | undefined,
+  tag: T
+): asserts value is Extract<SqValue, { tag: T }> {
+  if (!value) {
+    throw new Error("Undefined value");
+  }
+  if (value.tag !== tag) {
+    throw new Error(`Expected ${tag} value, got ${value.tag}`);
+  }
+}

packages/squiggle-lang/__tests__/library/mixture_test.ts

@@ -35,14 +35,14 @@ describe("mixture", () => {
         await testRun(
           "a = Sym.normal(0,1); m = mx(a, 3, [.999999,.00001]); stdev(a - m)"
         )
-      ).value
+      ).result.value
     ).toBeGreaterThan(1);
     expect(
       (
         await testRun(
           "a = normal(0,1); m = mx(a, 3, [.999999,.00001]); stdev(a - m)"
         )
-      ).value
+      ).result.value
     ).toBeCloseTo(0);
   });
 });

packages/squiggle-lang/__tests__/library/pointset_test.ts

@@ -35,16 +35,16 @@ describe("Mean of mixture is weighted average of means", () => {
         async (normalMean, normalStdev, betaA, betaB, x, y) => {
           // normaalize is due to https://github.com/quantified-uncertainty/squiggle/issues/1400 bug
           const squiggleString = `mean(mixture(Sym.normal(${normalMean},${normalStdev}), Sym.beta(${betaA},${betaB}), [${x}, ${y}])->normalize)`;
-          const res = await testRun(squiggleString);
+          const { result } = await testRun(squiggleString);
           const weightDenom = x + y;
           const normalWeight = x / weightDenom;
           const betaWeight = y / weightDenom;
           const betaMean = betaA / (betaA + betaB);
-          if (res.tag !== "Number") {
-            throw new Error(`Expected number result, got: ${res.tag}`);
+          if (result.tag !== "Number") {
+            throw new Error(`Expected number result, got: ${result.tag}`);
           }
           expectErrorToBeBounded(
-            res.value,
+            result.value,
             normalWeight * normalMean + betaWeight * betaMean,
             // this is a huge allowed error, but it's the highest precision we can achieve because of this bug: https://github.com/quantified-uncertainty/squiggle/issues/1414, even on relatively high \alpha and \beta values
             { epsilon: 0.7 }
@@ -65,11 +65,11 @@ describe("Discrete", () => {
 
   test("sample", async () => {
     for (let i = 0; i < 100; i++) {
-      const res = await testRun("mx(3,5) -> sample");
-      if (res.tag !== "Number") {
-        throw new Error(`Expected number result, got: ${res.tag}`);
+      const { result } = await testRun("mx(3,5) -> sample");
+      if (result.tag !== "Number") {
+        throw new Error(`Expected number result, got: ${result.tag}`);
       }
-      expect(res.value === 5 || res.value === 3).toBe(true);
+      expect(result.value === 5 || result.value === 3).toBe(true);
     }
   });
 });

packages/squiggle-lang/__tests__/library/sampleSet_test.ts

@@ -77,7 +77,7 @@ const arrayGen = () =>
 
 async function makeSampleSet(samples: number[]) {
   const sampleList = samples.map((x) => x.toFixed(20)).join(",");
-  const result = await testRun(`SampleSet.fromList([${sampleList}])`);
+  const { result } = await testRun(`SampleSet.fromList([${sampleList}])`);
   if (result.tag === "Dist") {
     return result.value;
   } else {
@@ -267,7 +267,7 @@ describe("fromSamples function", () => {
         const squiggleString = `x = fromSamples([${xsString}]); mean(x)`;
         const squiggleResult = await testRun(squiggleString);
         const mean = xs.reduce((a, b) => a + b, 0.0) / xs.length;
-        expect(squiggleResult.value).toBeCloseTo(mean, 4);
+        expect(squiggleResult.result.value).toBeCloseTo(mean, 4);
       })
     );
   });

packages/squiggle-lang/__tests__/library/scalars_test.ts

@@ -7,11 +7,11 @@ describe("Scalar manipulation is well-modeled by javascript math", () => {
     fc.assert(
       fc.asyncProperty(fc.nat(), async (x) => {
         const squiggleString = `log(${x})`;
-        const squiggleResult = await testRun(squiggleString);
+        const { result } = await testRun(squiggleString);
         if (x === 0) {
-          expect(squiggleResult.value).toEqual(-Infinity);
+          expect(result.value).toEqual(-Infinity);
         } else {
-          expect(squiggleResult.value).toEqual(Math.log(x));
+          expect(result.value).toEqual(Math.log(x));
         }
       })
     );
@@ -25,8 +25,8 @@ describe("Scalar manipulation is well-modeled by javascript math", () => {
         fc.integer(),
         async (x, y, z) => {
           const squiggleString = `x = ${x}; y = ${y}; z = ${z}; x + y + z`;
-          const squiggleResult = await testRun(squiggleString);
-          expect(squiggleResult.value).toBeCloseTo(x + y + z);
+          const { result } = await testRun(squiggleString);
+          expect(result.value).toBeCloseTo(x + y + z);
         }
       )
     );