Skip to content

Commit 55789a3

Browse files
committed
refact: clean code and some improvements
1 parent d23a49d commit 55789a3

File tree

3 files changed

+106
-105
lines changed

3 files changed

+106
-105
lines changed

src/__snapshots__/index.test.ts.snap

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
11
// Jest Snapshot v1, https://goo.gl/fbAQLP
22

33
exports[`transformer should compile 1`] = `
4-
"const input_1 = [1, 2, 3, 4];
4+
"declare function MACRO<T>(t: T): T;
5+
declare interface Array<T> {
6+
MAP: Array<T>[\\"map\\"];
7+
FILTER: Array<T>[\\"filter\\"];
8+
}
9+
const input_1 = [1, 2, 3, 4];
510
const length_1 = input_1.length;
611
const result_1 = [];
712
for (let i_1 = 0; i_1 < length_1; i_1++) {
@@ -11,15 +16,15 @@ for (let i_1 = 0; i_1 < length_1; i_1++) {
1116
}
1217
const input_2 = result_1;
1318
const length_2 = input_2.length;
14-
const result_2 = new Array(length_2);
19+
const result_2 = (new Array(length_2) as L[]);
1520
for (let i_2 = 0; i_2 < length_2; i_2++) {
1621
result_2[i_2] =
1722
input_2[i_2]
1823
+ 1;
1924
}
2025
const input_3 = result_2;
2126
const length_3 = input_3.length;
22-
const result_3 = new Array(length_3);
27+
const result_3 = (new Array(length_3) as L[]);
2328
for (let i_3 = 0; i_3 < length_3; i_3++) {
2429
result_3[i_3] =
2530
input_3[i_3].toString();

src/index.test.ts

Lines changed: 14 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -6,39 +6,22 @@ import { readFileSync } from "fs";
66
describe("transformer", () => {
77
it("should compile", () => {
88
const inputFile = resolve(__dirname, "__fixtures/input.ts");
9-
const result = compileCode(readFileSync(inputFile).toString());
9+
const result = transform(readFileSync(inputFile).toString());
1010
expect(result).toMatchSnapshot();
1111
});
1212
});
1313

14-
// function compile(file: string): string {
15-
// let content = "";
16-
// const program = ts.createProgram([file], {
17-
// target: ts.ScriptTarget.ESNext,
18-
// module: ts.ModuleKind.CommonJS
19-
// });
20-
// program.emit(
21-
// undefined,
22-
// (_, result) => (content = result),
23-
// undefined,
24-
// undefined,
25-
// {
26-
// after: [transformer(program)]
27-
// }
28-
// );
29-
// return content;
30-
// }
31-
32-
function compileCode(source: string): string {
33-
let result = ts.transpileModule(source, {
34-
transformers: {
35-
before: [transformer()]
36-
},
37-
compilerOptions: {
38-
target: ts.ScriptTarget.ESNext,
39-
module: ts.ModuleKind.CommonJS
40-
}
41-
});
42-
43-
return result.outputText;
14+
function transform(sourceText: string) {
15+
const source = ts.createSourceFile(
16+
"temp.ts",
17+
sourceText,
18+
ts.ScriptTarget.ESNext
19+
);
20+
const result = ts.transform<ts.SourceFile>(source, [transformer()], {});
21+
const printer = ts.createPrinter();
22+
return printer.printNode(
23+
ts.EmitHint.Unspecified,
24+
result.transformed[0],
25+
ts.createSourceFile("result.ts", "", ts.ScriptTarget.ESNext)
26+
);
4427
}

src/index.ts

Lines changed: 84 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
import * as ts from "typescript";
22

33
class 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

159144
const 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

164177
export default transformer;

0 commit comments

Comments
 (0)