[K/N] Add custom allocator prototype ^KT-55364

Enable custom allocator with -Xallocator=custom.

If the gc is cms (default):
 * a patched version of the memory manager is used that does not build a
   list of allocated objects.
 * a patched version of the cms is used that defers to the allocator for
   sweeping.

Otherwise, a warning is printed, and the allocator is used using the
standard api.

Design doc:
https://docs.google.com/document/d/15xMp-nE-DWL8OrtOc8DoXB80AHUphFICEGjj5K0aNFc

Co-authored-by: Troels Lund <troels@google.com>

Merge-request: KOTLIN-MR-546
Merged-by: Alexander Shabalin <alexander.shabalin@jetbrains.com>

compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2NativeCompilerArguments.kt

@@ -315,7 +315,7 @@ class K2NativeCompilerArguments : CommonCompilerArguments() {
     @Argument(value="-Xoverride-clang-options", valueDescription = "<arg1,arg2,...>", description = "Explicit list of Clang options")
     var clangOptions: Array<String>? = null
 
-    @Argument(value="-Xallocator", valueDescription = "std | mimalloc", description = "Allocator used in runtime")
+    @Argument(value="-Xallocator", valueDescription = "std | mimalloc | custom", description = "Allocator used in runtime")
     var allocator: String? = null
 
     @Argument(value = "-Xmetadata-klib", description = "Produce a klib that only contains the declarations metadata")

kotlin-native/backend.native/cli.bc/src/org/jetbrains/kotlin/cli/bc/K2Native.kt

@@ -324,8 +324,9 @@ class K2Native : CLICompiler<K2NativeCompilerArguments>() {
                     null -> null
                     "std" -> AllocationMode.STD
                     "mimalloc" -> AllocationMode.MIMALLOC
+                    "custom" -> AllocationMode.CUSTOM
                     else -> {
-                        configuration.report(ERROR, "Expected 'std' or 'mimalloc' for allocator")
+                        configuration.report(ERROR, "Expected 'std', 'mimalloc', or 'custom' for allocator")
                         AllocationMode.STD
                     }
                 })

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

@@ -7,5 +7,6 @@ package org.jetbrains.kotlin.backend.konan
 
 enum class AllocationMode {
     STD,
-    MIMALLOC
+    MIMALLOC,
-}
+    CUSTOM
+}

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

@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.library.metadata.resolver.TopologicalLibraryOrder
 import org.jetbrains.kotlin.util.removeSuffixIfPresent
 
 class KonanConfig(val project: Project, val configuration: CompilerConfiguration) {
-
     internal val distribution = run {
         val overridenProperties = mutableMapOf<String, String>().apply {
             configuration.get(KonanConfigKeys.OVERRIDE_KONAN_PROPERTIES)?.let(this::putAll)
@@ -240,6 +239,7 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
         get() = configuration.get(KonanConfigKeys.SHORT_MODULE_NAME)
 
     fun librariesWithDependencies(moduleDescriptor: ModuleDescriptor?): List<KonanLibrary> {
+
         if (moduleDescriptor == null) error("purgeUnneeded() only works correctly after resolve is over, and we have successfully marked package files as needed or not needed.")
         return resolvedLibraries.filterRoots { (!it.isDefault && !this.purgeUserLibs) || it.isNeededForLink }.getFullList(TopologicalLibraryOrder).map { it as KonanLibrary }
     }
@@ -267,6 +267,13 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
                     AllocationMode.STD
                 }
             }
+            AllocationMode.CUSTOM -> {
+                if (gc != GC.CONCURRENT_MARK_AND_SWEEP) {
+                    configuration.report(CompilerMessageSeverity.STRONG_WARNING,
+                            "Custom allocator is currently only integrated with concurrent mark and sweep gc. Performance will not be ideal with selected gc.")
+                }
+                AllocationMode.CUSTOM
+            }
         }
     }
 
@@ -292,7 +299,11 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
                         add("noop_gc.bc")
                     }
                     GC.CONCURRENT_MARK_AND_SWEEP -> {
-                        add("concurrent_ms_gc.bc")
+                        if (allocationMode == AllocationMode.CUSTOM) {
+                            add("concurrent_ms_gc_custom.bc")
+                        } else {
+                            add("concurrent_ms_gc.bc")
+                        }
                     }
                 }
             }
@@ -313,6 +324,9 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
             AllocationMode.STD -> {
                 add("std_alloc.bc")
             }
+            AllocationMode.CUSTOM -> {
+                add("custom_alloc.bc")
+            }
         }
     }.map {
         File(distribution.defaultNatives(target)).child(it).absolutePath

kotlin-native/runtime/build.gradle.kts

@@ -129,6 +129,13 @@ bitcode {
             headersDirs.from(files("src/main/cpp"))
         }
 
+        module("custom_alloc") {
+            headersDirs.from(files("src/main/cpp", "src/mm/cpp", "src/gc/common/cpp", "src/gc/cms/cpp"))
+            compilerArgs.add("-DCUSTOM_ALLOCATOR")
+            // Directly depends on cms which is only supported with threads.
+            onlyIf { targetSupportsThreads(target.name) }
+        }
+
         module("opt_alloc") {
             headersDirs.from(files("src/main/cpp"))
         }
@@ -191,10 +198,22 @@ bitcode {
             onlyIf { targetSupportsThreads(target.name) }
         }
 
+        module("concurrent_ms_gc_custom", file("src/gc/cms")) {
+            headersDirs.from(files("src/gc/cms/cpp", "src/gc/common/cpp", "src/mm/cpp", "src/main/cpp", "src/custom_alloc/cpp"))
+            compilerArgs.add("-DCUSTOM_ALLOCATOR")
+
+            onlyIf { targetSupportsThreads(target.name) }
+        }
+
         testsGroup("std_alloc_runtime_tests") {
             testedModules.addAll("main", "legacy_memory_manager", "strict", "std_alloc", "objc")
         }
 
+        testsGroup("custom_alloc_runtime_tests") {
+            testedModules.addAll("custom_alloc")
+            testSupportModules.addAll("main", "experimental_memory_manager", "common_gc", "concurrent_ms_gc", "objc")
+        }
+
         testsGroup("mimalloc_runtime_tests") {
             testedModules.addAll("main", "legacy_memory_manager", "strict", "mimalloc", "opt_alloc", "objc")
         }

kotlin-native/runtime/src/custom_alloc/cpp/Allocator.cpp

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "CustomLogging.hpp"
+
+// These functions are just stubs to make the existing object creation
+// infrastructure link correctly, but if they are ever called, something went
+// wrong.
+
+namespace kotlin {
+
+void* allocateInObjectPool(size_t size) noexcept {
+    CustomAllocWarning("static allocateInObjectPool(%zu) not supported", size);
+    return nullptr;
+}
+
+void freeInObjectPool(void* ptr) noexcept {
+    CustomAllocWarning("static freeInObjectPool(%p) not supported", ptr);
+}
+
+void initObjectPool() noexcept {}
+void compactObjectPoolInMainThread() noexcept {}
+void compactObjectPoolInCurrentThread() noexcept {}
+
+} // namespace kotlin

kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_ATOMICSTACK_HPP_
+#define CUSTOM_ALLOC_CPP_ATOMICSTACK_HPP_
+
+#include <atomic>
+
+#include "CustomLogging.hpp"
+#include "KAssert.h"
+
+namespace kotlin::alloc {
+
+template <class T>
+class AtomicStack {
+public:
+    // Pop() is not fully thread-safe, in that the returned page must not be
+    // immediately freed, if another thread might be simultaneously Popping
+    // from the same stack. As of writing this comment, this is handled by only
+    // freeing pages during STW.
+    T* Pop() noexcept {
+        T* elm = stack_.load(std::memory_order_acquire);
+        while (elm && !stack_.compare_exchange_weak(elm, elm->next_, std::memory_order_acq_rel)) {}
+        CustomAllocDebug("AtomicStack(%p)::Pop() = %p", this, elm);
+        return elm;
+    }
+
+    void Push(T* elm) noexcept {
+        T* head = nullptr;
+        do {
+            elm->next_ = head;
+        } while (!stack_.compare_exchange_weak(head, elm, std::memory_order_acq_rel));
+    }
+
+    void TransferAllFrom(AtomicStack<T>& other) noexcept {
+        // Clear out the `other` stack.
+        T* otherHead = nullptr;
+        while (!other.stack_.compare_exchange_weak(otherHead, nullptr, std::memory_order_acq_rel)) {}
+        // If the `other` stack was empty, do nothing.
+        if (!otherHead) return;
+        // Now find the tail of `other`. If no deletions are performed, this is safe.
+        T* otherTail = otherHead;
+        while (otherTail->next_) otherTail = otherTail->next_;
+        // Now make `otherTail->next_` point to the current head of `this` and
+        // simultaneously make `otherHead` the new current head.
+        T* thisHead = nullptr;
+        // can't be because of the loop above
+        RuntimeAssert(otherTail->next_ == nullptr, "otherTail->next_ must be a tail");
+        while (!stack_.compare_exchange_weak(thisHead, otherHead, std::memory_order_acq_rel)) {
+            otherTail->next_ = thisHead;
+        }
+    }
+
+    bool isEmpty() noexcept { return stack_.load(std::memory_order_relaxed) == nullptr; }
+
+    ~AtomicStack() noexcept {
+        RuntimeAssert(isEmpty(), "AtomicStack must be empty when destroyed");
+    }
+
+private:
+    std::atomic<T*> stack_{nullptr};
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/Cell.cpp

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "Cell.hpp"
+
+#include <cstdint>
+#include <cstring>
+
+#include "CustomLogging.hpp"
+#include "KAssert.h"
+
+namespace kotlin::alloc {
+
+Cell::Cell(uint32_t size) noexcept : isAllocated_(false), size_(size) {
+    CustomAllocDebug("Cell@%p::Cell(%u)", this, size);
+}
+
+uint8_t* Cell::TryAllocate(uint32_t cellsNeeded) noexcept {
+    CustomAllocDebug("Cell@%p{ allocated = %d, size = %u }::TryAllocate(%u)", this, isAllocated_, size_, cellsNeeded);
+    if (isAllocated_ || cellsNeeded > size_) {
+        CustomAllocDebug("Failed to allocate in Cell");
+        return nullptr;
+    }
+    uint32_t oldSize = size_;
+    uint32_t remainingSize = size_ - cellsNeeded;
+    Cell* newBlock = this + remainingSize;
+    size_ = remainingSize;
+    newBlock->isAllocated_ = true;
+    newBlock->size_ = cellsNeeded;
+    RuntimeAssert(remainingSize == 0 || size_ + newBlock->size_ == oldSize, "sizes don't add up");
+    return newBlock->data_; // Payload starts after header
+}
+
+void Cell::Deallocate() noexcept {
+    CustomAllocDebug("Cell@%p{ allocated = %d, size = %u }::Deallocate()", this, isAllocated_, size_);
+    RuntimeAssert(isAllocated_, "Cell is not currently allocated");
+    isAllocated_ = false;
+}
+
+Cell* Cell::Next() noexcept {
+    return this + size_;
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/Cell.hpp

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_CELL_HPP_
+#define CUSTOM_ALLOC_CPP_CELL_HPP_
+
+#include <cstdint>
+#include <cstring>
+
+namespace kotlin::alloc {
+
+// All allocations are whole units of cells.
+class Cell {
+public:
+    explicit Cell(uint32_t size) noexcept;
+
+    // Allocate `cellsNeeded` blocks at the end of this block, possibly the
+    // whole block, or null if it doesn't fit.
+    uint8_t* TryAllocate(uint32_t cellsNeeded) noexcept;
+
+    // Marks block as no longer allocated.
+    void Deallocate() noexcept;
+
+    // The next block.
+    Cell* Next() noexcept;
+
+private:
+    friend class MediumPage;
+
+    uint32_t isAllocated_;
+    uint32_t size_;
+    // The allocated data follows immediately after the header block
+    uint8_t data_[];
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocConstants.hpp

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_CUSTOMALLOCCONSTANTS_HPP_
+#define CUSTOM_ALLOC_CPP_CUSTOMALLOCCONSTANTS_HPP_
+
+#include <cstddef>
+#include <cstdint>
+
+#include "SmallPage.hpp"
+#include "MediumPage.hpp"
+
+inline constexpr const size_t KiB = 1024;
+
+inline constexpr const size_t SMALL_PAGE_SIZE = (256 * KiB);
+inline constexpr const int SMALL_PAGE_MAX_BLOCK_SIZE = 128;
+inline constexpr const size_t SMALL_PAGE_CELL_COUNT =
+        ((SMALL_PAGE_SIZE - sizeof(kotlin::alloc::SmallPage)) / sizeof(kotlin::alloc::SmallCell));
+
+inline constexpr const size_t MEDIUM_PAGE_SIZE = (256 * KiB);
+inline constexpr const size_t MEDIUM_PAGE_CELL_COUNT =
+        ((MEDIUM_PAGE_SIZE - sizeof(kotlin::alloc::MediumPage)) / sizeof(kotlin::alloc::Cell));
+
+inline constexpr const size_t LARGE_PAGE_SIZE_THRESHOLD = (MEDIUM_PAGE_CELL_COUNT - 1);
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocator.cpp

@@ -0,0 +1,135 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "CustomAllocator.hpp"
+
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <cinttypes>
+#include <new>
+
+#include "ConcurrentMarkAndSweep.hpp"
+#include "CustomLogging.hpp"
+#include "GCScheduler.hpp"
+#include "LargePage.hpp"
+#include "MediumPage.hpp"
+#include "SmallPage.hpp"
+
+namespace kotlin::alloc {
+
+using ObjectData = gc::ConcurrentMarkAndSweep::ObjectData;
+
+struct HeapObjHeader {
+    ObjectData gcData;
+    alignas(kObjectAlignment) ObjHeader object;
+};
+
+// Needs to be kept compatible with `HeapObjHeader` just like `ArrayHeader` is compatible
+// with `ObjHeader`: the former can always be casted to the other.
+struct HeapArrayHeader {
+    ObjectData gcData;
+    alignas(kObjectAlignment) ArrayHeader array;
+};
+
+size_t ObjectAllocatedDataSize(const TypeInfo* typeInfo) noexcept {
+    size_t membersSize = typeInfo->instanceSize_ - sizeof(ObjHeader);
+    return AlignUp(sizeof(HeapObjHeader) + membersSize, kObjectAlignment);
+}
+
+uint64_t ArrayAllocatedDataSize(const TypeInfo* typeInfo, uint32_t count) noexcept {
+    // -(int32_t min) * uint32_t max cannot overflow uint64_t. And are capped
+    // at about half of uint64_t max.
+    uint64_t membersSize = static_cast<uint64_t>(-typeInfo->instanceSize_) * count;
+    // Note: array body is aligned, but for size computation it is enough to align the sum.
+    return AlignUp<uint64_t>(sizeof(HeapArrayHeader) + membersSize, kObjectAlignment);
+}
+
+CustomAllocator::CustomAllocator(Heap& heap, gc::GCSchedulerThreadData& gcScheduler) noexcept :
+    heap_(heap), gcScheduler_(gcScheduler), mediumPage_(nullptr) {
+    CustomAllocInfo("CustomAllocator::CustomAllocator(heap)");
+    memset(smallPages_, 0, sizeof(smallPages_));
+}
+
+ObjHeader* CustomAllocator::CreateObject(const TypeInfo* typeInfo) noexcept {
+    RuntimeAssert(!typeInfo->IsArray(), "Must not be an array");
+    size_t allocSize = ObjectAllocatedDataSize(typeInfo);
+    auto* heapObject = new (Allocate(allocSize)) HeapObjHeader();
+    auto* object = &heapObject->object;
+    object->typeInfoOrMeta_ = const_cast<TypeInfo*>(typeInfo);
+    return object;
+}
+
+ArrayHeader* CustomAllocator::CreateArray(const TypeInfo* typeInfo, uint32_t count) noexcept {
+    RuntimeAssert(typeInfo->IsArray(), "Must be an array");
+    auto allocSize = ArrayAllocatedDataSize(typeInfo, count);
+    auto* heapArray = new (Allocate(allocSize)) HeapArrayHeader();
+    auto* array = &heapArray->array;
+    array->typeInfoOrMeta_ = const_cast<TypeInfo*>(typeInfo);
+    array->count_ = count;
+    return array;
+}
+
+void CustomAllocator::PrepareForGC() noexcept {
+    CustomAllocInfo("CustomAllocator@%p::PrepareForGC()", this);
+    mediumPage_ = nullptr;
+    memset(smallPages_, 0, sizeof(smallPages_));
+}
+
+uint8_t* CustomAllocator::Allocate(uint64_t size) noexcept {
+    gcScheduler_.OnSafePointAllocation(size);
+    CustomAllocDebug("CustomAllocator::Allocate(%" PRIu64 ")", size);
+    uint64_t cellCount = (size + sizeof(Cell) - 1) / sizeof(Cell);
+    uint8_t* ptr;
+    if (cellCount <= SMALL_PAGE_MAX_BLOCK_SIZE) {
+        ptr = AllocateInSmallPage(cellCount);
+    } else if (cellCount > LARGE_PAGE_SIZE_THRESHOLD) {
+        ptr = AllocateInLargePage(cellCount);
+    } else {
+        ptr = AllocateInMediumPage(cellCount);
+    }
+    memset(ptr, 0, size);
+    return ptr;
+}
+
+uint8_t* CustomAllocator::AllocateInLargePage(uint64_t cellCount) noexcept {
+    CustomAllocDebug("CustomAllocator::AllocateInLargePage(%" PRIu64 ")", cellCount);
+    uint8_t* block = heap_.GetLargePage(cellCount)->TryAllocate();
+    return block;
+}
+
+uint8_t* CustomAllocator::AllocateInMediumPage(uint32_t cellCount) noexcept {
+    CustomAllocDebug("CustomAllocator::AllocateInMediumPage(%u)", cellCount);
+    if (mediumPage_) {
+        uint8_t* block = mediumPage_->TryAllocate(cellCount);
+        if (block) return block;
+    }
+    CustomAllocDebug("Failed to allocate in curPage");
+    while (true) {
+        mediumPage_ = heap_.GetMediumPage(cellCount);
+        uint8_t* block = mediumPage_->TryAllocate(cellCount);
+        if (block) return block;
+    }
+}
+
+uint8_t* CustomAllocator::AllocateInSmallPage(uint32_t cellCount) noexcept {
+    CustomAllocDebug("CustomAllocator::AllocateInSmallPage(%u)", cellCount);
+    SmallPage* page = smallPages_[cellCount];
+    if (page) {
+        uint8_t* block = page->TryAllocate();
+        if (block) return block;
+    }
+    CustomAllocDebug("Failed to allocate in current SmallPage");
+    while ((page = heap_.GetSmallPage(cellCount))) {
+        uint8_t* block = page->TryAllocate();
+        if (block) {
+            smallPages_[cellCount] = page;
+            return block;
+        }
+    }
+    return nullptr;
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocator.hpp

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_ALLOCATOR_HPP_
+#define CUSTOM_ALLOC_CPP_ALLOCATOR_HPP_
+
+#include <atomic>
+#include <cstring>
+
+#include "GCScheduler.hpp"
+#include "Heap.hpp"
+#include "MediumPage.hpp"
+#include "Memory.h"
+#include "SmallPage.hpp"
+
+namespace kotlin::alloc {
+
+class CustomAllocator {
+public:
+    explicit CustomAllocator(Heap& heap, gc::GCSchedulerThreadData& gcScheduler) noexcept;
+
+    ObjHeader* CreateObject(const TypeInfo* typeInfo) noexcept;
+
+    ArrayHeader* CreateArray(const TypeInfo* typeInfo, uint32_t count) noexcept;
+
+    void PrepareForGC() noexcept;
+
+private:
+    uint8_t* Allocate(uint64_t cellCount) noexcept;
+    uint8_t* AllocateInLargePage(uint64_t cellCount) noexcept;
+    uint8_t* AllocateInMediumPage(uint32_t cellCount) noexcept;
+    uint8_t* AllocateInSmallPage(uint32_t cellCount) noexcept;
+
+    Heap& heap_;
+    gc::GCSchedulerThreadData& gcScheduler_;
+    MediumPage* mediumPage_;
+    SmallPage* smallPages_[SMALL_PAGE_MAX_BLOCK_SIZE + 1];
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocatorTest.cpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include <cstdint>
+#include <random>
+
+#include "CustomAllocator.hpp"
+#include "GCScheduler.hpp"
+#include "Memory.h"
+#include "gtest/gtest.h"
+#include "Heap.hpp"
+#include "SmallPage.hpp"
+#include "TypeInfo.h"
+
+namespace {
+
+using Heap = typename kotlin::alloc::Heap;
+using CustomAllocator = typename kotlin::alloc::CustomAllocator;
+
+#define MIN_BLOCK_SIZE 2
+
+TEST(CustomAllocTest, SmallAllocNonNull) {
+    const int N = 200;
+    TypeInfo fakeTypes[N];
+    for (int i = 1; i < N; ++i) {
+        fakeTypes[i] = {.instanceSize_ = 8 * i, .flags_ = 0};
+    }
+    Heap heap;
+    kotlin::gc::GCSchedulerConfig config;
+    kotlin::gc::GCSchedulerThreadData schedulerData(config, [](auto&) {});
+    CustomAllocator ca(heap, schedulerData);
+    ObjHeader* obj[N];
+    for (int i = 1; i < N; ++i) {
+        TypeInfo* type = fakeTypes + i;
+        obj[i] = ca.CreateObject(type);
+        EXPECT_TRUE(obj[i]);
+    }
+}
+
+TEST(CustomAllocTest, SmallAllocSameSmallPage) {
+    const int N = SMALL_PAGE_CELL_COUNT / SMALL_PAGE_MAX_BLOCK_SIZE;
+    for (int blocks = MIN_BLOCK_SIZE; blocks < SMALL_PAGE_MAX_BLOCK_SIZE; ++blocks) {
+        Heap heap;
+        kotlin::gc::GCSchedulerConfig config;
+        kotlin::gc::GCSchedulerThreadData schedulerData(config, [](auto&) {});
+        CustomAllocator ca(heap, schedulerData);
+        TypeInfo fakeType = {.instanceSize_ = 8 * blocks, .flags_ = 0};
+        uint8_t* first = reinterpret_cast<uint8_t*>(ca.CreateObject(&fakeType));
+        for (int i = 1; i < N; ++i) {
+            uint8_t* obj = reinterpret_cast<uint8_t*>(ca.CreateObject(&fakeType));
+            uint64_t dist = abs(obj - first);
+            EXPECT_TRUE(dist < SMALL_PAGE_SIZE);
+        }
+    }
+}
+
+TEST(CustomAllocTest, TwoAllocatorsDifferentPages) {
+    for (int blocks = MIN_BLOCK_SIZE; blocks < 2000; ++blocks) {
+        Heap heap;
+        kotlin::gc::GCScheduler scheduler;
+        kotlin::gc::GCSchedulerConfig config;
+        kotlin::gc::GCSchedulerThreadData schedulerData1(config, [](auto&) {});
+        kotlin::gc::GCSchedulerThreadData schedulerData2(config, [](auto&) {});
+        CustomAllocator ca1(heap, schedulerData1);
+        CustomAllocator ca2(heap, schedulerData2);
+        TypeInfo fakeType = {.instanceSize_ = 8 * blocks, .flags_ = 0};
+        uint8_t* obj1 = reinterpret_cast<uint8_t*>(ca1.CreateObject(&fakeType));
+        uint8_t* obj2 = reinterpret_cast<uint8_t*>(ca2.CreateObject(&fakeType));
+        uint64_t dist = abs(obj2 - obj1);
+        EXPECT_TRUE(dist >= SMALL_PAGE_SIZE);
+    }
+}
+
+#undef MIN_BLOCK_SIZE
+} // namespace

kotlin-native/runtime/src/custom_alloc/cpp/CustomLogging.hpp

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_CUSTOMLOGGING_HPP_
+#define CUSTOM_ALLOC_CPP_CUSTOMLOGGING_HPP_
+
+#include "Logging.hpp"
+#include "Porting.h"
+
+#define CustomAllocInfo(format, ...) RuntimeLogInfo({"alloc"}, "t%u " format, konan::currentThreadId(), ##__VA_ARGS__)
+#define CustomAllocDebug(format, ...) RuntimeLogDebug({"alloc"}, "t%u " format, konan::currentThreadId(), ##__VA_ARGS__)
+#define CustomAllocWarning(format, ...) RuntimeLogWarning({"alloc"}, "t%u " format, konan::currentThreadId(), ##__VA_ARGS__)
+#define CustomAllocError(format, ...) RuntimeLogError({"alloc"}, "t%u " format, konan::currentThreadId(), ##__VA_ARGS__)
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "GCApi.hpp"
+
+#include <limits>
+
+#include "ConcurrentMarkAndSweep.hpp"
+#include "CustomLogging.hpp"
+#include "ObjectFactory.hpp"
+
+namespace kotlin::alloc {
+
+bool TryResetMark(void* ptr) noexcept {
+    using Node = typename kotlin::mm::ObjectFactory<kotlin::gc::ConcurrentMarkAndSweep>::Storage::Node;
+    using NodeRef = typename kotlin::mm::ObjectFactory<kotlin::gc::ConcurrentMarkAndSweep>::NodeRef;
+    Node& node = Node::FromData(ptr);
+    NodeRef ref = NodeRef(node);
+    auto& objectData = ref.ObjectData();
+    if (!objectData.tryResetMark()) {
+        auto* objHeader = ref.GetObjHeader();
+        if (HasFinalizers(objHeader)) {
+            CustomAllocWarning("FINALIZER IGNORED");
+        }
+        return false;
+    }
+    return true;
+}
+
+void* SafeAlloc(uint64_t size) noexcept {
+    void* memory;
+    if (size > std::numeric_limits<size_t>::max() || !(memory = std_support::malloc(size))) {
+        konan::consoleErrorf("Out of memory trying to allocate %" PRIu64 "bytes. Aborting.\n", size);
+        konan::abort();
+    }
+    return memory;
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/GCApi.hpp

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_GCAPI_HPP_
+#define CUSTOM_ALLOC_CPP_GCAPI_HPP_
+
+#include <cstdint>
+#include <inttypes.h>
+#include <limits>
+#include <stdlib.h>
+
+namespace kotlin::alloc {
+
+bool TryResetMark(void* ptr) noexcept;
+
+void* SafeAlloc(uint64_t size) noexcept;
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/Heap.cpp

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "Heap.hpp"
+
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <cinttypes>
+#include <new>
+
+#include "CustomAllocConstants.hpp"
+#include "CustomLogging.hpp"
+#include "LargePage.hpp"
+#include "MediumPage.hpp"
+#include "SmallPage.hpp"
+#include "ThreadRegistry.hpp"
+#include "GCImpl.hpp"
+
+namespace kotlin::alloc {
+
+void Heap::PrepareForGC() noexcept {
+    CustomAllocDebug("Heap::PrepareForGC()");
+    for (auto& thread : kotlin::mm::ThreadRegistry::Instance().LockForIter()) {
+        thread.gc().impl().alloc().PrepareForGC();
+    }
+
+    mediumPages_.PrepareForGC();
+    largePages_.PrepareForGC();
+    for (size_t blockSize = 0; blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE; ++blockSize) {
+        smallPages_[blockSize].PrepareForGC();
+    }
+}
+
+void Heap::Sweep() noexcept {
+    CustomAllocDebug("Heap::Sweep()");
+    for (size_t blockSize = 0; blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE; ++blockSize) {
+        smallPages_[blockSize].Sweep();
+    }
+    mediumPages_.Sweep();
+    largePages_.Sweep();
+}
+
+MediumPage* Heap::GetMediumPage(uint32_t cellCount) noexcept {
+    CustomAllocDebug("Heap::GetMediumPage()");
+    return mediumPages_.GetPage(cellCount);
+}
+
+SmallPage* Heap::GetSmallPage(uint32_t cellCount) noexcept {
+    CustomAllocDebug("Heap::GetSmallPage()");
+    return smallPages_[cellCount].GetPage(cellCount);
+}
+
+LargePage* Heap::GetLargePage(uint64_t cellCount) noexcept {
+    CustomAllocInfo("CustomAllocator::AllocateInLargePage(%" PRIu64 ")", cellCount);
+    return largePages_.NewPage(cellCount);
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/Heap.hpp

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_HEAP_HPP_
+#define CUSTOM_ALLOC_CPP_HEAP_HPP_
+
+#include <atomic>
+#include <cstring>
+
+#include "CustomAllocConstants.hpp"
+#include "LargePage.hpp"
+#include "MediumPage.hpp"
+#include "PageStore.hpp"
+#include "SmallPage.hpp"
+
+namespace kotlin::alloc {
+
+class Heap {
+public:
+    // Called once by the GC thread after all mutators have been suspended
+    void PrepareForGC() noexcept;
+
+    // Sweep through all remaining pages, freeing those blocks where CanReclaim
+    // returns true. If multiple sweepers are active, each page will only be
+    // seen by one sweeper.
+    void Sweep() noexcept;
+
+    SmallPage* GetSmallPage(uint32_t cellCount) noexcept;
+    MediumPage* GetMediumPage(uint32_t cellCount) noexcept;
+    LargePage* GetLargePage(uint64_t cellCount) noexcept;
+
+private:
+    PageStore<SmallPage> smallPages_[SMALL_PAGE_MAX_BLOCK_SIZE + 1];
+    PageStore<MediumPage> mediumPages_;
+    PageStore<LargePage> largePages_;
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/HeapTest.cpp

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include <algorithm>
+#include <cstdint>
+#include <random>
+
+#include "LargePage.hpp"
+#include "gtest/gtest.h"
+#include "Heap.hpp"
+#include "SmallPage.hpp"
+
+namespace {
+
+using Heap = typename kotlin::alloc::Heap;
+using SmallPage = typename kotlin::alloc::SmallPage;
+using MediumPage = typename kotlin::alloc::MediumPage;
+using LargePage = typename kotlin::alloc::LargePage;
+
+#define MIN_BLOCK_SIZE 2
+
+void mark(void* obj) {
+    reinterpret_cast<uint64_t*>(obj)[0] = 1;
+}
+
+TEST(CustomAllocTest, HeapReuseSmallPages) {
+    Heap heap;
+    const int MIN = MIN_BLOCK_SIZE;
+    const int MAX = SMALL_PAGE_MAX_BLOCK_SIZE + 1;
+    SmallPage* pages[MAX];
+    for (int blocks = MIN; blocks < MAX; ++blocks) {
+        pages[blocks] = heap.GetSmallPage(blocks);
+        void* obj = pages[blocks]->TryAllocate();
+        mark(obj); // to make the page survive a sweep
+    }
+    heap.PrepareForGC();
+    heap.Sweep();
+    for (int blocks = MIN; blocks < MAX; ++blocks) {
+        EXPECT_EQ(pages[blocks], heap.GetSmallPage(blocks));
+    }
+}
+
+TEST(CustomAllocTest, HeapReuseMediumPages) {
+    Heap heap;
+    const uint32_t BLOCKSIZE = SMALL_PAGE_MAX_BLOCK_SIZE + 42;
+    MediumPage* page = heap.GetMediumPage(BLOCKSIZE);
+    void* obj = page->TryAllocate(BLOCKSIZE);
+    mark(obj); // to make the page survive a sweep
+    heap.PrepareForGC();
+    heap.Sweep();
+    EXPECT_EQ(page, heap.GetMediumPage(BLOCKSIZE));
+}
+
+} // namespace

kotlin-native/runtime/src/custom_alloc/cpp/LargePage.cpp

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "LargePage.hpp"
+
+#include <atomic>
+#include <cstdint>
+
+#include "CustomLogging.hpp"
+#include "CustomAllocConstants.hpp"
+#include "GCApi.hpp"
+
+namespace kotlin::alloc {
+
+LargePage* LargePage::Create(uint64_t cellCount) noexcept {
+    CustomAllocInfo("LargePage::Create(%" PRIu64 ")", cellCount);
+    RuntimeAssert(cellCount > LARGE_PAGE_SIZE_THRESHOLD, "blockSize too small for large page");
+    uint64_t size = sizeof(LargePage) + cellCount * sizeof(uint64_t);
+    return new (SafeAlloc(size)) LargePage();
+}
+
+void LargePage::Destroy() noexcept {
+    std_support::free(this);
+}
+
+uint8_t* LargePage::Data() noexcept {
+    return data_;
+}
+
+uint8_t* LargePage::TryAllocate() noexcept {
+    if (isAllocated_) return nullptr;
+    isAllocated_ = true;
+    return Data();
+}
+
+bool LargePage::Sweep() noexcept {
+    CustomAllocDebug("LargePage@%p::Sweep()", this);
+    if (!TryResetMark(Data())) {
+        isAllocated_ = false;
+        return false;
+    }
+    return true;
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/LargePage.hpp

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_LARGEPAGE_HPP_
+#define CUSTOM_ALLOC_CPP_LARGEPAGE_HPP_
+
+#include <atomic>
+#include <cstdint>
+
+#include "AtomicStack.hpp"
+
+namespace kotlin::alloc {
+
+class alignas(8) LargePage {
+public:
+    static LargePage* Create(uint64_t cellCount) noexcept;
+
+    void Destroy() noexcept;
+
+    uint8_t* TryAllocate() noexcept;
+
+    uint8_t* Data() noexcept;
+
+    bool Sweep() noexcept;
+
+private:
+    friend class AtomicStack<LargePage>;
+    LargePage* next_;
+    bool isAllocated_ = false;
+    struct alignas(8) {
+        uint8_t data_[];
+    };
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/LargePageTest.cpp

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include <cstdint>
+#include <random>
+
+#include "CustomAllocConstants.hpp"
+#include "gtest/gtest.h"
+#include "LargePage.hpp"
+#include "TypeInfo.h"
+
+namespace {
+
+using LargePage = typename kotlin::alloc::LargePage;
+
+TypeInfo fakeType = {.flags_ = 0}; // a type without a finalizer
+
+#define MIN_BLOCK_SIZE MEDIUM_PAGE_CELL_COUNT
+
+void mark(void* obj) {
+    reinterpret_cast<uint64_t*>(obj)[0] = 1;
+}
+
+LargePage* alloc(uint64_t blockSize) {
+    LargePage* page = LargePage::Create(blockSize);
+    uint64_t* ptr = reinterpret_cast<uint64_t*>(page->TryAllocate());
+    memset(ptr, 0, 8 * blockSize);
+    ptr[1] = reinterpret_cast<uint64_t>(&fakeType);
+    return page;
+}
+
+TEST(CustomAllocTest, LargePageSweepEmptyPage) {
+    LargePage* page = alloc(MIN_BLOCK_SIZE);
+    EXPECT_TRUE(page);
+    EXPECT_FALSE(page->Sweep());
+    page->Destroy();
+}
+
+TEST(CustomAllocTest, LargePageSweepFullPage) {
+    LargePage* page = alloc(MIN_BLOCK_SIZE);
+    EXPECT_TRUE(page);
+    EXPECT_TRUE(page->Data());
+    mark(page->Data());
+    EXPECT_TRUE(page->Sweep());
+    page->Destroy();
+}
+
+#undef MIN_BLOCK_SIZE
+} // namespace

kotlin-native/runtime/src/custom_alloc/cpp/MediumPage.cpp

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "MediumPage.hpp"
+
+#include <atomic>
+#include <cstdint>
+
+#include "CustomLogging.hpp"
+#include "CustomAllocConstants.hpp"
+#include "GCApi.hpp"
+
+namespace kotlin::alloc {
+
+MediumPage* MediumPage::Create(uint32_t cellCount) noexcept {
+    CustomAllocInfo("MediumPage::Create(%u)", cellCount);
+    RuntimeAssert(cellCount < MEDIUM_PAGE_CELL_COUNT, "cellCount is too large for medium page");
+    return new (SafeAlloc(MEDIUM_PAGE_SIZE)) MediumPage(cellCount);
+}
+
+void MediumPage::Destroy() noexcept {
+    std_support::free(this);
+}
+
+MediumPage::MediumPage(uint32_t cellCount) noexcept : curBlock_(cells_) {
+    cells_[0] = Cell(0); // Size 0 ensures any actual use would break
+    cells_[1] = Cell(MEDIUM_PAGE_CELL_COUNT - 1);
+}
+
+uint8_t* MediumPage::TryAllocate(uint32_t blockSize) noexcept {
+    CustomAllocDebug("MediumPage@%p::TryAllocate(%u)", this, blockSize);
+    // +1 accounts for header, since cell->size also includes header cell
+    uint32_t cellsNeeded = blockSize + 1;
+    uint8_t* block = curBlock_->TryAllocate(cellsNeeded);
+    if (block) return block;
+    UpdateCurBlock(cellsNeeded);
+    return curBlock_->TryAllocate(cellsNeeded);
+}
+
+bool MediumPage::Sweep() noexcept {
+    CustomAllocDebug("MediumPage@%p::Sweep()", this);
+    Cell* end = cells_ + MEDIUM_PAGE_CELL_COUNT;
+    bool alive = false;
+    for (Cell* block = cells_ + 1; block != end; block = block->Next()) {
+        if (block->isAllocated_) {
+            if (TryResetMark(block->data_)) {
+                alive = true;
+            } else {
+                block->Deallocate();
+            }
+        }
+    }
+    Cell* maxBlock = cells_; // size 0 block
+    for (Cell* block = cells_ + 1; block != end; block = block->Next()) {
+        if (block->isAllocated_) continue;
+        while (block->Next() != end && !block->Next()->isAllocated_) {
+            block->size_ += block->Next()->size_;
+        }
+        if (block->size_ > maxBlock->size_) maxBlock = block;
+    }
+    curBlock_ = maxBlock;
+    return alive;
+}
+
+void MediumPage::UpdateCurBlock(uint32_t cellsNeeded) noexcept {
+    CustomAllocDebug("MediumPage@%p::UpdateCurBlock(%u)", this, cellsNeeded);
+    if (curBlock_ == cells_) curBlock_ = cells_ + 1; // only used as a starting point
+    Cell* end = cells_ + MEDIUM_PAGE_CELL_COUNT;
+    Cell* maxBlock = cells_; // size 0 block
+    for (Cell* block = curBlock_; block != end; block = block->Next()) {
+        if (!block->isAllocated_ && block->size_ > maxBlock->size_) {
+            maxBlock = block;
+            if (block->size_ >= cellsNeeded) {
+                curBlock_ = maxBlock;
+                return;
+            }
+        }
+    }
+    CustomAllocDebug("MediumPage@%p::UpdateCurBlock: starting from beginning", this);
+    for (Cell* block = cells_ + 1; block != curBlock_; block = block->Next()) {
+        if (!block->isAllocated_ && block->size_ > maxBlock->size_) {
+            maxBlock = block;
+            if (block->size_ >= cellsNeeded) {
+                curBlock_ = maxBlock;
+                return;
+            }
+        }
+    }
+    curBlock_ = maxBlock;
+}
+
+bool MediumPage::CheckInvariants() noexcept {
+    if (curBlock_ < cells_ || curBlock_ >= cells_ + MEDIUM_PAGE_CELL_COUNT) return false;
+    for (Cell* cur = cells_ + 1;; cur = cur->Next()) {
+        if (cur->Next() <= cur) return false;
+        if (cur->Next() > cells_ + MEDIUM_PAGE_CELL_COUNT) return false;
+        if (cur->Next() == cells_ + MEDIUM_PAGE_CELL_COUNT) return true;
+    }
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/MediumPage.hpp

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_MEDIUMPAGE_HPP_
+#define CUSTOM_ALLOC_CPP_MEDIUMPAGE_HPP_
+
+#include <atomic>
+#include <cstdint>
+
+#include "AtomicStack.hpp"
+#include "Cell.hpp"
+
+namespace kotlin::alloc {
+
+class alignas(8) MediumPage {
+public:
+    static MediumPage* Create(uint32_t cellCount) noexcept;
+
+    void Destroy() noexcept;
+
+    // Tries to allocate in current page, returns null if no free block in page is big enough
+    uint8_t* TryAllocate(uint32_t blockSize) noexcept;
+
+    bool Sweep() noexcept;
+
+    // Testing method
+    bool CheckInvariants() noexcept;
+
+private:
+    MediumPage(uint32_t cellCount) noexcept;
+
+    // Looks for a block big enough to hold cellsNeeded. If none big enough is
+    // found, update to the largest one.
+    void UpdateCurBlock(uint32_t cellsNeeded) noexcept;
+
+    friend class AtomicStack<MediumPage>;
+    MediumPage* next_;
+
+    Cell* curBlock_;
+    Cell cells_[]; // cells_[0] is reserved for an empty block
+};
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/MediumPageTest.cpp

@@ -0,0 +1,122 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include <cstdint>
+#include <random>
+
+#include "Cell.hpp"
+#include "CustomAllocConstants.hpp"
+#include "gtest/gtest.h"
+#include "MediumPage.hpp"
+#include "TypeInfo.h"
+
+namespace {
+
+using MediumPage = typename kotlin::alloc::MediumPage;
+using Cell = typename kotlin::alloc::Cell;
+
+TypeInfo fakeType = {.flags_ = 0}; // a type without a finalizer
+
+#define MIN_BLOCK_SIZE (SMALL_PAGE_MAX_BLOCK_SIZE + 1)
+
+void mark(void* obj) {
+    reinterpret_cast<uint64_t*>(obj)[0] = 1;
+}
+
+uint8_t* alloc(MediumPage* page, uint32_t blockSize) {
+    uint8_t* ptr = page->TryAllocate(blockSize);
+    if (!page->CheckInvariants()) {
+        ADD_FAILURE();
+        return nullptr;
+    }
+    if (ptr == nullptr) return nullptr;
+    memset(ptr, 0, 8 * blockSize);
+    reinterpret_cast<uint64_t*>(ptr)[1] = reinterpret_cast<uint64_t>(&fakeType);
+    if (!page->CheckInvariants()) {
+        ADD_FAILURE();
+        return nullptr;
+    }
+    return ptr;
+}
+
+TEST(CustomAllocTest, MediumPageAlloc) {
+    MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+    uint8_t* p1 = alloc(page, MIN_BLOCK_SIZE);
+    uint8_t* p2 = alloc(page, MIN_BLOCK_SIZE);
+    uint64_t dist = abs(p1 - p2);
+    EXPECT_EQ(dist, (MIN_BLOCK_SIZE + 1) * sizeof(kotlin::alloc::Cell));
+    page->Destroy();
+}
+
+TEST(CustomAllocTest, MediumPageSweepEmptyPage) {
+    MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+    EXPECT_FALSE(page->Sweep());
+    page->Destroy();
+}
+
+TEST(CustomAllocTest, MediumPageSweepFullUnmarkedPage) {
+    for (uint32_t seed = 0xC0FFEE0; seed <= 0xC0FFEEF; ++seed) {
+        std::minstd_rand r(seed);
+        MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+        while (alloc(page, MIN_BLOCK_SIZE + r() % 100)) {}
+        EXPECT_FALSE(page->Sweep());
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, MediumPageSweepSingleMarked) {
+    MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+    mark(alloc(page, MIN_BLOCK_SIZE));
+    EXPECT_TRUE(page->Sweep());
+    page->Destroy();
+}
+
+TEST(CustomAllocTest, MediumPageSweepSingleReuse) {
+    for (uint32_t seed = 0xC0FFEE0; seed <= 0xC0FFEEF; ++seed) {
+        std::minstd_rand r(seed);
+        MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+        int count1 = 0;
+        while (alloc(page, MIN_BLOCK_SIZE + r() % 100)) ++count1;
+        EXPECT_FALSE(page->Sweep());
+        r.seed(seed);
+        int count2 = 0;
+        while (alloc(page, MIN_BLOCK_SIZE + r() % 100)) ++count2;
+        EXPECT_EQ(count1, count2);
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, MediumPageSweepReuse) {
+    for (uint32_t seed = 0xC0FFEE0; seed <= 0xC0FFEEF; ++seed) {
+        std::minstd_rand r(seed);
+        MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+        int unmarked = 0;
+        while (true) {
+            uint8_t* ptr = alloc(page, MIN_BLOCK_SIZE);
+            if (ptr == nullptr) break;
+            if (r() & 1) {
+                mark(ptr);
+            } else {
+                ++unmarked;
+            }
+        }
+        page->Sweep();
+        int freed = 0;
+        while (alloc(page, MIN_BLOCK_SIZE)) ++freed;
+        EXPECT_EQ(freed, unmarked);
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, MediumPageSweepCoallesce) {
+    MediumPage* page = MediumPage::Create(MIN_BLOCK_SIZE);
+    EXPECT_TRUE(alloc(page, (MEDIUM_PAGE_CELL_COUNT-1) / 2 - 1));
+    EXPECT_FALSE(page->Sweep());
+    EXPECT_TRUE(alloc(page, (MEDIUM_PAGE_CELL_COUNT-1) - 1));
+    page->Destroy();
+}
+
+#undef MIN_BLOCK_SIZE
+} // namespace

kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_PAGESTORE_HPP_
+#define CUSTOM_ALLOC_CPP_PAGESTORE_HPP_
+
+#include <atomic>
+
+#include "AtomicStack.hpp"
+
+namespace kotlin::alloc {
+
+template <class T>
+class PageStore {
+public:
+    void PrepareForGC() noexcept {
+        unswept_.TransferAllFrom(ready_);
+        unswept_.TransferAllFrom(used_);
+        T* page;
+        while ((page = empty_.Pop())) page->Destroy();
+    }
+
+    T* SweepAndFreeEmpty(AtomicStack<T>& from, AtomicStack<T>& to) noexcept {
+        T* page;
+        while ((page = from.Pop())) {
+            if (!page->Sweep()) {
+                empty_.Push(page);
+            } else {
+                to.Push(page);
+                return page;
+            }
+        }
+        return nullptr;
+    }
+
+    void Sweep() noexcept {
+        while (SweepAndFreeEmpty(unswept_, ready_)) {}
+    }
+
+    T* GetPage(uint32_t cellCount) noexcept {
+        T* page;
+        if ((page = SweepAndFreeEmpty(unswept_, used_))) {
+            return page;
+        }
+        if ((page = ready_.Pop())) {
+            used_.Push(page);
+            return page;
+        }
+        if ((page = empty_.Pop())) {
+            used_.Push(page);
+            return page;
+        }
+        return NewPage(cellCount);
+    }
+
+    T* NewPage(uint64_t cellCount) noexcept {
+        T* page = T::Create(cellCount);
+        used_.Push(page);
+        return page;
+    }
+
+    ~PageStore() noexcept {
+        T* page;
+        while ((page = empty_.Pop())) page->Destroy();
+        while ((page = ready_.Pop())) page->Destroy();
+        while ((page = used_.Pop())) page->Destroy();
+        while ((page = unswept_.Pop())) page->Destroy();
+    }
+
+private:
+    AtomicStack<T> empty_;
+    AtomicStack<T> ready_;
+    AtomicStack<T> used_;
+    AtomicStack<T> unswept_;
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/SmallPage.cpp

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "SmallPage.hpp"
+
+#include <atomic>
+#include <cstdint>
+#include <cstring>
+#include <random>
+
+#include "CustomLogging.hpp"
+#include "CustomAllocConstants.hpp"
+#include "GCApi.hpp"
+
+namespace kotlin::alloc {
+
+SmallPage* SmallPage::Create(uint32_t blockSize) noexcept {
+    CustomAllocInfo("SmallPage::Create(%u)", blockSize);
+    RuntimeAssert(blockSize <= SMALL_PAGE_MAX_BLOCK_SIZE, "blockSize too large for small page");
+    return new (SafeAlloc(SMALL_PAGE_SIZE)) SmallPage(blockSize);
+}
+
+void SmallPage::Destroy() noexcept {
+    std_support::free(this);
+}
+
+SmallPage::SmallPage(uint32_t blockSize) noexcept : blockSize_(blockSize) {
+    CustomAllocInfo("SmallPage(%p)::SmallPage(%u)", this, blockSize);
+    nextFree_ = cells_;
+    SmallCell* end = cells_ + (SMALL_PAGE_CELL_COUNT + 1 - blockSize_);
+    for (SmallCell* cell = cells_; cell < end; cell = cell->nextFree) {
+        cell->nextFree = cell + blockSize;
+    }
+}
+
+uint8_t* SmallPage::TryAllocate() noexcept {
+    SmallCell* end = cells_ + (SMALL_PAGE_CELL_COUNT + 1 - blockSize_);
+    SmallCell* freeBlock = nextFree_;
+    if (freeBlock >= end) {
+        return nullptr;
+    }
+    nextFree_ = freeBlock->nextFree;
+    CustomAllocDebug("SmallPage(%p){%u}::TryAllocate() = %p", this, blockSize_, freeBlock);
+    return freeBlock->data;
+}
+
+bool SmallPage::Sweep() noexcept {
+    CustomAllocInfo("SmallPage(%p)::Sweep()", this);
+    // `end` is after the last legal allocation of a block, but does not
+    // necessarily match an actual block starting point.
+    SmallCell* end = cells_ + (SMALL_PAGE_CELL_COUNT + 1 - blockSize_);
+    bool alive = false;
+    SmallCell** nextFree = &nextFree_;
+    for (SmallCell* cell = cells_; cell < end; cell += blockSize_) {
+        // If the current cell is free, move on.
+        if (cell == *nextFree) {
+            nextFree = &cell->nextFree;
+            continue;
+        }
+        // If the current cell was marked, it's alive, and the whole page is alive.
+        if (TryResetMark(cell)) {
+            alive = true;
+            continue;
+        }
+        // Free the current block and insert it into the free list.
+        cell->nextFree = *nextFree;
+        *nextFree = cell;
+        nextFree = &cell->nextFree;
+    }
+    return alive;
+}
+
+} // namespace kotlin::alloc

kotlin-native/runtime/src/custom_alloc/cpp/SmallPage.hpp

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef CUSTOM_ALLOC_CPP_SMALLPAGE_HPP_
+#define CUSTOM_ALLOC_CPP_SMALLPAGE_HPP_
+
+#include <atomic>
+#include <cstdint>
+
+#include "AtomicStack.hpp"
+
+namespace kotlin::alloc {
+
+struct alignas(8) SmallCell {
+    // The SmallCell either contains data or a pointer to the next free cell
+    union {
+        uint8_t data[];
+        SmallCell* nextFree;
+    };
+};
+
+class alignas(8) SmallPage {
+public:
+    static SmallPage* Create(uint32_t blockSize) noexcept;
+
+    void Destroy() noexcept;
+
+    // Tries to allocate in current page, returns null if no free block in page
+    uint8_t* TryAllocate() noexcept;
+
+    bool Sweep() noexcept;
+
+private:
+    friend class AtomicStack<SmallPage>;
+
+    explicit SmallPage(uint32_t blockSize) noexcept;
+
+    // Used for linking pages together in `pages` queue or in `unswept` queue.
+    SmallPage* next_;
+    uint32_t blockSize_;
+    SmallCell* nextFree_;
+    SmallCell cells_[];
+};
+
+} // namespace kotlin::alloc
+
+#endif

kotlin-native/runtime/src/custom_alloc/cpp/SmallPageTest.cpp

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include <cstdint>
+
+#include "Cell.hpp"
+#include "CustomAllocConstants.hpp"
+#include "gtest/gtest.h"
+#include "SmallPage.hpp"
+#include "TypeInfo.h"
+
+namespace {
+
+using SmallPage = typename kotlin::alloc::SmallPage;
+
+TypeInfo fakeType = {.flags_ = 0}; // a type without a finalizer
+
+void mark(void* obj) {
+    reinterpret_cast<uint64_t*>(obj)[0] = 1;
+}
+
+uint8_t* alloc(SmallPage* page, size_t blockSize) {
+    uint8_t* ptr = page->TryAllocate();
+    if (ptr) {
+        memset(ptr, 0, 8 * blockSize);
+        reinterpret_cast<uint64_t*>(ptr)[1] = reinterpret_cast<uint64_t>(&fakeType);
+    }
+    return ptr;
+}
+
+TEST(CustomAllocTest, SmallPageConsequtiveAlloc) {
+    for (uint32_t size = 2; size <= SMALL_PAGE_MAX_BLOCK_SIZE; ++size) {
+        SmallPage* page = SmallPage::Create(size);
+        uint8_t* prev = alloc(page, size);
+        uint8_t* cur;
+        while ((cur = alloc(page, size))) {
+            EXPECT_EQ(prev + sizeof(kotlin::alloc::Cell) * size, cur);
+            prev = cur;
+        }
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, SmallPageSweepEmptyPage) {
+    for (uint32_t size = 2; size <= SMALL_PAGE_MAX_BLOCK_SIZE; ++size) {
+        SmallPage* page = SmallPage::Create(size);
+        EXPECT_FALSE(page->Sweep());
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, SmallPageSweepFullUnmarkedPage) {
+    for (uint32_t size = 2; size <= SMALL_PAGE_MAX_BLOCK_SIZE; ++size) {
+        SmallPage* page = SmallPage::Create(size);
+        uint32_t count = 0;
+        while (alloc(page, size)) ++count;
+        EXPECT_EQ(count, SMALL_PAGE_CELL_COUNT / size);
+        EXPECT_FALSE(page->Sweep());
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, SmallPageSweepSingleMarked) {
+    for (uint32_t size = 2; size <= SMALL_PAGE_MAX_BLOCK_SIZE; ++size) {
+        SmallPage* page = SmallPage::Create(size);
+        uint8_t* ptr = alloc(page, size);
+        mark(ptr);
+        EXPECT_TRUE(page->Sweep());
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, SmallPageSweepSingleReuse) {
+    for (uint32_t size = 2; size <= SMALL_PAGE_MAX_BLOCK_SIZE; ++size) {
+        SmallPage* page = SmallPage::Create(size);
+        uint8_t* ptr = alloc(page, size);
+        EXPECT_FALSE(page->Sweep());
+        EXPECT_EQ(alloc(page, size), ptr);
+        page->Destroy();
+    }
+}
+
+TEST(CustomAllocTest, SmallPageSweepReuse) {
+    for (uint32_t size = 2; size <= SMALL_PAGE_MAX_BLOCK_SIZE; ++size) {
+        SmallPage* page = SmallPage::Create(size);
+        uint8_t* ptr;
+        for (int count = 0; (ptr = alloc(page, size)); ++count) {
+            if (count % 2 == 0) mark(ptr);
+        }
+        EXPECT_TRUE(page->Sweep());
+        uint32_t count = 0;
+        for (; (ptr = alloc(page, size)); ++count) {
+            if (count % 2 == 0) mark(ptr);
+        }
+        EXPECT_EQ(count, SMALL_PAGE_CELL_COUNT / size / 2);
+        page->Destroy();
+    }
+}
+} // namespace

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.cpp

@@ -22,6 +22,10 @@
 #include "FinalizerProcessor.hpp"
 #include "GCStatistics.hpp"
 
+#ifdef CUSTOM_ALLOCATOR
+#include "Heap.hpp"
+#endif
+
 using namespace kotlin;
 
 namespace {
@@ -94,7 +98,9 @@ NO_EXTERNAL_CALLS_CHECK void gc::ConcurrentMarkAndSweep::ThreadData::OnSuspendFo
 
 gc::ConcurrentMarkAndSweep::ConcurrentMarkAndSweep(
         mm::ObjectFactory<ConcurrentMarkAndSweep>& objectFactory, GCScheduler& gcScheduler) noexcept :
+#ifndef CUSTOM_ALLOCATOR
     objectFactory_(objectFactory),
+#endif
     gcScheduler_(gcScheduler),
     finalizerProcessor_(std_support::make_unique<FinalizerProcessor>([this](int64_t epoch) {
         state_.finalized(epoch);
@@ -154,6 +160,10 @@ bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     WaitForThreadsReadyToMark();
     gcHandle.threadsAreSuspended();
 
+#ifdef CUSTOM_ALLOCATOR
+    heap_.PrepareForGC();
+#endif
+
     auto& scheduler = gcScheduler_;
     scheduler.gcData().OnPerformFullGC();
 
@@ -171,15 +181,18 @@ bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
 
     gc::SweepExtraObjects<SweepTraits>(gcHandle, extraObjectDataFactory);
 
+#ifndef CUSTOM_ALLOCATOR
     auto objectFactoryIterable = objectFactory_.LockForIter();
-
     mm::ResumeThreads();
     gcHandle.threadsAreResumed();
-
     auto finalizerQueue = gc::Sweep<SweepTraits>(gcHandle, objectFactoryIterable);
-
+#else
+    mm::ResumeThreads();
+    gcHandle.threadsAreResumed();
+    SweepTraits::ObjectFactory::FinalizerQueue finalizerQueue;
+    heap_.Sweep();
+#endif
     kotlin::compactObjectPoolInMainThread();
-
     state_.finish(epoch);
     gcHandle.finalizersScheduled(finalizerQueue.size());
     gcHandle.finished();

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.hpp

@@ -14,18 +14,19 @@
 #include "MarkAndSweepUtils.hpp"
 #include "ObjectFactory.hpp"
 #include "ScopedThread.hpp"
+#include "ThreadData.hpp"
 #include "Types.h"
 #include "Utils.hpp"
 #include "GCState.hpp"
 #include "std_support/Memory.hpp"
 #include "GCStatistics.hpp"
 
+#ifdef CUSTOM_ALLOCATOR
+#include "CustomAllocator.hpp"
+#include "Heap.hpp"
+#endif
+
 namespace kotlin {
-
-namespace mm {
-class ThreadData;
-}
-
 namespace gc {
 
 class FinalizerProcessor;
@@ -72,6 +73,7 @@ public:
     class ThreadData : private Pinned {
     public:
         using ObjectData = ConcurrentMarkAndSweep::ObjectData;
+
         using Allocator = AllocatorWithGC<Allocator, ThreadData>;
 
         explicit ThreadData(ConcurrentMarkAndSweep& gc, mm::ThreadData& threadData, GCSchedulerThreadData& gcScheduler) noexcept :
@@ -111,11 +113,19 @@ public:
     void WaitForThreadsReadyToMark() noexcept;
     void CollectRootSetAndStartMarking(GCHandle gcHandle) noexcept;
 
+#ifdef CUSTOM_ALLOCATOR
+    alloc::Heap& heap() noexcept { return heap_; }
+#endif
+
 private:
     // Returns `true` if GC has happened, and `false` if not (because someone else has suspended the threads).
     bool PerformFullGC(int64_t epoch) noexcept;
 
+#ifndef CUSTOM_ALLOCATOR
     mm::ObjectFactory<ConcurrentMarkAndSweep>& objectFactory_;
+#else
+    alloc::Heap heap_;
+#endif
     GCScheduler& gcScheduler_;
 
     GCStateHolder state_;

kotlin-native/runtime/src/gc/cms/cpp/GCImpl.cpp

@@ -47,19 +47,31 @@ void gc::GC::ThreadData::ScheduleAndWaitFullGCWithFinalizers() noexcept {
 }
 
 void gc::GC::ThreadData::Publish() noexcept {
+#ifndef CUSTOM_ALLOCATOR
     impl_->objectFactoryThreadQueue().Publish();
+#endif
 }
 
 void gc::GC::ThreadData::ClearForTests() noexcept {
+#ifndef CUSTOM_ALLOCATOR
     impl_->objectFactoryThreadQueue().ClearForTests();
+#endif
 }
 
 ALWAYS_INLINE ObjHeader* gc::GC::ThreadData::CreateObject(const TypeInfo* typeInfo) noexcept {
+#ifndef CUSTOM_ALLOCATOR
     return impl_->objectFactoryThreadQueue().CreateObject(typeInfo);
+#else
+    return impl_->alloc().CreateObject(typeInfo);
+#endif
 }
 
 ALWAYS_INLINE ArrayHeader* gc::GC::ThreadData::CreateArray(const TypeInfo* typeInfo, uint32_t elements) noexcept {
+#ifndef CUSTOM_ALLOCATOR
     return impl_->objectFactoryThreadQueue().CreateArray(typeInfo, elements);
+#else
+    return impl_->alloc().CreateArray(typeInfo, elements);
+#endif
 }
 
 void gc::GC::ThreadData::OnStoppedForGC() noexcept {

kotlin-native/runtime/src/gc/cms/cpp/GCImpl.hpp

@@ -9,6 +9,10 @@
 
 #include "ConcurrentMarkAndSweep.hpp"
 
+#ifdef CUSTOM_ALLOCATOR
+#include "CustomAllocator.hpp"
+#endif
+
 namespace kotlin {
 namespace gc {
 
@@ -33,16 +37,28 @@ public:
     Impl(GC& gc, mm::ThreadData& threadData) noexcept :
         gcScheduler_(gc.impl_->gcScheduler().NewThreadData()),
         gc_(gc.impl_->gc(), threadData, gcScheduler_),
+#ifndef CUSTOM_ALLOCATOR
         objectFactoryThreadQueue_(gc.impl_->objectFactory(), gc_.CreateAllocator()) {}
+#else
+        alloc_(gc.impl_->gc().heap(), gcScheduler_) {}
+#endif
 
     GCSchedulerThreadData& gcScheduler() noexcept { return gcScheduler_; }
     GCImpl::ThreadData& gc() noexcept { return gc_; }
+#ifdef CUSTOM_ALLOCATOR
+    alloc::CustomAllocator& alloc() noexcept { return alloc_; }
+#else
     mm::ObjectFactory<GCImpl>::ThreadQueue& objectFactoryThreadQueue() noexcept { return objectFactoryThreadQueue_; }
+#endif
 
 private:
     GCSchedulerThreadData gcScheduler_;
     GCImpl::ThreadData gc_;
+#ifdef CUSTOM_ALLOCATOR
+    alloc::CustomAllocator alloc_;
+#else
     mm::ObjectFactory<GCImpl>::ThreadQueue objectFactoryThreadQueue_;
+#endif
 };
 
 } // namespace gc