From 2217be40035e5577513af5baaa413699f2b555a7 Mon Sep 17 00:00:00 2001
From: Svyatoslav Scherbina <svyatoslav.scherbina@jetbrains.com>
Date: Tue, 19 Sep 2017 11:53:27 +0300
Subject: [PATCH] Add `defer` support to `memScope`

---
 .../src/main/kotlin/kotlinx/cinterop/Utils.kt | 56 +++++++++++++++----
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/Interop/Runtime/src/main/kotlin/kotlinx/cinterop/Utils.kt b/Interop/Runtime/src/main/kotlin/kotlinx/cinterop/Utils.kt
index 7b7d193a25a..1de7f3d3bd7 100644
--- a/Interop/Runtime/src/main/kotlin/kotlinx/cinterop/Utils.kt
+++ b/Interop/Runtime/src/main/kotlin/kotlinx/cinterop/Utils.kt
@@ -36,12 +36,41 @@ object nativeHeap : NativeFreeablePlacement {
     override fun free(mem: NativePtr) = nativeMemUtils.free(mem)
 }
 
+private typealias Deferred = () -> Unit
+
+open class DeferScope {
+
+    @PublishedApi
+    internal var topDeferred: Deferred? = null
+
+    internal fun executeAllDeferred() {
+        topDeferred?.let {
+            it.invoke()
+            topDeferred = null
+        }
+    }
+
+    inline fun defer(crossinline block: () -> Unit) {
+        val currentTop = topDeferred
+        topDeferred = {
+            try {
+                block()
+            } finally {
+                // TODO: it is possible to implement chaining without recursion,
+                // but it would require using an anonymous object here
+                // which is not yet supported in Kotlin Native inliner.
+                currentTop?.invoke()
+            }
+        }
+    }
+}
+
 // TODO: implement optimally
-class Arena(private val parent: NativeFreeablePlacement = nativeHeap) : NativePlacement {
+open class ArenaBase(private val parent: NativeFreeablePlacement = nativeHeap) : NativePlacement, DeferScope() {
 
     private val allocatedChunks = ArrayList<NativePointed>()
 
-    override fun alloc(size: Long, align: Int): NativePointed {
+    final override fun alloc(size: Long, align: Int): NativePointed {
         val res = parent.alloc(size, align)
         try {
             allocatedChunks.add(res)
@@ -52,7 +81,10 @@ class Arena(private val parent: NativeFreeablePlacement = nativeHeap) : NativePl
         }
     }
 
-    fun clear() {
+    @PublishedApi
+    internal fun clearImpl() {
+        this.executeAllDeferred()
+
         allocatedChunks.forEach {
             parent.free(it)
         }
@@ -62,6 +94,10 @@ class Arena(private val parent: NativeFreeablePlacement = nativeHeap) : NativePl
 
 }
 
+class Arena(parent: NativeFreeablePlacement = nativeHeap) : ArenaBase(parent) {
+    fun clear() = this.clearImpl()
+}
+
 /**
  * Allocates variable of given type.
  *
@@ -355,18 +391,14 @@ fun CPointer<ByteVar>.toKString(): String {
     return decodeFromUtf8(bytes)
 }
 
-class MemScope : NativePlacement {
+class MemScope : ArenaBase() {
 
-    private val arena = Arena()
-
-    override fun alloc(size: Long, align: Int) = arena.alloc(size, align)
-
-    fun clear() = arena.clear()
-
-    val memScope: NativePlacement
+    val memScope: MemScope
         get() = this
 }
 
+// TODO: consider renaming `memScoped` because it now supports `defer`.
+
 /**
  * Runs given [block] providing allocation of memory
  * which will be automatically disposed at the end of this scope.
@@ -376,7 +408,7 @@ inline fun <R> memScoped(block: MemScope.()->R): R {
     try {
         return memScope.block()
     } finally {
-        memScope.clear()
+        memScope.clearImpl()
     }
 }