From a945a9824ee495435bc2d965b81e9b02527a7b33 Mon Sep 17 00:00:00 2001
From: xinhui pan <xinhui.pan@amd.com>
Date: Wed, 19 Apr 2023 10:57:37 +0800
Subject: [PATCH] thunk: Fix and optimise for pointer range search

Previous code might fail to get the correct ln node. And trigger extra
walk through of the tree. Fix it.

While walking through the tree, better to search from right to left as
the node->start likely close to *address*.

Change-Id: If86ddf73e59a1eb88225d1ea90797818e8165488
Signed-off-by: xinhui pan <xinhui.pan@amd.com>


[ROCm/ROCR-Runtime commit: 77761836aece988c60ee945e7b9309374e04a687]
---
 projects/rocr-runtime/src/fmm.c    | 48 ++++++++++++++++--------------
 projects/rocr-runtime/src/rbtree.c | 28 +++++++++++++++++
 projects/rocr-runtime/src/rbtree.h |  2 ++
 3 files changed, 56 insertions(+), 22 deletions(-)

diff --git a/projects/rocr-runtime/src/fmm.c b/projects/rocr-runtime/src/fmm.c
index ff02fa9394..2f088d2e51 100644
--- a/projects/rocr-runtime/src/fmm.c
+++ b/projects/rocr-runtime/src/fmm.c
@@ -470,17 +470,34 @@ static vm_object_t *vm_find_object_by_address_userptr_range(manageable_aperture_
 	vm_object_t *cur = NULL;
 	rbtree_t *tree = vm_object_tree(app, is_userptr);
 	rbtree_key_t key = rbtree_key((unsigned long)address, 0);
-	rbtree_node_t *ln = rbtree_lookup_nearest(tree, &key,
-			LKP_ALL, LEFT);
-	rbtree_node_t *rn = rbtree_lookup_nearest(tree, &key,
-			LKP_ALL, RIGHT);
+	rbtree_node_t *rn = rbtree_lookup_nearest(tree, &key, LKP_ALL, RIGHT);
+	rbtree_node_t *ln;
 	void *start;
 	uint64_t size;
-	int bad = 0;
 
-loop:
-	while (ln) {
-		cur = vm_object_entry(ln, is_userptr);
+	/* all nodes might sit on left side of *address*, in this case rn is NULL.
+	 * So pick up the rightest one as rn.
+	 */
+	if (!rn)
+		rn = rbtree_min_max(tree, RIGHT);
+
+	if (is_userptr) {
+		/* userptr might overlap. Need walk through the tree from right to left as only left nodes
+		 * can obtain the *address*
+		 */
+		ln = rbtree_min_max(tree, LEFT);
+	} else {
+		/* if key->size is -1, it match the node with start <= address.
+		 * if key->size is 0, it match the node with start < address.
+		 */
+		key = rbtree_key((unsigned long)address, -1);
+		ln = rbtree_lookup_nearest(tree, &key, LKP_ALL, LEFT);
+	}
+	if (!ln)
+		return NULL;
+
+	while (rn) {
+		cur = vm_object_entry(rn, is_userptr);
 		if (is_userptr == 0) {
 			start = cur->start;
 			size = cur->size;
@@ -498,20 +515,7 @@ loop:
 		if (ln == rn)
 			break;
 
-		ln = rbtree_next(tree, ln);
-	}
-
-	if (cur == NULL && bad == 0) {
-		/* As there is area overlap, say, (address, size) like
-		 * (0x100, 32), (0x108, 8), and the key.address is 0x118
-		 * The lookup above only find (0x108, 8), but the correct node should
-		 * be (0x100, 16). So try to walk though the tree to find the node.
-		 */
-		rn = ln;
-		key = rbtree_key(0, 0);
-		ln = rbtree_lookup_nearest(tree, &key, LKP_ALL, RIGHT);
-		bad = 1;
-		goto loop;
+		rn = rbtree_prev(tree, rn);
 	}
 
 	return cur; /* NULL if not found */
diff --git a/projects/rocr-runtime/src/rbtree.c b/projects/rocr-runtime/src/rbtree.c
index 86e2d59f7b..21d0d70d5a 100644
--- a/projects/rocr-runtime/src/rbtree.c
+++ b/projects/rocr-runtime/src/rbtree.c
@@ -372,3 +372,31 @@ rbtree_next(rbtree_t *tree, rbtree_node_t *node)
 		node = parent;
 	}
 }
+
+rbtree_node_t *
+rbtree_prev(rbtree_t *tree, rbtree_node_t *node)
+{
+	rbtree_node_t  *root, *sentinel, *parent;
+
+	sentinel = &tree->sentinel;
+
+	if (node->left != sentinel) {
+		return rbtree_max(node->left, sentinel);
+	}
+
+	root = tree->root;
+
+	for ( ;; ) {
+		parent = node->parent;
+
+		if (node == root) {
+			return NULL;
+		}
+
+		if (node == parent->right) {
+			return parent;
+		}
+
+		node = parent;
+	}
+}
diff --git a/projects/rocr-runtime/src/rbtree.h b/projects/rocr-runtime/src/rbtree.h
index 965c8fe773..fa13959369 100644
--- a/projects/rocr-runtime/src/rbtree.h
+++ b/projects/rocr-runtime/src/rbtree.h
@@ -64,6 +64,8 @@ struct rbtree_s {
 
 void rbtree_insert(rbtree_t *tree, rbtree_node_t *node);
 void rbtree_delete(rbtree_t *tree, rbtree_node_t *node);
+rbtree_node_t *rbtree_prev(rbtree_t *tree,
+		rbtree_node_t *node);
 rbtree_node_t *rbtree_next(rbtree_t *tree,
 		rbtree_node_t *node);