Skip to content

Commit

Permalink
fix: allocate collections in bump allocator (#57)
Browse files Browse the repository at this point in the history 
Fixes a memory leak where collections weren't being deallocated.

Co-authored-by: irbull irbull@users.noreply.github.com
  • Loading branch information
@dsherret
dsherret committed Feb 21, 2024
1 parent 89338b3 commit 928dfb6
Show file tree
Hide file tree
Showing 7 changed files with 198 additions and 156 deletions.
10 changes: 10 additions & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
[workspace]
resolver = "2"
members = [
"rs-lib",
]
Expand Down
69 changes: 44 additions & 25 deletions generation/generate/generate.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,18 @@ export function generate(analysisResult: AnalysisResult): string {
// hack to avoid yet another lifetime
writer.writeLine("let bump_ref = unsafe { mem::transmute::<&Bump, &'a Bump>(&bump_borrow) };");
writer.writeLine("let info_ref = unsafe { mem::transmute::<&ModuleInfo, &'a ModuleInfo<'a>>(&info) };");
writer.writeLine(`let ast_view = ${getViewForFunctionName("Module")}(info_ref, bump_ref);`);
writer.writeLine("let tokens = info_ref.tokens.map(|t| TokenContainer::new(t));");
writer.writeLine("let tokens_ref = tokens.as_ref().map(|t| unsafe { mem::transmute::<&TokenContainer, &'a TokenContainer<'a>>(t) });");
writer.write("let comments = info_ref.comments.map(|c| CommentContainer::new(");
writer.indent(() => {
writer.writeLine("c.leading,");
writer.writeLine("c.trailing,");
writer.writeLine(`tokens_ref.expect("Tokens must be provided when using comments."),`);
writer.writeLine(`info_ref.text_info.expect("Text info must be provided when using comments"),`);
});
writer.write("));").newLine();
writer.writeLine("let comments_ref = comments.as_ref().map(|t| unsafe { mem::transmute::<&CommentContainer, &'a CommentContainer<'a>>(t) });");
writer.writeLine(`let ast_view = ${getViewForFunctionName("Module")}(info_ref, tokens_ref, comments_ref, bump_ref);`);
writer.writeLine(`let result = with_view(ast_view);`);
writer.writeLine("bump_borrow.reset();");
writer.writeLine("result");
Expand All @@ -102,7 +113,18 @@ export function generate(analysisResult: AnalysisResult): string {
// hack to avoid yet another lifetime
writer.writeLine("let bump_ref = unsafe { mem::transmute::<&Bump, &'a Bump>(&bump_borrow) };");
writer.writeLine("let info_ref = unsafe { mem::transmute::<&ScriptInfo, &'a ScriptInfo<'a>>(&info) };");
writer.writeLine(`let ast_view = ${getViewForFunctionName("Script")}(info_ref, bump_ref);`);
writer.writeLine("let tokens = info_ref.tokens.map(|t| TokenContainer::new(t));");
writer.writeLine("let tokens_ref = tokens.as_ref().map(|t| unsafe { mem::transmute::<&TokenContainer, &'a TokenContainer<'a>>(t) });");
writer.write("let comments = info_ref.comments.map(|c| CommentContainer::new(");
writer.indent(() => {
writer.writeLine("c.leading,");
writer.writeLine("c.trailing,");
writer.writeLine(`tokens_ref.expect("Tokens must be provided when using comments."),`);
writer.writeLine(`info_ref.text_info.expect("Text info must be provided when using comments"),`);
});
writer.write("));").newLine();
writer.writeLine("let comments_ref = comments.as_ref().map(|t| unsafe { mem::transmute::<&CommentContainer, &'a CommentContainer<'a>>(t) });");
writer.writeLine(`let ast_view = ${getViewForFunctionName("Script")}(info_ref, tokens_ref, comments_ref, bump_ref);`);
writer.writeLine(`let result = with_view(ast_view);`);
writer.writeLine("bump_borrow.reset();");
writer.writeLine("result");
Expand Down Expand Up @@ -655,9 +677,13 @@ export function generate(analysisResult: AnalysisResult): string {
function writeStructFunctions() {
writer.write(`fn ${getViewForFunctionName(struct.name)}<'a>(`);
if (struct.name === "Module") {
writer.write(`source_file_info: &'a ModuleInfo<'a>`);
writer.write(`source_file_info: &'a ModuleInfo<'a>, `);
writer.write(`tokens: Option<&'a TokenContainer<'a>>, `);
writer.write(`comments: Option<&'a CommentContainer<'a>>`);
} else if (struct.name === "Script") {
writer.write(`source_file_info: &'a ScriptInfo<'a>`);
writer.write(`source_file_info: &'a ScriptInfo<'a>, `);
writer.write(`tokens: Option<&'a TokenContainer<'a>>, `);
writer.write(`comments: Option<&'a CommentContainer<'a>>`);
} else {
writer.write(`inner: &'a swc_ast::${struct.name}`);
}
Expand All @@ -669,19 +695,6 @@ export function generate(analysisResult: AnalysisResult): string {
} else if (struct.name === "Script") {
writer.writeLine("let inner = source_file_info.script;");
}
if (struct.name === "Module" || struct.name === "Script") {
writer.writeLine("let tokens = source_file_info.tokens.map(|t| &*bump.alloc(TokenContainer::new(t)));");
writer.writeLine(
`let comments = source_file_info.comments.map(|c| &*bump.alloc(CommentContainer::new(`,
);
writer.indent(() => {
writer.writeLine("c.leading,");
writer.writeLine("c.trailing,");
writer.writeLine(`tokens.expect("Tokens must be provided when using comments."),`);
writer.writeLine(`source_file_info.text_info.expect("Text info must be provided when using comments"),`);
});
writer.writeLine(")));");
}
// writer.writeLine(`println!("Entered ${struct.name}");`);

writer.write(`let node = bump.alloc(${struct.name} `).inlineBlock(() => {
Expand All @@ -696,10 +709,7 @@ export function generate(analysisResult: AnalysisResult): string {
}
for (const field of structFields) {
writer.write(`${field.name}: `);
writeGetViewTypeExpression(field.type, `&inner.${field.name}`);
if (isVecType(field.type)) {
writer.write(".collect()");
}
writeMaybeVecCollect(field.type, `&inner.${field.name}`);
writer.write(`,`).newLine();
}
}).write(");").newLine();
Expand Down Expand Up @@ -748,6 +758,18 @@ export function generate(analysisResult: AnalysisResult): string {
});
}

function writeMaybeVecCollect(type: TypeDefinition, name: string) {
if (isVecType(type)) {
writer.write("bump.alloc({");
writer.write(`let mut vec = allocator_api2::vec::Vec::with_capacity_in(${name.replace(/^&/, "")}.len(), bump);`);
writer.write("vec.extend(");
}
writeGetViewTypeExpression(type, name);
if (isVecType(type)) {
writer.write("); vec })");
}
}

function writeGetViewTypeExpression(type: TypeDefinition, name: string) {
if (type.kind === "Primitive") {
throw new Error("Primitive types not handled here.");
Expand All @@ -756,10 +778,7 @@ export function generate(analysisResult: AnalysisResult): string {
if (type.name === "Option") {
writer.write(`match ${name} `).inlineBlock(() => {
writer.write("Some(value) => Some(");
writeGetViewTypeExpression(type.genericArgs[0], "value");
if (isVecType(type.genericArgs[0])) {
writer.write(".collect()");
}
writeMaybeVecCollect(type.genericArgs[0], "value");
writer.write("),").newLine();
writer.writeLine("None => None,");
});
Expand Down
14 changes: 11 additions & 3 deletions generation/generate/helpers.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,18 @@ export function writeType(writer: Writer, analysisResult: AnalysisResult, type:
writer.write(path);
writer.write("<'a>");
} else if (type.genericArgs.length > 0) {
writer.write(path);
writer.write("<");
if (path === "Vec") {
writer.write("&'a [");
} else {
writer.write(path);
writer.write("<");
}
writer.write(type.genericArgs.map(type => writeType(writer, analysisResult, type, writeStructReference)).join(", "));
writer.write(">");
if (path === "Vec") {
writer.write("]");
} else {
writer.write(">");
}
} else {
writer.write(path);
}
Expand Down
3 changes: 2 additions & 1 deletion rs-lib/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@ sourcemap = ["swc_common/sourcemap"]
view = []

[dependencies]
bumpalo = "3.14.0"
allocator-api2 = "0.2.16"
bumpalo = { version = "3.14.0", features = ["collections", "allocator-api2"] }
num-bigint = "0.4"
rustc-hash = "1.1.0"
swc_atoms = "0.6.5"
Expand Down
Loading

0 comments on commit 928dfb6

Please # to comment.