fix: removeCSS not work with container (#211)

* fix: `removeCSS` not work with `container`

(cherry picked from commit 2fb097e)

* test: add unit test

* fixup! fix: `removeCSS` not work with `container`

* ⚡️improve perf

* Revert "fixup! fix: `removeCSS` not work with `container`"

This reverts commit bb54111.

* fixup! fix: `removeCSS` not work with `container`

---------

Co-authored-by: Derek <[email protected]>

Author: Wxh16144

src/hooks/useCSSVarRegister.ts

@@ -58,7 +58,7 @@ const useCSSVarRegister = <V, T extends Record<string, V>>(
     },
     ([, , styleId]) => {
       if (isClientSide) {
-        removeCSS(styleId, { mark: ATTR_MARK });
+        removeCSS(styleId, { mark: ATTR_MARK, attachTo: container });
       }
     },
     ([, cssVarsStr, styleId]) => {

src/hooks/useCacheToken.tsx

@@ -92,15 +92,13 @@ const TOKEN_THRESHOLD = 0;
 function cleanTokenStyle(tokenKey: string, instanceId: string) {
   tokenKeys.set(tokenKey, (tokenKeys.get(tokenKey) || 0) - 1);
 
-  const tokenKeyList = Array.from(tokenKeys.keys());
-  const cleanableKeyList = tokenKeyList.filter((key) => {
-    const count = tokenKeys.get(key) || 0;
-
-    return count <= 0;
-  });
+  const cleanableKeyList = new Set<string>();
+  tokenKeys.forEach((value, key) => {
+    if (value <= 0) cleanableKeyList.add(key);
+  })
 
   // Should keep tokens under threshold for not to insert style too often
-  if (tokenKeyList.length - cleanableKeyList.length > TOKEN_THRESHOLD) {
+  if (tokenKeys.size - cleanableKeyList.size > TOKEN_THRESHOLD) {
     cleanableKeyList.forEach((key) => {
       removeStyleTags(key, instanceId);
       tokenKeys.delete(key);

src/hooks/useStyleRegister.tsx

@@ -333,7 +333,6 @@ export const parseStyle = (
   if (!root) {
     styleStr = `{${styleStr}}`;
   } else if (layer) {
-
     // fixme: https://github.com/thysultan/stylis/pull/339
     if (styleStr) {
       styleStr = `@layer ${layer.name} {${styleStr}}`;
@@ -462,7 +461,7 @@ export default function useStyleRegister(
       // Remove cache if no need
       ([, , styleId], fromHMR) => {
         if ((fromHMR || autoClear) && isClientSide) {
-          removeCSS(styleId, { mark: ATTR_MARK });
+          removeCSS(styleId, { mark: ATTR_MARK, attachTo: container });
         }
       },
 

tests/index.spec.tsx

@@ -1,7 +1,8 @@
 import { render } from '@testing-library/react';
 import classNames from 'classnames';
+import type { ReactElement, ReactNode } from 'react';
 import * as React from 'react';
-import { ReactElement, ReactNode, StrictMode } from 'react';
+import { StrictMode } from 'react';
 import { describe, expect } from 'vitest';
 import type { CSSInterpolation, DerivativeFunc } from '../src';
 import {
@@ -426,6 +427,45 @@ describe('csssinjs', () => {
     expect(container.querySelectorAll('style')).toHaveLength(1);
   });
 
+  // https://github.com/ant-design/cssinjs/issues/189
+  it('should cleanup style when unmount', () => {
+    const container = document.createElement('div');
+
+    const CssVarBox = () => {
+      const [token] = useCacheToken<DerivativeToken>(
+        theme,
+        [{ primaryColor: 'red' }],
+        {
+          salt: 'test',
+        },
+      );
+
+      useCSSVarRegister(
+        {
+          key: 'color-2',
+          path: ['cssinjs-cleanup-style-when-unmount'],
+          token,
+        },
+        () => ({
+          token: token.primaryColor,
+        }),
+      );
+
+      return null;
+    };
+
+    const { unmount } = render(
+      <StyleProvider cache={createCache()} container={container} autoClear>
+        <Box />
+        <CssVarBox />
+      </StyleProvider>,
+    );
+
+    expect(container.querySelectorAll('style')).toHaveLength(2);
+    unmount();
+    expect(container.querySelectorAll('style')).toHaveLength(0);
+  });
+
   describe('nonce', () => {
     function test(name: string, nonce: string | (() => string)) {
       it(name, () => {