From c925f36c45ea8829e18abaaa6fc193ac3da6c8d4 Mon Sep 17 00:00:00 2001
From: Brijesh Bittu <brijeshb42@gmail.com>
Date: Thu, 24 Oct 2024 16:06:06 +0530
Subject: [PATCH] [react] Fix scoping issue with sx dynamic value transform
 (#286)

Co-authored-by: Brijesh Bittu <brijesh42@gmail.com>
---
 .../pigment-css-react/src/processors/sx.ts    | 10 +++++++
 .../src/utils/cssFnValueToVariable.ts         |  2 --
 .../src/utils/sxObjectExtractor.ts            | 27 ++++++++++++++-----
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/packages/pigment-css-react/src/processors/sx.ts b/packages/pigment-css-react/src/processors/sx.ts
index d5a81273..32fb3c40 100644
--- a/packages/pigment-css-react/src/processors/sx.ts
+++ b/packages/pigment-css-react/src/processors/sx.ts
@@ -99,6 +99,16 @@ export class SxProcessor extends BaseProcessor {
         ([variableId, expression, isUnitLess]) => {
           switch (expression.type) {
             case 'ArrowFunctionExpression': {
+              const fnBody = expression.body as Expression;
+              // try to deserialize AST if possible
+              if (fnBody.type === 'StringLiteral') {
+                try {
+                  const body = JSON.parse(fnBody.value);
+                  return t.objectProperty(t.stringLiteral(variableId), body);
+                } catch (ex) {
+                  // proceed as usual
+                }
+              }
               return t.objectProperty(
                 t.stringLiteral(variableId),
                 t.arrayExpression([expression.body as Expression, t.booleanLiteral(isUnitLess)]),
diff --git a/packages/pigment-css-react/src/utils/cssFnValueToVariable.ts b/packages/pigment-css-react/src/utils/cssFnValueToVariable.ts
index 459b2757..6befb70e 100644
--- a/packages/pigment-css-react/src/utils/cssFnValueToVariable.ts
+++ b/packages/pigment-css-react/src/utils/cssFnValueToVariable.ts
@@ -100,7 +100,6 @@ function parseAndWrapExpression(
 }
 
 function transformThemeKeysInFn(
-  styleKey: string,
   functionString: string,
   options: PluginCustomOptions,
   filename?: string,
@@ -182,7 +181,6 @@ function iterateAndReplaceFunctions(
     try {
       const fnString = value.toString();
       const expression = transformThemeKeysInFn(
-        key,
         fnString,
         options,
         filename,
diff --git a/packages/pigment-css-react/src/utils/sxObjectExtractor.ts b/packages/pigment-css-react/src/utils/sxObjectExtractor.ts
index 01cc360f..37521417 100644
--- a/packages/pigment-css-react/src/utils/sxObjectExtractor.ts
+++ b/packages/pigment-css-react/src/utils/sxObjectExtractor.ts
@@ -1,5 +1,11 @@
 import type { NodePath } from '@babel/core';
-import { arrowFunctionExpression, cloneNode } from '@babel/types';
+import {
+  arrayExpression,
+  arrowFunctionExpression,
+  booleanLiteral,
+  cloneNode,
+  stringLiteral,
+} from '@babel/types';
 import type {
   ArrowFunctionExpression,
   Expression,
@@ -9,6 +15,7 @@ import type {
 } from '@babel/types';
 import { findIdentifiers } from '@wyw-in-js/transform';
 import { isStaticObjectOrArrayExpression } from './checkStaticObjectOrArray';
+import { isUnitLess } from './isUnitLess';
 
 function validateObjectKey(
   keyPath: NodePath<PrivateName | Expression>,
@@ -66,7 +73,8 @@ function traverseObjectExpression(
   const properties = nodePath.get('properties');
   properties.forEach((property) => {
     if (property.isObjectProperty()) {
-      validateObjectKey(property.get('key'), parentCall);
+      const key = property.get('key');
+      validateObjectKey(key, parentCall);
 
       const value = property.get('value');
       if (!value.isExpression()) {
@@ -93,10 +101,17 @@ function traverseObjectExpression(
           localIdentifiers.push(id);
         });
         if (localIdentifiers.length) {
-          const arrowFn = arrowFunctionExpression(
-            localIdentifiers.map((i) => i.node),
-            cloneNode(value.node),
-          );
+          let cssKey = '';
+          if (key.isIdentifier()) {
+            cssKey = key.node.name;
+          } else if (key.isStringLiteral()) {
+            cssKey = key.node.value;
+          }
+          const unitLess = isUnitLess(cssKey);
+          const fnBody = arrayExpression([cloneNode(value.node), booleanLiteral(unitLess)]);
+          // Serialize the actual AST as a string
+          // which then gets deserialized in sx.ts
+          const arrowFn = arrowFunctionExpression([], stringLiteral(JSON.stringify(fnBody)));
           value.replaceWith(arrowFn);
         }
       }