fix: all tests passed

Signed-off-by: Akhilender Bongirwar <[email protected]>

Author: akhilender-bongirwar

kyo-offheap/native/src/main/java/lang/foreign/Arena.scala

@@ -9,7 +9,8 @@ import scala.scalanative.unsafe.*
   * An Arena is used to track off-heap allocations so that they can be automatically deallocated when the scope ends.
   */
 final class Arena extends AutoCloseable:
-    private val allocations = new ConcurrentLinkedQueue[Ptr[Byte]]()
+    private val allocations                  = new ConcurrentLinkedQueue[MemorySegment]()
+    @volatile private var _isClosed: Boolean = false
 
     /** Allocates a block of off-heap memory of the given size (in bytes).
       *
@@ -19,19 +20,24 @@ final class Arena extends AutoCloseable:
       *   A MemorySegment representing the allocated block.
       */
     def allocate(byteSize: Long): MemorySegment =
-        val segment = MemorySegment.allocate(byteSize)
-        allocations.add(segment.ptr)
+        if _isClosed then
+            throw new IllegalStateException("Arena is closed")
+        val segment = MemorySegment.allocate(byteSize, this)
+        allocations.add(segment)
         segment
     end allocate
 
     /** Frees all memory that was allocated through this Arena.
       */
     override def close(): Unit =
-        var ptr = allocations.poll()
-        while ptr != null do
-            free(ptr)
-            ptr = allocations.poll()
+        _isClosed = true
+        var seg = allocations.poll()
+        while seg != null do
+            free(seg.ptr)
+            seg = allocations.poll()
     end close
+
+    def isClosed: Boolean = _isClosed
 end Arena
 
 object Arena:

kyo-offheap/native/src/main/java/lang/foreign/MemorySegment.scala

@@ -10,7 +10,11 @@ import scala.scalanative.unsigned.*
   * This implementation wraps a pointer along with the allocated size in bytes. It also provides basic support for slicing and for
   * reading/writing primitive values, which are used by the Layout instances in the shared code.
   */
-final class MemorySegment private (private[foreign] val ptr: Ptr[Byte], val byteSize: Long):
+final class MemorySegment private (private[foreign] val ptr: Ptr[Byte], val byteSize: Long, private val arena: Arena):
+    private def checkOpen(): Unit =
+        if arena.isClosed then
+            throw new IllegalStateException("MemorySegment accessed after Arena was closed")
+
     /** Creates a new MemorySegment that is a slice of this segment.
       *
       * @param offset
@@ -21,70 +25,107 @@ final class MemorySegment private (private[foreign] val ptr: Ptr[Byte], val byte
       *   A new MemorySegment representing the slice.
       */
     def asSlice(offset: Long, newSize: Long): MemorySegment =
+        checkOpen()
         if offset < 0 || newSize < 0 || offset + newSize > byteSize then
             throw new IllegalArgumentException(s"Invalid slice parameters: byteSize=$byteSize, offset=$offset, newSize=$newSize")
         else
-            MemorySegment(ptr + offset, newSize)
+            MemorySegment(ptr + offset, newSize, arena)
+        end if
+    end asSlice
 
     /** Reads a value from memory using the provided layout.
       */
     def get(layout: ValueLayout.OfBoolean, offset: Long): Boolean =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[CBool]]
+    end get
     def get(layout: ValueLayout.OfByte, offset: Long): Byte =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Byte]]
+    end get
     def get(layout: ValueLayout.OfShort, offset: Long): Short =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Short]]
+    end get
     def get(layout: ValueLayout.OfInt, offset: Long): Int =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Int]]
+    end get
     def get(layout: ValueLayout.OfLong, offset: Long): Long =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Long]]
+    end get
     def get(layout: ValueLayout.OfFloat, offset: Long): Float =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Float]]
+    end get
     def get(layout: ValueLayout.OfDouble, offset: Long): Double =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Double]]
+    end get
     def get(layout: ValueLayout.OfChar, offset: Long): Char =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Char]]
+    end get
     def get(layout: AddressLayout, offset: Long): MemorySegment =
+        checkOpen()
         val newByteSize = byteSize - offset
         require(newByteSize >= 0)
-        new MemorySegment((ptr + offset).asInstanceOf[Ptr[Byte]], newByteSize)
+        new MemorySegment((ptr + offset).asInstanceOf[Ptr[Byte]], newByteSize, arena)
     end get
 
     /** Writes a value to memory using the provided layout.
       */
     def set(layout: ValueLayout.OfBoolean, offset: Long, value: Boolean): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[CBool]] = value
+    end set
     def set(layout: ValueLayout.OfByte, offset: Long, value: Byte): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Byte]] = value
+    end set
     def set(layout: ValueLayout.OfShort, offset: Long, value: Short): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Short]] = value
+    end set
     def set(layout: ValueLayout.OfInt, offset: Long, value: Int): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Int]] = value
+    end set
     def set(layout: ValueLayout.OfLong, offset: Long, value: Long): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Long]] = value
+    end set
     def set(layout: ValueLayout.OfFloat, offset: Long, value: Float): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Float]] = value
+    end set
     def set(layout: ValueLayout.OfDouble, offset: Long, value: Double): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Double]] = value
+    end set
     def set(layout: ValueLayout.OfChar, offset: Long, value: Char): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         !(ptr + offset).asInstanceOf[Ptr[Char]] = value
+    end set
     def set(layout: AddressLayout, offset: Long, value: MemorySegment): Unit =
+        checkOpen()
         require(offset + layout.byteSize <= byteSize)
         val _ = memcpy((ptr + offset).asInstanceOf[Ptr[Byte]], value.ptr, value.byteSize.toCSize)
     end set
@@ -93,10 +134,10 @@ end MemorySegment
 object MemorySegment:
 
     /** Allocates a new MemorySegment of the given byte size. */
-    private[foreign] def allocate(byteSize: Long): MemorySegment =
+    private[foreign] def allocate(byteSize: Long, arena: Arena): MemorySegment =
         val ptr = malloc(byteSize).asInstanceOf[Ptr[Byte]]
         if ptr == null then throw new RuntimeException("malloc returned null")
-        new MemorySegment(ptr, byteSize)
+        new MemorySegment(ptr, byteSize, arena)
     end allocate
 
     /** Copies a block of memory from the source segment to the destination segment.