From 9ea0ee9729fa1f419254a480ed247186262b57e0 Mon Sep 17 00:00:00 2001
From: Balint Gabor <127662+gbalint@users.noreply.github.com>
Date: Fri, 12 Jul 2024 17:55:32 +0200
Subject: [PATCH] More draft fixes

---
 .../ui-jsx-canvas-component-renderer.tsx      | 74 ++++++++++++++++---
 1 file changed, 62 insertions(+), 12 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
index 0607baa26b0c..9486ff47d134 100644
--- 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
@@ -102,15 +102,29 @@ export function createComponentRendererComponent(params: {
     const instancePath: ElementPath | null = tryToGetInstancePath(instancePathAny, pathsString)
 
     function shouldUpdate() {
-      return (
-        ElementsToRerenderGLOBAL.current === 'rerender-all-elements' ||
-        ElementsToRerenderGLOBAL.current.some((er) => {
-          return (
-            (instancePath != null &&
-              (EP.pathsEqual(er, instancePath) || EP.isParentComponentOf(instancePath, er))) ||
-            isElementInChildrenOrPropsTree(EP.toString(er), realPassedProps)
-          )
-        })
+      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
+      }
+
+      if (ElementsToRerenderGLOBAL.current.length === 1) {
+        return isElementInChildrenOrPropsTreeSingle(
+          EP.toString(ElementsToRerenderGLOBAL.current[0]),
+          realPassedProps,
+        )
+      }
+
+      return isElementInChildrenOrPropsTreeMulti(
+        ElementsToRerenderGLOBAL.current.map(EP.toString),
+        realPassedProps,
       )
     }
 
@@ -379,7 +393,11 @@ function isRenderProp(prop: any): prop is { props: { [UTOPIA_PATH_KEY]: string }
   )
 }
 
-function isElementInChildrenOrPropsTree(elementPath: string, props: any): boolean {
+function isElementInChildrenOrPropsTreeSingle(elementPath: string, props: any): boolean {
+  if (props.children == null || typeof props.children === 'string') {
+    return false
+  }
+
   const childrenArr = fastReactChildrenToArray(props.children)
   for (let c of childrenArr) {
     if ((c.props as any)[UTOPIA_PATH_KEY] === elementPath) {
@@ -394,17 +412,49 @@ function isElementInChildrenOrPropsTree(elementPath: string, props: any): boolea
   }
 
   for (let c of childrenArr) {
-    if (isElementInChildrenOrPropsTree(elementPath, c.props)) {
+    if (isElementInChildrenOrPropsTreeSingle(elementPath, c.props)) {
+      return true
+    }
+  }
+
+  for (let p in props) {
+    if (isRenderProp(p) && isElementInChildrenOrPropsTreeSingle(elementPath, p.props)) {
+      return true
+    }
+  }
+
+  return false
+}
+function isElementInChildrenOrPropsTreeMulti(elementPaths: Array<string>, props: any): boolean {
+  if (props.children == null || typeof props.children === 'string') {
+    return false
+  }
+
+  const childrenArr = fastReactChildrenToArray(props.children)
+
+  for (let c of childrenArr) {
+    if (elementPaths.includes((c.props as any)[UTOPIA_PATH_KEY])) {
       return true
     }
   }
 
   for (let p in props) {
-    if (isRenderProp(p) && isElementInChildrenOrPropsTree(elementPath, p.props)) {
+    if (React.isValidElement(p) && elementPaths.includes((p.props as any)[UTOPIA_PATH_KEY])) {
       return true
     }
   }
 
+  for (let c of childrenArr) {
+    if (isElementInChildrenOrPropsTreeMulti(elementPaths, c.props)) {
+      return true
+    }
+  }
+
+  for (let p in props) {
+    if (isRenderProp(p) && isElementInChildrenOrPropsTreeMulti(elementPaths, p.props)) {
+      return true
+    }
+  }
   return false
 }