Reapply "Fix/use editor state inside canvas" (#6105)

This reverts commit 4271d2e.

editor/src/components/canvas/canvas-component-entry.tsx

@@ -104,6 +104,7 @@ const CanvasComponentEntryInner = React.memo((props: CanvasComponentEntryProps)
     </>
   )
 })
+CanvasComponentEntryInner.displayName = 'CanvasComponentEntryInner'
 
 function CanvasInner({
   canvasProps,

editor/src/components/canvas/remix/utopia-remix-root-component.tsx

@@ -5,12 +5,7 @@ import type { DataRouteObject, Location, RouteObject } from 'react-router'
 import { createMemoryRouter, RouterProvider } from 'react-router'
 import { UTOPIA_PATH_KEY } from '../../../core/model/utopia-constants'
 import type { ElementPath } from '../../../core/shared/project-file-types'
-import {
-  useRefEditorState,
-  useEditorState,
-  Substores,
-  useSelectorWithCallback,
-} from '../../editor/store/store-hook'
+import { Substores, useSelectorWithCallback } from '../../editor/store/store-hook'
 import * as EP from '../../../core/shared/element-path'
 import { PathPropHOC } from './path-props-hoc'
 import { atom, useAtom, useSetAtom } from 'jotai'
@@ -23,6 +18,11 @@ import { AlwaysFalse, usePubSubAtomReadOnly } from '../../../core/shared/atom-wi
 import { CreateRemixDerivedDataRefsGLOBAL } from '../../editor/store/remix-derived-data'
 import { patchRoutesWithContext } from '../../../third-party/remix/create-remix-stub'
 import type { AppLoadContext } from '@remix-run/server-runtime'
+import {
+  useAllowRerenderOnlyOnAllElements,
+  useCanvasState,
+  useRefCanvasState,
+} from '../ui-jsx-canvas-renderer/canvas-state-hooks'
 
 type RouteModule = RouteModules[keyof RouteModules]
 type RouterType = ReturnType<typeof createMemoryRouter>
@@ -77,19 +77,20 @@ export function useRemixNavigationContext(
 }
 
 function useGetRouteModules(basePath: ElementPath) {
-  const remixDerivedDataRef = useRefEditorState((store) => store.derived.remixData)
-  const projectContentsRef = useRefEditorState((store) => store.editor.projectContents)
-  const fileBlobsRef = useRefEditorState((store) => store.editor.canvas.base64Blobs)
-  const hiddenInstancesRef = useRefEditorState((store) => store.editor.hiddenInstances)
-  const displayNoneInstancesRef = useRefEditorState((store) => store.editor.displayNoneInstances)
+  const remixDerivedDataRef = useRefCanvasState((store) => store.derived.remixData)
+  const projectContentsRef = useRefCanvasState((store) => store.editor.projectContents)
+  const fileBlobsRef = useRefCanvasState((store) => store.editor.canvas.base64Blobs)
+  const hiddenInstancesRef = useRefCanvasState((store) => store.editor.hiddenInstances)
+  const displayNoneInstancesRef = useRefCanvasState((store) => store.editor.displayNoneInstances)
 
   let metadataContext: UiJsxCanvasContextData = forceNotNull(
     `Missing UiJsxCanvasCtxAtom provider`,
     usePubSubAtomReadOnly(UiJsxCanvasCtxAtom, AlwaysFalse),
   )
 
-  const defaultExports = useEditorState(
+  const defaultExports = useCanvasState(
     Substores.derived,
+    useAllowRerenderOnlyOnAllElements(),
     (store) => {
       const routeModuleCreators = store.derived.remixData?.routeModuleCreators ?? {}
       return Object.values(routeModuleCreators).map(
@@ -171,17 +172,18 @@ export const RouteExportsForRouteObject: Array<keyof RouteObject> = [
 function useGetRoutes(
   getLoadContext?: (request: Request) => Promise<AppLoadContext> | AppLoadContext,
 ) {
-  const routes = useEditorState(
+  const routes = useCanvasState(
     Substores.derived,
+    useAllowRerenderOnlyOnAllElements(),
     (store) => store.derived.remixData?.routes ?? [],
     'UtopiaRemixRootComponent routes',
   )
 
-  const remixDerivedDataRef = useRefEditorState((store) => store.derived.remixData)
-  const projectContentsRef = useRefEditorState((store) => store.editor.projectContents)
-  const fileBlobsRef = useRefEditorState((store) => store.editor.canvas.base64Blobs)
-  const hiddenInstancesRef = useRefEditorState((store) => store.editor.hiddenInstances)
-  const displayNoneInstancesRef = useRefEditorState((store) => store.editor.displayNoneInstances)
+  const remixDerivedDataRef = useRefCanvasState((store) => store.derived.remixData)
+  const projectContentsRef = useRefCanvasState((store) => store.editor.projectContents)
+  const fileBlobsRef = useRefCanvasState((store) => store.editor.canvas.base64Blobs)
+  const hiddenInstancesRef = useRefCanvasState((store) => store.editor.hiddenInstances)
+  const displayNoneInstancesRef = useRefCanvasState((store) => store.editor.displayNoneInstances)
 
   let metadataContext: UiJsxCanvasContextData = forceNotNull(
     `Missing UiJsxCanvasCtxAtom provider`,
@@ -290,7 +292,7 @@ export interface UtopiaRemixRootComponentProps {
 }
 
 export const UtopiaRemixRootComponent = (props: UtopiaRemixRootComponentProps) => {
-  const remixDerivedDataRef = useRefEditorState((store) => store.derived.remixData)
+  const remixDerivedDataRef = useRefCanvasState((store) => store.derived.remixData)
 
   const routes = useRoutesWithIdEquality(props.getLoadContext)
 

editor/src/components/canvas/ui-jsx-canvas-renderer/canvas-state-hooks.tsx

@@ -0,0 +1,60 @@
+import React from 'react'
+import { shallowEqual } from '../../../core/shared/equality-utils'
+import type { EditorStorePatched } from '../../editor/store/editor-state'
+import type { StateSelector, Substores } from '../../editor/store/store-hook'
+import {
+  CanvasStateContext,
+  useEditorState,
+  useRefEditorState,
+} from '../../editor/store/store-hook'
+import type { StoreKey } from '../../editor/store/store-hook-substore-types'
+import { ElementsToRerenderGLOBAL } from '../ui-jsx-canvas'
+
+export function useAllowRerenderOnlyOnAllElements() {
+  return React.useCallback(() => ElementsToRerenderGLOBAL.current === 'rerender-all-elements', [])
+}
+
+export function useWrappedSelectorToPreventRerender<T, U>(
+  selector: StateSelector<T, U>,
+  allowRerender: () => boolean,
+): StateSelector<T, U> {
+  const previousResultRef = React.useRef<U | null>(null)
+  return React.useCallback(
+    (state: T) => {
+      if (allowRerender() || previousResultRef.current == null) {
+        // if allowRerender() is false, we will not recompute the selected state, preventing a re-render of the owner component
+        previousResultRef.current = selector(state)
+      }
+      return previousResultRef.current
+    },
+    [selector, allowRerender],
+  )
+}
+
+export const useCanvasState = <K extends StoreKey, S extends (typeof Substores)[K], U>(
+  storeKey: S,
+  allowRerender: () => boolean,
+  selector: StateSelector<Parameters<S>[0], U>,
+  selectorName: string,
+  equalityFn: (oldSlice: U, newSlice: U) => boolean = shallowEqual,
+) => {
+  return useEditorState(
+    storeKey,
+    useWrappedSelectorToPreventRerender(selector, allowRerender),
+    selectorName,
+    equalityFn,
+    CanvasStateContext,
+  )
+}
+
+/**
+ * Like useCanvasState, but DOES NOT TRIGGER A RE-RENDER
+ *
+ * ONLY USE IT IF YOU ARE SURE ABOUT WHAT YOU ARE DOING
+ */
+export const useRefCanvasState = <U,>(
+  selector: StateSelector<EditorStorePatched, U>,
+  explainMe = false,
+): { readonly current: U } => {
+  return useRefEditorState(selector, explainMe, CanvasStateContext)
+}

editor/src/components/canvas/ui-jsx-canvas-renderer/remix-scene-component.tsx

@@ -59,3 +59,4 @@ export const RemixSceneComponent = React.memo((props: React.PropsWithChildren<Re
     </div>
   )
 })
+RemixSceneComponent.displayName = 'RemixSceneComponent'

editor/src/components/canvas/ui-jsx-canvas.tsx

@@ -89,7 +89,7 @@ import {
   getListOfEvaluatedFiles,
 } from '../../core/shared/code-exec-utils'
 import { forceNotNull } from '../../core/shared/optional-utils'
-import { useRefEditorState } from '../editor/store/store-hook'
+import { EditorStateContext, useRefEditorState } from '../editor/store/store-hook'
 import { matchRoutes } from 'react-router'
 import { useAtom } from 'jotai'
 import { RemixNavigationAtom } from './remix/utopia-remix-root-component'
@@ -578,23 +578,26 @@ export const UiJsxCanvas = React.memo<UiJsxCanvasPropsWithErrorCallback>((props)
         all: 'initial',
       }}
     >
-      <Helmet>{parse(linkTags)}</Helmet>
-      <ElementsToRerenderContext.Provider
-        value={useKeepReferenceEqualityIfPossible(ElementsToRerenderGLOBAL.current)} // TODO this should either be moved to EditorState or the context should be moved to the root level
-      >
-        <RerenderUtopiaCtxAtom.Provider value={rerenderUtopiaContextValue}>
-          <UtopiaProjectCtxAtom.Provider value={utopiaProjectContextValue}>
-            <CanvasContainer
-              validRootPaths={rootValidPathsArray}
-              canvasRootElementElementPath={storyboardRootElementPath}
-            >
-              <SceneLevelUtopiaCtxAtom.Provider value={sceneLevelUtopiaContextValue}>
-                {StoryboardRoot}
-              </SceneLevelUtopiaCtxAtom.Provider>
-            </CanvasContainer>
-          </UtopiaProjectCtxAtom.Provider>
-        </RerenderUtopiaCtxAtom.Provider>
-      </ElementsToRerenderContext.Provider>
+      {/* deliberately breaking useEditorState and useRefEditorState to enforce the usage of useCanvasState */}
+      <EditorStateContext.Provider value={null}>
+        <Helmet>{parse(linkTags)}</Helmet>
+        <ElementsToRerenderContext.Provider
+          value={useKeepReferenceEqualityIfPossible(ElementsToRerenderGLOBAL.current)} // TODO this should either be moved to EditorState or the context should be moved to the root level
+        >
+          <RerenderUtopiaCtxAtom.Provider value={rerenderUtopiaContextValue}>
+            <UtopiaProjectCtxAtom.Provider value={utopiaProjectContextValue}>
+              <CanvasContainer
+                validRootPaths={rootValidPathsArray}
+                canvasRootElementElementPath={storyboardRootElementPath}
+              >
+                <SceneLevelUtopiaCtxAtom.Provider value={sceneLevelUtopiaContextValue}>
+                  {StoryboardRoot}
+                </SceneLevelUtopiaCtxAtom.Provider>
+              </CanvasContainer>
+            </UtopiaProjectCtxAtom.Provider>
+          </RerenderUtopiaCtxAtom.Provider>
+        </ElementsToRerenderContext.Provider>
+      </EditorStateContext.Provider>
     </div>
   )
 })

editor/src/components/editor/store/store-hook.ts

@@ -128,9 +128,10 @@ export const useEditorState = <K extends StoreKey, S extends (typeof Substores)[
   selector: StateSelector<Parameters<S>[0], U>,
   selectorName: string,
   equalityFn: (oldSlice: U, newSlice: U) => boolean = shallowEqual,
+  storeContext: React.Context<UtopiaStoreAPI | null> = EditorStateContext,
 ): U => {
   const storeKey: K = storeKey_.name as K
-  const context = React.useContext(EditorStateContext)
+  const context = React.useContext(storeContext)
 
   const wrappedSelector = useWrapSelectorInPerformanceMeasureBlock(storeKey, selector, selectorName)
 
@@ -221,8 +222,9 @@ export const useSelectorWithCallback = <K extends StoreKey, S extends (typeof Su
 export const useRefEditorState = <U>(
   selector: StateSelector<EditorStorePatched, U>,
   explainMe = false,
+  storeContext: React.Context<UtopiaStoreAPI | null> = EditorStateContext,
 ): { readonly current: U } => {
-  const context = React.useContext(EditorStateContext)
+  const context = React.useContext(storeContext)
   if (context == null) {
     throw new Error('useStore is missing from editor context')
   }