From 231748a5580ce16de4efd689d2227e96ce0b12e7 Mon Sep 17 00:00:00 2001
From: Brijesh Bittu <brijeshb42@gmail.com>
Date: Fri, 29 Mar 2024 17:04:09 +0530
Subject: [PATCH] [pigment-css][react] Fix sx prop transformation on Box

This current fix is very crude and does not handle the case when a user
might have their own wrapper component over Box and they'll be passing
sx prop through the component.

Also made `Box` as a default import instead of name import.

Fixes #41704
---
 .../src/app/material-ui/page.tsx              | 10 ++++---
 apps/pigment-css-vite-app/package.json        |  5 ++++
 apps/pigment-css-vite-app/src/main.tsx        |  2 +-
 .../src/pages/material-ui/index.tsx           | 20 +++++++------
 packages/pigment-css-react/src/Box.d.ts       |  2 +-
 packages/pigment-css-react/src/Box.jsx        |  4 ++-
 .../pigment-css-react/src/processors/sx.ts    | 29 +++++++++++++++++--
 .../pigment-css-react/tests/Box/box.test.tsx  | 13 +++++++++
 .../tests/Box/fixtures/box.input.js           | 20 +++++++++++++
 .../tests/Box/fixtures/box.output.css         |  9 ++++++
 .../tests/Box/fixtures/box.output.js          |  8 +++++
 11 files changed, 104 insertions(+), 18 deletions(-)
 create mode 100644 packages/pigment-css-react/tests/Box/box.test.tsx
 create mode 100644 packages/pigment-css-react/tests/Box/fixtures/box.input.js
 create mode 100644 packages/pigment-css-react/tests/Box/fixtures/box.output.css
 create mode 100644 packages/pigment-css-react/tests/Box/fixtures/box.output.js

diff --git a/apps/pigment-css-next-app/src/app/material-ui/page.tsx b/apps/pigment-css-next-app/src/app/material-ui/page.tsx
index 3a7c71710986c5..bcea6b4568eaef 100644
--- a/apps/pigment-css-next-app/src/app/material-ui/page.tsx
+++ b/apps/pigment-css-next-app/src/app/material-ui/page.tsx
@@ -4,6 +4,7 @@ import Link from 'next/link';
 import fs from 'fs/promises';
 import path from 'path';
 import { css } from '@pigment-css/react';
+import Box from '@pigment-css/react/Box';
 
 export default async function MaterialUIPage() {
   const rootPaths = await fs.readdir(path.join(process.cwd(), `src/app/material-ui`));
@@ -11,8 +12,9 @@ export default async function MaterialUIPage() {
     <div>
       <h1>Material UI Components</h1>
       <nav>
-        <ul
-          className={css({
+        <Box
+          as="ul"
+          sx={{
             margin: 0,
             marginBlock: '1rem',
             padding: 0,
@@ -20,7 +22,7 @@ export default async function MaterialUIPage() {
             display: 'flex',
             flexDirection: 'column',
             gap: '0.5rem',
-          })}
+          }}
         >
           {rootPaths
             .filter((item) => !item.match(/\.(js|tsx)$/))
@@ -37,7 +39,7 @@ export default async function MaterialUIPage() {
                 </Link>
               </li>
             ))}
-        </ul>
+        </Box>
       </nav>
     </div>
   );
diff --git a/apps/pigment-css-vite-app/package.json b/apps/pigment-css-vite-app/package.json
index 883667eb642e00..913eb80cd92bc8 100644
--- a/apps/pigment-css-vite-app/package.json
+++ b/apps/pigment-css-vite-app/package.json
@@ -42,6 +42,11 @@
         "dependsOn": [
           "^build"
         ]
+      },
+      "dev": {
+        "dependsOn": [
+          "^build"
+        ]
       }
     }
   }
diff --git a/apps/pigment-css-vite-app/src/main.tsx b/apps/pigment-css-vite-app/src/main.tsx
index 9760ba4a72aef3..15e708be89b719 100644
--- a/apps/pigment-css-vite-app/src/main.tsx
+++ b/apps/pigment-css-vite-app/src/main.tsx
@@ -5,7 +5,7 @@ import { ThemeProvider, createTheme } from '@mui/material/styles';
 import CircularProgress from '@mui/material/CircularProgress';
 import CssBaseline from '@mui/material/CssBaseline';
 import { css } from '@pigment-css/react';
-import { Box } from '@pigment-css/react/Box';
+import Box from '@pigment-css/react/Box';
 import { ErrorBoundary } from 'react-error-boundary';
 import routes from '~react-pages';
 import '@pigment-css/react/styles.css';
diff --git a/apps/pigment-css-vite-app/src/pages/material-ui/index.tsx b/apps/pigment-css-vite-app/src/pages/material-ui/index.tsx
index 4f78230508237e..9e7fadd4922b39 100644
--- a/apps/pigment-css-vite-app/src/pages/material-ui/index.tsx
+++ b/apps/pigment-css-vite-app/src/pages/material-ui/index.tsx
@@ -1,6 +1,6 @@
 import * as React from 'react';
 import { useLocation, matchRoutes, Link } from 'react-router-dom';
-import { css } from '@pigment-css/react';
+import Box from '@pigment-css/react/Box';
 import routes from '~react-pages';
 import Layout from '../../Layout';
 
@@ -17,8 +17,9 @@ export default function MaterialIndex() {
       <div>
         <h1>Material UI Components</h1>
         <nav>
-          <ul
-            className={css({
+          <Box
+            as="ul"
+            sx={{
               margin: 0,
               marginBlock: '1rem',
               padding: 0,
@@ -26,24 +27,25 @@ export default function MaterialIndex() {
               display: 'flex',
               flexDirection: 'column',
               gap: '0.5rem',
-            })}
+            }}
           >
             {childRoutes
               .filter((item) => !!item.path)
               .map((item) => (
                 <li key={item.path}>
-                  <Link
+                  <Box
+                    as={Link}
                     to={`/material-ui/${item.path}`}
-                    className={css({
+                    sx={{
                       textDecoration: 'underline',
                       fontSize: '17px',
-                    })}
+                    }}
                   >
                     {item.path}
-                  </Link>
+                  </Box>
                 </li>
               ))}
-          </ul>
+          </Box>
         </nav>
       </div>
     </Layout>
diff --git a/packages/pigment-css-react/src/Box.d.ts b/packages/pigment-css-react/src/Box.d.ts
index 8cfb27534aa6d4..0fe8b89db82c79 100644
--- a/packages/pigment-css-react/src/Box.d.ts
+++ b/packages/pigment-css-react/src/Box.d.ts
@@ -23,4 +23,4 @@ export interface PolymorphicComponent<BaseProps extends BaseDefaultProps>
 
 declare const Box: PolymorphicComponent<{}>;
 
-export { Box };
+export default Box;
diff --git a/packages/pigment-css-react/src/Box.jsx b/packages/pigment-css-react/src/Box.jsx
index 235cb67cf10210..de32424c204c3b 100644
--- a/packages/pigment-css-react/src/Box.jsx
+++ b/packages/pigment-css-react/src/Box.jsx
@@ -1,7 +1,7 @@
 /* eslint-disable react/prop-types */
 import * as React from 'react';
 
-export const Box = React.forwardRef(
+const Box = React.forwardRef(
   (
     {
       as = 'div',
@@ -50,3 +50,5 @@ export const Box = React.forwardRef(
     return <Component ref={ref} className={classes} style={styles} {...rest} />;
   },
 );
+
+export default Box;
diff --git a/packages/pigment-css-react/src/processors/sx.ts b/packages/pigment-css-react/src/processors/sx.ts
index b0293aa8c58f93..141c470ed8c1ae 100644
--- a/packages/pigment-css-react/src/processors/sx.ts
+++ b/packages/pigment-css-react/src/processors/sx.ts
@@ -11,6 +11,31 @@ import { processCssObject } from '../utils/processCssObject';
 import { cssFnValueToVariable } from '../utils/cssFnValueToVariable';
 import BaseProcessor from './base-processor';
 
+// @TODO: Maybe figure out a better way allow imports.
+const allowedSxTransformImports = [`${process.env.PACKAGE_NAME}/Box`];
+
+/**
+ * Specifically looks for whether the sx prop should be transformed or not.
+ * If it's a Pigment CSS styled component, the value of `argumentValue` will
+ * be a string className starting with "."
+ * In other cases, it explicitly checks if the import is allowed for sx transformation.
+ */
+function allowSxTransform(argumentExpression: ExpressionValue, argumentValue?: string) {
+  if (!argumentExpression) {
+    return false;
+  }
+  if (
+    argumentExpression.kind === ValueType.LAZY &&
+    argumentExpression.importedFrom?.some((i) => allowedSxTransformImports.includes(i))
+  ) {
+    return true;
+  }
+  if (argumentValue && argumentValue[0] === '.') {
+    return true;
+  }
+  return false;
+}
+
 export class SxProcessor extends BaseProcessor {
   sxArguments: ExpressionValue[] = [];
 
@@ -41,7 +66,7 @@ export class SxProcessor extends BaseProcessor {
       }
     }
 
-    if (!this.elementClassName || this.elementClassName[0] !== '.') {
+    if (!allowSxTransform(elementClassExpression, this.elementClassName)) {
       return;
     }
 
@@ -95,7 +120,7 @@ export class SxProcessor extends BaseProcessor {
   doRuntimeReplacement() {
     const t = this.astService;
     // do not replace if sx prop is not a Pigment styled component
-    if (!this.elementClassName) {
+    if (this.artifacts.length === 0) {
       return;
     }
     if (this.collectedVariables.length) {
diff --git a/packages/pigment-css-react/tests/Box/box.test.tsx b/packages/pigment-css-react/tests/Box/box.test.tsx
new file mode 100644
index 00000000000000..a5a33298b875b4
--- /dev/null
+++ b/packages/pigment-css-react/tests/Box/box.test.tsx
@@ -0,0 +1,13 @@
+import path from 'node:path';
+import { runTransformation, expect } from '../testUtils';
+
+describe('Pigment CSS - Box', () => {
+  it('should transform and render sx prop', async () => {
+    const { output, fixture } = await runTransformation(
+      path.join(__dirname, 'fixtures/box.input.js'),
+    );
+
+    expect(output.js).to.equal(fixture.js);
+    expect(output.css).to.equal(fixture.css);
+  });
+});
diff --git a/packages/pigment-css-react/tests/Box/fixtures/box.input.js b/packages/pigment-css-react/tests/Box/fixtures/box.input.js
new file mode 100644
index 00000000000000..b8548e7690dc31
--- /dev/null
+++ b/packages/pigment-css-react/tests/Box/fixtures/box.input.js
@@ -0,0 +1,20 @@
+import Box from '@pigment-css/react/Box';
+
+export function App() {
+  return (
+    <Box
+      as="ul"
+      sx={{
+        margin: 0,
+        marginBlock: '1rem',
+        padding: 0,
+        paddingLeft: '1.5rem',
+        display: 'flex',
+        flexDirection: 'column',
+        gap: '0.5rem',
+      }}
+    >
+      Hello Box
+    </Box>
+  );
+}
diff --git a/packages/pigment-css-react/tests/Box/fixtures/box.output.css b/packages/pigment-css-react/tests/Box/fixtures/box.output.css
new file mode 100644
index 00000000000000..a8ce65359e1e17
--- /dev/null
+++ b/packages/pigment-css-react/tests/Box/fixtures/box.output.css
@@ -0,0 +1,9 @@
+.bc1d15y {
+  margin: 0;
+  margin-block: 1rem;
+  padding: 0;
+  padding-left: 1.5rem;
+  display: flex;
+  flex-direction: column;
+  gap: 0.5rem;
+}
diff --git a/packages/pigment-css-react/tests/Box/fixtures/box.output.js b/packages/pigment-css-react/tests/Box/fixtures/box.output.js
new file mode 100644
index 00000000000000..a5deb916f15eff
--- /dev/null
+++ b/packages/pigment-css-react/tests/Box/fixtures/box.output.js
@@ -0,0 +1,8 @@
+import Box from '@pigment-css/react/Box';
+export function App() {
+  return (
+    <Box as="ul" sx={'bc1d15y'}>
+      Hello Box
+    </Box>
+  );
+}