Skip to content

Commit 87ea445

Browse files
zhiyuanzmjsxzz
zhiyuanzmj
and
sxzz
authored
feat(walk): add walkIdentifiers (#106)
Co-authored-by: 三咲智子 Kevin Deng <sxzz@sxzz.moe>
1 parent 4abdb92 commit 87ea445

File tree

5 files changed

+767
-4
lines changed

5 files changed

+767
-4
lines changed

‎src/check.ts‎

Lines changed: 97 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ export function isIdentifierOf(
8686
node: t.Node | undefined | null,
8787
test: string | string[] | ((id: string) => boolean),
8888
): node is t.Identifier {
89-
return !!node && node.type === 'Identifier' && match(node.name, test)
89+
return isIdentifier(node) && match(node.name, test)
9090
}
9191

9292
/**
@@ -388,4 +388,100 @@ export function isReferenced(
388388
return true
389389
}
390390

391+
export function isIdentifier(
392+
node?: t.Node | undefined | null,
393+
): node is t.Identifier {
394+
return !!node && (node.type === 'Identifier' || node.type === 'JSXIdentifier')
395+
}
396+
397+
export function isStaticProperty(
398+
node?: t.Node | undefined | null,
399+
): node is t.ObjectProperty {
400+
return (
401+
!!node &&
402+
(node.type === 'ObjectProperty' || node.type === 'ObjectMethod') &&
403+
!node.computed
404+
)
405+
}
406+
407+
export function isStaticPropertyKey(node: t.Node, parent: t.Node): boolean {
408+
return isStaticProperty(parent) && parent.key === node
409+
}
410+
411+
export function isForStatement(
412+
stmt: t.Node,
413+
): stmt is t.ForStatement | t.ForOfStatement | t.ForInStatement {
414+
return (
415+
stmt.type === 'ForOfStatement' ||
416+
stmt.type === 'ForInStatement' ||
417+
stmt.type === 'ForStatement'
418+
)
419+
}
420+
421+
export function isReferencedIdentifier(
422+
id: t.Identifier,
423+
parent: t.Node | null | undefined,
424+
parentStack: t.Node[],
425+
): boolean {
426+
if (!parent) {
427+
return true
428+
}
429+
430+
// is a special keyword but parsed as identifier
431+
if (id.name === 'arguments') {
432+
return false
433+
}
434+
435+
if (isReferenced(id, parent)) {
436+
return true
437+
}
438+
439+
// babel's isReferenced check returns false for ids being assigned to, so we
440+
// need to cover those cases here
441+
switch (parent.type) {
442+
case 'AssignmentExpression':
443+
case 'AssignmentPattern':
444+
return true
445+
case 'ObjectPattern':
446+
case 'ArrayPattern':
447+
return isInDestructureAssignment(parent, parentStack)
448+
}
449+
450+
return false
451+
}
452+
453+
export function isInDestructureAssignment(
454+
parent: t.Node,
455+
parentStack: t.Node[],
456+
): boolean {
457+
if (
458+
parent &&
459+
(parent.type === 'ObjectProperty' || parent.type === 'ArrayPattern')
460+
) {
461+
let i = parentStack.length
462+
while (i--) {
463+
const p = parentStack[i]
464+
if (p.type === 'AssignmentExpression') {
465+
return true
466+
} else if (p.type !== 'ObjectProperty' && !p.type.endsWith('Pattern')) {
467+
break
468+
}
469+
}
470+
}
471+
return false
472+
}
473+
474+
export function isInNewExpression(parentStack: t.Node[]): boolean {
475+
let i = parentStack.length
476+
while (i--) {
477+
const p = parentStack[i]
478+
if (p.type === 'NewExpression') {
479+
return true
480+
} else if (p.type !== 'MemberExpression') {
481+
break
482+
}
483+
}
484+
return false
485+
}
486+
391487
/* v8 ignore end */

‎src/extract.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,12 @@ export function extractIdentifiers(
1212
): t.Identifier[] {
1313
switch (node.type) {
1414
case 'Identifier':
15-
identifiers.push(node)
15+
case 'JSXIdentifier':
16+
identifiers.push(node as t.Identifier)
1617
break
1718

18-
case 'MemberExpression': {
19+
case 'MemberExpression':
20+
case 'JSXMemberExpression': {
1921
let object: any = node
2022
while (object.type === 'MemberExpression') {
2123
object = object.object

‎src/walk.ts‎

Lines changed: 184 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11
import { asyncWalk, walk } from 'estree-walker'
2-
import { isExpressionType } from './check'
2+
import {
3+
isExpressionType,
4+
isForStatement,
5+
isFunctionType,
6+
isIdentifier,
7+
isReferencedIdentifier,
8+
} from './check'
9+
import { extractIdentifiers } from './extract'
310
import { resolveString } from './resolve'
11+
import { TS_NODE_TYPES } from './utils'
412
import type { LiteralUnion } from './types'
513
import type * as t from '@babel/types'
614

@@ -219,3 +227,178 @@ export function walkExportDeclaration(
219227

220228
setExport()
221229
}
230+
231+
/**
232+
* Modified from https://github.com/vuejs/core/blob/main/packages/compiler-core/src/babelUtils.ts
233+
* To support browser environments and JSX.
234+
*
235+
* https://github.com/vuejs/core/blob/main/LICENSE
236+
*/
237+
238+
/**
239+
* Return value indicates whether the AST walked can be a constant
240+
*/
241+
export function walkIdentifiers(
242+
root: t.Node,
243+
onIdentifier: (
244+
node: t.Identifier,
245+
parent: t.Node | null | undefined,
246+
parentStack: t.Node[],
247+
isReference: boolean,
248+
isLocal: boolean,
249+
) => void,
250+
includeAll = false,
251+
parentStack: t.Node[] = [],
252+
knownIds: Record<string, number> = Object.create(null),
253+
): void {
254+
const rootExp =
255+
root.type === 'Program'
256+
? root.body[0].type === 'ExpressionStatement' && root.body[0].expression
257+
: root
258+
259+
walkAST<t.Node>(root, {
260+
enter(node: t.Node & { scopeIds?: Set<string> }, parent) {
261+
parent && parentStack.push(parent)
262+
if (
263+
parent &&
264+
parent.type.startsWith('TS') &&
265+
!TS_NODE_TYPES.includes(parent.type as any)
266+
) {
267+
return this.skip()
268+
}
269+
if (isIdentifier(node)) {
270+
const isLocal = !!knownIds[node.name]
271+
const isRefed = isReferencedIdentifier(node, parent, parentStack)
272+
if (includeAll || (isRefed && !isLocal)) {
273+
onIdentifier(node, parent, parentStack, isRefed, isLocal)
274+
}
275+
} else if (
276+
node.type === 'ObjectProperty' &&
277+
parent?.type === 'ObjectPattern'
278+
) {
279+
// mark property in destructure pattern
280+
;(node as any).inPattern = true
281+
} else if (isFunctionType(node)) {
282+
/* v8 ignore start */
283+
if (node.scopeIds) {
284+
node.scopeIds.forEach((id) => markKnownIds(id, knownIds))
285+
/* v8 ignore end */
286+
} else {
287+
// walk function expressions and add its arguments to known identifiers
288+
// so that we don't prefix them
289+
walkFunctionParams(node, (id) =>
290+
markScopeIdentifier(node, id, knownIds),
291+
)
292+
}
293+
} else if (node.type === 'BlockStatement') {
294+
/* v8 ignore start */
295+
if (node.scopeIds) {
296+
node.scopeIds.forEach((id) => markKnownIds(id, knownIds))
297+
/* v8 ignore end */
298+
} else {
299+
// #3445 record block-level local variables
300+
walkBlockDeclarations(node, (id) =>
301+
markScopeIdentifier(node, id, knownIds),
302+
)
303+
}
304+
} else if (node.type === 'CatchClause' && node.param) {
305+
for (const id of extractIdentifiers(node.param)) {
306+
markScopeIdentifier(node, id, knownIds)
307+
}
308+
} else if (isForStatement(node)) {
309+
walkForStatement(node, false, (id) =>
310+
markScopeIdentifier(node, id, knownIds),
311+
)
312+
}
313+
},
314+
leave(node: t.Node & { scopeIds?: Set<string> }, parent) {
315+
parent && parentStack.pop()
316+
if (node !== rootExp && node.scopeIds) {
317+
for (const id of node.scopeIds) {
318+
knownIds[id]--
319+
if (knownIds[id] === 0) {
320+
delete knownIds[id]
321+
}
322+
}
323+
}
324+
},
325+
})
326+
}
327+
328+
export function walkFunctionParams(
329+
node: t.Function,
330+
onIdent: (id: t.Identifier) => void,
331+
): void {
332+
for (const p of node.params) {
333+
for (const id of extractIdentifiers(p)) {
334+
onIdent(id)
335+
}
336+
}
337+
}
338+
339+
export function walkBlockDeclarations(
340+
block: t.BlockStatement | t.Program,
341+
onIdent: (node: t.Identifier) => void,
342+
): void {
343+
for (const stmt of block.body) {
344+
if (stmt.type === 'VariableDeclaration') {
345+
if (stmt.declare) continue
346+
for (const decl of stmt.declarations) {
347+
for (const id of extractIdentifiers(decl.id)) {
348+
onIdent(id)
349+
}
350+
}
351+
} else if (
352+
stmt.type === 'FunctionDeclaration' ||
353+
stmt.type === 'ClassDeclaration'
354+
) {
355+
/* v8 ignore next */
356+
if (stmt.declare || !stmt.id) continue
357+
onIdent(stmt.id)
358+
} else if (isForStatement(stmt)) {
359+
walkForStatement(stmt, true, onIdent)
360+
}
361+
}
362+
}
363+
364+
function walkForStatement(
365+
stmt: t.ForStatement | t.ForOfStatement | t.ForInStatement,
366+
isVar: boolean,
367+
onIdent: (id: t.Identifier) => void,
368+
) {
369+
const variable = stmt.type === 'ForStatement' ? stmt.init : stmt.left
370+
if (
371+
variable &&
372+
variable.type === 'VariableDeclaration' &&
373+
(variable.kind === 'var' ? isVar : !isVar)
374+
) {
375+
for (const decl of variable.declarations) {
376+
for (const id of extractIdentifiers(decl.id)) {
377+
onIdent(id)
378+
}
379+
}
380+
}
381+
}
382+
383+
function markKnownIds(name: string, knownIds: Record<string, number>) {
384+
if (name in knownIds) {
385+
knownIds[name]++
386+
} else {
387+
knownIds[name] = 1
388+
}
389+
}
390+
391+
function markScopeIdentifier(
392+
node: t.Node & { scopeIds?: Set<string> },
393+
child: t.Identifier,
394+
knownIds: Record<string, number>,
395+
) {
396+
const { name } = child
397+
/* v8 ignore start */
398+
if (node.scopeIds && node.scopeIds.has(name)) {
399+
return
400+
}
401+
/* v8 ignore end */
402+
markKnownIds(name, knownIds)
403+
;(node.scopeIds || (node.scopeIds = new Set())).add(name)
404+
}

0 commit comments

Comments
 (0)