Support parsing components with top level if statements (#4696)

* feat(editor) Support parsing components with top level if statements

* fix(parser-printer) Fix a bug when parsing if statements, plus tidying

* chore(parser-printer) Clarifying some code

* chore(editor) Removed stupid comment

* chore(tests) Adding a test for top level if statement parsing

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

@@ -46931,8 +46931,8 @@ exports[`UiJsxCanvas render renders fine with two components that reference each
       data-uid=\\"scene\\"
     >
       <div
-        data-uid=\\"aaa-unparsed-no-template-path\\"
-        data-path=\\"utopia-storyboard-uid/scene/app-entity:BBB\\"
+        data-uid=\\"aaa-root-true~~~1\\"
+        data-path=\\"utopia-storyboard-uid/scene/app-entity:BBB:bbb-root/bbb-root-false~~~1:aaa-root/aaa-root-false~~~1:bbb-root/bbb-root-false~~~1:aaa-root/aaa-root-false~~~1:bbb-root/bbb-root-false~~~1:aaa-root/aaa-root-true~~~1\\"
       >
         great
       </div>

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

@@ -1565,18 +1565,20 @@ import Utopia, {
 } from 'utopia-api'
 
 export var A = (props) => {
+  // @utopia/uid=aaa-root
   if (props.x === 0) {
-    return <div>great</div>
+    return <div data-uid='aaa-root-true'>great</div>
   } else {
-    return <B data-uid={'bbb-unparsed-no-template-path'} x={props.x - 1} />
+    return <B data-uid='aaa-root-false' x={props.x - 1} />
   }
 }
 
 export var B = (props) => {
+  // @utopia/uid=bbb-root
   if (props.x === 0) {
-    return <div>great</div>
+    return <div data-uid='bbb-root-true'>great</div>
   } else {
-    return <A data-uid={'aaa-unparsed-no-template-path'} x={props.x - 1} />
+    return <A data-uid='bbb-root-false' x={props.x - 1} />
   }
 }
 

editor/src/core/shared/element-template.ts

@@ -1323,6 +1323,20 @@ export type JSXElementChild =
   | JSXFragment
   | JSXConditionalExpression
 
+export function canBeRootElementOfComponent(element: JSXElementChild): boolean {
+  if (isJSXElement(element) || isJSXFragment(element) || isJSXConditionalExpression(element)) {
+    return true
+  }
+
+  if (isJSExpression(element)) {
+    if (hasElementsWithin(element)) {
+      return Object.keys(element.elementsWithin).length > 0
+    }
+  }
+
+  return false
+}
+
 export function isJSXElement(element: JSXElementChild): element is JSXElement {
   return element.type === 'JSX_ELEMENT'
 }

editor/src/core/workers/parser-printer/parser-printer-code-preservation.spec.ts

@@ -82,6 +82,7 @@ import { GithubPicker } from "react-color";
 function Picker() {
   const [color, setColor] = useThemeContext();
   const [visible, setVisible] = usePickerVisibilityContext();
+
   if (visible) {
     return <GithubPicker
       style={{ position: "absolute" }}

editor/src/core/workers/parser-printer/parser-printer-functional-components.spec.ts

@@ -167,6 +167,16 @@ export var whatever = (props) => {
 }
 `
 
+const codeWithComponentWithTopLevelIf = `import React from "react";
+import { View } from "utopia-api";
+export var whatever = (props) => {
+  if (props.showA) {
+    return <View data-uid={'aaa'} />
+  } else {
+    return <View data-uid={'bbb'} />
+  }
+}`
+
 describe('Parsing a function component with props', () => {
   it('Correctly parses a basic props object', () => {
     const actualResult = clearParseResultUniqueIDsAndEmptyBlocks(
@@ -903,6 +913,63 @@ describe('Parsing a function component with props', () => {
     )
     expect(actualResult).toEqual(expectedResult)
   })
+
+  it('Correctly parses a component with a top level if statement', () => {
+    const actualResult = clearParseResultUniqueIDsAndEmptyBlocks(
+      testParseCode(codeWithComponentWithTopLevelIf),
+    )
+    const viewA = clearJSXElementChildUniqueIDs(
+      jsxElement(
+        'View',
+        'aaa',
+        jsxAttributesFromMap({
+          'data-uid': jsExpressionValue('aaa', emptyComments),
+        }),
+        [],
+      ),
+    )
+    const viewB = clearJSXElementChildUniqueIDs(
+      jsxElement(
+        'View',
+        'bbb',
+        jsxAttributesFromMap({
+          'data-uid': jsExpressionValue('bbb', emptyComments),
+        }),
+        [],
+      ),
+    )
+    const exported = utopiaJSXComponent(
+      'whatever',
+      true,
+      'var',
+      'block',
+      defaultPropsParam,
+      ['showA'],
+      expect.objectContaining({
+        javascript: `if (props.showA) {\n    return <View data-uid={'aaa'} />\n  } else {\n    return <View data-uid={'bbb'} />\n  }`,
+        definedElsewhere: expect.arrayContaining(['props']),
+        elementsWithin: {
+          aaa: viewA,
+          bbb: viewB,
+        },
+      }),
+      expect.objectContaining({}),
+      false,
+      emptyComments,
+    )
+
+    const topLevelElements = [exported]
+    const expectedResult = parseSuccess(
+      JustImportViewAndReact,
+      expect.arrayContaining(topLevelElements),
+      expect.objectContaining({}),
+      null,
+      null,
+      [exportFunction('whatever')],
+      expect.objectContaining({}),
+    )
+    expect(actualResult).toEqual(expectedResult)
+  })
 })
 
 describe('Parsing, printing, reparsing a function component with props', () => {
@@ -996,4 +1063,8 @@ describe('Parsing, printing, reparsing a function component with props', () => {
   it('Correctly parses back and forth a component with a renamed function', () => {
     testParsePrintParse(codeWithARenamedFunction)
   })
+
+  it('Correctly parses back and forth a component with a top level if statement', () => {
+    testParsePrintParse(codeWithComponentWithTopLevelIf)
+  })
 })

editor/src/core/workers/parser-printer/parser-printer-parsing.ts

@@ -104,6 +104,7 @@ import {
   transpileJavascriptFromCode,
   transpileJavascript,
   insertDataUIDsIntoCode,
+  wrapAndTranspileJavascript,
 } from './parser-printer-transpiling'
 import * as PP from '../../shared/property-path'
 import type { ElementsWithinInPosition } from './parser-printer-utils'
@@ -824,6 +825,12 @@ function parseOtherJavaScript<E extends TS.Node, T extends { uid: string }>(
               addIfDefinedElsewhere(scope, node.condition, false)
               addIfDefinedElsewhere(scope, node.whenTrue, false)
               addIfDefinedElsewhere(scope, node.whenFalse, false)
+            } else if (TS.isIfStatement(node)) {
+              addIfDefinedElsewhere(scope, node.expression, false)
+              addIfDefinedElsewhere(scope, node.thenStatement, false)
+              if (node.elseStatement != null) {
+                addIfDefinedElsewhere(scope, node.elseStatement, false)
+              }
             } else if (TS.isDecorator(node)) {
               addIfDefinedElsewhere(scope, node.expression, false)
             } else if (TS.isDeleteExpression(node)) {
@@ -1053,15 +1060,16 @@ export function parseAttributeOtherJavaScript(
     (code, _, definedElsewhere, fileSourceNode, parsedElementsWithin, isList) => {
       const { code: codeFromFile, map } = fileSourceNode.toStringWithSourceMap({ file: filename })
       const rawMap = JSON.parse(map.toString())
-      const transpileEither = transpileJavascriptFromCode(
+
+      const transpileEither = wrapAndTranspileJavascript(
         sourceFile.fileName,
         sourceFile.text,
         codeFromFile,
         rawMap,
         parsedElementsWithin,
-        true,
         applySteganography,
       )
+
       return mapEither((transpileResult) => {
         const prependedWithReturn = prependToSourceString(
           sourceFile.fileName,
@@ -1183,13 +1191,12 @@ function parseJSExpression(
           sourceFile.fileName,
         )
         return flatMapEither((dataUIDFixResult) => {
-          const transpileEither = transpileJavascriptFromCode(
+          const transpileEither = wrapAndTranspileJavascript(
             sourceFile.fileName,
             sourceFile.text,
             dataUIDFixResult.code,
             dataUIDFixResult.sourceMap,
             parsedElementsWithin,
-            true,
             applySteganography,
           )
 
@@ -2889,7 +2896,6 @@ export function parseArbitraryNodes(
         // which resulted in the `data-uid` and `data-path` being created as if they were generated elements.
         // In those cases that is incorrect as they are just regularly used elements which were in a class component (for example).
         rootLevel ? [] : parsedElementsWithin,
-        false,
         applySteganography,
       )
       const dataUIDFixed = insertDataUIDsIntoCode(
@@ -3047,15 +3053,51 @@ export function parseOutFunctionContents(
         alreadyExistingUIDs,
         applySteganography,
       )
-      return mapEither((parsed) => {
-        highlightBounds = mergeHighlightBounds(highlightBounds, parsed.highlightBounds)
-        return withParserMetadata(
-          functionContents(jsBlock, 'block', parsed.value, returnStatementComments),
-          highlightBounds,
-          propsUsed.concat(parsed.propsUsed),
-          definedElsewhere.concat(parsed.definedElsewhere),
-        )
-      }, parsedElements)
+      return foldEither(
+        () => {
+          // If we aren't able to parse out the individual JSX elements (because they don't form part of a simple return statement)
+          // we attempt to parse the entire function body as arbitrary JS
+          const parsedAsArbitrary = parseAttributeOtherJavaScript(
+            sourceFile,
+            sourceText,
+            filename,
+            imports,
+            topLevelNames,
+            propsObjectName,
+            possibleElement,
+            highlightBounds,
+            alreadyExistingUIDs,
+            applySteganography,
+          )
+
+          return bimapEither(
+            (failure) => failure,
+            (success) => {
+              highlightBounds = mergeHighlightBounds(highlightBounds, success.highlightBounds)
+              const elem = successfullyParsedElement(sourceFile, possibleElement, success.value)
+              return withParserMetadata(
+                functionContents(jsBlock, 'block', [elem], returnStatementComments),
+                highlightBounds,
+                propsUsed.concat(success.propsUsed),
+                definedElsewhere.concat(success.definedElsewhere),
+              )
+            },
+            parsedAsArbitrary,
+          )
+        },
+        (parsed) => {
+          highlightBounds = mergeHighlightBounds(highlightBounds, parsed.highlightBounds)
+          return right(
+            withParserMetadata(
+              functionContents(jsBlock, 'block', parsed.value, returnStatementComments),
+              highlightBounds,
+              propsUsed.concat(parsed.propsUsed),
+              definedElsewhere.concat(parsed.definedElsewhere),
+            ),
+          )
+        },
+        parsedElements,
+      )
     }
   } else {
     const parsedElements = parseOutJSXElements(

editor/src/core/workers/parser-printer/parser-printer-transpiling.ts

@@ -5,7 +5,7 @@ import ReactSyntaxPlugin from 'babel-plugin-syntax-jsx'
 import ReactTransformPlugin from 'babel-plugin-transform-react-jsx'
 import type { SourceNode } from 'source-map'
 import type { Either } from '../../shared/either'
-import { left, right } from '../../shared/either'
+import { isRight, left, right } from '../../shared/either'
 import type { JSXElement } from '../../shared/element-template'
 import {
   getDefinedElsewhereFromElement,
@@ -16,7 +16,11 @@ import { fastForEach } from '../../shared/utils'
 import type { RawSourceMap } from '../ts/ts-typings/RawSourceMap'
 import infiniteLoopPrevention from './transform-prevent-infinite-loops'
 import type { ElementsWithinInPosition, CodeWithMap } from './parser-printer-utils'
-import { wrapCodeInParens, wrapCodeInParensWithMap } from './parser-printer-utils'
+import {
+  wrapCodeInAnonFunctionWithMap,
+  wrapCodeInParens,
+  wrapCodeInParensWithMap,
+} from './parser-printer-utils'
 import { JSX_CANVAS_LOOKUP_FUNCTION_NAME } from '../../shared/dom-utils'
 import type { SteganoTextData } from '../../shared/stegano-text'
 import { cleanSteganoTextData, encodeSteganoData } from '../../shared/stegano-text'
@@ -196,7 +200,6 @@ export function transpileJavascript(
   sourceFileText: string,
   fileSourceNode: typeof SourceNode,
   elementsWithin: ElementsWithinInPosition,
-  wrapInParens: boolean,
   applySteganography: SteganographyMode,
 ): Either<string, TranspileResult> {
   try {
@@ -208,7 +211,7 @@ export function transpileJavascript(
       code,
       rawMap,
       elementsWithin,
-      wrapInParens,
+      'do-not-wrap',
       applySteganography,
     )
   } catch (e: any) {
@@ -309,20 +312,53 @@ function applySteganographyPlugin(
   })
 }
 
+export function wrapAndTranspileJavascript(
+  sourceFileName: string,
+  sourceFileText: string,
+  code: string,
+  map: RawSourceMap,
+  elementsWithin: ElementsWithinInPosition,
+  applySteganography: SteganographyMode,
+): Either<string, TranspileResult> {
+  const wrappedInParensResult = transpileJavascriptFromCode(
+    sourceFileName,
+    sourceFileText,
+    code,
+    map,
+    elementsWithin,
+    'wrap-in-parens',
+    applySteganography,
+  )
+
+  if (isRight(wrappedInParensResult)) {
+    return wrappedInParensResult
+  } else {
+    return transpileJavascriptFromCode(
+      sourceFileName,
+      sourceFileText,
+      code,
+      map,
+      elementsWithin,
+      'wrap-in-anon-fn',
+      applySteganography,
+    )
+  }
+}
+
 export function transpileJavascriptFromCode(
   sourceFileName: string,
   sourceFileText: string,
   code: string,
   map: RawSourceMap,
   elementsWithin: ElementsWithinInPosition,
-  wrapInParens: boolean,
+  wrap: 'wrap-in-parens' | 'wrap-in-anon-fn' | 'do-not-wrap',
   applySteganography: SteganographyMode,
 ): Either<string, TranspileResult> {
   try {
     let codeToUse: string = code
     let mapToUse: RawSourceMap = map
 
-    if (wrapInParens) {
+    if (wrap === 'wrap-in-parens') {
       const wrappedInParens = wrapCodeInParensWithMap(
         sourceFileName,
         sourceFileText,
@@ -331,6 +367,15 @@ export function transpileJavascriptFromCode(
       )
       codeToUse = wrappedInParens.code
       mapToUse = wrappedInParens.sourceMap
+    } else if (wrap === 'wrap-in-anon-fn') {
+      const wrappedInAnonFunction = wrapCodeInAnonFunctionWithMap(
+        sourceFileName,
+        sourceFileText,
+        codeToUse,
+        mapToUse,
+      )
+      codeToUse = wrappedInAnonFunction.code
+      mapToUse = wrappedInAnonFunction.sourceMap
     }
 
     let plugins: Array<any> =