[v4] Only allow CSS variables in @theme (#13125)

* add `WalkAction` to indicate how the `walk` function should behave

You can `Continue` its execution (the default behaviour), `Skip` walking
the current node, or `Stop` walking entirely.

* walk the nodes of `@theme` directly

There is no need to walk the `@theme` itself or to create a temporary
array here.

* improvement: skip walking

* only allow CSS variables inside `@theme`

* update error message

packages/tailwindcss/src/ast.ts

@@ -42,28 +42,44 @@ export function comment(value: string): Comment {
   }
 }
 
+export enum WalkAction {
+  /** Continue walking, which is the default */
+  Continue,
+
+  /** Skip visiting the children of this node */
+  Skip,
+
+  /** Stop the walk entirely */
+  Stop,
+}
+
 export function walk(
   ast: AstNode[],
   visit: (
     node: AstNode,
     utils: {
       replaceWith(newNode: AstNode | AstNode[]): void
     },
-  ) => void | false,
+  ) => void | WalkAction,
 ) {
   for (let i = 0; i < ast.length; i++) {
     let node = ast[i]
-    let shouldContinue = visit(node, {
-      replaceWith(newNode) {
-        ast.splice(i, 1, ...(Array.isArray(newNode) ? newNode : [newNode]))
-        // We want to visit the newly replaced node(s), which start at the current
-        // index (i). By decrementing the index here, the next loop will process
-        // this position (containing the replaced node) again.
-        i--
-      },
-    })
-
-    if (shouldContinue === false) return
+    let status =
+      visit(node, {
+        replaceWith(newNode) {
+          ast.splice(i, 1, ...(Array.isArray(newNode) ? newNode : [newNode]))
+          // We want to visit the newly replaced node(s), which start at the
+          // current index (i). By decrementing the index here, the next loop
+          // will process this position (containing the replaced node) again.
+          i--
+        },
+      }) ?? WalkAction.Continue
+
+    // Stop the walk entirely
+    if (status === WalkAction.Stop) return
+
+    // Skip visiting the children of this node
+    if (status === WalkAction.Skip) continue
 
     if (node.kind === 'rule') {
       walk(node.nodes, visit)

packages/tailwindcss/src/index.test.ts

@@ -51,6 +51,32 @@ describe('compiling CSS', () => {
     `)
   })
 
+  test('that only CSS variables are allowed', () => {
+    expect(() =>
+      compileCss(
+        css`
+          @theme {
+            --color-primary: red;
+            .foo {
+              --color-primary: blue;
+            }
+          }
+          @tailwind utilities;
+        `,
+        ['bg-primary'],
+      ),
+    ).toThrowErrorMatchingInlineSnapshot(`
+      [Error: \`@theme\` blocks must only contain custom properties or \`@keyframes\`.
+
+        @theme {
+      >   .foo {
+      >     --color-primary: blue;
+      >   }
+        }
+        ]
+    `)
+  })
+
   test('`@tailwind utilities` is only processed once', () => {
     expect(
       compileCss(

packages/tailwindcss/src/index.ts

@@ -1,6 +1,6 @@
 import { Features, transform } from 'lightningcss'
 import { version } from '../package.json'
-import { comment, decl, rule, toCss, walk, type AstNode, type Rule } from './ast'
+import { WalkAction, comment, decl, rule, toCss, walk, type AstNode, type Rule } from './ast'
 import { compileCandidates } from './compile'
 import * as CSS from './css-parser'
 import { buildDesignSystem } from './design-system'
@@ -23,19 +23,29 @@ export function compile(css: string, rawCandidates: string[]) {
     if (node.selector !== '@theme') return
 
     // Record all custom properties in the `@theme` declaration
-    walk([node], (node, { replaceWith }) => {
+    walk(node.nodes, (node, { replaceWith }) => {
       // Collect `@keyframes` rules to re-insert with theme variables later,
       // since the `@theme` rule itself will be removed.
       if (node.kind === 'rule' && node.selector.startsWith('@keyframes ')) {
         keyframesRules.push(node)
         replaceWith([])
+        return WalkAction.Skip
+      }
+
+      if (node.kind === 'comment') return
+      if (node.kind === 'declaration' && node.property.startsWith('--')) {
+        theme.add(node.property, node.value)
         return
       }
 
-      if (node.kind !== 'declaration') return
-      if (!node.property.startsWith('--')) return
+      let snippet = toCss([rule('@theme', [node])])
+        .split('\n')
+        .map((line, idx, all) => `${idx === 0 || idx >= all.length - 2 ? ' ' : '>'} ${line}`)
+        .join('\n')
 
-      theme.add(node.property, node.value)
+      throw new Error(
+        `\`@theme\` blocks must only contain custom properties or \`@keyframes\`.\n\n${snippet}`,
+      )
     })
 
     // Keep a reference to the first `@theme` rule to update with the full theme
@@ -45,6 +55,7 @@ export function compile(css: string, rawCandidates: string[]) {
     } else {
       replaceWith([])
     }
+    return WalkAction.Skip
   })
 
   // Output final set of theme variables at the position of the first `@theme`
@@ -94,7 +105,7 @@ export function compile(css: string, rawCandidates: string[]) {
       // Stop walking after finding `@tailwind utilities` to avoid walking all
       // of the generated CSS. This means `@tailwind utilities` can only appear
       // once per file but that's the intended usage at this point in time.
-      return false
+      return WalkAction.Stop
     }
   })
 
@@ -145,6 +156,7 @@ export function compile(css: string, rawCandidates: string[]) {
     walk(ast, (node, { replaceWith }) => {
       if (node.kind === 'rule' && node.selector === '@media reference') {
         replaceWith([])
+        return WalkAction.Skip
       }
     })
   }