Rewrote escape analysis

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Context.kt

@@ -170,6 +170,9 @@ internal class Context(config: KonanConfig) : KonanBackendContext(config) {
     // But we have to wait until the code generation phase,
     // to dump this information into generated file.
     var serializedLinkData: LinkData? = null
+    val moduleEscapeAnalysisResult: ModuleEscapeAnalysisResult.ModuleEAResult.Builder by lazy {
+        ModuleEscapeAnalysisResult.ModuleEAResult.newBuilder()
+    }
 
     @Deprecated("")
     lateinit var psi2IrGeneratorContext: GeneratorContext

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/EscapeAnalysis.proto

@@ -0,0 +1,15 @@
+syntax = "proto2";
+package org.jetbrains.kotlin.backend.konan;
+
+option java_outer_classname = "ModuleEscapeAnalysisResult";
+option optimize_for = LITE_RUNTIME;
+
+message FunctionEAResult {
+    required string fqName = 1;
+    required int32 escapes = 2;
+    repeated int32 pointsTo = 3;
+}
+
+message ModuleEAResult {
+    repeated FunctionEAResult functionEAResults = 1;
+}

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/GraphAlgorithms.kt

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.backend.konan
+
+internal interface DirectedGraphNode<out K> {
+    val key: K
+    val directEdges: List<K>
+    val reversedEdges: List<K>
+}
+
+internal interface DirectedGraph<K, out N: DirectedGraphNode<K>> {
+    val nodes: Collection<N>
+    fun get(key: K): N
+}
+
+internal class DirectedGraphMultiNode<out K>(val nodes: Set<K>)
+
+internal class DirectedGraphCondensation<out K>(val topologicalOrder: List<DirectedGraphMultiNode<K>>)
+
+internal class DirectedGraphCondensationBuilder<K, out N: DirectedGraphNode<K>>(val graph: DirectedGraph<K, N>) {
+    private val visited = mutableSetOf<K>()
+    private val order = mutableListOf<N>()
+    private val nodeToMultiNodeMap = mutableMapOf<N, DirectedGraphMultiNode<K>>()
+    private val multiNodesOrder = mutableListOf<DirectedGraphMultiNode<K>>()
+
+    fun build(): DirectedGraphCondensation<K> {
+        // First phase.
+        graph.nodes.forEach {
+            if (!visited.contains(it.key))
+                findOrder(it)
+        }
+
+        // Second phase.
+        visited.clear()
+        val multiNodes = mutableListOf<DirectedGraphMultiNode<K>>()
+        order.reversed().forEach {
+            if (!visited.contains(it.key)) {
+                val nodes = mutableSetOf<K>()
+                paint(it, nodes)
+                multiNodes += DirectedGraphMultiNode(nodes)
+            }
+        }
+
+        // Topsort of built condensation.
+        multiNodes.forEach { multiNode ->
+            multiNode.nodes.forEach { nodeToMultiNodeMap.put(graph.get(it), multiNode) }
+        }
+        visited.clear()
+        multiNodes.forEach {
+            if (!visited.contains(it.nodes.first()))
+                findMultiNodesOrder(it)
+        }
+
+        return DirectedGraphCondensation(multiNodesOrder)
+    }
+
+    private fun findOrder(node: N) {
+        visited += node.key
+        node.directEdges.forEach {
+            if (!visited.contains(it))
+                findOrder(graph.get(it))
+        }
+        order += node
+    }
+
+    private fun paint(node: N, multiNode: MutableSet<K>) {
+        visited += node.key
+        multiNode += node.key
+        node.reversedEdges.forEach {
+            if (!visited.contains(it))
+                paint(graph.get(it), multiNode)
+        }
+    }
+
+    private fun findMultiNodesOrder(node: DirectedGraphMultiNode<K>) {
+        visited.addAll(node.nodes)
+        node.nodes.forEach {
+            graph.get(it).directEdges.forEach {
+                if (!visited.contains(it))
+                    findMultiNodesOrder(nodeToMultiNodeMap[graph.get(it)]!!)
+            }
+        }
+        multiNodesOrder += node
+    }
+}

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanPhases.kt

@@ -56,6 +56,7 @@ enum class KonanPhase(val description: String,
     /* ... ... */ AUTOBOX("Autoboxing of primitive types", BRIDGES_BUILDING, LOWER_COROUTINES),
     /* ... */ BITCODE("LLVM BitCode Generation"),
     /* ... ... */ RTTI("RTTI Generation"),
+    /* ... ... */ ESCAPE_ANALYSIS("Escape analysis"),
     /* ... ... */ CODEGEN("Code Generation"),
     /* ... ... */ BITCODE_LINKER("Bitcode linking"),
     /* */ LINK_STAGE("Link stage"),

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/library/KonanLibrary.kt

@@ -41,6 +41,8 @@ interface KonanLibraryLayout {
         get() = File(libDir, "linkdata")
     val moduleHeaderFile 
         get() = File(linkdataDir, "module")
+    val escapeAnalysisFile
+        get() = File(linkdataDir, "module_escape_analysis")
     fun packageFile(packageName: String)
         = File(linkdataDir, if (packageName == "") "root_package" else "package_$packageName")
 }

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/library/KonanLibraryReader.kt

@@ -23,6 +23,7 @@ interface KonanLibraryReader {
     val libraryName: String
     val bitcodePaths: List<String>
     val linkerOpts: List<String>
+    val escapeAnalysis: ByteArray?
     fun moduleDescriptor(specifics: LanguageVersionSettings): ModuleDescriptor
 }
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/library/KonanLibraryWriter.kt

@@ -23,6 +23,7 @@ interface KonanLibraryWriter {
     fun addNativeBitcode(library: String)
     fun addKotlinBitcode(llvmModule: LLVMModuleRef)
     fun addManifestAddend(path: String)
+    fun addEscapeAnalysis(escapeAnalysis: ByteArray)
     val mainBitcodeFileName: String
     fun commit()
 }

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/library/impl/KonanLibraryReaderImpl.kt

@@ -47,6 +47,7 @@ class LibraryReaderImpl(var libraryFile: File, val currentAbiVersion: Int, val t
         }
 
     val targetList = inPlace.targetsDir.listFiles.map{it.name}
+    override val escapeAnalysis: ByteArray? by lazy { inPlace.escapeAnalysisFile.let { if (it.exists) it.readBytes() else null } }
 
     override val libraryName 
         get() = inPlace.libraryName

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/library/impl/KonanLibraryWriterImpl.kt

@@ -82,6 +82,10 @@ class LibraryWriterImpl(override val libDir: File, currentAbiVersion: Int,
         println("manifest addend: ${properties.stringPropertyNames().joinToString(" ")}")
     }
 
+    override fun addEscapeAnalysis(escapeAnalysis: ByteArray) {
+        escapeAnalysisFile.writeBytes(escapeAnalysis)
+    }
+
     override fun commit() {
         manifestProperties.saveToFile(manifestFile)
         if (!nopack) {
@@ -99,7 +103,8 @@ internal fun buildLibrary(
     output: String, 
     llvmModule: LLVMModuleRef, 
     nopack: Boolean, 
-    manifest: String?): KonanLibraryWriter {
+    manifest: String?,
+    escapeAnalysis: ByteArray?): KonanLibraryWriter {
 
     val library = LibraryWriterImpl(output, abiVersion, target, nopack)
 
@@ -109,6 +114,7 @@ internal fun buildLibrary(
         library.addNativeBitcode(it)
     }
     manifest ?.let { library.addManifestAddend(it) }
+    escapeAnalysis?.let { library.addEscapeAnalysis(it) }
 
     library.commit()
     return library

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/CodeGenerator.kt

@@ -226,7 +226,9 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
     fun call(llvmFunction: LLVMValueRef, args: List<LLVMValueRef>,
              resultLifetime: Lifetime = Lifetime.IRRELEVANT,
              lazyLandingpad: () -> LLVMBasicBlockRef? = { null }): LLVMValueRef {
-        val callArgs = if (isObjectReturn(llvmFunction.type)) {
+        val callArgs = if (!isObjectReturn(llvmFunction.type)) {
+            args
+        } else {
             // If function returns an object - create slot for the returned value or give local arena.
             // This allows appropriate rootset accounting by just looking at the stack slots,
             // along with ability to allocate in appropriate arena.
@@ -235,15 +237,28 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
                     localAllocs++
                     arenaSlot!!
                 }
+
                 SlotType.RETURN -> returnSlot!!
-            // TODO: for RETURN_IF_ARENA choose between created slot and arenaSlot
+
-            // dynamically.
+                SlotType.ANONYMOUS -> vars.createAnonymousSlot()
-                SlotType.ANONYMOUS, SlotType.RETURN_IF_ARENA -> vars.createAnonymousSlot()
+
-                else -> throw Error("Incorrect slot type")
+                SlotType.RETURN_IF_ARENA -> returnSlot.let {
+                    if (it != null)
+                        call(context.llvm.getReturnSlotIfArenaFunction, listOf(it, vars.createAnonymousSlot()))
+                    else {
+                        // Return type is not an object type - can allocate locally.
+                        localAllocs++
+                        arenaSlot!!
+                    }
+                }
+
+                is SlotType.PARAM_IF_ARENA ->
+                    call(context.llvm.getParamSlotIfArenaFunction,
+                            listOf(vars.load(resultLifetime.slotType.parameter), vars.createAnonymousSlot()))
+
+                else -> throw Error("Incorrect slot type: ${resultLifetime.slotType}")
             }
             args + resultSlot
-        } else {
-            args
         }
         return callRaw(llvmFunction, callArgs, lazyLandingpad)
     }
@@ -385,7 +400,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
     internal fun prologue() {
         assert(returns.isEmpty())
         if (isObjectType(returnType!!)) {
-            this.returnSlot = LLVMGetParam(function, numParameters(function.type) - 1)
+            returnSlot = LLVMGetParam(function, numParameters(function.type) - 1)
         }
         positionAtEnd(localsInitBb)
         slotsPhi = phi(kObjHeaderPtrPtr)

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/ContextUtils.kt

@@ -32,40 +32,105 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 
-internal enum class SlotType {
+internal sealed class SlotType {
     // Frame local arena slot can be used.
-    ARENA,
+    class ARENA: SlotType()
     // Return slot can be used.
-    RETURN,
+    class RETURN: SlotType()
     // Return slot, if it is an arena, can be used.
-    RETURN_IF_ARENA,
+    class RETURN_IF_ARENA: SlotType()
+    // Return slot, if it is an arena, can be used.
+    class PARAM_IF_ARENA(val parameter: Int): SlotType()
     // Anonymous slot.
-    ANONYMOUS,
+    class ANONYMOUS: SlotType()
     // Unknown slot type.
-    UNKNOWN
+    class UNKNOWN: SlotType()
+
+    companion object {
+        val ARENA = ARENA()
+        val RETURN = RETURN()
+        val RETURN_IF_ARENA = RETURN_IF_ARENA()
+        val ANONYMOUS = ANONYMOUS()
+        val UNKNOWN = UNKNOWN()
+    }
 }
 
 // Lifetimes class of reference, computed by escape analysis.
-enum class Lifetime(val slotType: SlotType) {
+internal sealed class Lifetime(val slotType: SlotType) {
     // If reference is frame-local (only obtained from some call and never leaves).
-    LOCAL(SlotType.ARENA),
+    class LOCAL: Lifetime(SlotType.ARENA) {
+        override fun toString(): String {
+            return "LOCAL"
+        }
+    }
+
     // If reference is only returned.
-    RETURN_VALUE(SlotType.RETURN),
+    class RETURN_VALUE: Lifetime(SlotType.RETURN) {
+        override fun toString(): String {
+            return "RETURN_VALUE"
+        }
+    }
+
     // If reference is set as field of references of class RETURN_VALUE or INDIRECT_RETURN_VALUE.
-    INDIRECT_RETURN_VALUE(SlotType.RETURN_IF_ARENA),
+    class INDIRECT_RETURN_VALUE: Lifetime(SlotType.RETURN_IF_ARENA) {
+        override fun toString(): String {
+            return "INDIRECT_RETURN_VALUE"
+        }
+    }
+
     // If reference is stored to the field of an incoming parameters.
-    PARAMETER_FIELD(SlotType.ANONYMOUS),
+    class PARAMETER_FIELD(val parameter: Int): Lifetime(SlotType.PARAM_IF_ARENA(parameter)) {
+        override fun toString(): String {
+            return "PARAMETER_FIELD($parameter)"
+        }
+    }
+
     // If reference refers to the global (either global object or global variable).
-    GLOBAL(SlotType.ANONYMOUS),
+    class GLOBAL: Lifetime(SlotType.ANONYMOUS) {
+        override fun toString(): String {
+            return "GLOBAL"
+        }
+    }
+
     // If reference used to throw.
-    THROW(SlotType.ANONYMOUS),
+    class THROW: Lifetime(SlotType.ANONYMOUS) {
+        override fun toString(): String {
+            return "THROW"
+        }
+    }
+
     // If reference used as an argument of outgoing function. Class can be improved by escape analysis
     // of called function.
-    ARGUMENT(SlotType.ANONYMOUS),
+    class ARGUMENT: Lifetime(SlotType.ANONYMOUS) {
+        override fun toString(): String {
+            return "ARGUMENT"
+        }
+    }
+
     // If reference class is unknown.
-    UNKNOWN(SlotType.UNKNOWN),
+    class UNKNOWN: Lifetime(SlotType.UNKNOWN) {
+        override fun toString(): String {
+            return "UNKNOWN"
+        }
+    }
+
     // If reference class is irrelevant.
-    IRRELEVANT(SlotType.UNKNOWN)
+    class IRRELEVANT: Lifetime(SlotType.UNKNOWN) {
+        override fun toString(): String {
+            return "IRRELEVANT"
+        }
+    }
+
+    companion object {
+        val LOCAL = LOCAL()
+        val RETURN_VALUE = RETURN_VALUE()
+        val INDIRECT_RETURN_VALUE = INDIRECT_RETURN_VALUE()
+        val GLOBAL = GLOBAL()
+        val THROW = THROW()
+        val ARGUMENT = ARGUMENT()
+        val UNKNOWN = UNKNOWN()
+        val IRRELEVANT = IRRELEVANT()
+    }
 }
 
 /**
@@ -253,6 +318,8 @@ internal class Llvm(val context: Context, val llvmModule: LLVMModuleRef) {
     val setRefFunction = importRtFunction("SetRef")
     val updateRefFunction = importRtFunction("UpdateRef")
     val leaveFrameFunction = importRtFunction("LeaveFrame")
+    val getReturnSlotIfArenaFunction = importRtFunction("GetReturnSlotIfArena")
+    val getParamSlotIfArenaFunction = importRtFunction("GetParamSlotIfArena")
     val lookupOpenMethodFunction = importRtFunction("LookupOpenMethod")
     val isInstanceFunction = importRtFunction("IsInstance")
     val checkInstanceFunction = importRtFunction("CheckInstance")

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IrToBitcode.kt

@@ -71,8 +71,13 @@ internal fun emitLLVM(context: Context) {
 
         generateDebugInfoHeader(context)
 
+        val lifetimes = mutableMapOf<IrElement, Lifetime>()
+        phaser.phase(KonanPhase.ESCAPE_ANALYSIS) {
+            EscapeAnalysis.computeLifetimes(irModule, context, lifetimes)
+        }
+
         phaser.phase(KonanPhase.CODEGEN) {
-            irModule.acceptVoid(CodeGeneratorVisitor(context))
+            irModule.acceptVoid(CodeGeneratorVisitor(context, lifetimes))
         }
 
         if (context.shouldContainDebugInfo()) {
@@ -118,7 +123,8 @@ internal fun produceOutput(context: Context) {
                 libraryName, 
                 llvmModule,
                 nopack,
-                manifest)
+                manifest,
+                context.moduleEscapeAnalysisResult?.build()?.toByteArray())
 
             context.library = library
             context.bitcodeFileName = library.mainBitcodeFileName
@@ -232,10 +238,9 @@ internal interface CodeContext {
 
 //-------------------------------------------------------------------------//
 
-internal class CodeGeneratorVisitor(val context: Context) : IrElementVisitorVoid {
+internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrElement, Lifetime>) : IrElementVisitorVoid {
 
     val codegen = CodeGenerator(context)
-    val resultLifetimes = mutableMapOf<IrElement, Lifetime>()
 
     //-------------------------------------------------------------------------//
 
@@ -310,8 +315,6 @@ internal class CodeGeneratorVisitor(val context: Context) : IrElementVisitorVoid
     override fun visitModuleFragment(declaration: IrModuleFragment) {
         context.log{"visitModule                    : ${ir2string(declaration)}"}
 
-        // computeLifetimes(module, this.codegen, resultLifetimes)
-
         declaration.acceptChildrenVoid(this)
         appendLlvmUsed(context.llvm.usedFunctions)
         appendStaticInitializers(context.llvm.staticInitializers)
@@ -1847,7 +1850,7 @@ internal class CodeGeneratorVisitor(val context: Context) : IrElementVisitorVoid
 
     //-------------------------------------------------------------------------//
     private fun resultLifetime(callee: IrElement): Lifetime {
-        return resultLifetimes.getOrElse(callee) { /* TODO: make IRRELEVANT */ Lifetime.GLOBAL }
+        return lifetimes.getOrElse(callee) { /* TODO: make IRRELEVANT */ Lifetime.GLOBAL }
     }
 
     private fun evaluateConstructorCall(callee: IrCall, args: List<LLVMValueRef>): LLVMValueRef {

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/InlineConstructorsTransformation.kt

@@ -39,11 +39,14 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 
 //-----------------------------------------------------------------------------//
 
+private val inlineConstructor by lazy { FqName("konan.internal.InlineConstructor") }
+
+internal val FunctionDescriptor.isInlineConstructor get() = annotations.hasAnnotation(inlineConstructor)
+
 internal class InlineConstructorsTransformation(val context: Context): IrElementTransformerVoidWithContext() {
 
     private val deserializer      by lazy { DeserializerDriver(context) }
     private val substituteMap     by lazy { mutableMapOf<ValueDescriptor, IrExpression>() }
-    private val inlineConstructor by lazy { FqName("konan.internal.InlineConstructor") }
 
     //-------------------------------------------------------------------------//
 
@@ -55,7 +58,7 @@ internal class InlineConstructorsTransformation(val context: Context): IrElement
 
         val irCall = super.visitCall(expression) as IrCall
         val functionDescriptor = irCall.descriptor
-        if (!functionDescriptor.annotations.hasAnnotation(inlineConstructor)) return irCall                                // This call does not need inlining.
+        if (!functionDescriptor.isInlineConstructor) return irCall                        // This call does not need inlining.
 
         val functionDeclaration = getFunctionDeclaration(irCall)                            // Get declaration of the function to be inlined.
         if (functionDeclaration == null) {                                                  // We failed to get the declaration.

backend.native/tests/runtime/memory/escape0.kt

@@ -1,20 +1,13 @@
-fun foo1(arg: String) : String = arg
+//fun foo1(arg: String) : String = foo0(arg)
+fun foo1(arg: Any) : Any = foo0(arg)
 
-/*
+fun foo0(arg: Any) : Any = Any()
-fun foo2(arg: String) : String = arg + " foo2"
 
-fun bar(arg1: String, arg2: String) : String  = arg1 + " bar " + arg2
+var global : Any =  Any()
 
-fun zoo1() : String {
+fun foo0_escape(arg: Any) : Any{
-    var x = foo1("")
+    global = arg
-    var y = 4
+    return Any()
-    return x
-}
-
-fun zoo2() : String {
-    val x = foo1("")
-    var y = 5
-    return x
 }
 
 class Node(var previous: Node?)
@@ -38,10 +31,11 @@ fun zoo4(arg: Int) : Any {
         else -> c
     }
     return a
-} */
+}
 
-fun zoo5(arg: String) : String {
+fun zoo5(arg: Any) : Any{
-    return arg + foo1(arg)
+    foo1(arg)
+    return arg
 }
 
 fun zoo6(arg: Any) : Any {
@@ -53,6 +47,7 @@ fun zoo7(arg1: Any, arg2: Any, selector: Int) : Any {
 }
 
 fun main(args : Array<String>) {
-    val z = zoo7(Any(), Any(), 1)
+    //val z = zoo7(Any(), Any(), 1)
+    val x = zoo5(Any())
     //println(bar(foo1(foo2("")), foo2(foo1(""))))
 }

runtime/src/main/cpp/Arrays.cpp

@@ -137,8 +137,8 @@ void Kotlin_CharArray_set(KRef thiz, KInt index, KChar value) {
 
 OBJ_GETTER(Kotlin_CharArray_copyOf, KConstRef thiz, KInt newSize) {
   const ArrayHeader* array = thiz->array();
-  ArrayHeader* result = ArrayContainer(
+  ArrayHeader* result = AllocArrayInstance(
-      theCharArrayTypeInfo, newSize).GetPlace();
+      array->type_info(), newSize, OBJ_RESULT)->array();
   KInt toCopy = array->count_ < newSize ?  array->count_ : newSize;
   memcpy(
       PrimitiveArrayAddressOfElementAt<KChar>(result, 0),

runtime/src/main/cpp/Common.h

@@ -26,4 +26,6 @@
 #define THREAD_LOCAL_VARIABLE __thread
 #endif
 
+#define ARRAY_SIZE(a) (sizeof(a) / sizeof(a[0]))
+
 #endif // RUNTIME_COMMON_H

runtime/src/main/cpp/KString.h

@@ -52,6 +52,4 @@ int binarySearchRange(const T* array, int arrayLength, T needle) {
   return middle - (needle < value ? 1 : 0);
 }
 
-#define ARRAY_SIZE(a) (sizeof(a) / sizeof(a[0]))
-
 #endif // RUNTIME_STRING_H

runtime/src/main/cpp/Memory.cpp

@@ -29,7 +29,7 @@
 
 // If garbage collection algorithm for cyclic garbage to be used.
 #define USE_GC 1
-// Optmize management of cyclic garbage (increases memory footprint).
+// Optimize management of cyclic garbage (increases memory footprint).
 // Not recommended for low-end embedded targets.
 #define OPTIMIZE_GC 1
 // Define to 1 to print all memory operations.
@@ -467,6 +467,7 @@ bool ArenaContainer::allocContainer(container_size_t minSize) {
   RuntimeAssert(result != nullptr, "Cannot alloc memory");
   if (result == nullptr) return false;
   result->next = currentChunk_;
+  result->arena = this;
   result->asHeader()->refCount_ = (CONTAINER_TAG_STACK | CONTAINER_TAG_INCREMENT);
   currentChunk_ = result;
   current_ = reinterpret_cast<uint8_t*>(result->asHeader() + 1);
@@ -491,6 +492,16 @@ void* ArenaContainer::place(container_size_t size) {
   return result;
 }
 
+#define ARENA_SLOTS_CHUNK_SIZE 16
+
+ObjHeader** ArenaContainer::getSlot() {
+  if (slots_ == nullptr || slotsCount_ >= ARENA_SLOTS_CHUNK_SIZE) {
+    slots_ = PlaceArray(theArrayTypeInfo, ARENA_SLOTS_CHUNK_SIZE);
+    slotsCount_ = 0;
+  }
+  return ArrayAddressOfElementAt(slots_, slotsCount_++);
+}
+
 ObjHeader* ArenaContainer::PlaceObject(const TypeInfo* type_info) {
   RuntimeAssert(type_info->instanceSize_ >= 0, "must be an object");
   uint32_t size = type_info->instanceSize_ + sizeof(ObjHeader);
@@ -618,7 +629,7 @@ OBJ_GETTER(AllocInstance, const TypeInfo* type_info) {
     auto arena = initedArena(asArenaSlot(OBJ_RESULT));
     auto result = arena->PlaceObject(type_info);
 #if TRACE_MEMORY
-    fprintf(stderr, "instace %p in arena: %p\n", result, arena);
+    fprintf(stderr, "instance %p in arena: %p\n", result, arena);
 #endif
     return result;
   }
@@ -680,8 +691,30 @@ void SetRef(ObjHeader** location, const ObjHeader* object) {
   }
 }
 
+ObjHeader** GetReturnSlotIfArena(ObjHeader** returnSlot, ObjHeader** localSlot) {
+  return isArenaSlot(returnSlot) ? returnSlot : localSlot;
+}
+
+ObjHeader** GetParamSlotIfArena(ObjHeader* param, ObjHeader** localSlot) {
+  if (param == nullptr) return localSlot;
+  auto container = param->container();
+  if ((container->refCount_ & CONTAINER_TAG_MASK) != CONTAINER_TAG_STACK)
+    return localSlot;
+  auto chunk = reinterpret_cast<ContainerChunk*>(container) - 1;
+  return reinterpret_cast<ObjHeader**>(reinterpret_cast<uintptr_t>(&chunk->arena) | ARENA_BIT);
+}
+
 void UpdateReturnRef(ObjHeader** returnSlot, const ObjHeader* object) {
-  if (isArenaSlot(returnSlot)) return;
+  if (isArenaSlot(returnSlot)) {
+    return;
+    if (object == nullptr
+        || (object->container()->refCount_ & CONTAINER_TAG_MASK) > CONTAINER_TAG_NORMAL) {
+        // Not a subject of reference counting.
+        return;
+    }
+    auto arena = initedArena(asArenaSlot(returnSlot));
+    returnSlot = arena->getSlot();
+  }
   ObjHeader* old = *returnSlot;
 #if TRACE_MEMORY
   fprintf(stderr, "UpdateReturnRef *%p: %p -> %p\n", returnSlot, old, object);

runtime/src/main/cpp/Memory.h

@@ -36,7 +36,7 @@ typedef enum {
   CONTAINER_TAG_SEEN = 4,
   // Shift to get actual counter.
   CONTAINER_TAG_SHIFT = 3,
-  // Actual value to increment/decrement conatiner by. Tag is in lower bits.
+  // Actual value to increment/decrement container by. Tag is in lower bits.
   CONTAINER_TAG_INCREMENT = 1 << CONTAINER_TAG_SHIFT,
   // Mask for container type, disregard seen bit.
   CONTAINER_TAG_MASK = ((CONTAINER_TAG_INCREMENT >> 1) - 1)
@@ -239,6 +239,17 @@ class ArrayContainer : public Container {
 // Container is used for reference counting, and it is assumed that objects
 // with related placement will share container. Only
 // whole container can be freed, individual objects are not taken into account.
+class ArenaContainer;
+
+struct ContainerChunk {
+  ContainerChunk* next;
+  ArenaContainer* arena;
+  // Then we have ContainerHeader here.
+  ContainerHeader* asHeader() {
+    return reinterpret_cast<ContainerHeader*>(this + 1);
+  }
+};
+
 class ArenaContainer {
  public:
   void Init();
@@ -252,15 +263,9 @@ class ArenaContainer {
   // same operation could be used to place strings.
   ArrayHeader* PlaceArray(const TypeInfo* array_type_info, container_size_t count);
 
- private:
+  ObjHeader** getSlot();
-  struct ContainerChunk {
-    ContainerChunk* next;
-    // Then we have ContainerHeader here.
-    ContainerHeader* asHeader() {
-      return reinterpret_cast<ContainerHeader*>(this + 1);
-    }
-  };
 
+ private:
   void* place(container_size_t size);
   bool allocContainer(container_size_t minSize);
   void setMeta(ObjHeader* obj, const TypeInfo* type_info) {
@@ -272,6 +277,8 @@ class ArenaContainer {
   ContainerChunk* currentChunk_;
   uint8_t* current_;
   uint8_t* end_;
+  ArrayHeader* slots_;
+  uint32_t slotsCount_;
 };
 
 #ifdef __cplusplus
@@ -350,6 +357,10 @@ void UpdateReturnRef(ObjHeader** returnSlot, const ObjHeader* object) RUNTIME_NO
 void ReleaseRefs(ObjHeader** start, int count) RUNTIME_NOTHROW;
 // Called on frame leave, if it has object slots.
 void LeaveFrame(ObjHeader** start, int count) RUNTIME_NOTHROW;
+// Tries to use returnSlot's arena for allocation.
+ObjHeader** GetReturnSlotIfArena(ObjHeader** returnSlot, ObjHeader** localSlot) RUNTIME_NOTHROW;
+// Tries to use param's arena for allocation.
+ObjHeader** GetParamSlotIfArena(ObjHeader* param, ObjHeader** localSlot) RUNTIME_NOTHROW;
 // Collect garbage, which cannot be found by reference counting (cycles).
 void GarbageCollect() RUNTIME_NOTHROW;
 // Clears object subgraph references from memory subsystem, and optionally

runtime/src/main/kotlin/konan/Annotations.kt

@@ -89,3 +89,7 @@ public annotation class FixmeInline
  * Need to be fixed.
  */
 public annotation class Fixme
+
+public annotation class Escapes(val who: Int)
+
+public annotation class PointsTo(vararg val onWhom: Int)

runtime/src/main/kotlin/kotlin/Array.kt

@@ -38,9 +38,11 @@ public final class Array<T> {
         get() = getArrayLength()
 
     @SymbolName("Kotlin_Array_get")
+    @PointsTo(0b0100, 0, 0b0001) // <this> points to <return>, <return> points to <this>.
     external public operator fun get(index: Int): T
 
     @SymbolName("Kotlin_Array_set")
+    @PointsTo(0b0100) // <this> points to <value>.
     external public operator fun set(index: Int, value: T): Unit
 
     public operator fun iterator(): kotlin.collections.Iterator<T> {

runtime/src/main/kotlin/kotlin/collections/ArrayUtil.kt

@@ -191,6 +191,7 @@ internal fun <E> Array<E>.resetAt(index: Int) {
 }
 
 @SymbolName("Kotlin_Array_fillImpl")
+@PointsTo(0b01000) // <array> points to <value>.
 external private fun fillImpl(array: Array<Any>, fromIndex: Int, toIndex: Int, value: Any?)
 
 @SymbolName("Kotlin_IntArray_fillImpl")
@@ -212,6 +213,7 @@ internal fun IntArray.fill(fromIndex: Int, toIndex: Int, value: Int) {
 }
 
 @SymbolName("Kotlin_Array_copyImpl")
+@PointsTo(0, 0, 0b000001) // <destination> points to <array>.
 external private fun copyImpl(array: Array<Any>, fromIndex: Int,
                          destination: Array<Any>, toIndex: Int, count: Int)