From a5f2b38511ca9f938a4ae00f32ad8f797ca8dc65 Mon Sep 17 00:00:00 2001
From: siriwatknp <siriwatkunaporn@gmail.com>
Date: Tue, 30 Apr 2024 22:23:17 +0700
Subject: [PATCH] fix reassign forwardProp

---
 packages/pigment-css-react/src/styled.js           | 11 +++++++++++
 .../tests/styled/runtime-styled.test.js            | 14 ++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/packages/pigment-css-react/src/styled.js b/packages/pigment-css-react/src/styled.js
index 301a31d5..8ec2edf2 100644
--- a/packages/pigment-css-react/src/styled.js
+++ b/packages/pigment-css-react/src/styled.js
@@ -123,6 +123,17 @@ export default function styled(tag, componentMeta = {}) {
         getVariantClasses(inProps, variants),
       );
 
+      if (shouldUseAs && !shouldForwardProp) {
+        // Reassign `shouldForwardProp` based on the incoming `as` prop
+        if (isHtmlTag(Component)) {
+          finalShouldForwardProp = isPropValid;
+        } else if (slot === 'Root' || slot === 'root') {
+          finalShouldForwardProp = rootShouldForwardProp;
+        } else {
+          finalShouldForwardProp = slotShouldForwardProp;
+        }
+      }
+
       const newProps = {};
       // eslint-disable-next-line no-restricted-syntax
       for (const key in props) {
diff --git a/packages/pigment-css-react/tests/styled/runtime-styled.test.js b/packages/pigment-css-react/tests/styled/runtime-styled.test.js
index d00b1d04..a0a32538 100644
--- a/packages/pigment-css-react/tests/styled/runtime-styled.test.js
+++ b/packages/pigment-css-react/tests/styled/runtime-styled.test.js
@@ -252,6 +252,20 @@ describe('props filtering', () => {
       expect(container.firstChild).to.have.style('--foo', '300px');
     });
 
+    it('use component forward prop if provided `as` is a component', () => {
+      const StyledDiv = styled('div')({
+        classes: ['root'],
+      });
+
+      function Component({ TagComponent = 'span', ...props }) {
+        return <TagComponent {...props} />;
+      }
+
+      const { container } = render(<StyledDiv as={Component} TagComponent="button" disabled />);
+
+      expect(container.firstChild).to.have.tagName('button');
+    });
+
     it('should forward `as` prop', () => {
       // The components below is a simplified version of the `NativeSelect` component from Material UI.