From 482b5bc7af9b6f39baf5513c93bc9f16ac4eb601 Mon Sep 17 00:00:00 2001
From: Brad Cornes <bradlc41@gmail.com>
Date: Mon, 10 May 2021 13:41:48 +0100
Subject: [PATCH] =?UTF-8?q?fix=20@=E2=80=8Bapply=20in=20plugins=20in=20jit?=
 =?UTF-8?q?=20mode?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/util/jit.ts                           | 18 +++++------
 .../src/util/state.ts                         |  6 +++-
 src/server.ts                                 | 30 ++++++++++++++-----
 3 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/packages/tailwindcss-language-service/src/util/jit.ts b/packages/tailwindcss-language-service/src/util/jit.ts
index 311319d..482dcf8 100644
--- a/packages/tailwindcss-language-service/src/util/jit.ts
+++ b/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,
   }
 }
diff --git a/packages/tailwindcss-language-service/src/util/state.ts b/packages/tailwindcss-language-service/src/util/state.ts
index 8ff1a3c..826d50e 100644
--- a/packages/tailwindcss-language-service/src/util/state.ts
+++ b/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
diff --git a/src/server.ts b/src/server.ts
index 0a928f5..6f6e317 100644
--- a/src/server.ts
+++ b/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