Merge pull request #2979 from quantified-uncertainty/fix-viewerprovider

ViewerProvider should get viewerRef

packages/components/src/components/SquiggleOutputViewer/index.tsx

@@ -23,8 +23,10 @@ import { CodeEditorHandle } from "../CodeEditor/index.js";
 import { ErrorBoundary } from "../ErrorBoundary.js";
 import { PartialPlaygroundSettings } from "../PlaygroundSettings.js";
 import { SquiggleErrorAlert } from "../SquiggleErrorAlert.js";
-import { SquiggleViewerHandle } from "../SquiggleViewer/index.js";
-import { ViewerProvider } from "../SquiggleViewer/ViewerProvider.js";
+import {
+  SquiggleViewerHandle,
+  ViewerProvider,
+} from "../SquiggleViewer/ViewerProvider.js";
 import { Layout } from "./Layout.js";
 import { RenderingIndicator } from "./RenderingIndicator.js";
 
@@ -110,7 +112,11 @@ export const SquiggleOutputViewer = forwardRef<SquiggleViewerHandle, Props>(
     }
 
     return (
-      <ViewerProvider partialPlaygroundSettings={settings} editor={editor}>
+      <ViewerProvider
+        partialPlaygroundSettings={settings}
+        editor={editor}
+        ref={viewerRef}
+      >
         <Layout
           menu={
             <Dropdown

packages/components/src/components/SquigglePlayground/index.tsx

@@ -17,7 +17,7 @@ import {
   type PlaygroundSettings,
 } from "../PlaygroundSettings.js";
 import { SquiggleOutputViewer } from "../SquiggleOutputViewer/index.js";
-import { SquiggleViewerHandle } from "../SquiggleViewer/index.js";
+import { SquiggleViewerHandle } from "../SquiggleViewer/ViewerProvider.js";
 import {
   LeftPlaygroundPanel,
   LeftPlaygroundPanelHandle,

packages/components/src/components/SquiggleViewer/ViewerProvider.tsx

@@ -23,13 +23,16 @@ import {
   PartialPlaygroundSettings,
   PlaygroundSettings,
 } from "../PlaygroundSettings.js";
-import { SquiggleViewerHandle } from "./index.js";
 import {
   getChildrenValues,
   pathAsString,
   shouldBeginCollapsed,
 } from "./utils.js";
 
+export type SquiggleViewerHandle = {
+  viewValuePath(path: SqValuePath): void;
+};
+
 type ItemHandle = {
   element: HTMLDivElement;
   forceUpdate: () => void;
@@ -163,6 +166,7 @@ type ViewerContextShape = {
   editor?: CodeEditorHandle;
   itemStore: ItemStore;
   initialized: boolean;
+  handle: SquiggleViewerHandle;
 };
 
 export const ViewerContext = createContext<ViewerContextShape>({
@@ -171,6 +175,9 @@ export const ViewerContext = createContext<ViewerContextShape>({
   setFocused: () => undefined,
   editor: undefined,
   itemStore: new ItemStore(),
+  handle: {
+    viewValuePath: () => {},
+  },
   initialized: false,
 });
 
@@ -317,11 +324,13 @@ export const InnerViewerProvider = forwardRef<SquiggleViewerHandle, Props>(
       unstablePlaygroundSettings
     );
 
-    useImperativeHandle(ref, () => ({
+    const handle: SquiggleViewerHandle = {
       viewValuePath(path: SqValuePath) {
         itemStore.scrollToPath(path);
       },
-    }));
+    };
+
+    useImperativeHandle(ref, () => handle);
 
     const [focused, setFocused] = useState<SqValuePath | undefined>();
 
@@ -337,6 +346,7 @@ export const InnerViewerProvider = forwardRef<SquiggleViewerHandle, Props>(
           focused,
           setFocused,
           itemStore,
+          handle,
           initialized: true,
         }}
       >
@@ -347,15 +357,24 @@ export const InnerViewerProvider = forwardRef<SquiggleViewerHandle, Props>(
 );
 InnerViewerProvider.displayName = "InnerViewerProvider";
 
+const ProxyViewerProvider = forwardRef<SquiggleViewerHandle, Props>(
+  (props, ref) => {
+    const { handle } = useViewerContext();
+    useImperativeHandle(ref, () => handle);
+    return props.children; // TODO - props.settings will be ignored, what should we do?
+  }
+);
+ProxyViewerProvider.displayName = "ProxyViewerProvider";
+
 export const ViewerProvider = forwardRef<SquiggleViewerHandle, Props>(
   (props, ref) => {
     // `ViewerProvider` is a singleton, so if the context already exists, we don't initialize it again
     const { initialized } = useContext(ViewerContext);
     if (initialized) {
-      return props.children; // TODO: `ref` and settings will be ignored
+      return <ProxyViewerProvider ref={ref} {...props} />;
+    } else {
+      return <InnerViewerProvider ref={ref} {...props} />;
     }
-
-    return <InnerViewerProvider ref={ref} {...props} />;
   }
 );
 ViewerProvider.displayName = "ViewerProvider";

packages/components/src/components/SquiggleViewer/index.tsx

@@ -9,6 +9,7 @@ import { PartialPlaygroundSettings } from "../PlaygroundSettings.js";
 import { pathIsEqual, pathItemFormat, useGetSubvalueByPath } from "./utils.js";
 import { ValueViewer } from "./ValueViewer.js";
 import {
+  SquiggleViewerHandle,
   useFocus,
   useUnfocus,
   useViewerContext,
@@ -68,10 +69,6 @@ const FocusedNavigation: FC<{
   );
 };
 
-export type SquiggleViewerHandle = {
-  viewValuePath(path: SqValuePath): void;
-};
-
 export type SquiggleViewerProps = {
   value: SqValue;
   editor?: CodeEditorHandle;