Generalize the walk implementation

packages/@tailwindcss-node/src/urls.ts

@@ -5,8 +5,9 @@
 // Minor modifications have been made to work with the Tailwind CSS codebase
 
 import * as path from 'node:path'
-import { toCss, walk } from '../../tailwindcss/src/ast'
+import { toCss } from '../../tailwindcss/src/ast'
 import { parse } from '../../tailwindcss/src/css-parser'
+import { walk } from '../../tailwindcss/src/walk'
 import { normalizePath } from './normalize-path'
 
 const cssUrlRE =

packages/@tailwindcss-upgrade/src/codemods/css/migrate-preflight.ts

@@ -5,6 +5,7 @@ import type { Config } from '../../../../tailwindcss/src/compat/plugin-api'
 import type { DesignSystem } from '../../../../tailwindcss/src/design-system'
 import { toKeyPath } from '../../../../tailwindcss/src/utils/to-key-path'
 import * as ValueParser from '../../../../tailwindcss/src/value-parser'
+import { walk, WalkAction } from '../../../../tailwindcss/src/walk'
 import * as version from '../../utils/version'
 
 // Defaults in v4
@@ -117,7 +118,7 @@ function substituteFunctionsInValue(
   ast: ValueParser.ValueAstNode[],
   handle: (value: string, fallback?: string) => string | null,
 ) {
-  ValueParser.walk(ast, (node, { replaceWith }) => {
+  walk(ast, (node) => {
     if (node.kind === 'function' && node.value === 'theme') {
       if (node.nodes.length < 1) return
 
@@ -155,7 +156,7 @@ function substituteFunctionsInValue(
         fallbackValues.length > 0 ? handle(path, ValueParser.toCss(fallbackValues)) : handle(path)
       if (replacement === null) return
 
-      replaceWith(ValueParser.parse(replacement))
+      return WalkAction.Replace(ValueParser.parse(replacement))
     }
   })
 

packages/@tailwindcss-upgrade/src/codemods/template/migrate-automatic-var-injection.ts

@@ -1,9 +1,9 @@
-import { walk, WalkAction } from '../../../../tailwindcss/src/ast'
 import { cloneCandidate, type Candidate, type Variant } from '../../../../tailwindcss/src/candidate'
 import type { Config } from '../../../../tailwindcss/src/compat/plugin-api'
 import type { DesignSystem } from '../../../../tailwindcss/src/design-system'
 import type { Writable } from '../../../../tailwindcss/src/types'
 import * as ValueParser from '../../../../tailwindcss/src/value-parser'
+import { walk, WalkAction } from '../../../../tailwindcss/src/walk'
 
 export function migrateAutomaticVarInjection(
   designSystem: DesignSystem,

packages/@tailwindcss-upgrade/src/codemods/template/migrate-theme-to-var.ts

@@ -5,6 +5,7 @@ import { isValidSpacingMultiplier } from '../../../../tailwindcss/src/utils/infe
 import { segment } from '../../../../tailwindcss/src/utils/segment'
 import { toKeyPath } from '../../../../tailwindcss/src/utils/to-key-path'
 import * as ValueParser from '../../../../tailwindcss/src/value-parser'
+import { walk, WalkAction } from '../../../../tailwindcss/src/walk'
 
 export const enum Convert {
   All = 0,
@@ -27,27 +28,27 @@ export function createConverter(designSystem: DesignSystem, { prettyPrint = fals
     let themeModifierCount = 0
 
     // Analyze AST
-    ValueParser.walk(ast, (node) => {
+    walk(ast, (node) => {
       if (node.kind !== 'function') return
       if (node.value !== 'theme') return
 
       // We are only interested in the `theme` function
       themeUsageCount += 1
 
       // Figure out if a modifier is used
-      ValueParser.walk(node.nodes, (child) => {
+      walk(node.nodes, (child) => {
         // If we see a `,`, it means that we have a fallback value
         if (child.kind === 'separator' && child.value.includes(',')) {
-          return ValueParser.ValueWalkAction.Stop
+          return WalkAction.Stop
         }
 
         // If we see a `/`, we have a modifier
         else if (child.kind === 'word' && child.value === '/') {
           themeModifierCount += 1
-          return ValueParser.ValueWalkAction.Stop
+          return WalkAction.Stop
         }
 
-        return ValueParser.ValueWalkAction.Skip
+        return WalkAction.Skip
       })
     })
 
@@ -172,7 +173,7 @@ function substituteFunctionsInValue(
   ast: ValueParser.ValueAstNode[],
   handle: (value: string, fallback?: string) => string | null,
 ) {
-  ValueParser.walk(ast, (node, { parent, replaceWith }) => {
+  walk(ast, (node, ctx) => {
     if (node.kind === 'function' && node.value === 'theme') {
       if (node.nodes.length < 1) return
 
@@ -210,10 +211,10 @@ function substituteFunctionsInValue(
         fallbackValues.length > 0 ? handle(path, ValueParser.toCss(fallbackValues)) : handle(path)
       if (replacement === null) return
 
-      if (parent) {
-        let idx = parent.nodes.indexOf(node) - 1
+      if (ctx.parent) {
+        let idx = ctx.parent.nodes.indexOf(node) - 1
         while (idx !== -1) {
-          let previous = parent.nodes[idx]
+          let previous = ctx.parent.nodes[idx]
           // Skip the space separator
           if (previous.kind === 'separator' && previous.value.trim() === '') {
             idx -= 1
@@ -241,7 +242,7 @@ function substituteFunctionsInValue(
         }
       }
 
-      replaceWith(ValueParser.parse(replacement))
+      return WalkAction.Replace(ValueParser.parse(replacement))
     }
   })
 

packages/@tailwindcss-upgrade/src/codemods/template/migrate-variant-order.ts

@@ -1,7 +1,8 @@
-import { walk, type AstNode } from '../../../../tailwindcss/src/ast'
+import { type AstNode } from '../../../../tailwindcss/src/ast'
 import { type Variant } from '../../../../tailwindcss/src/candidate'
 import type { Config } from '../../../../tailwindcss/src/compat/plugin-api'
 import type { DesignSystem } from '../../../../tailwindcss/src/design-system'
+import { walk } from '../../../../tailwindcss/src/walk'
 import * as version from '../../utils/version'
 
 export function migrateVariantOrder(

packages/tailwindcss/src/apply.ts

@@ -1,10 +1,11 @@
 import { Features } from '.'
-import { cloneAstNode, rule, toCss, walk, WalkAction, type AstNode } from './ast'
+import { cloneAstNode, rule, toCss, type AstNode } from './ast'
 import { compileCandidates } from './compile'
 import type { DesignSystem } from './design-system'
 import type { SourceLocation } from './source-maps/source'
 import { DefaultMap } from './utils/default-map'
 import { segment } from './utils/segment'
+import { walk, WalkAction } from './walk'
 
 export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
   let features = Features.None
@@ -25,7 +26,7 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
   let definitions = new DefaultMap(() => new Set<AstNode>())
 
   // Collect all new `@utility` definitions and all `@apply` rules first
-  walk([root], (node, { parent, path }) => {
+  walk([root], (node, ctx) => {
     if (node.kind !== 'at-rule') return
 
     // Do not allow `@apply` rules inside `@keyframes` rules.
@@ -61,16 +62,15 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
     if (node.name === '@apply') {
       // `@apply` cannot be top-level, so we need to have a parent such that we
       // can replace the `@apply` node with the actual utility classes later.
-      if (parent === null) return
+      if (ctx.parent === null) return
 
       features |= Features.AtApply
 
-      parents.add(parent)
+      parents.add(ctx.parent)
 
       for (let dependency of resolveApplyDependencies(node, designSystem)) {
         // Mark every parent in the path as having a dependency to that utility.
-        for (let parent of path) {
-          if (parent === node) continue
+        for (let parent of ctx.path()) {
           if (!parents.has(parent)) continue
           dependencies.get(parent).add(dependency)
         }
@@ -158,7 +158,7 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
   for (let parent of sorted) {
     if (!('nodes' in parent)) continue
 
-    walk(parent.nodes, (child, { replaceWith }) => {
+    walk(parent.nodes, (child) => {
       if (child.kind !== 'at-rule' || child.name !== '@apply') return
 
       let parts = child.params.split(/(\s+)/g)
@@ -291,7 +291,7 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
           }
         }
 
-        replaceWith(newNodes)
+        return WalkAction.Replace(newNodes)
       }
     })
   }

packages/tailwindcss/src/ast.test.ts

@@ -2,17 +2,17 @@ import { expect, it } from 'vitest'
 import {
   atRule,
   context,
+  cssContext,
   decl,
   optimizeAst,
   styleRule,
   toCss,
-  walk,
-  WalkAction,
   type AstNode,
 } from './ast'
 import * as CSS from './css-parser'
 import { buildDesignSystem } from './design-system'
 import { Theme } from './theme'
+import { walk, WalkAction } from './walk'
 
 const css = String.raw
 const defaultDesignSystem = buildDesignSystem(new Theme())
@@ -31,7 +31,7 @@ it('should pretty print an AST', () => {
 })
 
 it('allows the placement of context nodes', () => {
-  const ast = [
+  let ast: AstNode[] = [
     styleRule('.foo', [decl('color', 'red')]),
     context({ context: 'a' }, [
       styleRule('.bar', [
@@ -48,17 +48,18 @@ it('allows the placement of context nodes', () => {
   let blueContext
   let greenContext
 
-  walk(ast, (node, { context }) => {
+  walk(ast, (node, _ctx) => {
     if (node.kind !== 'declaration') return
+    let ctx = cssContext(_ctx)
     switch (node.value) {
       case 'red':
-        redContext = context
+        redContext = ctx.context
         break
       case 'blue':
-        blueContext = context
+        blueContext = ctx.context
         break
       case 'green':
-        greenContext = context
+        greenContext = ctx.context
         break
     }
   })
@@ -292,25 +293,25 @@ it('should not emit exact duplicate declarations in the same rule', () => {
 it('should only visit children once when calling `replaceWith` with single element array', () => {
   let visited = new Set()
 
-  let ast = [
+  let ast: AstNode[] = [
     atRule('@media', '', [styleRule('.foo', [decl('color', 'blue')])]),
     styleRule('.bar', [decl('color', 'blue')]),
   ]
 
-  walk(ast, (node, { replaceWith }) => {
+  walk(ast, (node) => {
     if (visited.has(node)) {
       throw new Error('Visited node twice')
     }
     visited.add(node)
 
-    if (node.kind === 'at-rule') replaceWith(node.nodes)
+    if (node.kind === 'at-rule') return WalkAction.Replace(node.nodes)
   })
 })
 
 it('should only visit children once when calling `replaceWith` with multi-element array', () => {
   let visited = new Set()
 
-  let ast = [
+  let ast: AstNode[] = [
     atRule('@media', '', [
       context({}, [
         styleRule('.foo', [decl('color', 'red')]),
@@ -320,19 +321,20 @@ it('should only visit children once when calling `replaceWith` with multi-elemen
     styleRule('.bar', [decl('color', 'green')]),
   ]
 
-  walk(ast, (node, { replaceWith }) => {
+  walk(ast, (node) => {
     let key = id(node)
     if (visited.has(key)) {
       throw new Error('Visited node twice')
     }
     visited.add(key)
 
-    if (node.kind === 'at-rule') replaceWith(node.nodes)
+    if (node.kind === 'at-rule') return WalkAction.Replace(node.nodes)
   })
 
   expect(visited).toMatchInlineSnapshot(`
     Set {
       "@media ",
+      "<context>",
       ".foo",
       "color: red",
       ".baz",
@@ -348,14 +350,13 @@ it('should never visit children when calling `replaceWith` with `WalkAction.Skip
 
   let inner = styleRule('.foo', [decl('color', 'blue')])
 
-  let ast = [atRule('@media', '', [inner]), styleRule('.bar', [decl('color', 'blue')])]
+  let ast: AstNode[] = [atRule('@media', '', [inner]), styleRule('.bar', [decl('color', 'blue')])]
 
-  walk(ast, (node, { replaceWith }) => {
+  walk(ast, (node) => {
     visited.add(node)
 
     if (node.kind === 'at-rule') {
-      replaceWith(node.nodes)
-      return WalkAction.Skip
+      return WalkAction.ReplaceSkip(node.nodes)
     }
   })
 
@@ -413,11 +414,10 @@ it('should skip the correct number of children based on the replaced children no
       decl('--index', '4'),
     ]
     let visited: string[] = []
-    walk(ast, (node, { replaceWith }) => {
+    walk(ast, (node) => {
       visited.push(id(node))
       if (node.kind === 'declaration' && node.value === '2') {
-        replaceWith([])
-        return WalkAction.Skip
+        return WalkAction.ReplaceSkip([])
       }
     })
 
@@ -441,11 +441,10 @@ it('should skip the correct number of children based on the replaced children no
       decl('--index', '4'),
     ]
     let visited: string[] = []
-    walk(ast, (node, { replaceWith }) => {
+    walk(ast, (node) => {
       visited.push(id(node))
       if (node.kind === 'declaration' && node.value === '2') {
-        replaceWith([])
-        return WalkAction.Continue
+        return WalkAction.Replace([])
       }
     })
 
@@ -469,11 +468,10 @@ it('should skip the correct number of children based on the replaced children no
       decl('--index', '4'),
     ]
     let visited: string[] = []
-    walk(ast, (node, { replaceWith }) => {
+    walk(ast, (node) => {
       visited.push(id(node))
       if (node.kind === 'declaration' && node.value === '2') {
-        replaceWith([decl('--index', '2.1')])
-        return WalkAction.Skip
+        return WalkAction.ReplaceSkip([decl('--index', '2.1')])
       }
     })
 
@@ -497,11 +495,10 @@ it('should skip the correct number of children based on the replaced children no
       decl('--index', '4'),
     ]
     let visited: string[] = []
-    walk(ast, (node, { replaceWith }) => {
+    walk(ast, (node) => {
       visited.push(id(node))
       if (node.kind === 'declaration' && node.value === '2') {
-        replaceWith([decl('--index', '2.1')])
-        return WalkAction.Continue
+        return WalkAction.Replace([decl('--index', '2.1')])
       }
     })
 
@@ -526,11 +523,10 @@ it('should skip the correct number of children based on the replaced children no
       decl('--index', '4'),
     ]
     let visited: string[] = []
-    walk(ast, (node, { replaceWith }) => {
+    walk(ast, (node) => {
       visited.push(id(node))
       if (node.kind === 'declaration' && node.value === '2') {
-        replaceWith([decl('--index', '2.1'), decl('--index', '2.2')])
-        return WalkAction.Skip
+        return WalkAction.ReplaceSkip([decl('--index', '2.1'), decl('--index', '2.2')])
       }
     })
 
@@ -554,11 +550,10 @@ it('should skip the correct number of children based on the replaced children no
       decl('--index', '4'),
     ]
     let visited: string[] = []
-    walk(ast, (node, { replaceWith }) => {
+    walk(ast, (node) => {
       visited.push(id(node))
       if (node.kind === 'declaration' && node.value === '2') {
-        replaceWith([decl('--index', '2.1'), decl('--index', '2.2')])
-        return WalkAction.Continue
+        return WalkAction.Replace([decl('--index', '2.1'), decl('--index', '2.2')])
       }
     })