Fix/memoized should update (#6098)

**Problem:** <img width="1402" alt="image" src="https://github.com/user-attachments/assets/b83471ac-1eee-4e11-be7f-584c3a441e84"> There's a `shouldUpdate` function in ComponentRendererComponent which is not cheap, and it is memoized, and it's called 8 times for each canvas update, even for components where shouldUpdate will eventually be false. **Fix:** Turning shouldUpdate into a memoized function, only recalculating the value when any of its inputs change: <img width="1044" alt="image" src="https://github.com/user-attachments/assets/df866148-825f-4ac0-9230-78daa698f7c0"> The time spent goes from ~1.5ms to sub 0.2ms. This is doubly relevant because I have a follow-up PR which introduces shouldUpdate-style callbacks to even more places.
concrete-utopia · Jul 19, 2024 · b45b071 · b45b071
1 parent 95d88b8
commit b45b071
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 45 deletions.
diff --git a/editor/src/components/canvas/ui-jsx-canvas-renderer/ui-jsx-canvas-component-renderer.tsx b/editor/src/components/canvas/ui-jsx-canvas-renderer/ui-jsx-canvas-component-renderer.tsx
@@ -19,13 +19,13 @@ import type {
   UiJsxCanvasContextData,
   VariableData,
 } from '../ui-jsx-canvas'
-import {
-  DomWalkerInvalidatePathsCtxAtom,
-  UiJsxCanvasCtxAtom,
-  ElementsToRerenderGLOBAL,
-} from '../ui-jsx-canvas'
+import { DomWalkerInvalidatePathsCtxAtom, UiJsxCanvasCtxAtom } from '../ui-jsx-canvas'
 import type { MutableUtopiaCtxRefData } from './ui-jsx-canvas-contexts'
-import { RerenderUtopiaCtxAtom, SceneLevelUtopiaCtxAtom } from './ui-jsx-canvas-contexts'
+import {
+  ElementsToRerenderContext,
+  RerenderUtopiaCtxAtom,
+  SceneLevelUtopiaCtxAtom,
+} from './ui-jsx-canvas-contexts'
 import { applyPropsParamToPassedProps } from './ui-jsx-canvas-props-utils'
 import { runBlockUpdatingScope } from './ui-jsx-canvas-scope-utils'
 import * as EP from '../../../core/shared/element-path'
@@ -50,6 +50,7 @@ import { mapArrayToDictionary } from '../../../core/shared/array-utils'
 import { assertNever } from '../../../core/shared/utils'
 import { addFakeSpyEntry } from './ui-jsx-canvas-spy-wrapper'
 import type { FilePathMappings } from '../../../core/model/project-file-utils'
+import type { ElementsToRerender } from '../../editor/store/editor-state'
 
 function tryToGetInstancePath(
   maybePath: ElementPath | null,
@@ -101,25 +102,10 @@ export function createComponentRendererComponent(params: {
 
     const instancePath: ElementPath | null = tryToGetInstancePath(instancePathAny, pathsString)
 
-    function shouldUpdate() {
-      if (ElementsToRerenderGLOBAL.current === 'rerender-all-elements') {
-        return true
-      }
-
-      if (
-        instancePath != null &&
-        ElementsToRerenderGLOBAL.current.some(
-          (er) => EP.pathsEqual(er, instancePath) || EP.isParentComponentOf(instancePath, er),
-        )
-      ) {
-        return true
-      }
-
-      return areElementsInChildrenOrPropsTree(
-        ElementsToRerenderGLOBAL.current.map(EP.toString),
-        realPassedProps,
-      )
-    }
+    const shouldUpdate = useAllowRerenderForPath(
+      instancePath ?? 'rerender-all-elements',
+      realPassedProps,
+    )
 
     const rerenderUtopiaContext = usePubSubAtomReadOnly(RerenderUtopiaCtxAtom, shouldUpdate)
 
@@ -286,6 +272,7 @@ export function createComponentRendererComponent(params: {
       )
       applyBlockReturnFunctions(scope)
 
+      // possibly problematic: this should ONLY run if shouldUpdate() is true
       const arbitraryBlockResult = runBlockUpdatingScope(
         rootElementPath,
         params.filePath,
@@ -427,3 +414,36 @@ function fastReactChildrenToArray(children: any) {
   }
   return React.Children.toArray(children).filter(React.isValidElement)
 }
+
+export function useAllowRerenderForPath(
+  elementPath: ElementPath | 'rerender-all-elements',
+  realPassedProps: any,
+) {
+  const elementsToRerender = React.useContext(ElementsToRerenderContext)
+  const shouldUpdate = React.useMemo(
+    () => shouldUpdateInner(elementPath, realPassedProps, elementsToRerender),
+    [elementPath, realPassedProps, elementsToRerender],
+  )
+  return React.useCallback(() => shouldUpdate, [shouldUpdate])
+}
+
+function shouldUpdateInner(
+  elementPath: ElementPath | 'rerender-all-elements',
+  realPassedProps: any,
+  elementsToRerender: ElementsToRerender,
+) {
+  if (elementsToRerender === 'rerender-all-elements') {
+    return true
+  }
+
+  if (
+    elementPath != 'rerender-all-elements' &&
+    elementsToRerender.some(
+      (er) => EP.pathsEqual(er, elementPath) || EP.isParentComponentOf(elementPath, er),
+    )
+  ) {
+    return true
+  }
+
+  return areElementsInChildrenOrPropsTree(elementsToRerender.map(EP.toString), realPassedProps)
+}
diff --git a/editor/src/components/canvas/ui-jsx-canvas-renderer/ui-jsx-canvas-contexts.tsx b/editor/src/components/canvas/ui-jsx-canvas-renderer/ui-jsx-canvas-contexts.tsx
@@ -1,12 +1,16 @@
-import type React from 'react'
+import React from 'react'
 import { emptySet } from '../../../core/shared/set-utils'
 import type { MapLike } from 'typescript'
 import { atomWithPubSub } from '../../../core/shared/atom-with-pub-sub'
 import type { Either } from '../../../core/shared/either'
 import { left } from '../../../core/shared/either'
 import type { ElementPath } from '../../../core/shared/project-file-types'
 import type { ProjectContentTreeRoot } from '../../assets'
-import type { TransientFilesState, UIFileBase64Blobs } from '../../editor/store/editor-state'
+import type {
+  ElementsToRerender,
+  TransientFilesState,
+  UIFileBase64Blobs,
+} from '../../editor/store/editor-state'
 import type { VariableData } from '../ui-jsx-canvas'
 import type { FilePathMappings } from '../../../core/model/project-file-utils'
 
@@ -78,3 +82,8 @@ export const SceneLevelUtopiaCtxAtom = atomWithPubSub<SceneLevelContextProps>({
     validPaths: new Set(),
   },
 })
+
+type ElementsToRerenderContextValue = ElementsToRerender
+
+export const ElementsToRerenderContext =
+  React.createContext<ElementsToRerenderContextValue>('rerender-all-elements')
diff --git a/editor/src/components/canvas/ui-jsx-canvas.tsx b/editor/src/components/canvas/ui-jsx-canvas.tsx
@@ -54,6 +54,7 @@ import parse from 'html-react-parser'
 import type { ComponentRendererComponent } from './ui-jsx-canvas-renderer/component-renderer-component'
 import type { MutableUtopiaCtxRefData } from './ui-jsx-canvas-renderer/ui-jsx-canvas-contexts'
 import {
+  ElementsToRerenderContext,
   RerenderUtopiaCtxAtom,
   SceneLevelUtopiaCtxAtom,
   UtopiaProjectCtxAtom,
@@ -559,18 +560,22 @@ export const UiJsxCanvas = React.memo<UiJsxCanvasPropsWithErrorCallback>((props)
       }}
     >
       <Helmet>{parse(linkTags)}</Helmet>
-      <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
+        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>
     </div>
   )
 })

diff --git a/editor/src/core/performance/__snapshots__/performance-regression-tests.spec.tsx.snap b/editor/src/core/performance/__snapshots__/performance-regression-tests.spec.tsx.snap
@@ -1273,8 +1273,9 @@ Array [
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//Symbol(react.forward_ref)(CanvasContainer)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
+  "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedExoticType(Symbol(react.provider))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//div",
-  "//div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
+  "/UtopiaSpiedExoticType(Symbol(react.provider))/Symbol(react.provider)/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/Symbol(react.forward_ref)(CanvasContainer)/div:data-testid='canvas-container'",
   "/Symbol(react.forward_ref)(CanvasContainer)/div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
@@ -1571,8 +1572,9 @@ Array [
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//Symbol(react.forward_ref)(CanvasContainer)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
+  "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedExoticType(Symbol(react.provider))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//div",
-  "//div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
+  "/UtopiaSpiedExoticType(Symbol(react.provider))/Symbol(react.provider)/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/Symbol(react.forward_ref)(CanvasContainer)/div:data-testid='canvas-container'",
   "/Symbol(react.forward_ref)(CanvasContainer)/div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
@@ -2413,8 +2415,9 @@ Array [
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//Symbol(react.forward_ref)(CanvasContainer)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
+  "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedExoticType(Symbol(react.provider))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//div",
-  "//div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
+  "/UtopiaSpiedExoticType(Symbol(react.provider))/Symbol(react.provider)/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/Symbol(react.forward_ref)(CanvasContainer)/div:data-testid='canvas-container'",
   "/Symbol(react.forward_ref)(CanvasContainer)/div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
@@ -2699,8 +2702,9 @@ Array [
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//Symbol(react.forward_ref)(CanvasContainer)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedFunctionComponent(Provider)",
+  "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//UtopiaSpiedExoticType(Symbol(react.provider))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))//div",
-  "//div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
+  "/UtopiaSpiedExoticType(Symbol(react.provider))/Symbol(react.provider)/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",
   "/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))/Symbol(react.forward_ref)(CanvasContainer)/div:data-testid='canvas-container'",
   "/Symbol(react.forward_ref)(CanvasContainer)/div/UtopiaSpiedFunctionComponent(Provider)/UtopiaSpiedExoticType(Symbol(react.fragment))",

diff --git a/editor/src/core/performance/performance-regression-tests.spec.tsx b/editor/src/core/performance/performance-regression-tests.spec.tsx
@@ -65,7 +65,7 @@ describe('React Render Count Tests -', () => {
 
     const renderCountAfter = renderResult.getNumberOfRenders()
     // if this breaks, GREAT NEWS but update the test please :)
-    expect(renderCountAfter - renderCountBefore).toMatchInlineSnapshot(`787`)
+    expect(renderCountAfter - renderCountBefore).toMatchInlineSnapshot(`789`)
     expect(renderResult.getRenderInfo()).toMatchSnapshot()
   })
 
@@ -127,7 +127,7 @@ describe('React Render Count Tests -', () => {
 
     const renderCountAfter = renderResult.getNumberOfRenders()
     // if this breaks, GREAT NEWS but update the test please :)
-    expect(renderCountAfter - renderCountBefore).toMatchInlineSnapshot(`1135`)
+    expect(renderCountAfter - renderCountBefore).toMatchInlineSnapshot(`1137`)
     expect(renderResult.getRenderInfo()).toMatchSnapshot()
   })