Repository: tailwindlabs/tailwindcss (ts)

Issue:
JIT x CSS Modules x `@layer` directive results in no CSS module class emitted
### What version of Tailwind CSS are you using?

v2.1.1

### What version of Node.js are you using?

v12.14.1

### What browser are you using?

Chrome

### What operating system are you using?

macOS

### Reproduction repository

repro steps below

# The Bug

HTML elements aren't being assigned classnames when the Tailwind JIT engine is used with CSS modules and `@layer components`. This happens in both dev and prod

I think most people can remove `@layer` but I'm reporting this anyway because this seems like a regression

# Reproduction
1. `yarn create next-app --example with-tailwindcss`, then `cd` into the project dir
2. apply this patch: 
```patch
diff --git a/package.json b/package.json
index 46f0c5d..8132319 100644
--- a/package.json
+++ b/package.json
@@ -13,9 +13,8 @@
     "react-dom": "^17.0.1"
   },
   "devDependencies": {
-    "@tailwindcss/jit": "0.1.3",
     "autoprefixer": "^10.0.4",
     "postcss": "^8.1.10",
-    "tailwindcss": "^2.0.2"
+    "tailwindcss": "^2.1.1"
   }
 }
diff --git a/pages/index.js b/pages/index.js
index 5cb31bc..5084680 100644
--- a/pages/index.js
+++ b/pages/index.js
@@ -1,8 +1,10 @@
 import Head from 'next/head'
 
+import styles from "./index.module.css";
+
 export default function Home() {
   return (
-    <div className="flex flex-col items-center justify-center min-h-screen py-2">
+    <div className={styles.wrapper}>
       <Head>
         <title>Create Next App</title>
         <link rel="icon" href="/favicon.ico" />
diff --git a/pages/index.module.css b/pages/index.module.css
new file mode 100644
index 0000000..a062ca2
--- /dev/null
+++ b/pages/index.module.css
@@ -0,0 +1,5 @@
+@layer components {
+    .wrapper {
+        @apply flex flex-col items-center justify-center min-h-screen py-2;
+    }
+}
\ No newline at end of file
diff --git a/postcss.config.js b/postcss.config.js
index fc58cf3..2d5c5c0 100644
--- a/postcss.config.js
+++ b/postcss.config.js
@@ -2,7 +2,7 @@
 // https://tailwindcss.com/docs/using-with-preprocessors
 module.exports = {
   plugins: {
-    '@tailwindcss/jit': {},
+    'tailwindcss': {},
     autoprefixer: {},
   },
 }
diff --git a/tailwind.config.js b/tailwind.config.js
index d7f0874..32f0b66 100644
--- a/tailwind.config.js
+++ b/tailwind.config.js
@@ -1,4 +1,5 @@
 module.exports = {
+  mode: "jit",
   purge: ['./pages/**/*.{js,ts,jsx,tsx}', './components/**/*.{js,ts,jsx,tsx}'],
   darkMode: false, // or 'media' or 'class'
   theme: {
diff --git a/yarn.lock b/yarn.lock
index 4088c52..267ddf3 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -131,19 +131,6 @@
   resolved "https://registry.yarnpkg.com/@opentelemetry/context-base/-/context-base-0.14.0.tgz#c67fc20a4d891447ca1a855d7d70fa79a3533001"
   integrity sha512-sDOAZcYwynHFTbLo6n8kIbLiVF3a3BLkrmehJUyEbT9F+Smbi47kLGS2gG2g0fjBLR/Lr1InPD7kXL7FaTqEkw==
 
-"@tailwindcss/jit@0.1.3":
-  version "0.1.3"
-  resolved "https://registry.yarnpkg.com/@tailwindcss/jit/-/jit-0.1.3.tgz#50390d9ac95fee78ed23a7e2320ef20bd4a35354"
-  integrity sha512-7VAvHKNLJxbGWRKxo2Z+beiodag7vWPx8b/+Egd5fve4zFihsngeNt6RwQFnll+almjppRYefRC5Py5v5K+6vg==
-  dependencies:
-    chokidar "^3.5.1"
-    dlv "^1.1.3"
-    fast-glob "^3.2.5"
-    lodash.topath "^4.5.2"
-    object-hash "^2.1.1"
-    postcss-selector-parser "^6.0.4"
-    quick-lru "^5.1.1"
-
 "@types/node@*":
   version "14.14.37"
   resolved "https://registry.yarnpkg.com/@types/node/-/node-14.14.37.tgz#a3dd8da4eb84a996c36e331df98d82abd76b516e"
@@ -2213,7 +2200,7 @@ supports-color@^8.0.0:
   dependencies:
     has-flag "^4.0.0"
 
-tailwindcss@^2.0.2:
+tailwindcss@^2.1.1:
   version "2.1.1"
   resolved "https://registry.yarnpkg.com/tailwindcss/-/tailwindcss-2.1.1.tgz#642f6038c9283a8e1454da34585b8b7c1a1e8877"
   integrity sha512-zZ6axGqpSZOCBS7wITm/WNHkBzDt5CIZlDlx0eCVldwTxFPELCVGbgh7Xpb3/kZp3cUxOmK7bZUjqhuMrbN6xQ==
```
3. `yarn run dev`, notice that the root div has no class

This issue occurs: 
- if you switch to webpack 5

This issue does not occur: 
- if you disable JIT
- if you remove `@layer components`
- if you use `@apply` in a `@layer` directive in a global stylesheet, not in a CSS module (I think)
- if you use `@tailwindcss/jit` (I think)

Pull request description:
Error when `@layer` used without matching `@tailwind` directive
Also refactor to only detect `@tailwind` directives once per build to improve performance.

Resolves #4008

Patch:
diff --git a/src/jit/index.js b/src/jit/index.js
index 0c8f63c2..d155f1ee 100644
--- a/src/jit/index.js
+++ b/src/jit/index.js
@@ -30,9 +30,9 @@ export default function (configOrPath = {}) {
         })
       }
 
-      rewriteTailwindImports(root)
+      let tailwindDirectives = rewriteTailwindImports(root)
 
-      let context = setupContext(configOrPath)(result, root)
+      let context = setupContext(configOrPath, tailwindDirectives)(result, root)
 
       if (!env.TAILWIND_DISABLE_TOUCH) {
         if (context.configPath !== null) {
@@ -41,8 +41,8 @@ export default function (configOrPath = {}) {
       }
 
       return postcss([
-        removeLayerAtRules(context),
-        expandTailwindAtRules(context, registerDependency),
+        removeLayerAtRules(context, tailwindDirectives),
+        expandTailwindAtRules(context, registerDependency, tailwindDirectives),
         expandApplyAtRules(context),
         evaluateTailwindFunctions(context.tailwindConfig),
         substituteScreenAtRules(context.tailwindConfig),
diff --git a/src/jit/lib/expandTailwindAtRules.js b/src/jit/lib/expandTailwindAtRules.js
index cd528f1c..6b421d8a 100644
--- a/src/jit/lib/expandTailwindAtRules.js
+++ b/src/jit/lib/expandTailwindAtRules.js
@@ -111,9 +111,12 @@ function buildStylesheet(rules, context) {
   return returnValue
 }
 
-export default function expandTailwindAtRules(context, registerDependency) {
+export default function expandTailwindAtRules(context, registerDependency, tailwindDirectives) {
   return (root) => {
-    let foundTailwind = false
+    if (tailwindDirectives.size === 0) {
+      return root
+    }
+
     let layerNodes = {
       base: null,
       components: null,
@@ -126,8 +129,6 @@ export default function expandTailwindAtRules(context, registerDependency) {
     // file as a dependency since the output of this CSS does not depend on
     // the source of any templates. Think Vue <style> blocks for example.
     root.walkAtRules('tailwind', (rule) => {
-      foundTailwind = true
-
       if (rule.params === 'base') {
         layerNodes.base = rule
       }
@@ -145,10 +146,6 @@ export default function expandTailwindAtRules(context, registerDependency) {
       }
     })
 
-    if (!foundTailwind) {
-      return root
-    }
-
     // ---
 
     if (sharedState.env.TAILWIND_DISABLE_TOUCH) {
diff --git a/src/jit/lib/removeLayerAtRules.js b/src/jit/lib/removeLayerAtRules.js
index 009ae200..77ca8441 100644
--- a/src/jit/lib/removeLayerAtRules.js
+++ b/src/jit/lib/removeLayerAtRules.js
@@ -1,7 +1,22 @@
-export default function removeLayerAtRules() {
+export default function removeLayerAtRules(_context, tailwindDirectives) {
   return (root) => {
     root.walkAtRules((rule) => {
-      if (['layer', 'responsive', 'variants'].includes(rule.name)) {
+      if (rule.name === 'layer' && ['base', 'components', 'utilities'].includes(rule.params)) {
+        if (!tailwindDirectives.has(rule.params)) {
+          throw rule.error(
+            `\`@layer ${rule.params}\` is used but no matching \`@tailwind ${rule.params}\` directive is present.`
+          )
+        }
+        rule.remove()
+      } else if (rule.name === 'responsive') {
+        if (!tailwindDirectives.has('utilities')) {
+          throw rule.error('`@responsive` is used but `@tailwind utilities` is missing.')
+        }
+        rule.remove()
+      } else if (rule.name === 'variants') {
+        if (!tailwindDirectives.has('utilities')) {
+          throw rule.error('`@variants` is used but `@tailwind utilities` is missing.')
+        }
         rule.remove()
       }
     })
diff --git a/src/jit/lib/rewriteTailwindImports.js b/src/jit/lib/rewriteTailwindImports.js
index b022ebd6..51cf06aa 100644
--- a/src/jit/lib/rewriteTailwindImports.js
+++ b/src/jit/lib/rewriteTailwindImports.js
@@ -23,4 +23,12 @@ export default function rewriteTailwindImports(root) {
       atRule.params = 'screens'
     }
   })
+
+  let tailwindDirectives = new Set()
+
+  root.walkAtRules('tailwind', (rule) => {
+    tailwindDirectives.add(rule.params)
+  })
+
+  return tailwindDirectives
 }
diff --git a/src/jit/lib/setupContext.js b/src/jit/lib/setupContext.js
index 99534988..2fa79d63 100644
--- a/src/jit/lib/setupContext.js
+++ b/src/jit/lib/setupContext.js
@@ -314,6 +314,7 @@ function rebootWatcher(context) {
           touch(context.configPath)
         } else {
           context.changedFiles.add(path.resolve('.', file))
+          console.log(context.touchFile)
           touch(context.touchFile)
         }
       })
@@ -686,14 +687,8 @@ function cleanupContext(context) {
 // Retrieve an existing context from cache if possible (since contexts are unique per
 // source path), or set up a new one (including setting up watchers and registering
 // plugins) then return it
-export default function setupContext(configOrPath) {
+export default function setupContext(configOrPath, tailwindDirectives) {
   return (result, root) => {
-    let foundTailwind = false
-
-    root.walkAtRules('tailwind', () => {
-      foundTailwind = true
-    })
-
     let sourcePath = result.opts.from
     let [
       tailwindConfig,
@@ -712,7 +707,7 @@ export default function setupContext(configOrPath) {
     // We may want to think about `@layer` being part of this trigger too, but it's tough
     // because it's impossible for a layer in one file to end up in the actual @tailwind rule
     // in another file since independent sources are effectively isolated.
-    if (foundTailwind) {
+    if (tailwindDirectives.size > 0) {
       contextDependencies.add(sourcePath)
       for (let message of result.messages) {
         if (message.type === 'dependency') {


Test patch:
diff --git a/tests/jit/layer-without-tailwind.test.css b/tests/jit/layer-without-tailwind.test.css
new file mode 100644
index 00000000..28dfe0bb
--- /dev/null
+++ b/tests/jit/layer-without-tailwind.test.css
@@ -0,0 +1,5 @@
+@layer components {
+  .foo {
+    color: black;
+  }
+}
diff --git a/tests/jit/layer-without-tailwind.test.html b/tests/jit/layer-without-tailwind.test.html
new file mode 100644
index 00000000..a6d7176e
--- /dev/null
+++ b/tests/jit/layer-without-tailwind.test.html
@@ -0,0 +1,1 @@
+<div class="foo"></div>
diff --git a/tests/jit/layer-without-tailwind.test.js b/tests/jit/layer-without-tailwind.test.js
new file mode 100644
index 00000000..07d2d9f6
--- /dev/null
+++ b/tests/jit/layer-without-tailwind.test.js
@@ -0,0 +1,101 @@
+import postcss from 'postcss'
+import path from 'path'
+import tailwind from '../../src/jit/index.js'
+
+function run(input, config = {}) {
+  const { currentTestName } = expect.getState()
+
+  return postcss(tailwind(config)).process(input, {
+    from: `${path.resolve(__filename)}?test=${currentTestName}`,
+  })
+}
+
+test('using @layer without @tailwind', async () => {
+  let config = {
+    purge: [path.resolve(__dirname, './layer-without-tailwind.test.html')],
+    mode: 'jit',
+    theme: {},
+    plugins: [],
+  }
+
+  let css = `
+    @layer components {
+      .foo {
+        color: black;
+      }
+    }
+  `
+
+  await expect(run(css, config)).rejects.toThrowError(
+    '`@layer components` is used but no matching `@tailwind components` directive is present.'
+  )
+})
+
+test('using @responsive without @tailwind', async () => {
+  let config = {
+    purge: [path.resolve(__dirname, './layer-without-tailwind.test.html')],
+    mode: 'jit',
+    theme: {},
+    plugins: [],
+  }
+
+  let css = `
+    @responsive {
+      .foo {
+        color: black;
+      }
+    }
+  `
+
+  await expect(run(css, config)).rejects.toThrowError(
+    '`@responsive` is used but `@tailwind utilities` is missing.'
+  )
+})
+
+test('using @variants without @tailwind', async () => {
+  let config = {
+    purge: [path.resolve(__dirname, './layer-without-tailwind.test.html')],
+    mode: 'jit',
+    theme: {},
+    plugins: [],
+  }
+
+  let css = `
+    @variants hover {
+      .foo {
+        color: black;
+      }
+    }
+  `
+
+  await expect(run(css, config)).rejects.toThrowError(
+    '`@variants` is used but `@tailwind utilities` is missing.'
+  )
+})
+
+test('non-Tailwind @layer rules are okay', async () => {
+  let config = {
+    purge: [path.resolve(__dirname, './layer-without-tailwind.test.html')],
+    mode: 'jit',
+    theme: {},
+    plugins: [],
+  }
+
+  let css = `
+    @layer custom {
+      .foo {
+        color: black;
+      }
+    }
+  `
+
+  return run(css, config).then((result) => {
+    expect(result.css).toMatchFormattedCss(`
+      @layer custom {
+        .foo {
+          color: black;
+        }
+      }
+    `)
+  })
+})
diff --git a/tests/jit/prefix.test.js b/tests/jit/prefix.test.js
index 553a1ce7..3d21bfa2 100644
--- a/tests/jit/prefix.test.js
+++ b/tests/jit/prefix.test.js
@@ -66,7 +66,7 @@ test('prefix', () => {
       }
     }
     @tailwind utilities;
-    @layer utilites {
+    @layer utilities {
       .custom-utility {
         foo: bar;
       }