Rework apply to use local postcss root

We were storing user CSS in the context so we could use it with apply. The problem is that this CSS does not get updated on save unless it has a tailwind directive in it resulting in stale apply caches. This could result in either stale generation or errors about missing classes.

src/lib/expandApplyAtRules.js

@@ -5,6 +5,8 @@ import { resolveMatches } from './generateRules'
 import bigSign from '../util/bigSign'
 import escapeClassName from '../util/escapeClassName'
 
+/** @typedef {Map<string, [any, import('postcss').Rule[]]>} ApplyCache */
+
 function extractClasses(node) {
   let classes = new Set()
   let container = postcss.root({ nodes: [node.clone()] })
@@ -35,6 +37,83 @@ function prefix(context, selector) {
   return typeof prefix === 'function' ? prefix(selector) : prefix + selector
 }
 
+function* pathToRoot(node) {
+  yield node
+  while (node.parent) {
+    yield node.parent
+    node = node.parent
+  }
+}
+
+/**
+ * @param {import('postcss').Node} node
+ */
+function structuralCloneOfNode(node) {
+  for (let parent of pathToRoot(node)) {
+    if (node === parent) {
+      continue
+    }
+
+    if (parent.type === 'root') {
+      break
+    }
+
+    node = parent.clone({
+      nodes: [node],
+    })
+  }
+
+  return node
+}
+
+/**
+ * @param {import('postcss').Root} root
+ */
+function buildLocalApplyCache(root, context) {
+  /** @type {ApplyCache} */
+  let cache = new Map()
+
+  let reservedBits = 0n
+  let tmp = context.layerOrder.utilities >> 3n
+  while (tmp > 1n) {
+    tmp = tmp >> 1n
+    reservedBits++
+  }
+
+  let highestOffset = 1n << reservedBits
+
+  root.walkRules((rule, idx) => {
+    // Ignore rules generated by Tailwind
+    for (let node of pathToRoot(rule)) {
+      if (node.raws.tailwind?.layer !== undefined) {
+        return
+      }
+    }
+
+    // Walk to the top of the rule
+    let container = structuralCloneOfNode(rule)
+
+    for (let className of extractClasses(rule)) {
+      let list = cache.get(className) || []
+      cache.set(className, list)
+
+      list.push([
+        {
+          layer: 'user',
+          sort: BigInt(idx) + highestOffset,
+          important: false,
+        },
+        container,
+      ])
+    }
+  })
+
+  return cache
+}
+
+/**
+ * @returns {ApplyCache}
+ */
 function buildApplyCache(applyCandidates, context) {
   for (let candidate of applyCandidates) {
     if (context.notClassCache.has(candidate) || context.applyClassCache.has(candidate)) {
@@ -62,6 +141,17 @@ function buildApplyCache(applyCandidates, context) {
   return context.applyClassCache
 }
 
+/**
+ * @param {ApplyCache[]} caches
+ * @returns {ApplyCache}
+ */
+function combineCaches(caches) {
+  return {
+    get: (name) => caches.flatMap((cache) => cache.get(name) || []),
+    has: (name) => caches.some((cache) => cache.has(name)),
+  }
+}
+
 function extractApplyCandidates(params) {
   let candidates = params.split(/[\s\t\n]+/g)
 
@@ -72,7 +162,7 @@ function extractApplyCandidates(params) {
   return [candidates, false]
 }
 
-function processApply(root, context) {
+function processApply(root, context, localCache) {
   let applyCandidates = new Set()
 
   // Collect all @apply rules and candidates
@@ -90,7 +180,7 @@ function processApply(root, context) {
   // Start the @apply process if we have rules with @apply in them
   if (applies.length > 0) {
     // Fill up some caches!
-    let applyClassCache = buildApplyCache(applyCandidates, context)
+    let applyClassCache = combineCaches([localCache, buildApplyCache(applyCandidates, context)])
 
     /**
      * When we have an apply like this:
@@ -296,12 +386,14 @@ function processApply(root, context) {
     }
 
     // Do it again, in case we have other `@apply` rules
-    processApply(root, context)
+    processApply(root, context, localCache)
   }
 }
 
 export default function expandApplyAtRules(context) {
   return (root) => {
-    processApply(root, context)
+    let localCache = buildLocalApplyCache(root, context)
+
+    processApply(root, context, localCache)
   }
 }

src/lib/expandTailwindAtRules.js

@@ -204,17 +204,29 @@ export default function expandTailwindAtRules(context) {
     // Replace any Tailwind directives with generated CSS
 
     if (layerNodes.base) {
-      layerNodes.base.before(cloneNodes([...baseNodes, ...defaultNodes], layerNodes.base.source))
+      layerNodes.base.before(
+        cloneNodes([...baseNodes, ...defaultNodes], layerNodes.base.source, {
+          layer: 'base',
+        })
+      )
       layerNodes.base.remove()
     }
 
     if (layerNodes.components) {
-      layerNodes.components.before(cloneNodes([...componentNodes], layerNodes.components.source))
+      layerNodes.components.before(
+        cloneNodes([...componentNodes], layerNodes.components.source, {
+          layer: 'components',
+        })
+      )
       layerNodes.components.remove()
     }
 
     if (layerNodes.utilities) {
-      layerNodes.utilities.before(cloneNodes([...utilityNodes], layerNodes.utilities.source))
+      layerNodes.utilities.before(
+        cloneNodes([...utilityNodes], layerNodes.utilities.source, {
+          layer: 'utilities',
+        })
+      )
       layerNodes.utilities.remove()
     }
 
@@ -234,10 +246,18 @@ export default function expandTailwindAtRules(context) {
     })
 
     if (layerNodes.variants) {
-      layerNodes.variants.before(cloneNodes(variantNodes, layerNodes.variants.source))
+      layerNodes.variants.before(
+        cloneNodes(variantNodes, layerNodes.variants.source, {
+          layer: 'variants',
+        })
+      )
       layerNodes.variants.remove()
     } else if (variantNodes.length > 0) {
-      root.append(cloneNodes(variantNodes, root.source))
+      root.append(
+        cloneNodes(variantNodes, root.source, {
+          layer: 'variants',
+        })
+      )
     }
 
     // If we've got a utility layer and no utilities are generated there's likely something wrong

src/lib/setupContextUtils.js

@@ -230,17 +230,6 @@ function buildPluginApi(tailwindConfig, context, { variantList, variantMap, offs
       // Preserved for backwards compatibility but not used in v3.0+
       return []
     },
-    addUserCss(userCss) {
-      for (let [identifier, rule] of withIdentifiers(userCss)) {
-        let offset = offsets.user++
-
-        if (!context.candidateRuleMap.has(identifier)) {
-          context.candidateRuleMap.set(identifier, [])
-        }
-
-        context.candidateRuleMap.get(identifier).push([{ sort: offset, layer: 'user' }, rule])
-      }
-    },
     addBase(base) {
       for (let [identifier, rule] of withIdentifiers(base)) {
         let prefixedIdentifier = prefixIdentifier(identifier, {})
@@ -521,15 +510,6 @@ function collectLayerPlugins(root) {
     }
   })
 
-  root.walkRules((rule) => {
-    // At this point it is safe to include all the left-over css from the
-    // user's css file. This is because the `@tailwind` and `@layer` directives
-    // will already be handled and will be removed from the css tree.
-    layerPlugins.push(function ({ addUserCss }) {
-      addUserCss(rule, { respectPrefix: false })
-    })
-  })
-
   return layerPlugins
 }
 

src/util/cloneNodes.js

@@ -1,11 +1,18 @@
-export default function cloneNodes(nodes, source) {
+export default function cloneNodes(nodes, source = undefined, raws = undefined) {
   return nodes.map((node) => {
     let cloned = node.clone()
 
     if (source !== undefined) {
       cloned.source = source
     }
 
+    if (raws !== undefined) {
+      cloned.raws.tailwind = {
+        ...cloned.raws.tailwind,
+        ...raws,
+      }
+    }
+
     return cloned
   })
 }

tests/apply.test.js

@@ -1,15 +1,8 @@
 import fs from 'fs'
 import path from 'path'
-import * as sharedState from '../src/lib/sharedState.js'
 
 import { run, html, css, defaults } from './util/run'
 
-beforeEach(() => {
-  sharedState.contextMap.clear()
-  sharedState.configContextMap.clear()
-  sharedState.contextSourcesMap.clear()
-})
-
 test('@apply', () => {
   let config = {
     darkMode: 'class',
@@ -1338,3 +1331,54 @@ it('should be possible to use apply in plugins', async () => {
     `)
   })
 })
+
+it('The apply class cache is invalidated when rules change', async () => {
+  let config = {
+    content: [{ raw: html`<div></div>` }],
+    plugins: [],
+  }
+
+  let inputBefore = css`
+    .foo {
+      color: green;
+    }
+
+    .bar {
+      @apply foo;
+    }
+  `
+
+  let inputAfter = css`
+    .foo {
+      color: red;
+    }
+
+    .bar {
+      @apply foo;
+    }
+  `
+
+  let result = await run(inputBefore, config)
+
+  expect(result.css).toMatchFormattedCss(css`
+    .foo {
+      color: green;
+    }
+
+    .bar {
+      color: green;
+    }
+  `)
+
+  result = await run(inputAfter, config)
+
+  expect(result.css).toMatchFormattedCss(css`
+    .foo {
+      color: red;
+    }
+
+    .bar {
+      color: red;
+    }
+  `)
+})

tests/context-reuse.test.js

@@ -83,5 +83,5 @@ it('a build re-uses the context across multiple files with the same config', asy
   expect(dependencies[3]).toEqual([path.resolve(__dirname, 'context-reuse.tailwind.config.js')])
 
   // And none of this should have resulted in multiple contexts being created
-  // expect(sharedState.contextSourcesMap.size).toBe(1)
+  expect(sharedState.contextSourcesMap.size).toBe(1)
 })