Very Popular
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359/** * TypeScript AST Transformer for react refresh. * @link https://github.com/facebook/react/issues/16604#issuecomment-528663101 * @link https://github.com/facebook/react/blob/master/packages/react-refresh/src/ReactFreshBabelPlugin.js */
import ts from 'https://esm.sh/typescript'
const f = ts.factory
type TSFunctionLike = ts.FunctionDeclaration | ts.FunctionExpression | ts.ArrowFunction
export class RefreshTransformer { #sf: ts.SourceFile
static refreshSig = '$RefreshSig$' static refreshReg = '$RefreshReg$'
constructor(sf: ts.SourceFile) { this.#sf = sf }
transform() { const statements: ts.Statement[] = [] const components: ts.Identifier[] = [] const signatures: ts.Identifier[] = [] const hookCalls: WeakMap<TSFunctionLike, { id: ts.Identifier, key: string, customHooks: string[] }> = new WeakMap() const seenHooks: Set<string> = new Set()
this.#sf.statements.forEach(node => { if (ts.isFunctionDeclaration(node)) { if (node.name && isComponentishName(node.name.text)) { components.push(node.name) } if (node.name && isHookName(node.name.text)) { seenHooks.add(node.name.text) } const hookCallsSignature = this._getHookCallsSignature(node) if (hookCallsSignature) { const id = f.createUniqueName('_s', ts.GeneratedIdentifierFlags.Optimistic) signatures.push(id) hookCalls.set(node, { id, ...hookCallsSignature }) } } else if (ts.isVariableStatement(node)) { node.declarationList.declarations.forEach(({ name, initializer, modifiers }) => { if ( initializer && ts.isIdentifier(name) && (ts.isFunctionExpression(initializer) || ts.isArrowFunction(initializer)) ) { const hookCallsSignature = this._getHookCallsSignature(initializer) if (isComponentishName(name.text)) { components.push(name) } if (isHookName(name.text)) { seenHooks.add(name.text) } if (hookCallsSignature) { const id = f.createUniqueName('_s', ts.GeneratedIdentifierFlags.Optimistic) signatures.push(id) hookCalls.set(initializer, { id, ...hookCallsSignature }) } } }) } else if (ts.isImportDeclaration(node)) { const name = node.importClause?.name const namedBindings = node.importClause?.namedBindings if (name && isHookName(name.text)) { seenHooks.add(name.text) } if (namedBindings) { namedBindings.forEachChild(node => { if (ts.isImportSpecifier(node) && isHookName(node.name.text)) { seenHooks.add(node.name.text) } }) } } statements.push(node) })
components.forEach(name => { statements.push(f.createExpressionStatement( ts.createCall( f.createIdentifier(RefreshTransformer.refreshReg), undefined, [ name, f.createStringLiteralFromNode(name) ] ) )) })
if (signatures.length > 0) { statements.unshift(f.createVariableStatement( undefined, f.createVariableDeclarationList(signatures.map(id => { return f.createVariableDeclaration( id, undefined, undefined, ts.createCall(f.createIdentifier(RefreshTransformer.refreshSig), undefined, undefined) ) }), ts.NodeFlags.Const) )) }
return ts.updateSourceFileNode( this.#sf, ts.setTextRange( f.createNodeArray( statements.map(node => { if (ts.isFunctionDeclaration(node) && hookCalls.has(node)) { const { id, key, customHooks } = hookCalls.get(node)! const _customHooks = customHooks.filter(name => seenHooks.has(name)) const forceResetComment = !!ts.getLeadingCommentRanges(this.#sf.text, node.pos)?.filter(({ pos, end }) => this.#sf.text.substring(pos, end).includes('@refresh reset')).length; return this._sign(node, node.name!.text, id, key, forceResetComment || _customHooks.length !== customHooks.length, _customHooks)! as ts.Statement[] } if (ts.isVariableStatement(node)) { const forceResetComment = !!ts.getLeadingCommentRanges(this.#sf.text, node.pos)?.filter(({ pos, end }) => this.#sf.text.substring(pos, end).includes('@refresh reset')).length; const _ss: ts.Statement[] = [] const vs = f.createVariableStatement( node.modifiers, f.createVariableDeclarationList( node.declarationList.declarations.map(decl => { const { name, initializer } = decl if ( initializer && ts.isIdentifier(name) && (ts.isFunctionExpression(initializer) || ts.isArrowFunction(initializer)) && hookCalls.has(initializer) ) { const { id, key, customHooks } = hookCalls.get(initializer)! const _customHooks = customHooks.filter(name => seenHooks.has(name)) const [_initializer, _s] = this._sign(initializer, name.text, id, key, forceResetComment || _customHooks.length !== customHooks.length, _customHooks)! _ss.push(_s as ts.Statement) return f.createVariableDeclaration( name, decl.exclamationToken, decl.type, _initializer as ts.ArrowFunction ) } return decl }), node.declarationList.flags ) ) return [vs, ..._ss] } return node }).flat() ), this.#sf.statements ) ) }
private _getHookCallsSignature(fnNode: ts.FunctionDeclaration | ts.FunctionExpression | ts.ArrowFunction) { const hookCalls: { name: string, key: string }[] = [] if (fnNode.body && ts.isBlock(fnNode.body)) { fnNode.body.statements.forEach(s => { if (ts.isVariableStatement(s)) { s.declarationList.declarations.forEach(({ initializer }) => { if ( initializer && ts.isCallExpression(initializer) ) { const sig = this._getHookCallSignature(initializer) if (sig) { hookCalls.push(sig) } } }) } else if ( ts.isExpressionStatement(s) && ts.isCallExpression(s.expression) ) { const sig = this._getHookCallSignature(s.expression) if (sig) { hookCalls.push(sig) } } }) } if (hookCalls.length === 0) { return null } return { key: hookCalls.map(call => call.name + '{' + call.key + '}').join('\n'), customHooks: hookCalls .filter(call => !isBuiltinHook(call.name)) .map(call => call.name), } }
private _getHookCallSignature(ctx: ts.CallExpression) { let name: string const { expression, arguments: args } = ctx if (ts.isIdentifier(expression)) { name = expression.text } else if (ts.isPropertyAccessExpression(expression)) { name = expression.name.text } else { return null } if (!isHookName(name)) { return null }
let key = '' if (ts.isVariableDeclaration(ctx.parent)) { // TODO: if there is no LHS, consider some other heuristic. key = ctx.parent.name.getText() }
// Some built-in Hooks reset on edits to arguments. if (name === 'useState' && args.length > 0) { // useState second argument is initial state. key += '(' + args[0].getText() + ')' } else if (name === 'useReducer' && args.length > 1) { // useReducer second argument is initial state. key += '(' + args[1].getText() + ')' }
return { name, key, } }
private _sign(fnNode: TSFunctionLike, fnName: string, sigId: ts.Identifier, key: string, forceReset: boolean, customHooks: string[]) { if (fnNode.body && ts.isBlock(fnNode.body)) { const _s = f.createExpressionStatement(ts.createCall( sigId, undefined, [ f.createIdentifier(fnName), f.createStringLiteral(key), ...(forceReset || customHooks.length > 0 ? [ forceReset ? f.createTrue() : f.createFalse(), ...(customHooks.length > 0 ? [ f.createFunctionExpression( undefined, undefined, undefined, undefined, undefined, undefined, f.createBlock([ ts.createReturn( ts.createArrayLiteral( customHooks.map(name => f.createIdentifier(name)), false ) ) ], true) ) ] : []) ] : []) ] )) if (ts.isFunctionDeclaration(fnNode)) { return [ f.createFunctionDeclaration( fnNode.decorators, fnNode.modifiers, fnNode.asteriskToken, fnNode.name, fnNode.typeParameters, fnNode.parameters, fnNode.type, f.createBlock([ f.createExpressionStatement(ts.createCall(sigId, undefined, undefined)), ...fnNode.body.statements ], true) ), _s ] } else if (ts.isFunctionExpression(fnNode)) { return [ f.createFunctionExpression( fnNode.modifiers, fnNode.asteriskToken, fnNode.name, fnNode.typeParameters, fnNode.parameters, fnNode.type, f.createBlock([ f.createExpressionStatement(ts.createCall(sigId, undefined, undefined)), ...fnNode.body.statements ], true) ), _s ] } else if (ts.isArrowFunction(fnNode)) { return [ f.createArrowFunction( fnNode.modifiers, fnNode.typeParameters, fnNode.parameters, fnNode.type, fnNode.equalsGreaterThanToken, f.createBlock([ f.createExpressionStatement(ts.createCall(sigId, undefined, undefined)), ...fnNode.body.statements ], true) ), _s ] } } }}
function isComponentishName(name: string) { const c = name.charAt(0) return c >= 'A' && c <= 'Z'}
export function isHookName(name: string) { let c: string return name.startsWith('use') && (c = name.charAt(3)) && c >= 'A' && c <= 'Z'}
function isBuiltinHook(hookName: string) { switch (hookName) { case 'useState': case 'React.useState': case 'useReducer': case 'React.useReducer': case 'useEffect': case 'React.useEffect': case 'useLayoutEffect': case 'React.useLayoutEffect': case 'useMemo': case 'React.useMemo': case 'useCallback': case 'React.useCallback': case 'useRef': case 'React.useRef': case 'useContext': case 'React.useContext': case 'useImperativeMethods': case 'React.useImperativeMethods': case 'useDebugValue': case 'React.useDebugValue': return true; default: return false; }}
export default function transformReactRefresh(ctx: ts.TransformationContext, sf: ts.SourceFile): ts.SourceFile { const t = new RefreshTransformer(sf) return t.transform()}