11import * as ts from "typescript" ;
22
33class Transformer {
4- rootMacros : Record < string , ts . Expression > = { } ;
4+ rootMacros = new Map < string , ts . Expression > ( ) ;
55 constructor ( public context : ts . TransformationContext ) { }
66 transform ( node : ts . Node ) : ts . Node {
7- const postExtract = ts . visitNode ( node , this . extractMacros ) ;
8- return ts . visitNode ( postExtract , this . resolveMacros ) ;
7+ return ts . visitNode (
8+ ts . visitNode ( node , this . extractMacros ) ,
9+ this . resolveMacros
10+ ) ;
911 }
10- // Removes macro definition from code and save them for later
11- private extractMacros = ( node : ts . Node ) : ts . Node | undefined => {
12+ extractMacros = ( node : ts . Node ) : ts . Node | undefined => {
1213 if ( ts . isVariableStatement ( node ) ) {
1314 const firstDeclaration = node . declarationList . declarations [ 0 ] ; // TODO maybe check for more
1415 if (
@@ -24,16 +25,15 @@ class Transformer {
2425 ) ;
2526 }
2627 const value = firstDeclaration . initializer . arguments [ 0 ] ;
27- this . rootMacros [ name . text ] = value ;
28+ this . rootMacros . set ( name . text , value ) ;
2829 return undefined ;
2930 }
3031 }
3132 return ts . visitEachChild ( node , this . extractMacros , this . context ) ;
3233 } ;
33- // Search for macros calls and replace them with the macros
34- private resolveMacros = ( node : ts . Node ) : ts . Node | undefined => {
34+ resolveMacros = ( node : ts . Node ) : ts . Node | undefined => {
3535 if ( ts . isBlock ( node ) || ts . isSourceFile ( node ) ) {
36- const newBlock = this . replaceMacros ( node , this . rootMacros ) ;
36+ const newBlock = this . replaceMacros ( node . statements , this . rootMacros ) ;
3737 if ( ts . isBlock ( node ) ) {
3838 return ts . visitEachChild (
3939 ts . updateBlock ( node , newBlock ) ,
@@ -50,10 +50,7 @@ class Transformer {
5050 }
5151 return ts . visitEachChild ( node , this . resolveMacros , this . context ) ;
5252 } ;
53- // Prefix macros variables to avoid name collision, returns "return" expression
54- private fixMacros = (
55- node : ts . Block
56- ) : [ ts . Expression | undefined , ts . Block ] => {
53+ cleanMacro = < T extends ts . Node > ( node : T ) : [ ts . Expression | undefined , T ] => {
5754 const visit = ( node : ts . Node ) : ts . Node | undefined => {
5855 if ( ts . isReturnStatement ( node ) ) {
5956 if ( ! node . expression ) throw new Error ( "Expected macro to return value" ) ;
@@ -67,88 +64,76 @@ class Transformer {
6764 ) ;
6865 }
6966 if ( ts . isVariableDeclaration ( node ) && ts . isIdentifier ( node . name ) ) {
70- variableMap [ node . name . text ] = ts . createUniqueName ( node . name . text ) ;
67+ variableMap . set ( node . name . text , ts . createUniqueName ( node . name . text ) ) ;
7168 }
72- if ( ts . isIdentifier ( node ) && variableMap [ node . text ] ) {
73- return variableMap [ node . text ] ;
69+ if ( ts . isIdentifier ( node ) && variableMap . has ( node . text ) ) {
70+ return variableMap . get ( node . text ) ! ;
7471 }
7572 return ts . visitEachChild ( node , visit , this . context ) ;
7673 } ;
77- const variableMap : Record < string , ts . Identifier > = { } ;
74+ const variableMap = new Map < string , ts . Identifier > ( ) ;
7875 let result : ts . Expression | undefined = undefined ;
7976 const resultNode = ts . visitNode ( node , visit ) ;
8077 return [ result , resultNode ] ;
8178 } ;
82- // Actually replace the macros in the code
83- private replaceMacros = (
84- block : ts . BlockLike ,
85- macros : Record < string , ts . Expression >
79+ replaceMacros = (
80+ statements : ts . NodeArray < ts . Statement > ,
81+ macros : Map < string , ts . Expression >
8682 ) : ts . Statement [ ] => {
87- const visit = ( child : ts . Node ) : ts . Node => {
88- if ( ts . isBlock ( child ) ) {
89- return ts . createBlock ( this . replaceMacros ( child , macros ) ) ;
83+ const visit = ( node : ts . Node ) : ts . Node | undefined => {
84+ if (
85+ [
86+ ts . SyntaxKind . InterfaceDeclaration ,
87+ ts . SyntaxKind . PropertySignature
88+ ] . includes ( node . kind )
89+ ) {
90+ return node ;
91+ }
92+ if ( ts . isBlock ( node ) ) {
93+ return ts . createBlock ( this . replaceMacros ( node . statements , macros ) ) ;
94+ }
95+ if ( ts . isIdentifier ( node ) && macros . has ( node . text ) ) {
96+ return macros . get ( node . text ) ! ;
9097 }
9198 if (
92- ts . isIdentifier ( child ) &&
93- ( macros as Object ) . hasOwnProperty ( child . text )
99+ ts . isCallExpression ( node ) &&
100+ ts . isPropertyAccessExpression ( node . expression ) &&
101+ macros . has ( node . expression . name . text )
94102 ) {
95- return macros [ child . text ] ;
103+ return ts . visitNode (
104+ ts . updateCall ( node , node . expression . name , node . typeArguments , [
105+ node . expression . expression ,
106+ ...node . arguments
107+ ] ) ,
108+ visit
109+ ) ;
96110 }
97111 if (
98- ts . isCallExpression ( child ) &&
99- ts . isIdentifier ( child . expression ) &&
100- ( macros as Object ) . hasOwnProperty ( child . expression . text )
112+ ts . isCallExpression ( node ) &&
113+ ts . isIdentifier ( node . expression ) &&
114+ macros . has ( node . expression . text )
101115 ) {
102- const value = macros [ child . expression . text ] ;
103- if ( ! ts . isArrowFunction ( value ) ) {
116+ const value = macros . get ( node . expression . text ) ! ;
117+ if ( ! ts . isArrowFunction ( value ) && ! ts . isFunctionExpression ( value ) ) {
104118 throw new Error ( "Expected function expression for macro value" ) ;
105119 }
106- const newMacros = { ...macros } ;
107- for (
108- let i = 0 ;
109- i < child . arguments . length && i < value . parameters . length ;
110- i ++
111- ) {
112- const argName = value . parameters [ i ] . name ;
113- if ( ! ts . isIdentifier ( argName ) ) {
114- throw new Error ( "Expected identifier in macro function definition" ) ;
115- }
116- const argValue = child . arguments [ i ] ;
117- newMacros [ argName . text ] = argValue ;
118- }
119-
120- const block = ts . isBlock ( value . body )
121- ? value . body
122- : ts . createBlock ( [ ts . createReturn ( value . body ) ] ) ;
123- const [ resultName , resultBlock ] = this . fixMacros (
120+ const newMacros = new Map ( [
121+ ...macros . entries ( ) ,
122+ ...getNameValueMap ( node . arguments , value . parameters ) . entries ( )
123+ ] ) ;
124+ const [ resultName , resultBlock ] = this . cleanMacro (
124125 ts . visitNode (
125- ts . createBlock ( this . replaceMacros ( block , newMacros ) ) ,
126+ ts . createBlock ( this . replaceMacros ( getStatements ( value ) , newMacros ) ) ,
126127 visit
127128 )
128129 ) ;
129130 result = result . concat ( resultBlock . statements ) ;
130- if ( ! resultName ) {
131- throw new Error ( "Macro should return value" ) ;
132- }
133131 return resultName ;
134132 }
135- if (
136- ts . isCallExpression ( child ) &&
137- ts . isPropertyAccessExpression ( child . expression ) &&
138- ( macros as Object ) . hasOwnProperty ( child . expression . name . text )
139- ) {
140- return ts . visitNode (
141- ts . updateCall ( child , child . expression . name , child . typeArguments , [
142- child . expression . expression ,
143- ...child . arguments
144- ] ) ,
145- visit
146- ) ;
147- }
148- return ts . visitEachChild ( child , visit , this . context ) ;
133+ return ts . visitEachChild ( node , visit , this . context ) ;
149134 } ;
150135 let result : ts . Statement [ ] = [ ] ;
151- for ( const statement of block . statements ) {
136+ for ( const statement of statements ) {
152137 const newStatement = ts . visitNode ( statement , visit ) ;
153138 result . push ( newStatement ) ;
154139 }
@@ -158,7 +143,35 @@ class Transformer {
158143
159144const transformer = (
160145 _program ?: ts . Program
161- ) : ts . TransformerFactory < any > => context => node =>
162- new Transformer ( context ) . transform ( node ) ;
146+ ) : ts . TransformerFactory < any > => context => {
147+ return node => {
148+ return new Transformer ( context ) . transform ( node ) ;
149+ } ;
150+ } ;
151+
152+ function getStatements (
153+ node : ts . FunctionExpression | ts . ArrowFunction
154+ ) : ts . NodeArray < ts . Statement > {
155+ if ( ts . isBlock ( node . body ) ) {
156+ return node . body . statements ;
157+ }
158+ return ts . createNodeArray ( [ ts . createReturn ( node . body ) ] ) ;
159+ }
160+
161+ function getNameValueMap (
162+ values : ts . NodeArray < ts . Expression > ,
163+ args : ts . NodeArray < ts . ParameterDeclaration >
164+ ) {
165+ const map = new Map < string , ts . Expression > ( ) ;
166+ for ( let i = 0 ; i < values . length && i < args . length ; i ++ ) {
167+ const argName = args [ i ] . name ;
168+ if ( ! ts . isIdentifier ( argName ) ) {
169+ throw new Error ( "Expected identifier in macro function definition" ) ;
170+ }
171+ const argValue = values [ i ] ;
172+ map . set ( argName . text , argValue ) ;
173+ }
174+ return map ;
175+ }
163176
164177export default transformer ;
0 commit comments