Respond to PR comments

packages/components/README.md

@@ -20,7 +20,7 @@ Add to `App.js`:
 ```jsx
 import { SquiggleEditor } from "@quri/squiggle-components";
 <SquiggleEditor
-  squiggleString="x = beta($alpha, 10); x + $shift"
+  defaultCode="x = beta($alpha, 10); x + $shift"
   jsImports={{ alpha: 3, shift: 20 }}
 />;
 ```
@@ -50,7 +50,7 @@ export function DynamicSquiggleChart({ squiggleString }) {
   } else {
     return (
         <SquiggleChart
-          squiggleString={squiggleString}
+          defaultCode={squiggleString}
           width={445}
           height={200}
           showSummary={true}

packages/components/src/components/SquiggleEditor.tsx

@@ -1,10 +1,10 @@
-import React, { useState } from "react";
+import React from "react";
 import { CodeEditor } from "./CodeEditor";
 import { environment, bindings, jsImports } from "@quri/squiggle-lang";
 import { defaultImports, defaultBindings } from "@quri/squiggle-lang";
 import { SquiggleContainer } from "./SquiggleContainer";
 import { SquiggleChart, SquiggleChartProps } from "./SquiggleChart";
-import { useSquigglePartial } from "../lib/hooks";
+import { useSquigglePartial, useMaybeControlledValue } from "../lib/hooks";
 import { SquiggleErrorAlert } from "./SquiggleErrorAlert";
 
 const WrappedCodeEditor: React.FC<{
@@ -28,21 +28,16 @@ export type SquiggleEditorProps = SquiggleChartProps & {
 };
 
 export const SquiggleEditor: React.FC<SquiggleEditorProps> = (props) => {
-  let defaultCode = props.defaultCode ?? "";
-  const [uncontrolledCode, setCode] = useState(defaultCode);
-  let code = props.code ?? uncontrolledCode;
+  const [code, setCode] = useMaybeControlledValue({
+    value: props.code,
+    defaultValue: props.defaultCode,
+    onChange: props.onCodeChange,
+  });
 
   let chartProps = { ...props, code };
   return (
     <SquiggleContainer>
-      <WrappedCodeEditor
-        code={code}
-        setCode={(code) => {
-          if (props.onCodeChange) props.onCodeChange(code);
-
-          if (props.code === undefined) setCode(code);
-        }}
-      />
+      <WrappedCodeEditor code={code} setCode={setCode} />
       <SquiggleChart {...chartProps} />
     </SquiggleContainer>
   );
@@ -66,7 +61,7 @@ export interface SquigglePartialProps {
 }
 
 export const SquigglePartial: React.FC<SquigglePartialProps> = ({
-  code,
+  code: controlledCode,
   defaultCode = "",
   onChange,
   onCodeChange,
@@ -74,11 +69,14 @@ export const SquigglePartial: React.FC<SquigglePartialProps> = ({
   environment,
   jsImports = defaultImports,
 }: SquigglePartialProps) => {
-  const [uncontrolledCode, setCode] = useState(defaultCode);
-  let codeProp = code ?? uncontrolledCode;
+  const [code, setCode] = useMaybeControlledValue<string>({
+    value: controlledCode,
+    defaultValue: defaultCode,
+    onChange: onCodeChange,
+  });
 
   const result = useSquigglePartial({
-    code: codeProp,
+    code,
     bindings,
     environment,
     jsImports,
@@ -87,14 +85,7 @@ export const SquigglePartial: React.FC<SquigglePartialProps> = ({
 
   return (
     <SquiggleContainer>
-      <WrappedCodeEditor
-        code={codeProp}
-        setCode={(code) => {
-          if (onCodeChange) onCodeChange(code);
-
-          if (code === undefined) setCode(code);
-        }}
-      />
+      <WrappedCodeEditor code={code} setCode={setCode} />
       {result.tag !== "Ok" ? <SquiggleErrorAlert error={result.value} /> : null}
     </SquiggleContainer>
   );

packages/components/src/components/SquigglePlayground.tsx

@@ -1,6 +1,7 @@
 import React, { FC, Fragment, useState, useEffect } from "react";
 import { Path, useForm, UseFormRegister, useWatch } from "react-hook-form";
 import * as yup from "yup";
+import { useMaybeControlledValue } from "../lib/hooks";
 import { yupResolver } from "@hookform/resolvers/yup";
 import { Tab } from "@headlessui/react";
 import {
@@ -216,7 +217,11 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
   onSettingsChange,
   showEditor = true,
 }) => {
-  const [uncontrolledCode, setUncontrolledCode] = useState(defaultCode);
+  const [code, setCode] = useMaybeControlledValue({
+    value: controlledCode,
+    defaultValue: defaultCode,
+    onChange: onCodeChange,
+  });
   const [importString, setImportString] = useState("{}");
   const [imports, setImports] = useState({});
   const [importsAreValid, setImportsAreValid] = useState(true);
@@ -261,8 +266,6 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
     }
   };
 
-  const code = controlledCode ?? uncontrolledCode;
-
   const samplingSettings = (
     <div className="space-y-6 p-3 max-w-xl">
       <div>
@@ -434,14 +437,8 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
   const firstTab = vars.showEditor ? (
     <div className="border border-slate-200">
       <CodeEditor
-        value={code}
-        onChange={(newCode) => {
-          if (controlledCode === undefined) {
-            // uncontrolled mode
-            setUncontrolledCode(newCode);
-          }
-          onCodeChange?.(newCode);
-        }}
+        value={code ?? ""}
+        onChange={setCode}
         oneLine={false}
         showGutter={true}
         height={height - 1}

packages/components/src/lib/hooks.ts

@@ -5,7 +5,7 @@ import {
   run,
   runPartial,
 } from "@quri/squiggle-lang";
-import { useEffect, useMemo } from "react";
+import { useEffect, useMemo, useState } from "react";
 
 type SquiggleArgs<T extends ReturnType<typeof run | typeof runPartial>> = {
   code: string;
@@ -42,3 +42,23 @@ export const useSquigglePartial = (
 export const useSquiggle = (args: SquiggleArgs<ReturnType<typeof run>>) => {
   return useSquiggleAny(args, run);
 };
+
+type ControlledValueArgs<T> = {
+  value?: T;
+  defaultValue: T;
+  onChange?: (x: T) => void;
+};
+export function useMaybeControlledValue<T>(
+  args: ControlledValueArgs<T>
+): [T, (x: T) => void] {
+  let [uncontrolledValue, setUncontrolledValue] = useState(args.defaultValue);
+  let value = args.value ?? uncontrolledValue;
+  let onChange = (newValue: T) => {
+    if (args.value === undefined) {
+      // uncontrolled mode
+      setUncontrolledValue(newValue);
+    }
+    args.onChange?.(newValue);
+  };
+  return [value, onChange];
+}

packages/components/src/stories/SquiggleEditor.stories.mdx

@@ -21,6 +21,19 @@ the distribution.
   </Story>
 </Canvas>
 
+It's also possible to create a controlled version of the same component
+
+<Canvas>
+  <Story
+    name="Controlled"
+    args={{
+      code: "normal(5,2)",
+    }}
+  >
+    {Template.bind({})}
+  </Story>
+</Canvas>
+
 You can also name variables like so:
 
 <Canvas>

packages/components/src/stories/SquigglePlayground.stories.mdx

@@ -14,7 +14,7 @@ including sampling settings, in squiggle.
   <Story
     name="Normal"
     args={{
-      initialSquiggleString: "normal(5,2)",
+      defaultCode: "normal(5,2)",
       height: 800,
     }}
   >