Skip to content

Commit acb27ef

Browse files
authored
Generalize the walk implementation (#19126)
This PR generalizes the `walk` implementations we have. What's important here is that we currently have multiple `walk` implementations, one for the AST, one for the `SelectorParser`, one for the `ValueParser`. Sometimes, we also need to go up the tree in a depth-first manner. For that, we have `walkDepth` implementations. The funny thing is, all these implementations are very very similar, even the kinds of trees are very similar. They are just objects with `nodes: []` as children. So this PR introduces a generic `walk` function that can work on all of these trees. There are also some situations where you need to go down and back up the tree. For this reason, we added an `enter` and `exit` phase: ```ts walk(ast, { enter(node, ctx) {}, exit(node, ctx) {}, }) ``` This means that you don't need to `walk(ast)` and later `walkDepth(ast)` in case you wanted to do something _after_ visiting all nodes. The API of these walk functions also slightly changed to fix some problems we've had before. One is the `replaceWith` function. You could technically call it multiple times, but that doesn't make sense so instead you always have to return an explicit `WalkAction`. The possibilities are: ```ts // The ones we already had WalkAction.Continue // Continue walking as normal, the default behavior WalkAction.Skip // Skip walking the `nodes` of the current node WalkAction.Stop // Stop the entire walk // The new ones WalkAction.Replace(newNode) // Replace the current node, and continue walking the new node(s) WalkAction.ReplaceSkip(newNode) // Replace the current node, but don't walk the new node(s) WalkAction.ReplaceStop(newNode) // Replace the current node, but stop the entire walk ``` To make sure that we can walk in both directions, and to make sure we have proper control over when to walk which nodes, the `walk` function is implemented in an iterative manner using a stack instead of recursion. This also means that a `WalkAction.Stop` or `WalkAction.ReplaceStop` will immediately stop the walk, without unwinding the entire call stack. Some notes: - The CSS AST does have `context` nodes, for this we can build up the context lazily when we need it. I added a `cssContext(ctx)` that gives you an enhanced context including the `context` object that you can read information from. - The second argument of the `walk` function can still be a normal function, which is equivalent to `{ enter: fn }`. Let's also take a look at some numbers. With this new implementation, each `walk` is roughly ~1.3-1.5x faster than before. If you look at the memory usage (especially in Bun) we go from `~2.2GB` peak memory usage, to `~300mb` peak memory usage. Some benchmarks on small and big trees (M1 Max): <img width="2062" height="1438" alt="image" src="https://github.com/user-attachments/assets/5ec8c22a-9de8-4e08-869a-18c0d30eb7e8" /> <img width="2062" height="1246" alt="image" src="https://github.com/user-attachments/assets/e89d4b8e-29ca-4aee-8fd2-b7c043d3bbf4" /> We also ran some benchmarks on @thecrypticace's M3 Max: <img width="1598" height="1452" alt="image" src="https://github.com/user-attachments/assets/3b06b6fe-2497-4f24-a428-1a0e2af3896a" /> In node the memory difference isn't that big, but the performance itself is still better: <img width="2034" height="1586" alt="image" src="https://github.com/user-attachments/assets/ef28ae14-b53e-4912-9621-531f3b02898f" /> In summary: 1. Single `walk` implementation for multiple use cases 2. Support for `enter` and `exit` phases 3. New `WalkAction` possibilities for better control 4. Overall better performance 5. ... and lower memory usage ## Test plan 1. All tests still pass (but had to adjust some of the APIs if `walk` was used inside tests). 2. Added new tests for the `walk` implementation 3. Ran local benchmarks to verify the performance improvements
1 parent fc63ce7 commit acb27ef

29 files changed

+2174
-731
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,9 @@
55
// Minor modifications have been made to work with the Tailwind CSS codebase
66

77
import * as path from 'node:path'
8-
import { toCss, walk } from '../../tailwindcss/src/ast'
8+
import { toCss } from '../../tailwindcss/src/ast'
99
import { parse } from '../../tailwindcss/src/css-parser'
10+
import { walk } from '../../tailwindcss/src/walk'
1011
import { normalizePath } from './normalize-path'
1112

1213
const cssUrlRE =

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import type { Config } from '../../../../tailwindcss/src/compat/plugin-api'
55
import type { DesignSystem } from '../../../../tailwindcss/src/design-system'
66
import { toKeyPath } from '../../../../tailwindcss/src/utils/to-key-path'
77
import * as ValueParser from '../../../../tailwindcss/src/value-parser'
8+
import { walk, WalkAction } from '../../../../tailwindcss/src/walk'
89
import * as version from '../../utils/version'
910

1011
// Defaults in v4
@@ -117,7 +118,7 @@ function substituteFunctionsInValue(
117118
ast: ValueParser.ValueAstNode[],
118119
handle: (value: string, fallback?: string) => string | null,
119120
) {
120-
ValueParser.walk(ast, (node, { replaceWith }) => {
121+
walk(ast, (node) => {
121122
if (node.kind === 'function' && node.value === 'theme') {
122123
if (node.nodes.length < 1) return
123124

@@ -155,7 +156,7 @@ function substituteFunctionsInValue(
155156
fallbackValues.length > 0 ? handle(path, ValueParser.toCss(fallbackValues)) : handle(path)
156157
if (replacement === null) return
157158

158-
replaceWith(ValueParser.parse(replacement))
159+
return WalkAction.Replace(ValueParser.parse(replacement))
159160
}
160161
})
161162

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
import { walk, WalkAction } from '../../../../tailwindcss/src/ast'
21
import { cloneCandidate, type Candidate, type Variant } from '../../../../tailwindcss/src/candidate'
32
import type { Config } from '../../../../tailwindcss/src/compat/plugin-api'
43
import type { DesignSystem } from '../../../../tailwindcss/src/design-system'
54
import type { Writable } from '../../../../tailwindcss/src/types'
65
import * as ValueParser from '../../../../tailwindcss/src/value-parser'
6+
import { walk, WalkAction } from '../../../../tailwindcss/src/walk'
77

88
export function migrateAutomaticVarInjection(
99
designSystem: DesignSystem,

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

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { isValidSpacingMultiplier } from '../../../../tailwindcss/src/utils/infe
55
import { segment } from '../../../../tailwindcss/src/utils/segment'
66
import { toKeyPath } from '../../../../tailwindcss/src/utils/to-key-path'
77
import * as ValueParser from '../../../../tailwindcss/src/value-parser'
8+
import { walk, WalkAction } from '../../../../tailwindcss/src/walk'
89

910
export const enum Convert {
1011
All = 0,
@@ -27,27 +28,27 @@ export function createConverter(designSystem: DesignSystem, { prettyPrint = fals
2728
let themeModifierCount = 0
2829

2930
// Analyze AST
30-
ValueParser.walk(ast, (node) => {
31+
walk(ast, (node) => {
3132
if (node.kind !== 'function') return
3233
if (node.value !== 'theme') return
3334

3435
// We are only interested in the `theme` function
3536
themeUsageCount += 1
3637

3738
// Figure out if a modifier is used
38-
ValueParser.walk(node.nodes, (child) => {
39+
walk(node.nodes, (child) => {
3940
// If we see a `,`, it means that we have a fallback value
4041
if (child.kind === 'separator' && child.value.includes(',')) {
41-
return ValueParser.ValueWalkAction.Stop
42+
return WalkAction.Stop
4243
}
4344

4445
// If we see a `/`, we have a modifier
4546
else if (child.kind === 'word' && child.value === '/') {
4647
themeModifierCount += 1
47-
return ValueParser.ValueWalkAction.Stop
48+
return WalkAction.Stop
4849
}
4950

50-
return ValueParser.ValueWalkAction.Skip
51+
return WalkAction.Skip
5152
})
5253
})
5354

@@ -172,7 +173,7 @@ function substituteFunctionsInValue(
172173
ast: ValueParser.ValueAstNode[],
173174
handle: (value: string, fallback?: string) => string | null,
174175
) {
175-
ValueParser.walk(ast, (node, { parent, replaceWith }) => {
176+
walk(ast, (node, ctx) => {
176177
if (node.kind === 'function' && node.value === 'theme') {
177178
if (node.nodes.length < 1) return
178179

@@ -210,10 +211,10 @@ function substituteFunctionsInValue(
210211
fallbackValues.length > 0 ? handle(path, ValueParser.toCss(fallbackValues)) : handle(path)
211212
if (replacement === null) return
212213

213-
if (parent) {
214-
let idx = parent.nodes.indexOf(node) - 1
214+
if (ctx.parent) {
215+
let idx = ctx.parent.nodes.indexOf(node) - 1
215216
while (idx !== -1) {
216-
let previous = parent.nodes[idx]
217+
let previous = ctx.parent.nodes[idx]
217218
// Skip the space separator
218219
if (previous.kind === 'separator' && previous.value.trim() === '') {
219220
idx -= 1
@@ -241,7 +242,7 @@ function substituteFunctionsInValue(
241242
}
242243
}
243244

244-
replaceWith(ValueParser.parse(replacement))
245+
return WalkAction.Replace(ValueParser.parse(replacement))
245246
}
246247
})
247248

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
import { walk, type AstNode } from '../../../../tailwindcss/src/ast'
1+
import { type AstNode } from '../../../../tailwindcss/src/ast'
22
import { type Variant } from '../../../../tailwindcss/src/candidate'
33
import type { Config } from '../../../../tailwindcss/src/compat/plugin-api'
44
import type { DesignSystem } from '../../../../tailwindcss/src/design-system'
5+
import { walk } from '../../../../tailwindcss/src/walk'
56
import * as version from '../../utils/version'
67

78
export function migrateVariantOrder(

packages/tailwindcss/src/apply.ts

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import { Features } from '.'
2-
import { cloneAstNode, rule, toCss, walk, WalkAction, type AstNode } from './ast'
2+
import { cloneAstNode, rule, toCss, type AstNode } from './ast'
33
import { compileCandidates } from './compile'
44
import type { DesignSystem } from './design-system'
55
import type { SourceLocation } from './source-maps/source'
66
import { DefaultMap } from './utils/default-map'
77
import { segment } from './utils/segment'
8+
import { walk, WalkAction } from './walk'
89

910
export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
1011
let features = Features.None
@@ -25,7 +26,7 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
2526
let definitions = new DefaultMap(() => new Set<AstNode>())
2627

2728
// Collect all new `@utility` definitions and all `@apply` rules first
28-
walk([root], (node, { parent, path }) => {
29+
walk([root], (node, ctx) => {
2930
if (node.kind !== 'at-rule') return
3031

3132
// Do not allow `@apply` rules inside `@keyframes` rules.
@@ -61,16 +62,15 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
6162
if (node.name === '@apply') {
6263
// `@apply` cannot be top-level, so we need to have a parent such that we
6364
// can replace the `@apply` node with the actual utility classes later.
64-
if (parent === null) return
65+
if (ctx.parent === null) return
6566

6667
features |= Features.AtApply
6768

68-
parents.add(parent)
69+
parents.add(ctx.parent)
6970

7071
for (let dependency of resolveApplyDependencies(node, designSystem)) {
7172
// Mark every parent in the path as having a dependency to that utility.
72-
for (let parent of path) {
73-
if (parent === node) continue
73+
for (let parent of ctx.path()) {
7474
if (!parents.has(parent)) continue
7575
dependencies.get(parent).add(dependency)
7676
}
@@ -158,7 +158,7 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
158158
for (let parent of sorted) {
159159
if (!('nodes' in parent)) continue
160160

161-
walk(parent.nodes, (child, { replaceWith }) => {
161+
walk(parent.nodes, (child) => {
162162
if (child.kind !== 'at-rule' || child.name !== '@apply') return
163163

164164
let parts = child.params.split(/(\s+)/g)
@@ -291,7 +291,7 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
291291
}
292292
}
293293

294-
replaceWith(newNodes)
294+
return WalkAction.Replace(newNodes)
295295
}
296296
})
297297
}

packages/tailwindcss/src/ast.test.ts

Lines changed: 30 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,17 @@ import { expect, it } from 'vitest'
22
import {
33
atRule,
44
context,
5+
cssContext,
56
decl,
67
optimizeAst,
78
styleRule,
89
toCss,
9-
walk,
10-
WalkAction,
1110
type AstNode,
1211
} from './ast'
1312
import * as CSS from './css-parser'
1413
import { buildDesignSystem } from './design-system'
1514
import { Theme } from './theme'
15+
import { walk, WalkAction } from './walk'
1616

1717
const css = String.raw
1818
const defaultDesignSystem = buildDesignSystem(new Theme())
@@ -31,7 +31,7 @@ it('should pretty print an AST', () => {
3131
})
3232

3333
it('allows the placement of context nodes', () => {
34-
const ast = [
34+
let ast: AstNode[] = [
3535
styleRule('.foo', [decl('color', 'red')]),
3636
context({ context: 'a' }, [
3737
styleRule('.bar', [
@@ -48,17 +48,18 @@ it('allows the placement of context nodes', () => {
4848
let blueContext
4949
let greenContext
5050

51-
walk(ast, (node, { context }) => {
51+
walk(ast, (node, _ctx) => {
5252
if (node.kind !== 'declaration') return
53+
let ctx = cssContext(_ctx)
5354
switch (node.value) {
5455
case 'red':
55-
redContext = context
56+
redContext = ctx.context
5657
break
5758
case 'blue':
58-
blueContext = context
59+
blueContext = ctx.context
5960
break
6061
case 'green':
61-
greenContext = context
62+
greenContext = ctx.context
6263
break
6364
}
6465
})
@@ -292,25 +293,25 @@ it('should not emit exact duplicate declarations in the same rule', () => {
292293
it('should only visit children once when calling `replaceWith` with single element array', () => {
293294
let visited = new Set()
294295

295-
let ast = [
296+
let ast: AstNode[] = [
296297
atRule('@media', '', [styleRule('.foo', [decl('color', 'blue')])]),
297298
styleRule('.bar', [decl('color', 'blue')]),
298299
]
299300

300-
walk(ast, (node, { replaceWith }) => {
301+
walk(ast, (node) => {
301302
if (visited.has(node)) {
302303
throw new Error('Visited node twice')
303304
}
304305
visited.add(node)
305306

306-
if (node.kind === 'at-rule') replaceWith(node.nodes)
307+
if (node.kind === 'at-rule') return WalkAction.Replace(node.nodes)
307308
})
308309
})
309310

310311
it('should only visit children once when calling `replaceWith` with multi-element array', () => {
311312
let visited = new Set()
312313

313-
let ast = [
314+
let ast: AstNode[] = [
314315
atRule('@media', '', [
315316
context({}, [
316317
styleRule('.foo', [decl('color', 'red')]),
@@ -320,19 +321,20 @@ it('should only visit children once when calling `replaceWith` with multi-elemen
320321
styleRule('.bar', [decl('color', 'green')]),
321322
]
322323

323-
walk(ast, (node, { replaceWith }) => {
324+
walk(ast, (node) => {
324325
let key = id(node)
325326
if (visited.has(key)) {
326327
throw new Error('Visited node twice')
327328
}
328329
visited.add(key)
329330

330-
if (node.kind === 'at-rule') replaceWith(node.nodes)
331+
if (node.kind === 'at-rule') return WalkAction.Replace(node.nodes)
331332
})
332333

333334
expect(visited).toMatchInlineSnapshot(`
334335
Set {
335336
"@media ",
337+
"<context>",
336338
".foo",
337339
"color: red",
338340
".baz",
@@ -348,14 +350,13 @@ it('should never visit children when calling `replaceWith` with `WalkAction.Skip
348350

349351
let inner = styleRule('.foo', [decl('color', 'blue')])
350352

351-
let ast = [atRule('@media', '', [inner]), styleRule('.bar', [decl('color', 'blue')])]
353+
let ast: AstNode[] = [atRule('@media', '', [inner]), styleRule('.bar', [decl('color', 'blue')])]
352354

353-
walk(ast, (node, { replaceWith }) => {
355+
walk(ast, (node) => {
354356
visited.add(node)
355357

356358
if (node.kind === 'at-rule') {
357-
replaceWith(node.nodes)
358-
return WalkAction.Skip
359+
return WalkAction.ReplaceSkip(node.nodes)
359360
}
360361
})
361362

@@ -413,11 +414,10 @@ it('should skip the correct number of children based on the replaced children no
413414
decl('--index', '4'),
414415
]
415416
let visited: string[] = []
416-
walk(ast, (node, { replaceWith }) => {
417+
walk(ast, (node) => {
417418
visited.push(id(node))
418419
if (node.kind === 'declaration' && node.value === '2') {
419-
replaceWith([])
420-
return WalkAction.Skip
420+
return WalkAction.ReplaceSkip([])
421421
}
422422
})
423423

@@ -441,11 +441,10 @@ it('should skip the correct number of children based on the replaced children no
441441
decl('--index', '4'),
442442
]
443443
let visited: string[] = []
444-
walk(ast, (node, { replaceWith }) => {
444+
walk(ast, (node) => {
445445
visited.push(id(node))
446446
if (node.kind === 'declaration' && node.value === '2') {
447-
replaceWith([])
448-
return WalkAction.Continue
447+
return WalkAction.Replace([])
449448
}
450449
})
451450

@@ -469,11 +468,10 @@ it('should skip the correct number of children based on the replaced children no
469468
decl('--index', '4'),
470469
]
471470
let visited: string[] = []
472-
walk(ast, (node, { replaceWith }) => {
471+
walk(ast, (node) => {
473472
visited.push(id(node))
474473
if (node.kind === 'declaration' && node.value === '2') {
475-
replaceWith([decl('--index', '2.1')])
476-
return WalkAction.Skip
474+
return WalkAction.ReplaceSkip([decl('--index', '2.1')])
477475
}
478476
})
479477

@@ -497,11 +495,10 @@ it('should skip the correct number of children based on the replaced children no
497495
decl('--index', '4'),
498496
]
499497
let visited: string[] = []
500-
walk(ast, (node, { replaceWith }) => {
498+
walk(ast, (node) => {
501499
visited.push(id(node))
502500
if (node.kind === 'declaration' && node.value === '2') {
503-
replaceWith([decl('--index', '2.1')])
504-
return WalkAction.Continue
501+
return WalkAction.Replace([decl('--index', '2.1')])
505502
}
506503
})
507504

@@ -526,11 +523,10 @@ it('should skip the correct number of children based on the replaced children no
526523
decl('--index', '4'),
527524
]
528525
let visited: string[] = []
529-
walk(ast, (node, { replaceWith }) => {
526+
walk(ast, (node) => {
530527
visited.push(id(node))
531528
if (node.kind === 'declaration' && node.value === '2') {
532-
replaceWith([decl('--index', '2.1'), decl('--index', '2.2')])
533-
return WalkAction.Skip
529+
return WalkAction.ReplaceSkip([decl('--index', '2.1'), decl('--index', '2.2')])
534530
}
535531
})
536532

@@ -554,11 +550,10 @@ it('should skip the correct number of children based on the replaced children no
554550
decl('--index', '4'),
555551
]
556552
let visited: string[] = []
557-
walk(ast, (node, { replaceWith }) => {
553+
walk(ast, (node) => {
558554
visited.push(id(node))
559555
if (node.kind === 'declaration' && node.value === '2') {
560-
replaceWith([decl('--index', '2.1'), decl('--index', '2.2')])
561-
return WalkAction.Continue
556+
return WalkAction.Replace([decl('--index', '2.1'), decl('--index', '2.2')])
562557
}
563558
})
564559

0 commit comments

Comments
 (0)