implement different approach

Instead of trying to traverse ASTs and tracking unused variables in
various spots, let's reverse the algorithm.

ALWAYS emit everything, this has the added benefit that emitting CSS
variables and keyframes is only done once (in the `compile(…)` step, not
in each `build(…)` step).

However, the removal of the unused CSS variables is part of the
`optimizeAst` step. This has a few benefits:

1. All the logic is done in a single spot, instead of tracking all over
   the codebase.
2. We are already traversing for the optimization step, so we can hook
   in without additional walks of the entire AST.
3. It will be more correct, e.g.: if you have a valid utility, but
   invalid variant, since the utility is handled first it could be that
   we mark variables as used even though the no CSS will be generated
   due to the invalid variant.

   Since this the `optimizeAst` step never even sees the thrown-away
   variant+utility, it means that we don't even have to worry about
   this.

Had to add a context node, to different between `:root {}` coming from
`@theme` or coming from user CSS (which should stay untouched).

Author: RobinMalfait

packages/tailwindcss/src/ast.ts

@@ -1,4 +1,8 @@
 import { parseAtRule } from './css-parser'
+import type { DesignSystem } from './design-system'
+import { ThemeOptions } from './theme'
+import { DefaultMap } from './utils/default-map'
+import * as ValueParser from './value-parser'
 
 const AT_SIGN = 0x40
 
@@ -237,9 +241,13 @@ export function walkDepth(
 
 // Optimize the AST for printing where all the special nodes that require custom
 // handling are handled such that the printing is a 1-to-1 transformation.
-export function optimizeAst(ast: AstNode[]) {
+export function optimizeAst(ast: AstNode[], designSystem: DesignSystem) {
   let atRoots: AstNode[] = []
   let seenAtProperties = new Set<string>()
+  let cssThemeVariables = new DefaultMap<
+    Extract<AstNode, { nodes: AstNode[] }>['nodes'],
+    Set<Declaration>
+  >(() => new Set())
 
   function transform(
     node: AstNode,
@@ -251,6 +259,22 @@ export function optimizeAst(ast: AstNode[]) {
       if (node.property === '--tw-sort' || node.value === undefined || node.value === null) {
         return
       }
+
+      // Track used CSS variables
+      if (node.value.includes('var(')) {
+        ValueParser.walk(ValueParser.parse(node.value), (node) => {
+          if (node.kind !== 'function' || node.value !== 'var') return
+
+          ValueParser.walk(node.nodes, (child) => {
+            if (child.kind !== 'word' || child.value[0] !== '-' || child.value[1] !== '-') return
+
+            designSystem.theme.markUsedVariable(child.value)
+          })
+
+          return ValueParser.ValueWalkAction.Skip
+        })
+      }
+
       parent.push(node)
     }
 
@@ -331,6 +355,15 @@ export function optimizeAst(ast: AstNode[]) {
         return
       }
 
+      if (node.context.theme) {
+        let declarations = cssThemeVariables.get(parent)
+        for (let child of node.nodes) {
+          if (child.kind === 'declaration') {
+            declarations.add(child)
+          }
+        }
+      }
+
       for (let child of node.nodes) {
         transform(child, parent, depth)
       }
@@ -352,6 +385,35 @@ export function optimizeAst(ast: AstNode[]) {
     transform(node, newAst, 0)
   }
 
+  // Remove unused theme variables
+  next: for (let [parent, declarations] of cssThemeVariables) {
+    for (let declaration of declarations) {
+      let options = designSystem.theme.getOptions(declaration.property)
+      if (options & ThemeOptions.USED) continue
+
+      // Remove the declaration (from its parent)
+      let idx = parent.indexOf(declaration)
+      parent.splice(idx, 1)
+
+      // If the parent is now empty, remove it from the AST
+      if (parent.length === 0) {
+        for (let [idx, node] of newAst.entries()) {
+          // Assumption, but right now the `@theme` must be top-level, so we
+          // don't need to traverse the entire AST to find the parent.
+          //
+          // Checking for `rule`, because at this stage the `@theme` is already
+          // converted to a normal style rule `:root, :host`
+          if (node.kind === 'rule' && node.nodes === parent) {
+            newAst.splice(idx, 1)
+            break
+          }
+        }
+
+        continue next
+      }
+    }
+  }
+
   return newAst.concat(atRoots)
 }
 

packages/tailwindcss/src/compile.ts

@@ -168,15 +168,6 @@ export function compileAstNodes(candidate: Candidate, designSystem: DesignSystem
       if (result === null) return []
     }
 
-    // Mark CSS variables as used. Variables resolved internally will already be
-    // marked as used. This is purely for arbitrary values and properties that use
-    // variables. E.g.: `[--color:var(--color-red-500)]`
-    //
-    //
-    if (candidate.raw.includes('var(')) {
-      designSystem.theme.trackUsedVariables([node])
-    }
-
     rules.push({
       node,
       propertySort,

packages/tailwindcss/src/design-system.ts

@@ -65,7 +65,7 @@ export function buildDesignSystem(theme: Theme): DesignSystem {
           },
         })
 
-        astNodes = optimizeAst(astNodes)
+        astNodes = optimizeAst(astNodes, designSystem)
 
         if (astNodes.length === 0 || wasInvalid) {
           result.push(null)

packages/tailwindcss/src/index.ts

@@ -4,6 +4,7 @@ import {
   atRoot,
   atRule,
   comment,
+  context,
   context as contextNode,
   decl,
   optimizeAst,
@@ -532,6 +533,39 @@ async function parseCss(
     customUtility(designSystem)
   }
 
+  // Output final set of theme variables at the position of the first
+  // `@theme` rule.
+  if (firstThemeRule) {
+    let nodes = []
+
+    for (let [key, value] of designSystem.theme.entries()) {
+      if (value.options & ThemeOptions.REFERENCE) continue
+
+      nodes.push(decl(escape(key), value.value))
+    }
+
+    let keyframesRules = designSystem.theme.getKeyframes()
+    if (keyframesRules.length > 0) {
+      let animationParts = [...designSystem.theme.namespace('--animate').values()].flatMap(
+        (animation) => animation.split(/\s+/),
+      )
+
+      for (let keyframesRule of keyframesRules) {
+        // Remove any keyframes that aren't used by an animation variable.
+        let keyframesName = keyframesRule.params
+        if (!animationParts.includes(keyframesName)) {
+          continue
+        }
+
+        // Wrap `@keyframes` in `AtRoot` so they are hoisted out of `:root` when
+        // printing.
+        nodes.push(atRoot([keyframesRule]))
+      }
+    }
+
+    firstThemeRule.nodes = [context({ theme: true }, nodes)]
+  }
+
   // Replace the `@tailwind utilities` node with a context since it prints
   // children directly.
   if (utilitiesNode) {
@@ -567,10 +601,6 @@ async function parseCss(
   features |= substituteFunctions(ast, designSystem)
   features |= substituteAtApply(ast, designSystem)
 
-  // Mark CSS variables as used. Right now they can only be used in
-  // declarations, because `@media` and `@container` don't support them.
-  designSystem.theme.trackUsedVariables(ast)
-
   // Remove `@utility`, we couldn't replace it before yet because we had to
   // handle the nested `@apply` at-rules first.
   walk(ast, (node, { replaceWith }) => {
@@ -592,7 +622,6 @@ async function parseCss(
     root,
     utilitiesNode,
     features,
-    firstThemeRule,
   }
 }
 
@@ -605,10 +634,7 @@ export async function compileAst(
   features: Features
   build(candidates: string[]): AstNode[]
 }> {
-  let { designSystem, ast, globs, root, utilitiesNode, features, firstThemeRule } = await parseCss(
-    input,
-    opts,
-  )
+  let { designSystem, ast, globs, root, utilitiesNode, features } = await parseCss(input, opts)
 
   if (process.env.NODE_ENV !== 'test') {
     ast.unshift(comment(`! tailwindcss v${version} | MIT License | https://tailwindcss.com `))
@@ -636,7 +662,7 @@ export async function compileAst(
       }
 
       if (!utilitiesNode) {
-        compiled ??= optimizeAst(ast)
+        compiled ??= optimizeAst(ast, designSystem)
         return compiled
       }
 
@@ -648,7 +674,7 @@ export async function compileAst(
       for (let candidate of newRawCandidates) {
         if (!designSystem.invalidCandidates.has(candidate)) {
           if (candidate[0] === '-' && candidate[1] === '-') {
-            emitNewCssVariables = designSystem.theme.use(candidate)
+            emitNewCssVariables = designSystem.theme.markUsedVariable(candidate)
             didChange ||= emitNewCssVariables
           } else {
             allValidCandidates.add(candidate)
@@ -660,61 +686,27 @@ export async function compileAst(
       // If no new candidates were added, we can return the original CSS. This
       // currently assumes that we only add new candidates and never remove any.
       if (!didChange) {
-        compiled ??= optimizeAst(ast)
+        compiled ??= optimizeAst(ast, designSystem)
         return compiled
       }
 
       let newNodes = compileCandidates(allValidCandidates, designSystem, {
         onInvalidCandidate,
       }).astNodes
 
-      // Output final set of theme variables at the position of the first
-      // `@theme` rule.
-      if (firstThemeRule) {
-        let nodes = []
-
-        for (let [key, value] of designSystem.theme.entries()) {
-          if (value.options & ThemeOptions.REFERENCE) continue
-          if (!(value.options & ThemeOptions.USED)) continue
-
-          nodes.push(decl(escape(key), value.value))
-        }
-
-        let keyframesRules = designSystem.theme.getKeyframes()
-        if (keyframesRules.length > 0) {
-          let animationParts = [...designSystem.theme.namespace('--animate').values()].flatMap(
-            (animation) => animation.split(/\s+/),
-          )
-
-          for (let keyframesRule of keyframesRules) {
-            // Remove any keyframes that aren't used by an animation variable.
-            let keyframesName = keyframesRule.params
-            if (!animationParts.includes(keyframesName)) {
-              continue
-            }
-
-            // Wrap `@keyframes` in `AtRoot` so they are hoisted out of `:root` when
-            // printing.
-            nodes.push(atRoot([keyframesRule]))
-          }
-        }
-
-        firstThemeRule.nodes = nodes
-      }
-
       // If no new ast nodes were generated, then we can return the original
       // CSS. This currently assumes that we only add new ast nodes and never
       // remove any.
       if (previousAstNodeCount === newNodes.length && !emitNewCssVariables) {
-        compiled ??= optimizeAst(ast)
+        compiled ??= optimizeAst(ast, designSystem)
         return compiled
       }
 
       previousAstNodeCount = newNodes.length
 
       utilitiesNode.nodes = newNodes
 
-      compiled = optimizeAst(ast)
+      compiled = optimizeAst(ast, designSystem)
       return compiled
     },
   }

packages/tailwindcss/src/theme.ts

@@ -1,6 +1,5 @@
-import { walk, WalkAction, type AstNode, type AtRule } from './ast'
+import { type AtRule } from './ast'
 import { escape } from './utils/escape'
-import * as ValueParser from './value-parser'
 
 export const enum ThemeOptions {
   NONE = 0,
@@ -108,7 +107,7 @@ export class Theme {
   }
 
   getOptions(key: string) {
-    return this.values.get(key)?.options ?? ThemeOptions.NONE
+    return this.values.get(this.#unprefixKey(key))?.options ?? ThemeOptions.NONE
   }
 
   entries() {
@@ -125,6 +124,11 @@ export class Theme {
     return `--${this.prefix}-${key.slice(2)}`
   }
 
+  #unprefixKey(key: string) {
+    if (!this.prefix) return key
+    return `--${key.slice(3 + this.prefix.length)}`
+  }
+
   clearNamespace(namespace: string, clearOptions: ThemeOptions) {
     let ignored = ignoredThemeKeyMap.get(namespace) ?? []
 
@@ -174,42 +178,13 @@ export class Theme {
       return null
     }
 
-    this.use(themeKey)
-
     return `var(${escape(this.#prefixKey(themeKey))})`
   }
 
-  trackUsedVariables(ast: AstNode[]) {
-    walk(ast, (node) => {
-      // Variables used in `@utility` and `@custom-variant` at-rules will be
-      // handled separately, because we only want to mark them as used if the
-      // utility or variant is used.
-      if (
-        node.kind === 'at-rule' &&
-        (node.name === '@utility' || node.name === '@custom-variant')
-      ) {
-        return WalkAction.Skip
-      }
-
-      if (node.kind !== 'declaration') return
-      if (!node.value?.includes('var(')) return
-
-      ValueParser.walk(ValueParser.parse(node.value), (node) => {
-        if (node.kind !== 'function' || node.value !== 'var') return
-
-        ValueParser.walk(node.nodes, (child) => {
-          if (child.kind !== 'word' || child.value[0] !== '-' || child.value[1] !== '-') return
-
-          this.use(child.value)
-        })
-      })
-    })
-  }
-
-  use(themeKey: string) {
-    let value = this.values.get(themeKey)
-    if (!value) return false // Unknown
-    if (value.options & ThemeOptions.USED) return false // Already used
+  markUsedVariable(themeKey: string) {
+    let value = this.values.get(this.#unprefixKey(themeKey))
+    if (!value) return false // Unknown variable
+    if (value.options & ThemeOptions.USED) return false // Variable already used
 
     value.options |= ThemeOptions.USED
     return true

packages/tailwindcss/src/utilities.ts

@@ -4824,8 +4824,6 @@ export function createCssUtility(node: AtRule) {
           }
         }
 
-        designSystem.theme.trackUsedVariables(atRule.nodes)
-
         return atRule.nodes
       })
 
@@ -4846,11 +4844,7 @@ export function createCssUtility(node: AtRule) {
 
   if (IS_VALID_STATIC_UTILITY_NAME.test(name)) {
     return (designSystem: DesignSystem) => {
-      designSystem.utilities.static(name, () => {
-        let ast = structuredClone(node.nodes)
-        designSystem.theme.trackUsedVariables(ast)
-        return ast
-      })
+      designSystem.utilities.static(name, () => structuredClone(node.nodes))
     }
   }