progress towards SwiftArena which manages swift values and heap objects

Author: ktoso

Samples/SwiftKitSampleApp/src/test/java/org/swift/swiftkit/SwiftArenaTests.java

@@ -0,0 +1,55 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package org.swift.swiftkit;
+
+import com.example.swift.generated.MySwiftClass;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledOnOs;
+import org.junit.jupiter.api.condition.OS;
+
+import java.lang.foreign.Arena;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class SwiftArenaTests {
+
+    @BeforeAll
+    static void beforeAll() {
+        System.out.printf("java.library.path = %s\n", System.getProperty("java.library.path"));
+
+        System.loadLibrary("swiftCore");
+        System.loadLibrary("ExampleSwiftLibrary");
+
+        System.setProperty("jextract.trace.downcalls", "true");
+    }
+
+    @Test
+    @DisabledOnOs(OS.LINUX) // FIXME: enable on Linux when we get new compiler with mangled names in swift interfaces
+    void arena_releaseClassOnClose() {
+        try (var arena = SwiftArena.ofConfined()) {
+            var obj = new MySwiftClass(1, 2);
+
+            assertEquals(1, SwiftKit.retainCount(obj.$memorySegment()));
+            // TODO: test directly on SwiftHeapObject inheriting obj
+
+            SwiftKit.retain(obj.$memorySegment());
+            assertEquals(2, SwiftKit.retainCount(obj.$memorySegment()));
+
+            SwiftKit.release(obj.$memorySegment());
+            assertEquals(1, SwiftKit.retainCount(obj.$memorySegment()));
+        }
+    }
+}

Sources/JExtractSwift/Swift2JavaTranslator+Printing.swift

@@ -168,7 +168,7 @@ extension Swift2JavaTranslator {
   }
 
   public func printClass(_ printer: inout CodePrinter, _ decl: ImportedNominalType, body: (inout CodePrinter) -> Void) {
-    printer.printTypeDecl("public final class \(decl.javaClassName)") { printer in
+    printer.printTypeDecl("public final class \(decl.javaClassName) implements SwiftHeapObject") { printer in
       // ==== Storage of the class
       // FIXME: implement the self storage for the memory address and accessors
       printClassSelfProperty(&printer, decl)

SwiftKit/src/main/java/org/swift/swiftkit/SwiftArena.java

@@ -16,50 +16,134 @@
 
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.concurrent.ConcurrentSkipListSet;
-
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * A Swift arena manages Swift allocated memory for classes, structs, enums etc.
+ * When an arena is closed, it will destroy all managed swift objects in a way appropriate to their type.
+ * <p>
+ * A confined arena has an associated owner thread that confines some operations to
+ * associated owner thread such as {@link #close()}.
+ */
 public interface SwiftArena extends Arena {
 
-    void release(MemorySegment segment);
+    static SwiftArena ofConfined() {
+        return new ConfinedSwiftMemorySession(Thread.currentThread());
+    }
 
-    void retain(MemorySegment segment);
+    /**
+     * Register a struct, enum or other non-reference counted Swift object.
+     * Its memory should be considered managed by this arena, and be destroyed when the arena is closed.
+     *
+     * @param object
+     */
+    void register(SwiftHeapObject object);
 
-    long retainCount(MemorySegment segment);
+    void register(SwiftValue value);
 
 }
 
-final class AutoSwiftArena implements SwiftArena {
-    Arena underlying;
+final class ConfinedSwiftMemorySession implements SwiftArena {
+
+    final Arena underlying;
+    final Thread owner;
+    final SwiftResourceList resources;
 
-    ConcurrentSkipListSet<MemorySegment> managedMemorySegments;
+    // TODO: just int and volatile updates
+    final int CLOSED = 0;
+    final int ACTIVE = 1;
+    final AtomicInteger state;
+
+    public ConfinedSwiftMemorySession(Thread owner) {
+        this.owner = owner;
+        underlying = Arena.ofConfined();
+        resources = new ConfinedResourceList();
+        state = new AtomicInteger(ACTIVE);
+    }
 
     @Override
     public MemorySegment allocate(long byteSize, long byteAlignment) {
-        return null;
+        return underlying.allocate(byteSize, byteAlignment);
     }
 
     @Override
     public MemorySegment.Scope scope() {
-        return null;
+        return underlying.scope();
     }
 
-    @Override
-    public void close() {
-
+    public void checkValid() throws RuntimeException {
+        if (this.owner != null && this.owner != Thread.currentThread()) {
+            throw new WrongThreadException("ConfinedSwift arena is confined to %s but was closed from %s!".formatted(this.owner, Thread.currentThread()));
+        } else if (this.state.get() < ACTIVE) {
+            throw new RuntimeException("Arena is already closed!");
+        }
     }
 
     @Override
-    public void release(MemorySegment segment) {
-        SwiftKit.release(segment);
+    public void register(SwiftHeapObject object) {
+        this.resources.add(new SwiftHeapObjectCleanup(object.$memorySegment()));
     }
 
     @Override
-    public void retain(MemorySegment segment) {
-        SwiftKit.retain(segment);
+    public void register(SwiftValue value) {
+        this.resources.add(new SwiftHeapObjectCleanup(value.$memorySegment()));
     }
 
     @Override
-    public long retainCount(MemorySegment segment) {
-        return SwiftKit.retainCount(segment);
+    public void close() {
+        checkValid();
+
+        // Cleanup all resources
+        if (this.state.compareAndExchange(ACTIVE, CLOSED) == ACTIVE) {
+            this.resources.cleanup();
+        } // else, was already closed; do nothing
+
+
+        // Those the underlying arena
+        this.underlying.close();
+
+        // After this method returns normally, the scope must be not alive anymore
+        assert (!this.scope().isAlive());
+    }
+
+    /**
+     * Represents a list of resources that need a cleanup, e.g. allocated classes/structs.
+     */
+    static abstract class SwiftResourceList implements Runnable {
+        // TODO: Could use intrusive linked list to avoid one indirection here
+        final List<SwiftMemoryResourceCleanup> resourceCleanups = new LinkedList<>();
+
+        abstract void add(SwiftMemoryResourceCleanup cleanup);
+
+        public abstract void cleanup();
+
+        public final void run() {
+            cleanup(); // cleaner interop
+        }
+    }
+
+    static final class ConfinedResourceList extends SwiftResourceList {
+        @Override
+        void add(SwiftMemoryResourceCleanup cleanup) {
+            resourceCleanups.add(cleanup);
+        }
+
+        @Override
+        public void cleanup() {
+            for (SwiftMemoryResourceCleanup cleanup : resourceCleanups) {
+                cleanup.run();
+            }
+        }
+    }
+}
+
+final class UnexpectedRetainCountException extends RuntimeException {
+    public UnexpectedRetainCountException(MemorySegment resource, long retainCount, int expectedRetainCount) {
+        super(("Attempting to cleanup managed memory segment %s, but it's retain count was different than [%d] (was %d)! " +
+                "This would result in destroying a swift object that is still retained by other code somewhere."
+        ).formatted(resource, expectedRetainCount, retainCount));
     }
 }

SwiftKit/src/main/java/org/swift/swiftkit/SwiftHeapObject.java

@@ -16,6 +16,8 @@
 
 import java.lang.foreign.MemorySegment;
 
-public interface SwiftHeapObject {
-    MemorySegment $self();
+/**
+ * Represents a wrapper around a Swift heap object, e.g. a {@code class} or an {@code actor}.
+ */
+public interface SwiftHeapObject extends SwiftMemoryResource {
 }