Skip to content

Commit 20e9014

Browse files
erbesharaterbesharat
committedDec 8, 2024·
Initial setup for concurrent workers
1 parent 0cc7191 commit 20e9014

File tree

6 files changed

+250
-236
lines changed

6 files changed

+250
-236
lines changed
 

‎package-lock.json

+1-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/cli/src/api/split.ts

+19-12
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,19 @@
11
import fs from "fs";
22
import path from "path";
3+
import { z } from "zod";
34
import DependencyTreeManager from "../dependencyManager/dependencyManager";
5+
import { Group } from "../dependencyManager/types";
46
import { cleanupOutputDir, createOutputDir } from "../helper/file";
57
import SplitRunner from "../splitRunner/splitRunner";
68
import { splitSchema } from "./helpers/validation";
7-
import { z } from "zod";
8-
import { Group } from "../dependencyManager/types";
99

1010
export function split(payload: z.infer<typeof splitSchema>) {
1111
console.time("split command");
1212
const groupMap: Record<number, Group> = {};
1313

1414
// Get the dependency tree
1515
const dependencyTreeManager = new DependencyTreeManager(
16-
payload.entrypointPath,
16+
payload.entrypointPath
1717
);
1818

1919
const outputDir = payload.outputDir || path.dirname(payload.entrypointPath);
@@ -33,15 +33,22 @@ export function split(payload: z.infer<typeof splitSchema>) {
3333
const targetDir = path.dirname(payload.entrypointPath);
3434
const annotationDirectory = path.join(outputDir, index.toString());
3535

36-
files.forEach((file) => {
37-
const relativeFileNamePath = path.relative(targetDir, file.path);
38-
const destinationPath = path.join(
39-
annotationDirectory,
40-
relativeFileNamePath,
41-
);
42-
fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
43-
fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
44-
});
36+
files
37+
.then((files) => {
38+
files.forEach((file) => {
39+
const relativeFileNamePath = path.relative(targetDir, file.path);
40+
const destinationPath = path.join(
41+
annotationDirectory,
42+
relativeFileNamePath
43+
);
44+
fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
45+
fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
46+
});
47+
})
48+
.catch((error) => {
49+
console.error(error);
50+
throw error;
51+
});
4552
});
4653

4754
// Store the processed annotations in the output directory

‎packages/cli/src/commands/split.ts

+19-12
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1-
import path from "path";
21
import fs from "fs";
2+
import path from "path";
33
import DependencyTreeManager from "../dependencyManager/dependencyManager";
4+
import { Group } from "../dependencyManager/types";
45
import { cleanupOutputDir, createOutputDir } from "../helper/file";
56
import SplitRunner from "../splitRunner/splitRunner";
6-
import { Group } from "../dependencyManager/types";
77

88
export default function splitCommandHandler(
99
entrypointPath: string, // Path to the entrypoint file
10-
outputDir: string, // Path to the output directory
10+
outputDir: string // Path to the output directory
1111
) {
1212
const groupMap: Record<number, Group> = {};
1313

@@ -27,15 +27,22 @@ export default function splitCommandHandler(
2727
const targetDir = path.dirname(entrypointPath);
2828
const annotationDirectory = path.join(outputDir, index.toString());
2929

30-
files.forEach((file) => {
31-
const relativeFileNamePath = path.relative(targetDir, file.path);
32-
const destinationPath = path.join(
33-
annotationDirectory,
34-
relativeFileNamePath,
35-
);
36-
fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
37-
fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
38-
});
30+
files
31+
.then((files) => {
32+
files.forEach((file) => {
33+
const relativeFileNamePath = path.relative(targetDir, file.path);
34+
const destinationPath = path.join(
35+
annotationDirectory,
36+
relativeFileNamePath
37+
);
38+
fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
39+
fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
40+
});
41+
})
42+
.catch((error) => {
43+
console.error(error);
44+
throw error;
45+
});
3946
});
4047

4148
// Store the processed annotations in the output directory

‎packages/cli/src/languagesPlugins/types.ts

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import { Group } from "../dependencyManager/types";
21
import Parser from "tree-sitter";
2+
import { Group } from "../dependencyManager/types";
33

44
export interface DepImportIdentifier {
55
// Specific to each programing languages. Used by the language plugins.
@@ -61,7 +61,7 @@ export interface LanguagePlugin {
6161

6262
removeAnnotationFromOtherGroups(
6363
sourceCode: string,
64-
groupToKeep: Group,
64+
groupToKeep: Group
6565
): string;
6666

6767
getImports(filePath: string, node: Parser.SyntaxNode): DepImport[];
@@ -71,7 +71,7 @@ export interface LanguagePlugin {
7171
cleanupInvalidImports(
7272
filePath: string,
7373
sourceCode: string,
74-
exportMap: Map<string, DepExport[]>,
74+
exportMap: Map<string, DepExport[]>
7575
): string;
7676

7777
cleanupUnusedImports(filePath: string, sourceCode: string): string;

‎packages/cli/src/splitRunner/splitRunner.ts

+20-208
Original file line numberDiff line numberDiff line change
@@ -1,231 +1,43 @@
1-
import { Group } from "../dependencyManager/types";
2-
import { removeIndexesFromSourceCode } from "../helper/file";
1+
import path from "path";
2+
import { Worker } from "worker_threads";
33
import DependencyTreeManager from "../dependencyManager/dependencyManager";
4+
import { Group } from "../dependencyManager/types";
45
import { File } from "./types";
5-
import Parser from "tree-sitter";
6-
import assert from "assert";
7-
import { getLanguagePlugin } from "../languagesPlugins";
8-
import { DepExport } from "../languagesPlugins/types";
96

107
class SplitRunner {
118
private dependencyTreeManager: DependencyTreeManager;
12-
private entrypointPath: string;
139
private group: Group;
14-
private files: File[];
1510

1611
constructor(dependencyTreeManager: DependencyTreeManager, group: Group) {
1712
this.dependencyTreeManager = dependencyTreeManager;
18-
this.entrypointPath = dependencyTreeManager.dependencyTree.path;
1913
this.group = group;
20-
this.files = dependencyTreeManager.getFiles();
21-
}
22-
23-
#removeAnnotationFromOtherGroups() {
24-
this.files = this.files.map((file) => {
25-
const languagePlugin = getLanguagePlugin(this.entrypointPath, file.path);
26-
27-
const updatedSourceCode = languagePlugin.removeAnnotationFromOtherGroups(
28-
file.sourceCode,
29-
this.group,
30-
);
31-
return { ...file, sourceCode: updatedSourceCode };
32-
});
33-
}
34-
35-
#getExportMap() {
36-
const exportMap = new Map<string, DepExport[]>();
37-
38-
this.files.forEach((file) => {
39-
const languagePlugin = getLanguagePlugin(this.entrypointPath, file.path);
40-
41-
const tree = languagePlugin.parser.parse(file.sourceCode);
42-
43-
const exports = languagePlugin.getExports(tree.rootNode);
44-
45-
exportMap.set(file.path, exports);
46-
});
47-
48-
return exportMap;
49-
}
50-
51-
#removeInvalidImportsAndUsages(exportMap: Map<string, DepExport[]>) {
52-
this.files = this.files.map((file) => {
53-
const languagePlugin = getLanguagePlugin(this.entrypointPath, file.path);
54-
55-
const updatedSourceCode = languagePlugin.cleanupInvalidImports(
56-
file.path,
57-
file.sourceCode,
58-
exportMap,
59-
);
60-
61-
return { ...file, sourceCode: updatedSourceCode };
62-
});
6314
}
6415

65-
#removeUnusedImports() {
66-
this.files = this.files.map((file) => {
67-
const languagePlugin = getLanguagePlugin(this.entrypointPath, file.path);
16+
async run(): Promise<File[]> {
17+
console.time("\nSplitting");
6818

69-
const updatedSourceCode = languagePlugin.cleanupUnusedImports(
70-
file.path,
71-
file.sourceCode,
72-
);
73-
74-
return { ...file, sourceCode: updatedSourceCode };
19+
const worker = new Worker(path.resolve(__dirname, "worker.js"), {
20+
workerData: {
21+
entrypointPath: this.dependencyTreeManager.dependencyTree.path,
22+
group: this.group,
23+
files: this.dependencyTreeManager.getFiles(),
24+
},
7525
});
76-
}
77-
78-
#removeUnusedFiles() {
79-
let fileRemoved = true;
80-
while (fileRemoved) {
81-
fileRemoved = false;
82-
83-
// We always want to keep the entrypoint file.
84-
// It will never be imported anywhere, so we add it now.
85-
const filesToKeep = new Set<string>();
86-
filesToKeep.add(this.dependencyTreeManager.dependencyTree.path);
87-
88-
this.files.forEach((file) => {
89-
const languagePlugin = getLanguagePlugin(
90-
this.entrypointPath,
91-
file.path,
92-
);
93-
94-
const tree = languagePlugin.parser.parse(file.sourceCode);
95-
96-
const imports = languagePlugin.getImports(file.path, tree.rootNode);
97-
98-
imports.forEach((depImport) => {
99-
if (depImport.isExternal || !depImport.source) {
100-
// Ignore external dependencies
101-
return;
102-
}
103-
104-
filesToKeep.add(depImport.source);
105-
});
106-
});
107-
108-
const previousFilesLength = this.files.length;
10926

110-
this.files = this.files.filter((file) => {
111-
return filesToKeep.has(file.path);
27+
return new Promise<File[]>((resolve, reject) => {
28+
worker.on("message", (updatedFiles: File[]) => {
29+
console.timeEnd("Splitting");
30+
resolve(updatedFiles);
11231
});
11332

114-
if (this.files.length !== previousFilesLength) {
115-
fileRemoved = true;
116-
}
117-
}
118-
}
119-
120-
#removeUnusedExports(exportMap: Map<string, DepExport[]>) {
121-
let exportDeleted = true;
122-
while (exportDeleted) {
123-
exportDeleted = false;
124-
125-
// const usedExportMap = new Map<string, Export>();
126-
127-
this.files = this.files.map((file) => {
128-
const languagePlugin = getLanguagePlugin(
129-
this.entrypointPath,
130-
file.path,
131-
);
132-
133-
const tree = languagePlugin.parser.parse(file.sourceCode);
134-
135-
const imports = languagePlugin.getImports(file.path, tree.rootNode);
136-
137-
imports.forEach((depImport) => {
138-
if (depImport.isExternal || !depImport.source) {
139-
// Ignore external dependencies
140-
return;
141-
}
142-
143-
// for each import, reconstruct the export map
144-
const depExport = exportMap.get(depImport.source);
145-
if (!depExport) {
146-
throw new Error("Export not found");
147-
}
148-
149-
// check named imports
150-
});
151-
152-
return file;
153-
});
154-
}
155-
// TODO
156-
// Step 1, create variable to track which export is used
157-
// Step 2, iterate over all file imports. If the import is used, mark the export as used
158-
// Step 3, iterate over each file, and remove the unused exports
159-
160-
// Repeat above step until no more unused exports are found
161-
assert(exportMap);
162-
}
163-
164-
#removeErrors() {
165-
this.files = this.files.map((file) => {
166-
const languagePlugin = getLanguagePlugin(this.entrypointPath, file.path);
167-
168-
const tree = languagePlugin.parser.parse(file.sourceCode);
169-
170-
const indexesToRemove: { startIndex: number; endIndex: number }[] = [];
171-
172-
const query = new Parser.Query(
173-
languagePlugin.parser.getLanguage(),
174-
"(ERROR) @error",
175-
);
176-
const errorCaptures = query.captures(tree.rootNode);
177-
errorCaptures.forEach((capture) => {
178-
indexesToRemove.push({
179-
startIndex: capture.node.startIndex,
180-
endIndex: capture.node.endIndex,
181-
});
33+
worker.on("error", reject);
34+
worker.on("exit", (code) => {
35+
if (code !== 0) {
36+
reject(new Error(`Worker stopped with exit code ${code}`));
37+
}
18238
});
183-
184-
const updatedSourceCode = removeIndexesFromSourceCode(
185-
file.sourceCode,
186-
indexesToRemove,
187-
);
188-
189-
return { ...file, sourceCode: updatedSourceCode };
19039
});
19140
}
192-
193-
run() {
194-
console.info("\n");
195-
console.time("Splitting");
196-
197-
console.time("remove annotation from other groups");
198-
this.#removeAnnotationFromOtherGroups();
199-
console.timeEnd("remove annotation from other groups");
200-
201-
console.time("Get export map");
202-
const exportMap = this.#getExportMap();
203-
console.timeEnd("Get export map");
204-
205-
console.time("Remove invalid imports and usages");
206-
this.#removeInvalidImportsAndUsages(exportMap);
207-
console.timeEnd("Remove invalid imports and usages");
208-
209-
console.time("Remove unused imports");
210-
this.#removeUnusedImports();
211-
console.timeEnd("Remove unused imports");
212-
213-
console.time("Remove unused files");
214-
this.#removeUnusedFiles();
215-
console.timeEnd("Remove unused files");
216-
217-
console.time("Remove unused exports");
218-
this.#removeUnusedExports(exportMap);
219-
console.timeEnd("Remove unused exports");
220-
221-
console.time("Remove errors");
222-
this.#removeErrors();
223-
console.timeEnd("Remove errors");
224-
225-
console.timeEnd("Splitting");
226-
227-
return this.files;
228-
}
22941
}
23042

23143
export default SplitRunner;

‎packages/cli/src/splitRunner/worker.ts

+188
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
import assert from "assert";
2+
import Parser from "tree-sitter";
3+
import { parentPort, workerData } from "worker_threads";
4+
import { Group } from "../dependencyManager/types";
5+
import { removeIndexesFromSourceCode } from "../helper/file";
6+
import { getLanguagePlugin } from "../languagesPlugins";
7+
import { DepExport } from "../languagesPlugins/types";
8+
import { File } from "./types";
9+
10+
const {
11+
entrypointPath,
12+
group,
13+
files,
14+
}: {
15+
entrypointPath: string;
16+
group: Group;
17+
files: File[];
18+
} = workerData;
19+
20+
let currentFiles = files;
21+
22+
function removeAnnotationFromOtherGroups() {
23+
currentFiles = currentFiles.map((file) => {
24+
const languagePlugin = getLanguagePlugin(entrypointPath, file.path);
25+
26+
const updatedSourceCode = languagePlugin.removeAnnotationFromOtherGroups(
27+
file.sourceCode,
28+
group
29+
);
30+
return { ...file, sourceCode: updatedSourceCode };
31+
});
32+
}
33+
34+
function getExportMap() {
35+
const exportMap = new Map<string, DepExport[]>();
36+
37+
currentFiles.forEach((file) => {
38+
const languagePlugin = getLanguagePlugin(entrypointPath, file.path);
39+
const tree = languagePlugin.parser.parse(file.sourceCode);
40+
const exports = languagePlugin.getExports(tree.rootNode);
41+
exportMap.set(file.path, exports);
42+
});
43+
44+
return exportMap;
45+
}
46+
47+
function removeInvalidImportsAndUsages(exportMap: Map<string, DepExport[]>) {
48+
currentFiles = currentFiles.map((file) => {
49+
const languagePlugin = getLanguagePlugin(entrypointPath, file.path);
50+
51+
const updatedSourceCode = languagePlugin.cleanupInvalidImports(
52+
file.path,
53+
file.sourceCode,
54+
exportMap
55+
);
56+
57+
return { ...file, sourceCode: updatedSourceCode };
58+
});
59+
}
60+
61+
function removeUnusedImports() {
62+
currentFiles = currentFiles.map((file) => {
63+
const languagePlugin = getLanguagePlugin(entrypointPath, file.path);
64+
65+
const updatedSourceCode = languagePlugin.cleanupUnusedImports(
66+
file.path,
67+
file.sourceCode
68+
);
69+
70+
return { ...file, sourceCode: updatedSourceCode };
71+
});
72+
}
73+
74+
function removeUnusedFiles() {
75+
let fileRemoved = true;
76+
while (fileRemoved) {
77+
fileRemoved = false;
78+
79+
// We always want to keep the entrypoint file.
80+
// It will never be imported anywhere, so we add it now.
81+
const filesToKeep = new Set<string>();
82+
filesToKeep.add(entrypointPath);
83+
84+
currentFiles.forEach((file) => {
85+
const languagePlugin = getLanguagePlugin(entrypointPath, file.path);
86+
87+
const tree = languagePlugin.parser.parse(file.sourceCode);
88+
89+
const imports = languagePlugin.getImports(file.path, tree.rootNode);
90+
91+
imports.forEach((depImport) => {
92+
if (depImport.isExternal || !depImport.source) {
93+
// Ignore external dependencies
94+
return;
95+
}
96+
97+
filesToKeep.add(depImport.source);
98+
});
99+
});
100+
101+
const previousFilesLength = currentFiles.length;
102+
103+
currentFiles = currentFiles.filter((file) => {
104+
return filesToKeep.has(file.path);
105+
});
106+
107+
if (currentFiles.length !== previousFilesLength) {
108+
fileRemoved = true;
109+
}
110+
}
111+
}
112+
113+
function removeUnusedExports(exportMap: Map<string, DepExport[]>) {
114+
let exportDeleted = true;
115+
while (exportDeleted) {
116+
exportDeleted = false;
117+
118+
// TODO: Implement logic if needed. For now, we just assert exportMap.
119+
assert(exportMap);
120+
}
121+
// TODO steps (left as comments):
122+
// Step 1, create variable to track which export is used
123+
// Step 2, iterate over all file imports. If the import is used, mark the export as used
124+
// Step 3, iterate over each file, and remove the unused exports
125+
// Repeat above step until no more unused exports are found
126+
}
127+
128+
function removeErrors() {
129+
currentFiles = currentFiles.map((file) => {
130+
const languagePlugin = getLanguagePlugin(entrypointPath, file.path);
131+
132+
const tree = languagePlugin.parser.parse(file.sourceCode);
133+
134+
const indexesToRemove: { startIndex: number; endIndex: number }[] = [];
135+
136+
const query = new Parser.Query(
137+
languagePlugin.parser.getLanguage(),
138+
"(ERROR) @error"
139+
);
140+
const errorCaptures = query.captures(tree.rootNode);
141+
errorCaptures.forEach((capture) => {
142+
indexesToRemove.push({
143+
startIndex: capture.node.startIndex,
144+
endIndex: capture.node.endIndex,
145+
});
146+
});
147+
148+
const updatedSourceCode = removeIndexesFromSourceCode(
149+
file.sourceCode,
150+
indexesToRemove
151+
);
152+
153+
return { ...file, sourceCode: updatedSourceCode };
154+
});
155+
}
156+
157+
(async () => {
158+
console.time("remove annotation from other groups");
159+
removeAnnotationFromOtherGroups();
160+
console.timeEnd("remove annotation from other groups");
161+
162+
console.time("Get export map");
163+
const exportMap = getExportMap();
164+
console.timeEnd("Get export map");
165+
166+
console.time("Remove invalid imports and usages");
167+
removeInvalidImportsAndUsages(exportMap);
168+
console.timeEnd("Remove invalid imports and usages");
169+
170+
console.time("Remove unused imports");
171+
removeUnusedImports();
172+
console.timeEnd("Remove unused imports");
173+
174+
console.time("Remove unused files");
175+
removeUnusedFiles();
176+
console.timeEnd("Remove unused files");
177+
178+
console.time("Remove unused exports");
179+
removeUnusedExports(exportMap);
180+
console.timeEnd("Remove unused exports");
181+
182+
console.time("Remove errors");
183+
removeErrors();
184+
console.timeEnd("Remove errors");
185+
186+
// Send updated files back to the parent
187+
parentPort?.postMessage(currentFiles);
188+
})();

0 commit comments

Comments
 (0)
Please sign in to comment.