extends ctx of CSS AstNode's

The CSS AST has `context` nodes that you can access while walking the
tree, where they build up information over time.

Additionally, when looking at the `path` or the `parent` all the
`context` nodes are hidden.

This adds a small wrapper for the context object when walking over
AstNode's and have to access the `.context`, `.parent` or `.path()`

Author: RobinMalfait

packages/tailwindcss/src/ast.test.ts

@@ -1,5 +1,14 @@
 import { expect, it } from 'vitest'
-import { atRule, context, decl, optimizeAst, styleRule, toCss, type AstNode } from './ast'
+import {
+  atRule,
+  context,
+  cssContext,
+  decl,
+  optimizeAst,
+  styleRule,
+  toCss,
+  type AstNode,
+} from './ast'
 import * as CSS from './css-parser'
 import { buildDesignSystem } from './design-system'
 import { Theme } from './theme'
@@ -39,23 +48,18 @@ it('allows the placement of context nodes', () => {
   let blueContext
   let greenContext
 
-  let walkContext: Record<string, string | boolean> = {}
-  walk(ast, (node) => {
-    if (node.kind === 'context') {
-      walkContext = { ...walkContext, ...node.context }
-      return
-    }
-
+  walk(ast, (node, _ctx) => {
     if (node.kind !== 'declaration') return
+    let ctx = cssContext(_ctx)
     switch (node.value) {
       case 'red':
-        redContext = walkContext
+        redContext = ctx.context
         break
       case 'blue':
-        blueContext = walkContext
+        blueContext = ctx.context
         break
       case 'green':
-        greenContext = walkContext
+        greenContext = ctx.context
         break
     }
   })

packages/tailwindcss/src/ast.ts

@@ -6,7 +6,7 @@ import { Theme, ThemeOptions } from './theme'
 import { DefaultMap } from './utils/default-map'
 import { extractUsedVariables } from './utils/variables'
 import * as ValueParser from './value-parser'
-import { walk, WalkAction } from './walk'
+import { walk, WalkAction, type VisitContext } from './walk'
 
 const AT_SIGN = 0x40
 
@@ -185,6 +185,36 @@ export function cloneAstNode<T extends AstNode>(node: T): T {
   }
 }
 
+export function cssContext(
+  ctx: VisitContext<AstNode>,
+): VisitContext<AstNode> & { context: Record<string, string | boolean> } {
+  return {
+    depth: ctx.depth,
+    get context() {
+      let context: Record<string, string | boolean> = {}
+      for (let child of ctx.path()) {
+        if (child.kind === 'context') {
+          Object.assign(context, child.context)
+        }
+      }
+
+      // Once computed, we never need to compute this again
+      Object.defineProperty(this, 'context', { value: context })
+      return context
+    },
+    get parent() {
+      let parent = (this.path().pop() as Extract<AstNode, { nodes: AstNode[] }>) ?? null
+
+      // Once computed, we never need to compute this again
+      Object.defineProperty(this, 'parent', { value: parent })
+      return parent
+    },
+    path() {
+      return ctx.path().filter((n) => n.kind !== 'context')
+    },
+  }
+}
+
 // 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(

packages/tailwindcss/src/compat/apply-compat-hooks.ts

@@ -1,9 +1,9 @@
 import { Features } from '..'
-import { styleRule, toCss, walk, WalkAction, type AstNode } from '../ast'
+import { cssContext, styleRule, toCss, type AstNode } from '../ast'
 import type { DesignSystem } from '../design-system'
 import type { SourceLocation } from '../source-maps/source'
 import { segment } from '../utils/segment'
-import { WalkAction } from '../walk'
+import { walk, WalkAction } from '../walk'
 import { applyConfigToTheme } from './apply-config-to-theme'
 import { applyKeyframesToTheme } from './apply-keyframes-to-theme'
 import { createCompatConfig } from './config/create-compat-config'
@@ -51,8 +51,9 @@ export async function applyCompatibilityHooks({
     src: SourceLocation | undefined
   }[] = []
 
-  walk(ast, (node, ctx) => {
+  walk(ast, (node, _ctx) => {
     if (node.kind !== 'at-rule') return
+    let ctx = cssContext(_ctx)
 
     // Collect paths from `@plugin` at-rules
     if (node.name === '@plugin') {
@@ -385,10 +386,12 @@ function upgradeToFullPluginSupport({
   if (typeof resolvedConfig.important === 'string') {
     let wrappingSelector = resolvedConfig.important
 
-    walk(ast, (node, ctx) => {
+    walk(ast, (node, _ctx) => {
       if (node.kind !== 'at-rule') return
       if (node.name !== '@tailwind' || node.params !== 'utilities') return
 
+      let ctx = cssContext(_ctx)
+
       // The AST node was already manually wrapped so there's nothing to do
       if (ctx.parent?.kind === 'rule' && ctx.parent.selector === wrappingSelector) {
         return WalkAction.Stop

packages/tailwindcss/src/index.ts

@@ -6,6 +6,7 @@ import {
   comment,
   context,
   context as contextNode,
+  cssContext,
   decl,
   optimizeAst,
   rule,
@@ -165,8 +166,9 @@ async function parseCss(
   let root = null as Root
 
   // Handle at-rules
-  walk(ast, (node, ctx) => {
+  walk(ast, (node, _ctx) => {
     if (node.kind !== 'at-rule') return
+    let ctx = cssContext(_ctx)
 
     // Find `@tailwind utilities` so that we can later replace it with the
     // actual generated utility class CSS.
@@ -460,7 +462,9 @@ async function parseCss(
 
             if (child.name === '@tailwind' && child.params === 'utilities') {
               child.params += ` source(${path})`
-              return WalkAction.ReplaceStop([contextNode({ sourceBase: context.base }, [child])])
+              return WalkAction.ReplaceStop([
+                contextNode({ sourceBase: ctx.context.base }, [child]),
+              ])
             }
           })
         }
@@ -475,6 +479,7 @@ async function parseCss(
           let hasReference = themeParams.includes('reference')
 
           walk(node.nodes, (child) => {
+            if (child.kind === 'context') return
             if (child.kind !== 'at-rule') {
               if (hasReference) {
                 throw new Error(

packages/tailwindcss/src/walk.ts

@@ -33,7 +33,7 @@ type ExitResult<T> = Exclude<WalkResult<T>, { kind: WalkKind.Skip }>
 
 type Parent<T> = T & { nodes: T[] }
 
-interface VisitContext<N> {
+export interface VisitContext<N> {
   parent: Parent<N> | null
   depth: number
   path: () => N[]