fix @​apply in plugins in jit mode

Author: bradlc

packages/tailwindcss-language-service/src/util/jit.ts

@@ -13,20 +13,16 @@ export function generateRules(state: State, classNames: string[]): { root: Root;
     .module(new Set(classNames), state.jitContext)
     .sort(([a], [z]) => bigSign(a - z))
 
-  let actualRules: Rule[] = []
+  let root = state.modules.postcss.module.root({ nodes: rules.map(([, rule]) => rule) })
+  state.modules.jit.expandApplyAtRules.module(state.jitContext)(root)
 
-  for (let [, rule] of rules) {
-    if (rule.type === 'rule') {
-      actualRules.push(rule)
-    } else if (rule.walkRules) {
-      rule.walkRules((subRule) => {
-        actualRules.push(subRule)
-      })
-    }
-  }
+  let actualRules: Rule[] = []
+  root.walkRules((subRule) => {
+    actualRules.push(subRule)
+  })
 
   return {
-    root: state.modules.postcss.module.root({ nodes: rules.map(([, rule]) => rule) }),
+    root,
     rules: actualRules,
   }
 }

packages/tailwindcss-language-service/src/util/state.ts

@@ -76,7 +76,11 @@ export interface State {
     postcss?: { version: string; module: Postcss }
     postcssSelectorParser?: { module: any }
     resolveConfig?: { module: any }
-    jit?: { generateRules: { module: any }; setupContext: { module: any } }
+    jit?: {
+      generateRules: { module: any }
+      setupContext: { module: any }
+      expandApplyAtRules: { module: any }
+    }
   }
   browserslist?: string[]
   featureFlags?: FeatureFlags

src/server.ts

@@ -445,6 +445,11 @@ async function createProjectService(
               resolveFrom(configDir, 'tailwindcss/lib/jit/lib/setupContext')
             ).default,
           },
+          expandApplyAtRules: {
+            module: __non_webpack_require__(
+              resolveFrom(configDir, 'tailwindcss/lib/jit/lib/expandApplyAtRules')
+            ).default,
+          },
         }
       } catch (_) {
         try {
@@ -459,6 +464,11 @@ async function createProjectService(
                 resolveFrom(configDir, 'tailwindcss/jit/lib/setupContext')
               ),
             },
+            expandApplyAtRules: {
+              module: __non_webpack_require__(
+                resolveFrom(configDir, 'tailwindcss/jit/lib/expandApplyAtRules')
+              ),
+            },
           }
         } catch (_) {}
       }
@@ -491,6 +501,10 @@ async function createProjectService(
     if (applyComplexClasses && !applyComplexClasses.default.__patched) {
       let _applyComplexClasses = applyComplexClasses.default
       applyComplexClasses.default = (config, ...args) => {
+        if (state.jit) {
+          return state.modules.jit.expandApplyAtRules.module(state.jitContext)
+        }
+
         let configClone = klona(config)
         configClone.separator = typeof state.separator === 'undefined' ? ':' : state.separator
 
@@ -629,6 +643,15 @@ async function createProjectService(
       throw error
     }
 
+    if (state.jit) {
+      state.jitContext = state.modules.jit.setupContext.module(state.configPath)(
+        { opts: {}, messages: [] },
+        state.modules.postcss.module.root()
+      )
+      state.jitContext.tailwindConfig.separator =
+        typeof userSeperator === 'undefined' ? ':' : userSeperator
+    }
+
     let postcssResult: Result
     try {
       postcssResult = await postcss
@@ -690,13 +713,6 @@ async function createProjectService(
     state.plugins = await getPlugins(config)
     state.classNames = (await extractClassNames(postcssResult.root)) as ClassNames
 
-    if (state.jit) {
-      state.jitContext = state.modules.jit.setupContext.module(state.configPath)(
-        { opts: {}, messages: [] },
-        state.modules.postcss.module.root()
-      )
-    }
-
     state.variants = getVariants(state)
 
     state.enabled = true