Merge pull request #3268 from quantified-uncertainty/profiler

Profiler

.changeset/fifty-queens-begin.md

@@ -0,0 +1,6 @@
+---
+"@quri/squiggle-lang": patch
+"@quri/squiggle-components": patch
+---
+
+Profiler mode

.changeset/silent-hats-scream.md

@@ -0,0 +1,6 @@
+---
+"@quri/squiggle-lang": patch
+"@quri/squiggle-components": patch
+---
+
+Normalized AST serialization - significantly reduces the produced bundle size, in some cases

.gitignore

@@ -9,3 +9,5 @@ node_modules
 /.vscode/settings.json
 .turbo
 *.tsbuildinfo
+
+*.tgz

packages/components/src/components/CodeEditor/profilerExtension.ts

@@ -0,0 +1,58 @@
+import { RangeSet, RangeSetBuilder } from "@codemirror/state";
+import { Decoration, EditorView } from "@codemirror/view";
+
+import { simulationField } from "./fields.js";
+
+const makeMark = (intensity: number) => {
+  const min = 100;
+  const color = 255 - min + min * (1 - intensity);
+  return Decoration.mark({
+    attributes: {
+      style: `background-color: rgb(255,${color},${color},50%)`,
+    },
+  });
+};
+
+export function profilerExtension() {
+  const empty = RangeSet.of<Decoration>([]);
+
+  const decorations = EditorView.decorations.compute(
+    ["doc", simulationField.field],
+    (state) => {
+      const simulation = state.field(simulationField.field);
+      if (!simulation) {
+        return empty;
+      }
+
+      if (simulation.code !== state.doc.toString()) {
+        return empty;
+      }
+
+      const { output } = simulation;
+      if (!output.ok) {
+        return empty;
+      }
+
+      const profile = output.value.raw.runOutput.profile;
+      if (!profile) {
+        return empty;
+      }
+
+      const times = profile.times;
+      const totalTimes = times.reduce((a, b) => a + b, 0);
+      const scaledTimes = times.map((t) => t / totalTimes);
+      const maxTime = Math.max(...scaledTimes);
+      const normalizedTimes = scaledTimes.map((t) => t / maxTime);
+
+      const builder = new RangeSetBuilder<Decoration>();
+      const docString = state.doc.toString();
+      for (let pos = 0; pos < docString.length; pos++) {
+        const mark = makeMark(normalizedTimes[pos]);
+        builder.add(pos, pos + 1, mark);
+      }
+      return builder.finish();
+    }
+  );
+
+  return [decorations];
+}

packages/components/src/components/CodeEditor/useSquiggleEditorView.ts

@@ -34,6 +34,7 @@ import {
 import { onFocusByPathField, simulationField } from "./fields.js";
 import { CodeEditorProps } from "./index.js";
 import { lightThemeHighlightingStyle } from "./languageSupport/highlightingStyle.js";
+import { profilerExtension } from "./profilerExtension.js";
 import { useErrorsExtension } from "./useErrorsExtension.js";
 import { useFormatSquiggleExtension } from "./useFormatSquiggleExtension.js";
 import { useGutterExtension } from "./useGutterExtension.js";
@@ -136,6 +137,7 @@ export function useSquiggleEditorExtensions(
     formatExtension,
     errorsExtension,
     tooltipsExtension,
+    profilerExtension(),
     squiggleLanguageExtension,
   ];
 

packages/components/src/components/PlaygroundSettings.tsx

@@ -1,4 +1,5 @@
 import { FC } from "react";
+import { useWatch } from "react-hook-form";
 import { z } from "zod";
 
 import { defaultRunnerName, RunnerName, SqScale } from "@quri/squiggle-lang";
@@ -15,12 +16,12 @@ import { SAMPLE_COUNT_MAX, SAMPLE_COUNT_MIN } from "../lib/constants.js";
 import { functionChartDefaults } from "../widgets/LambdaWidget/FunctionChart/utils.js";
 import { FormComment } from "./ui/FormComment.js";
 import { FormSection } from "./ui/FormSection.js";
-import { useWatch } from "react-hook-form";
 
 export const environmentSchema = z.object({
   sampleCount: z.number().int().gte(SAMPLE_COUNT_MIN).lte(SAMPLE_COUNT_MAX),
   xyPointLength: z.number().int().gte(10).lte(10000),
   seed: z.string(),
+  profile: z.boolean(),
 });
 
 const runnerSchema = z.union([
@@ -92,6 +93,7 @@ export const defaultPlaygroundSettings: PlaygroundSettings = {
     sampleCount: 1000,
     xyPointLength: 1000,
     seed: "default_seed",
+    profile: false,
   },
   runner: defaultRunnerName,
   functionChartSettings: {
@@ -226,6 +228,11 @@ export const RenderingSettingsForm: FC = () => {
         label="Coordinate Count (For PointSet Shapes)"
         description="When distributions are converted into PointSet shapes, we need to know how many coordinates to use."
       />
+      <CheckboxFormField<PlaygroundSettings>
+        name="environment.profile"
+        label="Performance Profiler"
+        description="When enabled, source code will be highlighted according to how much time was spent. Enabling this option will slow down the execution of the model."
+      />
       <SelectRunnerField />
     </div>
   );

packages/components/src/components/SquiggleEditor.tsx

@@ -72,6 +72,7 @@ export const SquiggleEditor: FC<SquiggleEditorProps> = ({
           showGutter={false}
           errors={errors}
           project={project}
+          simulation={simulation}
           sourceId={sourceId}
           onSubmit={runSimulation}
         />

packages/components/src/stories/SquigglePlayground.stories.tsx

@@ -397,3 +397,14 @@ export const Specification: Story = {
   `,
   },
 };
+
+export const WithProfiler: Story = {
+  args: {
+    defaultCode: `x = 5 // fast
+y = List.upTo(1, 1000) -> map({|v| v + 1 }) -> List.length // slow
+`,
+    environment: {
+      profile: true,
+    },
+  },
+};

packages/squiggle-lang/__tests__/SqProject/SqProject_imports_test.ts

@@ -39,7 +39,8 @@ describe("Imports", () => {
     });
 
     test("getImports", () => {
-      expect(project.getImports("main")).toEqual({
+      const imports = project.getImports("main");
+      expect(imports).toMatchObject({
         ok: true,
         value: [
           { type: "named", variable: "common", sourceId: "./common" },

packages/squiggle-lang/__tests__/ast/parse_test.ts

@@ -1,4 +1,5 @@
-import { ASTNode, parse } from "../../src/ast/parse.js";
+import { parse } from "../../src/ast/parse.js";
+import { ASTNode } from "../../src/ast/types.js";
 import {
   testEvalError,
   testEvalToBe,

packages/squiggle-lang/src/ast/operators.ts

@@ -0,0 +1,27 @@
+export const infixFunctions = {
+  "+": "add",
+  "-": "subtract",
+  "!=": "unequal",
+  ".-": "Dist.dotSubtract",
+  ".*": "Dist.dotMultiply",
+  "./": "Dist.dotDivide",
+  ".^": "Dist.dotPow",
+  ".+": "Dist.dotAdd",
+  "*": "multiply",
+  "/": "divide",
+  "&&": "and",
+  "^": "pow", // or xor
+  "<": "smaller",
+  "<=": "smallerEq",
+  "==": "equal",
+  ">": "larger",
+  ">=": "largerEq",
+  "||": "or",
+  to: "to",
+};
+
+export const unaryFunctions = {
+  "-": "unaryMinus",
+  "!": "not",
+  ".-": "unaryDotMinus",
+};

packages/squiggle-lang/src/ast/parse.ts

@@ -2,47 +2,23 @@ import { ICompileError } from "../errors/IError.js";
 import * as Result from "../utility/result.js";
 import { result } from "../utility/result.js";
 import { SExpr, SExprPrintOptions, sExprToString } from "../utility/sExpr.js";
-import { type ASTCommentNode, type ASTNode } from "./peggyHelpers.js";
 import {
   parse as peggyParse,
   SyntaxError as PeggySyntaxError,
 } from "./peggyParser.js";
-
-export { type ASTNode } from "./peggyHelpers.js";
-
-// Types copy-pasted from Peggy, but converted from interface to type.
-// We need a type because interfaces don't match JsonValue type that we use for serialization.
-
-/** Provides information pointing to a location within a source. */
-export type Location = {
-  /** Line in the parsed source (1-based). */
-  line: number;
-  /** Column in the parsed source (1-based). */
-  column: number;
-  /** Offset in the parsed source (0-based). */
-  offset: number;
-};
-
-/** The `start` and `end` position's of an object within the source. */
-export type LocationRange = {
-  /** Unlike in Peggy, this must be a string. */
-  source: string;
-  /** Position at the beginning of the expression. */
-  start: Location;
-  /** Position after the end of the expression. */
-  end: Location;
-};
+import {
+  AST,
+  type ASTCommentNode,
+  type ASTNode,
+  LocationRange,
+} from "./types.js";
 
 export type ParseError = {
   type: "SyntaxError";
   location: LocationRange;
   message: string;
 };
 
-export type AST = Extract<ASTNode, { type: "Program" }> & {
-  comments: ASTCommentNode[];
-};
-
 type ParseResult = result<AST, ICompileError>;
 
 export function parse(expr: string, source: string): ParseResult {