From 94c69fff2bc162bd7bc61a913897e033b73df7f3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 21 Aug 2019 10:44:40 -0700
Subject: [PATCH] feat(runtime/alloc): auxiliary functions for new runtime

---
 src/runtime/alloc.cpp | 27 +++++++++++++++++++++------
 src/runtime/alloc.h   |  2 ++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/runtime/alloc.cpp b/src/runtime/alloc.cpp
index 73fb67cfa7..e989e1ed58 100644
--- a/src/runtime/alloc.cpp
+++ b/src/runtime/alloc.cpp
@@ -51,6 +51,7 @@ struct page_header {
     page *           m_next;
     page *           m_prev;
     void *           m_free_list;
+    unsigned         m_obj_size;
     unsigned         m_max_free;
     unsigned         m_num_free;
     unsigned         m_slot_idx;
@@ -312,6 +313,7 @@ static page * alloc_page(heap * h, unsigned obj_size) {
     p->m_header.m_heap       = h;
     page_list_insert(h->m_curr_page[slot_idx], p);
     p->m_header.m_slot_idx   = slot_idx;
+    p->m_header.m_obj_size   = obj_size;
     char * curr_free         = p->m_data;
     set_next_obj(curr_free, nullptr);
     char * end               = p->m_data + (LEAN_PAGE_SIZE - sizeof(page_header));
@@ -400,12 +402,7 @@ void * alloc(size_t sz) {
     return r;
 }
 
-void dealloc(void * o, size_t sz) {
-    LEAN_RUNTIME_STAT_CODE(g_num_dealloc++);
-    sz = align(sz, LEAN_OBJECT_SIZE_DELTA);
-    if (LEAN_UNLIKELY(sz > LEAN_MAX_SMALL_OBJECT_SIZE)) {
-        return free(o);
-    }
+static inline void dealloc_small_core(void * o) {
     LEAN_RUNTIME_STAT_CODE(g_num_small_dealloc++);
     if (LEAN_UNLIKELY(g_heap == nullptr)) {
         init_heap(false);
@@ -425,6 +422,24 @@ void dealloc(void * o, size_t sz) {
     }
 }
 
+void dealloc(void * o, size_t sz) {
+    LEAN_RUNTIME_STAT_CODE(g_num_dealloc++);
+    sz = align(sz, LEAN_OBJECT_SIZE_DELTA);
+    if (LEAN_UNLIKELY(sz > LEAN_MAX_SMALL_OBJECT_SIZE)) {
+        return free(o);
+    }
+    dealloc_small_core(o);
+}
+
+void dealloc_small(void * o) {
+    dealloc_small_core(o);
+}
+
+unsigned get_small_object_size(void * o) {
+    page * p = get_page_of(o);
+    return p->m_header.m_obj_size;
+}
+
 void initialize_alloc() {
 #ifdef LEAN_SMALL_ALLOCATOR
     g_heap_manager = new heap_manager();
diff --git a/src/runtime/alloc.h b/src/runtime/alloc.h
index 75276dff44..216200ffb8 100644
--- a/src/runtime/alloc.h
+++ b/src/runtime/alloc.h
@@ -11,6 +11,8 @@ namespace lean {
 void init_thread_heap();
 void * alloc(size_t sz);
 void dealloc(void * o, size_t sz);
+void dealloc_small(void * o);
+unsigned get_small_object_size(void * o);
 void initialize_alloc();
 void finalize_alloc();
 }