WIP: another version

Author: sliemeobn

Makefile

@@ -19,7 +19,7 @@ unittest:
 	swift package --swift-sdk "$(SWIFT_SDK_ID)" \
 	    $(TRACING_ARGS) \
 	    --disable-sandbox \
-	    js test --prelude ./Tests/prelude.mjs -Xnode --expose-gc
+	    js test --prelude ./Tests/prelude.mjs -Xnode --expose-gc --verbose
 
 .PHONY: regenerate_swiftpm_resources
 regenerate_swiftpm_resources:

Plugins/PackageToJS/Templates/runtime.mjs

@@ -256,13 +256,18 @@ class JSObjectSpace {
             this._refCounts[id]++;
             return id;
         }
-        const newId = this._freeSlotStack.length > 0 ? this._freeSlotStack.pop() : this._values.length;
+        const newId = this._freeSlotStack.length > 0
+            ? this._freeSlotStack.pop()
+            : this._values.length;
         this._values[newId] = value;
         this._refCounts[newId] = 1;
         this._valueRefMap.set(value, newId);
         return newId;
     }
     retainByRef(ref) {
+        if (this._refCounts[ref] === 0) {
+            throw new ReferenceError("Attempted to retain invalid reference " + ref);
+        }
         this._refCounts[ref]++;
         return ref;
     }

Runtime/bench/_version2.ts

@@ -0,0 +1,67 @@
+import { globalVariable } from "../src/find-global.js";
+import { ref } from "../src/types.js";
+
+export class JSObjectSpace_v2 {
+    private _entryByValue: Map<any, { id: number; rc: number }>;
+    private _values: (any | undefined)[];
+    private _refCounts: number[];
+    private _nextRef: number;
+
+    constructor() {
+        this._values = [];
+        this._values[0] = undefined;
+        this._values[1] = globalVariable;
+
+        this._entryByValue = new Map();
+        this._entryByValue.set(globalVariable, { id: 1, rc: 1 });
+
+        this._refCounts = [];
+        this._refCounts[0] = 0;
+        this._refCounts[1] = 1;
+
+        // 0 is invalid, 1 is globalThis.
+        this._nextRef = 2;
+    }
+
+    retain(value: any) {
+        const entry = this._entryByValue.get(value);
+        if (entry) {
+            entry.rc++;
+            this._refCounts[entry.id]++;
+            return entry.id;
+        }
+
+        const id = this._nextRef++;
+        this._values[id] = value;
+        this._refCounts[id] = 1;
+        this._entryByValue.set(value, { id, rc: 1 });
+        return id;
+    }
+
+    retainByRef(ref: ref) {
+        return this.retain(this.getObject(ref));
+    }
+
+    release(ref: ref) {
+        const value = this._values[ref];
+        const entry = this._entryByValue.get(value)!;
+        entry.rc--;
+        this._refCounts[ref]--;
+        if (entry.rc != 0) return;
+
+        this._entryByValue.delete(value);
+        // Keep IDs monotonic; clear slot and leave possible holes.
+        this._values[ref] = undefined;
+        this._refCounts[ref] = 0;
+    }
+
+    getObject(ref: ref) {
+        const value = this._values[ref];
+        if (value === undefined) {
+            throw new ReferenceError(
+                "Attempted to read invalid reference " + ref,
+            );
+        }
+        return value;
+    }
+}

Runtime/bench/bench-runner.ts

@@ -5,6 +5,7 @@
 
 import { JSObjectSpace } from "../src/object-heap.js";
 import { JSObjectSpaceOriginal } from "./_original.js";
+import { JSObjectSpace_v2 } from "./_version2.js";
 
 export interface HeapLike {
     retain(value: unknown): number;
@@ -94,7 +95,8 @@ function runBenchmark(
 function main() {
     const implementations: Array<{ name: string; Heap: new () => HeapLike }> = [
         { name: "JSObjectSpaceOriginal", Heap: JSObjectSpaceOriginal },
-        { name: "JSObjectSpace (current)", Heap: JSObjectSpace },
+        { name: "JSObjectSpace_v2 (ref++, single map)", Heap: JSObjectSpace_v2 },
+        { name: "JSObjectSpace (current)", Heap: JSObjectSpace }
     ];
 
     console.log("JSObjectSpace benchmark");

Runtime/src/object-heap.ts

@@ -40,6 +40,12 @@ export class JSObjectSpace {
     }
 
     retainByRef(ref: ref) {
+        if (this._refCounts[ref] === 0) {
+            throw new ReferenceError(
+                "Attempted to retain invalid reference " + ref,
+            );
+        }
+
         this._refCounts[ref]++;
         return ref;
     }

Tests/JavaScriptKitTests/JSClosureTests.swift

@@ -92,62 +92,48 @@ class JSClosureTests: XCTestCase {
             throw XCTSkip("Missing --expose-gc flag")
         }
 
-        // Step 1: Create many JSClosure instances
+        // Step 1: Create many source closures and keep only JS references alive.
+        // These closures must remain callable even after heavy finalizer churn.
         let obj = JSObject()
-        var closurePointers: Set<UInt32> = []
         let numberOfSourceClosures = 10_000
 
         do {
             var closures: [JSClosure] = []
             for i in 0..<numberOfSourceClosures {
-                let closure = JSClosure { _ in .undefined }
+                let closure = JSClosure { _ in .number(Double(i)) }
                 obj["c\(i)"] = closure.jsValue
                 closures.append(closure)
-                // Store
-                closurePointers.insert(UInt32(UInt(bitPattern: Unmanaged.passUnretained(closure).toOpaque())))
-
-                // To avoid all JSClosures having a common address diffs, randomly allocate a new object.
-                if Bool.random() {
-                    _ = JSObject()
-                }
-            }
-        }
-
-        // Step 2: Create many JSObject to make JSObject.id close to Swift heap object address
-        let minClosurePointer = closurePointers.min() ?? 0
-        let maxClosurePointer = closurePointers.max() ?? 0
-        while true {
-            let obj = JSObject()
-            if minClosurePointer == obj.id {
-                break
             }
         }
 
-        // Step 3: Create JSClosure instances and find the one with JSClosure.id == &closurePointers[x]
+        // Step 2: Create many temporary objects/closures to stress ID reuse and finalizer paths.
+        // Under the optimized object heap, IDs are aggressively reused, so this should exercise
+        // the same misdeallocation surface without relying on monotonic ID growth.
         do {
-            while true {
-                let c = JSClosure { _ in .undefined }
-                if closurePointers.contains(c.id) || c.id > maxClosurePointer {
-                    break
+            let numberOfProbeClosures = 50_000
+            for i in 0..<numberOfProbeClosures {
+                let tempClosure = JSClosure { _ in .number(Double(i)) }
+                if i % 3 == 0 {
+                    let tempObject = JSObject()
+                    tempObject["probe"] = tempClosure.jsValue
                 }
-                // To avoid all JSClosures having a common JSObject.id diffs, randomly allocate a new JS object.
-                if Bool.random() {
+                if i % 7 == 0 {
                     _ = JSObject()
                 }
             }
         }
 
-        // Step 4: Trigger garbage collection to call the finalizer of the conflicting JSClosure instance
+        // Step 3: Trigger garbage collection to run finalizers for temporary closures.
         for _ in 0..<100 {
             gc()
             // Tick the event loop to allow the garbage collector to run finalizers
             // registered by FinalizationRegistry.
             try await Task.sleep(for: .milliseconds(0))
         }
 
-        // Step 5: Verify that the JSClosure instances are still alive and can be called
+        // Step 4: Verify source closures are still alive and correct.
         for i in 0..<numberOfSourceClosures {
-            _ = obj["c\(i)"].function!()
+            XCTAssertEqual(obj["c\(i)"].function!(), .number(Double(i)))
         }
     }
 }