Skip to content

Refactor closure capture analysis and minor cleanups #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BlobMaster41 merged 2 commits into from
Dec 12, 2025
Merged

Refactor closure capture analysis and minor cleanups #5

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80aa6ab
Refactor environment slot allocation logic
BlobMaster41 Dec 12, 2025
bc725a4
Refactor closure capture analysis and minor cleanups
BlobMaster41 Dec 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 30 additions & 43 deletions src/compiler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1526,7 +1526,7 @@ export class Compiler extends DiagnosticEmitter {

private ensureEnumToString(enumElement: Enum, reportNode: Node): string | null {
if (enumElement.toStringFunctionName) return enumElement.toStringFunctionName;

if (!this.compileEnum(enumElement)) return null;
if (enumElement.is(CommonFlags.Const)) {
this.errorRelated(
Expand Down Expand Up @@ -1782,8 +1782,7 @@ export class Compiler extends DiagnosticEmitter {
let preCapturedNames = instance.preCapturedNames;
if (preCapturedNames && preCapturedNames.size > 0) {
// Check if any parameters are captured
let parameterTypes = instance.signature.parameterTypes;
for (let i = 0, k = parameterTypes.length; i < k; i++) {
for (let i = 0, k = instance.signature.parameterTypes.length; i < k; i++) {
let paramName = instance.getParameterName(i);
if (preCapturedNames.has(paramName)) {
let local = flow.lookupLocal(paramName);
Expand All @@ -1798,24 +1797,21 @@ export class Compiler extends DiagnosticEmitter {
if (!capturedLocals.has(local)) {
// Calculate proper byte offset with alignment
// Reserve slot 0 for parent environment pointer (4 or 8 bytes depending on wasm32/64)
let ptrSize = this.options.usizeType.byteSize;
let currentOffset = ptrSize; // Start after parent pointer slot
let currentOffset = this.options.usizeType.byteSize; // Start after parent pointer slot
for (let _keys = Map_keys(capturedLocals), j = 0, m = _keys.length; j < m; ++j) {
let existingLocal = _keys[j];
let endOfSlot = existingLocal.envSlotIndex + existingLocal.type.byteSize;
if (endOfSlot > currentOffset) currentOffset = endOfSlot;
}
// Align to the type's natural alignment
let typeSize = local.type.byteSize;
let align = typeSize;
let align = local.type.byteSize;
currentOffset = (currentOffset + align - 1) & ~(align - 1);
local.envSlotIndex = currentOffset;
local.envOwner = instance; // Track which function owns this capture
capturedLocals.set(local, local.envSlotIndex);
}
if (!instance.envLocal) {
let envLocal = flow.addScopedLocal("$env", this.options.usizeType);
instance.envLocal = envLocal;
instance.envLocal = flow.addScopedLocal("$env", this.options.usizeType);
}
}
}
Expand All @@ -1832,23 +1828,20 @@ export class Compiler extends DiagnosticEmitter {
instance.capturedLocals = capturedLocals;
}
if (!capturedLocals.has(thisLocal)) {
let ptrSize = this.options.usizeType.byteSize;
let currentOffset = ptrSize;
let currentOffset = this.options.usizeType.byteSize;
for (let _keys = Map_keys(capturedLocals), j = 0, m = _keys.length; j < m; ++j) {
let existingLocal = _keys[j];
let endOfSlot = existingLocal.envSlotIndex + existingLocal.type.byteSize;
if (endOfSlot > currentOffset) currentOffset = endOfSlot;
}
let typeSize = thisLocal.type.byteSize;
let align = typeSize;
let align = thisLocal.type.byteSize;
currentOffset = (currentOffset + align - 1) & ~(align - 1);
thisLocal.envSlotIndex = currentOffset;
thisLocal.envOwner = instance;
capturedLocals.set(thisLocal, thisLocal.envSlotIndex);
}
if (!instance.envLocal) {
let envLocal = flow.addScopedLocal("$env", this.options.usizeType);
instance.envLocal = envLocal;
instance.envLocal = flow.addScopedLocal("$env", this.options.usizeType);
}
}
}
Expand All @@ -1858,8 +1851,7 @@ export class Compiler extends DiagnosticEmitter {
// the environment pointer. This is needed because indirect calls to other closures
// can overwrite the global $~lib/__closure_env.
if (instance.outerFunction && !instance.closureEnvLocal) {
let closureEnvLocal = flow.addScopedLocal("$closureEnv", this.options.usizeType);
instance.closureEnvLocal = closureEnvLocal;
instance.closureEnvLocal = flow.addScopedLocal("$closureEnv", this.options.usizeType);
}

// compile statements
Expand Down Expand Up @@ -2005,8 +1997,7 @@ export class Compiler extends DiagnosticEmitter {
valueTypeRef, property.memoryOffset
);
let flowBefore = this.currentFlow;
let flow = getterInstance.flow;
this.currentFlow = flow;
this.currentFlow = getterInstance.flow;
if (property.is(CommonFlags.DefinitelyAssigned) && valueType.isReference && !valueType.isNullableReference) {
body = this.makeRuntimeNonNullCheck(body, valueType, getterInstance.identifierNode);
}
Expand Down Expand Up @@ -2704,7 +2695,7 @@ export class Compiler extends DiagnosticEmitter {
// (then │ │ (body) │
// (?block $continue │ │ if loops: (incrementor) ─────┘
// (body) │ │ recompile body?
// ) ├◄┘
// ) ├◄┘
// (incrementor) ┌◄┘
// (br $loop)
// )
Expand Down Expand Up @@ -2993,17 +2984,17 @@ export class Compiler extends DiagnosticEmitter {
// Compile the condition (always executes)
let condExpr = this.compileExpression(statement.condition, Type.auto);
let condType = this.currentType;

// Shortcut if there are no cases
if (!numCases) return module.drop(condExpr);

// Assign the condition to a temporary local as we compare it multiple times
let outerFlow = this.currentFlow;
let tempLocal = outerFlow.getTempLocal(condType);
let tempLocalIndex = tempLocal.index;
let breaks = new Array<ExpressionRef>(1 + numCases);
breaks[0] = module.local_set(tempLocalIndex, condExpr, condType.isManaged);

// Make one br_if per labeled case and leave it to Binaryen to optimize the
// sequence of br_ifs to a br_table according to optimization levels
let breakIndex = 1;
Expand All @@ -3015,7 +3006,7 @@ export class Compiler extends DiagnosticEmitter {
defaultIndex = i;
continue;
}

// Compile the equality expression for this case
const left = statement.condition;
const leftExpr = module.local_get(tempLocalIndex, condType.toRef());
Expand All @@ -3030,7 +3021,7 @@ export class Compiler extends DiagnosticEmitter {
condType,
statement
);

// Add it to the list of breaks
breaks[breakIndex++] = module.br(`case${i}|${label}`, equalityExpr);
}
Expand Down Expand Up @@ -4021,7 +4012,7 @@ export class Compiler extends DiagnosticEmitter {
expression: BinaryExpression,
contextualType: Type,
): ExpressionRef {

const left = expression.left;
const leftExpr = this.compileExpression(left, contextualType);
const leftType = this.currentType;
Expand All @@ -4039,9 +4030,9 @@ export class Compiler extends DiagnosticEmitter {
);
}

/**
/**
* compile `==` `===` `!=` `!==` BinaryExpression, from previously compiled left and right expressions.
*
*
* This is split from `compileCommutativeCompareBinaryExpression` so that the logic can be reused
* for switch cases in `compileSwitchStatement`, where the left expression only should be compiled once.
*/
Expand All @@ -4059,15 +4050,15 @@ export class Compiler extends DiagnosticEmitter {

let module = this.module;
let operatorString = operatorTokenToString(operator);

// check operator overload
const operatorKind = OperatorKind.fromBinaryToken(operator);
const leftOverload = leftType.lookupOverload(operatorKind, this.program);
const rightOverload = rightType.lookupOverload(operatorKind, this.program);
if (leftOverload && rightOverload && leftOverload != rightOverload) {
this.error(
DiagnosticCode.Ambiguous_operator_overload_0_conflicting_overloads_1_and_2,
reportNode.range,
reportNode.range,
operatorString,
leftOverload.internalName,
rightOverload.internalName
Expand Down Expand Up @@ -4157,7 +4148,7 @@ export class Compiler extends DiagnosticEmitter {

leftExpr = this.compileExpression(left, contextualType);
leftType = this.currentType;

// check operator overload
const operatorKind = OperatorKind.fromBinaryToken(operator);
const leftOverload = leftType.lookupOverload(operatorKind, this.program);
Expand Down Expand Up @@ -4228,7 +4219,7 @@ export class Compiler extends DiagnosticEmitter {
return this.compileNonCommutativeCompareBinaryExpression(expression, contextualType);
}
case Token.Equals_Equals_Equals:
case Token.Equals_Equals:
case Token.Equals_Equals:
case Token.Exclamation_Equals_Equals:
case Token.Exclamation_Equals: {
return this.compileCommutativeCompareBinaryExpression(expression, contextualType);
Expand Down Expand Up @@ -6546,13 +6537,13 @@ export class Compiler extends DiagnosticEmitter {
if (numArguments < numParams) {
return argumentExpressions;
}

// make an array literal expression from the rest args
let elements = argumentExpressions.slice(numParams - 1);
let range = new Range(elements[0].range.start, elements[elements.length - 1].range.end);
range.source = reportNode.range.source;
let arrExpr = new ArrayLiteralExpression(elements, range);

// return the original args, but replace the rest args with the array
const exprs = argumentExpressions.slice(0, numParams - 1);
exprs.push(arrExpr);
Expand Down Expand Up @@ -8342,8 +8333,7 @@ export class Compiler extends DiagnosticEmitter {
case NodeKind.Function: {
// Found a function expression - analyze its captures
let funcExpr = <FunctionExpression>current;
let declaration = funcExpr.declaration;
let capturedNames = this.analyzeCapturedVariablesWithDeclared(declaration, flow, instance, declaredVars);
let capturedNames = this.analyzeCapturedVariablesWithDeclared(funcExpr.declaration, flow, declaredVars);
if (capturedNames.size > 0) {
// Check if closures feature is enabled
if (!this.options.hasFeature(Feature.Closures)) {
Expand Down Expand Up @@ -8593,7 +8583,6 @@ export class Compiler extends DiagnosticEmitter {
private analyzeCapturedVariablesWithDeclared(
declaration: FunctionDeclaration,
outerFlow: Flow,
outerFunc: Function,
declaredVars: Map<string, Type | null>
): Set<string> {
// For prescan, we just collect variable NAMES that are captured
Expand Down Expand Up @@ -8639,8 +8628,7 @@ export class Compiler extends DiagnosticEmitter {
if (endOfSlot > maxEnd) maxEnd = endOfSlot;
}
// Ensure total size is aligned to pointer size
let size = (maxEnd + usizeSize - 1) & ~(usizeSize - 1);
return size;
return (maxEnd + usizeSize - 1) & ~(usizeSize - 1);
}

/** Ensures a closure environment is set up for the outer function. */
Expand All @@ -8657,7 +8645,7 @@ export class Compiler extends DiagnosticEmitter {
for (let _keys = Map_keys(captures), i = 0, k = _keys.length; i < k; i++) {
let local = _keys[i];
if (!existingCaptures.has(local)) {
existingCaptures.set(local, captures.get(local)!);
existingCaptures.set(local, captures.get(local) as i32);
}
}
}
Expand All @@ -8666,8 +8654,7 @@ export class Compiler extends DiagnosticEmitter {

// Create a new environment local for the outer function
let usizeType = this.options.usizeType;
let envLocal = flow.addScopedLocal("$env", usizeType);
outerFunc.envLocal = envLocal;
outerFunc.envLocal = flow.addScopedLocal("$env", usizeType);;
outerFunc.capturedLocals = captures;

// Compute the environment size
Expand Down Expand Up @@ -9504,7 +9491,7 @@ export class Compiler extends DiagnosticEmitter {
stmts.length = 1;
stmts.push(
module.i32(1)
);
);
module.removeFunction(name);
module.addFunction(name, sizeType, TypeRef.I32, [ TypeRef.I32 ], module.block(null, stmts, TypeRef.I32));
}
Expand Down