Added more rbtrees in dynarec managment, to speedup FindDynablockFromNativeAddress function

3 files changed, 118 insertions, 62 deletions

diff --git a/src/custommem.c b/src/custommem.c
index d9804f34..ec45500a 100644
--- a/src/custommem.c
+++ b/src/custommem.c
@@ -35,6 +35,7 @@
 
 // init inside dynablocks.c
 static mmaplist_t          *mmaplist = NULL;
+static rbtree_t            *rbt_dynmem = NULL;
 static uint64_t jmptbl_allocated = 0, jmptbl_allocated1 = 0, jmptbl_allocated2 = 0, jmptbl_allocated3 = 0;
 #ifdef JMPTABL_SHIFT4
 static uint64_t jmptbl_allocated4 = 0;
@@ -794,8 +795,12 @@ size_t customGetUsableSize(void* p)
 
 #ifdef DYNAREC
 #define NCHUNK          64
+typedef struct mapchunk_s {
+    blocklist_t         chunk;
+    rbtree_t*           tree;
+} mapchunk_t;
 typedef struct mmaplist_s {
-    blocklist_t         chunks[NCHUNK];
+    mapchunk_t          chunks[NCHUNK];
     mmaplist_t*         next;
 } mmaplist_t;
 
@@ -806,31 +811,9 @@ dynablock_t* FindDynablockFromNativeAddress(void* p)
     
     uintptr_t addr = (uintptr_t)p;
 
-    int i= 0;
-    mmaplist_t* list = mmaplist;
-    if(!list)
-        return NULL;
-    while(list) {
-        if ((addr>(uintptr_t)list->chunks[i].block) 
-         && (addr<((uintptr_t)list->chunks[i].block+list->chunks[i].size))) {
-            blockmark_t* sub = (blockmark_t*)list->chunks[i].block;
-            while((uintptr_t)sub<addr) {
-                blockmark_t* n = NEXT_BLOCK(sub);
-                if((uintptr_t)n>addr) {
-                    // found it!
-                    // self is the field of a block
-                    return *(dynablock_t**)((uintptr_t)sub+sizeof(blockmark_t));
-                }
-                sub = n;
-            }
-            return NULL;
-        }
-        ++i;
-        if(i==NCHUNK) {
-            i = 0;
-            list = list->next;
-        }
-    }
+    mapchunk_t* bl = (mapchunk_t*)rb_get_64(rbt_dynmem, (uintptr_t)p);
+    if(bl)
+        return *(dynablock_t**)rb_get_64(bl->tree, (uintptr_t)p);
     return NULL;
 }
 
@@ -877,19 +860,20 @@ uintptr_t AllocDynarecMap(size_t size)
     int i = 0;
     uintptr_t sz = size + 2*sizeof(blockmark_t);
     while(1) {
-        if(list->chunks[i].maxfree>=size) {
+        if(list->chunks[i].chunk.maxfree>=size) {
             // looks free, try to alloc!
             size_t rsize = 0;
-            void* sub = getFirstBlock(list->chunks[i].block, size, &rsize, list->chunks[i].first);
+            void* sub = getFirstBlock(list->chunks[i].chunk.block, size, &rsize, list->chunks[i].chunk.first);
             if(sub) {
-                void* ret = allocBlock(list->chunks[i].block, sub, size, &list->chunks[i].first);
-                if(rsize==list->chunks[i].maxfree)
-                    list->chunks[i].maxfree = getMaxFreeBlock(list->chunks[i].block, list->chunks[i].size, list->chunks[i].first);
+                void* ret = allocBlock(list->chunks[i].chunk.block, sub, size, &list->chunks[i].chunk.first);
+                if(rsize==list->chunks[i].chunk.maxfree)
+                    list->chunks[i].chunk.maxfree = getMaxFreeBlock(list->chunks[i].chunk.block, list->chunks[i].chunk.size, list->chunks[i].chunk.first);
+                rb_set_64(list->chunks[i].tree, (uintptr_t)ret, (uintptr_t)ret+size, (uintptr_t)ret);
                 return (uintptr_t)ret;
             }
         }
         // check if new
-        if(!list->chunks[i].size) {
+        if(!list->chunks[i].chunk.size) {
             // alloc a new block, aversized or not, we are at the end of the list
             size_t allocsize = (sz>DYNMMAPSZ)?sz:DYNMMAPSZ;
             // allign sz with pagesize
@@ -923,9 +907,11 @@ uintptr_t AllocDynarecMap(size_t size)
 #endif
             setProtection((uintptr_t)p, allocsize, PROT_READ | PROT_WRITE | PROT_EXEC);
 
-            list->chunks[i].block = p;
-            list->chunks[i].first = p;
-            list->chunks[i].size = allocsize;
+            list->chunks[i].chunk.block = p;
+            list->chunks[i].chunk.first = p;
+            list->chunks[i].chunk.size = allocsize;
+            list->chunks[i].tree = rbtree_init("dynamap");
+            rb_set_64(rbt_dynmem, (uintptr_t)p, (uintptr_t)p+allocsize, (uintptr_t)&list->chunks[i]);
             // setup marks
             blockmark_t* m = (blockmark_t*)p;
             m->prev.x32 = 0;
@@ -935,10 +921,11 @@ uintptr_t AllocDynarecMap(size_t size)
             n->next.x32 = 0;
             n->prev.x32 = m->next.x32;
             // alloc 1st block
-            void* ret  = allocBlock(list->chunks[i].block, p, size, &list->chunks[i].first);
-            list->chunks[i].maxfree = getMaxFreeBlock(list->chunks[i].block, list->chunks[i].size, list->chunks[i].first);
-            if(list->chunks[i].maxfree)
-                list->chunks[i].first = getNextFreeBlock(m);
+            void* ret  = allocBlock(list->chunks[i].chunk.block, p, size, &list->chunks[i].chunk.first);
+            list->chunks[i].chunk.maxfree = getMaxFreeBlock(list->chunks[i].chunk.block, list->chunks[i].chunk.size, list->chunks[i].chunk.first);
+            if(list->chunks[i].chunk.maxfree)
+                list->chunks[i].chunk.first = getNextFreeBlock(m);
+            rb_set_64(list->chunks[i].tree, (uintptr_t)ret, (uintptr_t)ret+size, (uintptr_t)ret);
             return (uintptr_t)ret;
         }
         // next chunk...
@@ -957,23 +944,16 @@ void FreeDynarecMap(uintptr_t addr)
     if(!addr)
         return;
     
-    int i= 0;
-    mmaplist_t* list = mmaplist;
 
-    while(list) {
-        if ((addr>(uintptr_t)list->chunks[i].block) 
-         && (addr<((uintptr_t)list->chunks[i].block+list->chunks[i].size))) {
-            void* sub = (void*)(addr-sizeof(blockmark_t));
-            size_t newfree = freeBlock(list->chunks[i].block, list->chunks[i].size, sub, &list->chunks[i].first);
-            if(list->chunks[i].maxfree < newfree)
-                list->chunks[i].maxfree = newfree;
-            return;
-        }
-        ++i;
-        if(i==NCHUNK) {
-            i = 0;
-            list = list->next;
-        }
+    mapchunk_t* bl = (mapchunk_t*)rb_get_64(rbt_dynmem, addr);
+
+    if(bl) {
+        void* sub = (void*)(addr-sizeof(blockmark_t));
+        size_t newfree = freeBlock(bl->chunk.block, bl->chunk.size, sub, &bl->chunk.first);
+        if(bl->chunk.maxfree < newfree)
+            bl->chunk.maxfree = newfree;
+        rb_unset(bl->tree, addr, addr+newfree);
+        return;
     }
 }
 
@@ -2029,6 +2009,7 @@ void init_custommem_helper(box64context_t* ctx)
             box64_jmptbldefault0[i] = (uintptr_t)native_next;
     }
     lockaddress = kh_init(lockaddress);
+    rbt_dynmem = rbtree_init("rbt_dynmem");
 #endif
     pthread_atfork(NULL, NULL, atfork_child_custommem);
     // init mapallmem list
@@ -2086,8 +2067,10 @@ void fini_custommem_helper(box64context_t *ctx)
         mmaplist = NULL;
         while(head) {
             for (int i=0; i<NCHUNK; ++i) {
-                if(head->chunks[i].block)
-                    internal_munmap(head->chunks[i].block, head->chunks[i].size);
+                if(head->chunks[i].chunk.block)
+                    internal_munmap(head->chunks[i].chunk.block, head->chunks[i].chunk.size);
+                if(head->chunks[i].tree)
+                    rbtree_delete(head->chunks[i].tree);
             }
             mmaplist_t *old = head;
             head = head->next;
@@ -2120,6 +2103,8 @@ void fini_custommem_helper(box64context_t *ctx)
     }
     kh_destroy(lockaddress, lockaddress);
     lockaddress = NULL;
+    rbtree_delete(rbt_dynmem);
+    rbt_dynmem = NULL;
 #endif
     rbtree_delete(memprot);
     memprot = NULL;
diff --git a/src/include/rbtree.h b/src/include/rbtree.h
index 00a7a87a..83eddb78 100644
--- a/src/include/rbtree.h
+++ b/src/include/rbtree.h
@@ -58,6 +58,16 @@ void rbtree_delete(rbtree_t* tree);
 uint32_t rb_get(rbtree_t* tree, uintptr_t addr);
 
 /** 
+ * rb_get_64() - Retrieves data associated with a specific address in the red-black tree.
+ * @tree: Pointer to the red-black tree from which data is to be retrieved.
+ * @addr: The memory address used as a key to find the corresponding node in the tree.
+ *
+ * This function searches the red-black tree for a node that corresponds to the specified address.
+ * Return: The 64bits data associated with the address if found; otherwise, 0.
+ */
+uint64_t rb_get_64(rbtree_t* tree, uintptr_t addr);
+
+/** 
  * rb_get_end() - Searches for a node within a specified address range in a red-black tree and retrieves its data and end address.
  * @tree: Pointer to the red-black tree to be searched.
  * @addr: The address to search for within the nodes of the red-black tree.
@@ -71,6 +81,20 @@ uint32_t rb_get(rbtree_t* tree, uintptr_t addr);
  */
 int rb_get_end(rbtree_t* tree, uintptr_t addr, uint32_t* val, uintptr_t* end);
 
+/** 
+ * rb_get_end_64() - Searches for a node within a specified address range in a red-black tree and retrieves its data and end address.
+ * @tree: Pointer to the red-black tree to be searched.
+ * @addr: The address to search for within the nodes of the red-black tree.
+ * @val: Pointer to store the data of the node that contains the address if found.
+ * @end: Pointer to store the end address of the node that contains the address, or the start of the next node if not found, or UINTPTR_MAX if no next node exists.
+ *
+ * This function traverses the red-black tree starting from the root, searching for a node where the 'addr' falls between the node's start and end addresses (exclusive of end).
+ * If such a node is found, the function stores the node's data in '*val' and the node's end address in '*end', then returns 1 to indicate success.
+ * If no such node is found, the function sets '*val' to 0 and '*end' to the start address of the next node in the tree structure or to UINTPTR_MAX if there is no subsequent node.
+ * Return: 1 if a node containing the address is found, otherwise 0.
+ */
+int rb_get_end_64(rbtree_t* tree, uintptr_t addr, uint64_t* val, uintptr_t* end);
+
 /**
  * rb_set() - Set an address range in a red-black tree.
  * @tree: Pointer to the red-black tree where the address range will be set.
@@ -188,6 +212,19 @@ int rb_get_end(rbtree_t* tree, uintptr_t addr, uint32_t* val, uintptr_t* end);
  */
 int rb_set(rbtree_t* tree, uintptr_t start, uintptr_t end, uint32_t data);
 
+
+/**
+ * rb_set_64() - Set an address range in a red-black tree.
+ * @tree: Pointer to the red-black tree where the address range will be set.
+ * @start: The starting address of the range to be set.
+ * @end: The ending address of the range to be set.
+ * @data: The data value to associate with the address range.
+ *
+ * This function adds a new address range with associated data to the red-black tree. 
+ */
+
+int rb_set_64(rbtree_t* tree, uintptr_t start, uintptr_t end, uint64_t data);
+
 /**
  * rb_unset() - Removes a range of values from the red-black tree.
  * @tree: Pointer to the red-black tree.
diff --git a/src/tools/rbtree.c b/src/tools/rbtree.c
index 0d5ea61f..daa01f35 100644
--- a/src/tools/rbtree.c
+++ b/src/tools/rbtree.c
@@ -24,7 +24,7 @@ static void rbtree_print(const rbtree_t* tree);
 typedef struct rbnode {
     struct rbnode *left, *right, *parent;
     uintptr_t start, end;
-    uint32_t data;
+    uint64_t data;
     uint8_t meta;
 } rbnode;
 
@@ -58,7 +58,7 @@ void rbtree_delete(rbtree_t *tree) {
 #define IS_BLACK 0x2
 
 // Make sure prev is either the rightmost node before start or the leftmost range after start
-static int add_range_next_to(rbtree_t *tree, rbnode *prev, uintptr_t start, uintptr_t end, uint32_t data) {
+static int add_range_next_to(rbtree_t *tree, rbnode *prev, uintptr_t start, uintptr_t end, uint64_t data) {
 // printf("Adding %lx-%lx:%hhx next to %p\n", start, end, data, prev);
     rbnode *node = rbtreeMalloc(sizeof(*node));
     if (!node) return -1;
@@ -233,7 +233,7 @@ static int add_range_next_to(rbtree_t *tree, rbnode *prev, uintptr_t start, uint
     return -1; // unreachable
 }
 
-static int add_range(rbtree_t *tree, uintptr_t start, uintptr_t end, uint32_t data) {
+static int add_range(rbtree_t *tree, uintptr_t start, uintptr_t end, uint64_t data) {
 // printf("add_range\n");
     rbnode *cur = tree->root, *prev = NULL;
     while (cur) {
@@ -565,6 +565,12 @@ uint32_t rb_get(rbtree_t *tree, uintptr_t addr) {
     else return 0;
 }
 
+uint64_t rb_get_64(rbtree_t *tree, uintptr_t addr) {
+    rbnode *node = find_addr(tree, addr);
+    if (node) return node->data;
+    else return 0;
+}
+
 int rb_get_end(rbtree_t* tree, uintptr_t addr, uint32_t* val, uintptr_t* end) {
     rbnode *node = tree->root, *next = NULL;
     while (node) {
@@ -589,7 +595,31 @@ int rb_get_end(rbtree_t* tree, uintptr_t addr, uint32_t* val, uintptr_t* end) {
     return 0;
 }
 
-int rb_set(rbtree_t *tree, uintptr_t start, uintptr_t end, uint32_t data) {
+int rb_get_end_64(rbtree_t* tree, uintptr_t addr, uint64_t* val, uintptr_t* end) {
+    rbnode *node = tree->root, *next = NULL;
+    while (node) {
+        if ((node->start <= addr) && (node->end > addr)) {
+            *val = node->data;
+            *end = node->end;
+            return 1;
+        }
+        if (node->end <= addr) {
+            node = node->right;
+        } else {
+            next = node;
+            node = node->left;
+        }
+    }
+    *val = 0;
+    if (next) {
+        *end = next->start;
+    } else {
+        *end = (uintptr_t)-1;
+    }
+    return 0;
+}
+
+int rb_set_64(rbtree_t *tree, uintptr_t start, uintptr_t end, uint64_t data) {
 // printf("rb_set( "); rbtree_print(tree); printf(" , 0x%lx, 0x%lx, %hhu);\n", start, end, data); fflush(stdout);
 dynarec_log(LOG_DEBUG, "set %s: 0x%lx, 0x%lx, 0x%x\n", tree->name, start, end, data);
     if (!tree->root) {
@@ -722,6 +752,10 @@ dynarec_log(LOG_DEBUG, "set %s: 0x%lx, 0x%lx, 0x%x\n", tree->name, start, end, d
     return add_range_next_to(tree, last->left ? pred_node(last) : last, start, end, data);
 }
 
+int rb_set(rbtree_t *tree, uintptr_t start, uintptr_t end, uint32_t data) {
+    return rb_set_64(tree, start, end, data);
+}
+
 int rb_unset(rbtree_t *tree, uintptr_t start, uintptr_t end) {
 // printf("rb_unset( "); rbtree_print(tree); printf(" , 0x%lx, 0x%lx);\n", start, end); fflush(stdout);
 dynarec_log(LOG_DEBUG, "unset: %s 0x%lx, 0x%lx);\n", tree->name, start, end);
@@ -825,7 +859,7 @@ static void print_rbnode(const rbnode *node, unsigned depth, uintptr_t minstart,
     } else {
         print_rbnode(node->left, depth + ((node->meta & IS_BLACK) ? 1 : 0), minstart, node->start, bdepth);
     }
-    printf(", (%c/%p) %lx-%lx: %hhu, ", node->meta & IS_BLACK ? 'B' : 'R', node, node->start, node->end, node->data);
+    printf(", (%c/%p) %lx-%lx: %llu, ", node->meta & IS_BLACK ? 'B' : 'R', node, node->start, node->end, node->data);
     if (node->right && (node->right->meta & IS_LEFT)) {
         printf("<invalid meta>");
     } else if (node->right && (node->right->parent != node)) {