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.

Author: thecrypticace

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/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, {})
@@ -500,36 +489,39 @@ function collectLayerPlugins(root) {
     if (layerRule.params === 'base') {
       for (let node of layerRule.nodes) {
         layerPlugins.push(function ({ addBase }) {
+          node.raws.tailwind = {
+            ...node.raws.tailwind,
+            layer: layerRule.params,
+          }
           addBase(node, { respectPrefix: false })
         })
       }
       layerRule.remove()
     } else if (layerRule.params === 'components') {
       for (let node of layerRule.nodes) {
         layerPlugins.push(function ({ addComponents }) {
+          node.raws.tailwind = {
+            ...node.raws.tailwind,
+            layer: layerRule.params,
+          }
           addComponents(node, { respectPrefix: false })
         })
       }
       layerRule.remove()
     } else if (layerRule.params === 'utilities') {
       for (let node of layerRule.nodes) {
         layerPlugins.push(function ({ addUtilities }) {
+          node.raws.tailwind = {
+            ...node.raws.tailwind,
+            layer: layerRule.params,
+          }
           addUtilities(node, { respectPrefix: false })
         })
       }
       layerRule.remove()
     }
   })
 
-  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
 }
 

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)
 })