Handle Some Function Wrapped Components (#5808)

- Added `functionWrapping` to `UtopiaJSXComponent`, which is an array of
`FunctionWrap` elements.
- Added utility functions to clear `FunctionWrap` values of unique IDs
and source maps.
- `looksLikeCanvasElements` has been modified to capture the functions
wrapping
  the components alongside identifying the parameters.
- `createExecutionScope` now handles the wrapping of components for
execution.
- `printUtopiaJSXComponent` now handles printing out the components with
  their function wrappers.
- Reordered the logic in `createExecutionScope` and extracted out the
logic for
adding a component to the scope. The new function is used for both
arbitrary
blocks to make components available to their scope and also for their
original
  purpose of adding the components otherwise when created.

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

@@ -12,7 +12,6 @@ type ExtendedSceneProps = SceneProps & { [UTOPIA_SCENE_ID_KEY]: string }
 export const SceneComponent = React.memo(
   (props: React.PropsWithChildren<ExtendedSceneProps>) => {
     const colorTheme = useColorTheme()
-    const canvasIsLive = false
     const updateInvalidatedPaths = usePubSubAtomReadOnly(
       DomWalkerInvalidatePathsCtxAtom,
       AlwaysTrue,

editor/src/components/canvas/ui-jsx-canvas-renderer/ui-jsx-canvas-element-renderer-utils.tsx

@@ -957,7 +957,7 @@ export function utopiaCanvasJSXLookup(
   }
 }
 
-function runJSExpression(
+export function runJSExpression(
   block: JSExpression,
   elementPath: ElementPath | null,
   currentScope: MapLike<any>,

editor/src/components/canvas/ui-jsx-canvas-renderer/ui-jsx-canvas-execution-scope.tsx

@@ -19,7 +19,12 @@ import {
   updateMutableUtopiaCtxRefWithNewProps,
   UtopiaProjectCtxAtom,
 } from './ui-jsx-canvas-contexts'
-import { createLookupRender, utopiaCanvasJSXLookup } from './ui-jsx-canvas-element-renderer-utils'
+import type { RenderContext } from './ui-jsx-canvas-element-renderer-utils'
+import {
+  createLookupRender,
+  runJSExpression,
+  utopiaCanvasJSXLookup,
+} from './ui-jsx-canvas-element-renderer-utils'
 import { runBlockUpdatingScope } from './ui-jsx-canvas-scope-utils'
 import * as EP from '../../../core/shared/element-path'
 import type {
@@ -70,16 +75,6 @@ export function createExecutionScope(
   topLevelJsxComponents: Map<string | null, UtopiaJSXComponent>
   requireResult: MapLike<any>
 } {
-  const filePathMappings = getFilePathMappings(projectContents)
-  if (!(filePath in topLevelComponentRendererComponents.current)) {
-    // we make sure that the ref has an entry for this filepath
-    topLevelComponentRendererComponents.current[filePath] = {}
-  }
-  let topLevelComponentRendererComponentsForFile =
-    topLevelComponentRendererComponents.current[filePath]
-
-  const fileBlobsForFile = defaultIfNull(emptyFileBlobs, fileBlobs[filePath])
-
   const { topLevelElements, imports, jsxFactoryFunction, combinedTopLevelArbitraryBlock } =
     getParseSuccessForFilePath(filePath, projectContents)
   const requireResult: MapLike<any> = importResultFromImports(filePath, imports, customRequire)
@@ -105,57 +100,97 @@ export function createExecutionScope(
     process: process,
     ...requireResult,
   }
+  const filePathMappings = getFilePathMappings(projectContents)
+  if (!(filePath in topLevelComponentRendererComponents.current)) {
+    // we make sure that the ref has an entry for this filepath
+    topLevelComponentRendererComponents.current[filePath] = {}
+  }
+  let topLevelComponentRendererComponentsForFile =
+    topLevelComponentRendererComponents.current[filePath]
+
+  function addComponentToScope(component: UtopiaJSXComponent): void {
+    const elementName = component.name ?? 'default'
+    topLevelJsxComponents.set(component.name, component)
+    if (elementName in topLevelComponentRendererComponentsForFile) {
+      executionScope[elementName] = topLevelComponentRendererComponentsForFile[elementName]
+    } else {
+      const baseComponent = createComponentRendererComponent({
+        topLevelElementName: component.name,
+        mutableContextRef: mutableContextRef,
+        filePath: filePath,
+      })
+      const wrappedComponent = component.functionWrapping.reduceRight((workingComponent, wrap) => {
+        switch (wrap.type) {
+          case 'SIMPLE_FUNCTION_WRAP':
+            const functionExpression = runJSExpression(
+              wrap.functionExpression,
+              null,
+              executionScope,
+              renderContext,
+              wrap.functionExpression.uid,
+              null,
+              null,
+            )
+            return functionExpression(workingComponent)
+        }
+      }, baseComponent)
+      topLevelComponentRendererComponentsForFile[elementName] = wrappedComponent
+      executionScope[elementName] = wrappedComponent
+    }
+  }
+
+  const fileBlobsForFile = defaultIfNull(emptyFileBlobs, fileBlobs[filePath])
+
   // TODO All of this is run on every interaction o_O
 
   let topLevelJsxComponents: Map<string | null, UtopiaJSXComponent> = new Map()
 
-  // Make sure there is something in scope for all of the top level components
-  fastForEach(topLevelElements, (topLevelElement) => {
-    if (isUtopiaJSXComponent(topLevelElement)) {
-      topLevelJsxComponents.set(topLevelElement.name, topLevelElement)
-      const elementName = topLevelElement.name ?? 'default'
-      if (!(elementName in topLevelComponentRendererComponentsForFile)) {
-        topLevelComponentRendererComponentsForFile[elementName] = createComponentRendererComponent({
-          topLevelElementName: topLevelElement.name,
-          mutableContextRef: mutableContextRef,
-          filePath: filePath,
-        })
-      }
-    }
-  })
-
-  executionScope = {
-    ...executionScope,
-    ...topLevelComponentRendererComponentsForFile,
+  const renderContext: RenderContext = {
+    rootScope: executionScope,
+    parentComponentInputProps: {},
+    requireResult: requireResult,
+    hiddenInstances: hiddenInstances,
+    displayNoneInstances: displayNoneInstances,
+    fileBlobs: fileBlobsForFile,
+    validPaths: new Set(),
+    reactChildren: undefined,
+    metadataContext: metadataContext,
+    updateInvalidatedPaths: updateInvalidatedPaths,
+    jsxFactoryFunctionName: jsxFactoryFunction,
+    shouldIncludeCanvasRootInTheSpy: shouldIncludeCanvasRootInTheSpy,
+    filePath: filePath,
+    imports: imports,
+    code: '',
+    highlightBounds: {},
+    editedText: editedText,
+    variablesInScope: {},
+    filePathMappings: filePathMappings,
   }
 
   let spiedVariablesInRoot: VariableData = {}
 
   // First make sure everything is in scope
   if (combinedTopLevelArbitraryBlock != null && openStoryboardFileNameKILLME != null) {
+    for (const definedElsewhereValue of combinedTopLevelArbitraryBlock.definedElsewhere) {
+      const componentFromThisFile = topLevelElements.find(
+        (topLevelElement): topLevelElement is UtopiaJSXComponent => {
+          return (
+            isUtopiaJSXComponent(topLevelElement) && topLevelElement.name === definedElsewhereValue
+          )
+        },
+      )
+      if (componentFromThisFile != null) {
+        addComponentToScope(componentFromThisFile)
+      }
+    }
     const { highlightBounds, code } = getCodeAndHighlightBoundsForFile(filePath, projectContents)
     const propertiesFromParams = propertiesExposedByParams(combinedTopLevelArbitraryBlock.params)
     const lookupRenderer = createLookupRender(
       EP.emptyElementPath,
       {
-        rootScope: executionScope,
-        parentComponentInputProps: {},
-        requireResult: requireResult,
-        hiddenInstances: hiddenInstances,
-        displayNoneInstances: displayNoneInstances,
-        fileBlobs: fileBlobsForFile,
-        validPaths: new Set(),
-        reactChildren: undefined,
-        metadataContext: metadataContext,
-        updateInvalidatedPaths: updateInvalidatedPaths,
-        jsxFactoryFunctionName: jsxFactoryFunction,
-        shouldIncludeCanvasRootInTheSpy: shouldIncludeCanvasRootInTheSpy,
-        filePath: openStoryboardFileNameKILLME,
-        imports: imports,
+        ...renderContext,
         code: code,
         highlightBounds: highlightBounds,
-        editedText: editedText,
-        variablesInScope: {},
         filePathMappings: filePathMappings,
       },
       null,
@@ -179,6 +214,14 @@ export function createExecutionScope(
     )
     spiedVariablesInRoot = spiedVariablesDeclaredWithinBlock
   }
+
+  // Make sure there is something in scope for all of the top level components
+  fastForEach(topLevelElements, (topLevelElement) => {
+    if (isUtopiaJSXComponent(topLevelElement)) {
+      addComponentToScope(topLevelElement)
+    }
+  })
+
   // WARNING: mutating the mutableContextRef
   updateMutableUtopiaCtxRefWithNewProps(mutableContextRef, {
     ...mutableContextRef.current,

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

@@ -156,6 +156,91 @@ export { ToBeDefaultExported as default }`,
     )
   })
 
+  it('renders a canvas with a component wrapped in React.memo', () => {
+    testCanvasRender(
+      null,
+      `
+      import * as React from 'react'
+      import { Storyboard, Scene } from 'utopia-api'
+      
+      export var App = React.memo((props) => {
+        return (
+          <div
+            style={{ ...props.style, backgroundColor: '#FFFFFF' }}
+            data-uid={'aaa'}
+          >
+            <span style={{position: 'absolute'}} data-uid={'bbb'}>hi</span>
+          </div>
+        )
+      })
+      export var ${BakedInStoryboardVariableName} = (props) => {
+        return (
+          <Storyboard data-uid={'${BakedInStoryboardUID}'}>
+            <Scene
+              style={{ position: 'absolute', height: 200, left: 59, width: 200, top: 79 }}
+              data-uid={'${TestSceneUID}'}
+            >
+              <App
+                data-uid='${TestAppUID}'
+                style={{ position: 'absolute', height: '100%', width: '100%' }}
+                title={'Hi there!'}
+              />
+            </Scene>
+          </Storyboard>
+        )
+      }
+      `,
+    )
+  })
+
+  it('renders a canvas with a component wrapped in a function that wraps the component with more content', () => {
+    testCanvasRender(
+      null,
+      `
+      import * as React from 'react'
+      import { Storyboard, Scene } from 'utopia-api'
+
+      function wrapComponent(Component) {
+        return () => {
+          return <div>
+            <div>
+              <span>wrapComponent</span>
+            </div>
+            <div><Component /></div>
+          </div>
+        }
+      }
+      
+      export var App = wrapComponent(() => {
+        return (
+          <div
+            style={{ backgroundColor: '#FFFFFF' }}
+            data-uid={'aaa'}
+          >
+            <span style={{position: 'absolute'}} data-uid={'bbb'}>hi</span>
+          </div>
+        )
+      })
+      export var ${BakedInStoryboardVariableName} = (props) => {
+        return (
+          <Storyboard data-uid={'${BakedInStoryboardUID}'}>
+            <Scene
+              style={{ position: 'absolute', height: 200, left: 59, width: 200, top: 79 }}
+              data-uid={'${TestSceneUID}'}
+            >
+              <App
+                data-uid='${TestAppUID}'
+                style={{ position: 'absolute', height: '100%', width: '100%' }}
+                title={'Hi there!'}
+              />
+            </Scene>
+          </Storyboard>
+        )
+      }
+      `,
+    )
+  })
+
   it('renders a canvas defined by a utopia storyboard component', () => {
     testCanvasRender(
       null,

editor/src/components/editor/actions/actions.spec.tsx

@@ -154,6 +154,7 @@ function storyboardComponent(numberOfScenes: number): UtopiaJSXComponent {
     false,
     'var',
     'block',
+    [],
     null,
     [],
     jsxElement(
@@ -186,6 +187,7 @@ const originalModel = deepFreeze(
         true,
         'var',
         'block',
+        [],
         defaultPropsParam,
         [],
         jsxElement(

editor/src/components/editor/store/store-deep-equality-instances-2.spec.ts

@@ -788,6 +788,7 @@ describe('UtopiaJSXComponentKeepDeepEquality', () => {
     isFunction: true,
     declarationSyntax: 'const',
     blockOrExpression: 'block',
+    functionWrapping: [],
     param: {
       type: 'PARAM',
       dotDotDotToken: false,
@@ -837,6 +838,7 @@ describe('UtopiaJSXComponentKeepDeepEquality', () => {
     isFunction: true,
     declarationSyntax: 'const',
     blockOrExpression: 'block',
+    functionWrapping: [],
     param: {
       type: 'PARAM',
       dotDotDotToken: false,
@@ -886,6 +888,7 @@ describe('UtopiaJSXComponentKeepDeepEquality', () => {
     isFunction: true,
     declarationSyntax: 'const',
     blockOrExpression: 'block',
+    functionWrapping: [],
     param: {
       type: 'PARAM',
       dotDotDotToken: false,

editor/src/components/editor/store/store-deep-equality-instances.ts

@@ -604,6 +604,8 @@ import type { OnlineState } from '../online-status'
 import { onlineState } from '../online-status'
 import type { NavigatorRow } from '../../navigator/navigator-row'
 import { condensedNavigatorRow, regularNavigatorRow } from '../../navigator/navigator-row'
+import type { SimpleFunctionWrap, FunctionWrap } from 'utopia-shared/src/types'
+import { simpleFunctionWrap, isSimpleFunctionWrap } from 'utopia-shared/src/types'
 
 export function ElementPropertyPathKeepDeepEquality(): KeepDeepEqualityCall<ElementPropertyPath> {
   return combine2EqualityCalls(
@@ -1570,8 +1572,29 @@ export function JSXElementChildKeepDeepEquality(): KeepDeepEqualityCall<JSXEleme
   }
 }
 
+export const SimpleFunctionWrapKeepDeepEquality: KeepDeepEqualityCall<SimpleFunctionWrap> =
+  combine1EqualityCall(
+    (wrap) => wrap.functionExpression,
+    JSExpressionKeepDeepEqualityCall,
+    simpleFunctionWrap,
+  )
+
+export const FunctionWrapKeepDeepEquality: KeepDeepEqualityCall<FunctionWrap> = (
+  oldValue,
+  newValue,
+) => {
+  switch (oldValue.type) {
+    case 'SIMPLE_FUNCTION_WRAP':
+      if (isSimpleFunctionWrap(newValue)) {
+        return SimpleFunctionWrapKeepDeepEquality(oldValue, newValue)
+      }
+      break
+  }
+  return keepDeepEqualityResult(newValue, false)
+}
+
 export const UtopiaJSXComponentKeepDeepEquality: KeepDeepEqualityCall<UtopiaJSXComponent> =
-  combine10EqualityCalls(
+  combine11EqualityCalls(
     (component) => component.name,
     createCallWithTripleEquals(),
     (component) => component.isFunction,
@@ -1580,6 +1603,8 @@ export const UtopiaJSXComponentKeepDeepEquality: KeepDeepEqualityCall<UtopiaJSXC
     createCallWithTripleEquals(),
     (component) => component.blockOrExpression,
     createCallWithTripleEquals(),
+    (component) => component.functionWrapping,
+    arrayDeepEquality(FunctionWrapKeepDeepEquality),
     (component) => component.param,
     nullableDeepEquality(ParamKeepDeepEquality()),
     (component) => component.propsUsed,

editor/src/core/model/__snapshots__/storyboard-utils.spec.ts.snap

@@ -12,6 +12,7 @@ Array [
     "arbitraryJSBlock": null,
     "blockOrExpression": "block",
     "declarationSyntax": "var",
+    "functionWrapping": Array [],
     "isFunction": false,
     "name": "storyboard",
     "param": null,

editor/src/core/model/element-template-utils.spec.tsx

@@ -192,6 +192,7 @@ describe('removeJSXElementChild', () => {
       true,
       'var',
       'block',
+      [],
       defaultPropsParam,
       [],
       jsxElement(
@@ -212,6 +213,7 @@ describe('removeJSXElementChild', () => {
       true,
       'var',
       'block',
+      [],
       defaultPropsParam,
       [],
       jsxElement(