Fix closure environment slot calculation and test coverage

Corrects the calculation of environment slot offsets for captured variables in closures, ensuring proper byte offset handling and consistent environment setup. Updates test WAT files to reflect the new closure environment layout and stack management, improving correctness and coverage for closure, function expression, return, ternary, and typealias scenarios.

Author: BlobMaster41

src/compiler.ts

@@ -1779,6 +1779,40 @@ export class Compiler extends DiagnosticEmitter {
     // This ensures we know which locals need to be stored in the environment
     this.prescanForClosures(bodyNode, instance, flow);
 
+    // Mark captured parameters and set up environment if needed
+    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++) {
+        let paramName = instance.getParameterName(i);
+        if (preCapturedNames.has(paramName)) {
+          let local = flow.lookupLocal(paramName);
+          if (local && !local.isCaptured) {
+            local.isCaptured = true;
+            // Ensure environment structures are set up
+            if (!instance.capturedLocals) {
+              instance.capturedLocals = new Map();
+            }
+            if (!instance.capturedLocals.has(local)) {
+              // Calculate proper byte offset
+              let currentOffset = 0;
+              for (let _keys = Map_keys(instance.capturedLocals), j = 0, m = _keys.length; j < m; ++j) {
+                let existingLocal = _keys[j];
+                currentOffset += existingLocal.type.byteSize;
+              }
+              local.envSlotIndex = currentOffset;
+              instance.capturedLocals.set(local, local.envSlotIndex);
+            }
+            if (!instance.envLocal) {
+              let envLocal = flow.addScopedLocal("$env", this.options.usizeType);
+              instance.envLocal = envLocal;
+            }
+          }
+        }
+      }
+    }
+
     // compile statements
     if (bodyNode.kind == NodeKind.Block) {
       stmts = this.compileStatements((<BlockStatement>bodyNode).statements, stmts);
@@ -3239,12 +3273,12 @@ export class Compiler extends DiagnosticEmitter {
           if (!sourceFunc.capturedLocals.has(local)) {
             // Calculate proper byte offset based on current environment size
             let currentOffset = 0;
-            for (let _values = Map_values(sourceFunc.capturedLocals), i = 0, k = _values.length; i < k; ++i) {
-              let existingLocal = _values[i];
+            for (let _keys = Map_keys(sourceFunc.capturedLocals), i = 0, k = _keys.length; i < k; ++i) {
+              let existingLocal = _keys[i];
               currentOffset += existingLocal.type.byteSize;
             }
             local.envSlotIndex = currentOffset;
-            sourceFunc.capturedLocals.set(local, local);
+            sourceFunc.capturedLocals.set(local, local.envSlotIndex);
           }
           // Ensure we have an environment local
           if (!sourceFunc.envLocal) {
@@ -7507,7 +7541,13 @@ export class Compiler extends DiagnosticEmitter {
         }
         if (local && !captures.has(local)) {
           local.isCaptured = true;
-          local.envSlotIndex = captures.size;
+          // Calculate proper byte offset based on existing captures
+          let currentOffset = 0;
+          for (let _keys = Map_keys(captures), idx = 0, cnt = _keys.length; idx < cnt; ++idx) {
+            let existingLocal = _keys[idx];
+            currentOffset += existingLocal.type.byteSize;
+          }
+          local.envSlotIndex = currentOffset;
           captures.set(local, local.envSlotIndex);
         }
         break;
@@ -7810,7 +7850,10 @@ export class Compiler extends DiagnosticEmitter {
         this.collectDeclaredVariables(forStmt.body, vars);
         break;
       }
-      // Function expressions don't add to outer scope
+      case NodeKind.FunctionDeclaration:
+      case NodeKind.Function:
+        // Function declarations and expressions don't add their variables to outer scope
+        break;
     }
   }
 
@@ -8022,6 +8065,19 @@ export class Compiler extends DiagnosticEmitter {
         }
         break;
       }
+      case NodeKind.Variable: {
+        // Add declared variables to inner names so they're not considered captures
+        let varStmt = <VariableStatement>node;
+        for (let i = 0, k = varStmt.declarations.length; i < k; i++) {
+          let decl = varStmt.declarations[i];
+          innerFunctionNames.add(decl.name.text);
+          // Scan initializers for captures
+          if (decl.initializer) {
+            this.collectCapturedNames(decl.initializer, innerFunctionNames, outerFlow, declaredVars, capturedNames);
+          }
+        }
+        break;
+      }
       case NodeKind.Binary: {
         let binary = <BinaryExpression>node;
         this.collectCapturedNames(binary.left, innerFunctionNames, outerFlow, declaredVars, capturedNames);

tests/compiler/closure.debug.wat

@@ -2352,28 +2352,28 @@
   global.get $$~lib/__closure_env
   i32.load
   global.get $$~lib/__closure_env
-  i32.load offset=1
+  i32.load offset=4
   i32.add
   global.get $$~lib/__closure_env
-  i32.load offset=2
+  i32.load offset=8
   i32.add
  )
  (func $closure/testMultipleCaptures (param $a i32) (param $b i32) (result i32)
-  (local $c i32)
   (local $$env i32)
+  (local $c i32)
   (local $4 i32)
   i32.const 12
   call $~lib/rt/tlsf/__alloc
   local.set $$env
   local.get $$env
   local.get $a
-  i32.store offset=4
+  i32.store
   local.get $$env
   local.get $b
-  i32.store offset=8
+  i32.store offset=4
   local.get $$env
   i32.const 10
-  i32.store
+  i32.store offset=8
   i32.const 8
   i32.const 4
   call $~lib/rt/itcms/__new

tests/compiler/closure.release.wat

@@ -1549,11 +1549,11 @@
  )
  (func $closure/testMultipleCaptures~anonymous|0 (result i32)
   global.get $$~lib/__closure_env
-  i32.load offset=2
+  i32.load offset=8
   global.get $$~lib/__closure_env
   i32.load
   global.get $$~lib/__closure_env
-  i32.load offset=1
+  i32.load offset=4
   i32.add
   i32.add
  )
@@ -1853,13 +1853,13 @@
    i32.add
    local.tee $1
    i32.const 1
-   i32.store offset=4
+   i32.store
    local.get $0
    i32.const 2
-   i32.store offset=12
+   i32.store offset=8
    local.get $0
    i32.const 10
-   i32.store offset=4
+   i32.store offset=12
    i32.const 4
    call $~lib/rt/itcms/__new
    local.tee $0