diff --git a/drivers/base/regmap/regcache-rbtree.c b/drivers/base/regmap/regcache-rbtree.c
index 56486d9..353f602 100644
--- a/drivers/base/regmap/regcache-rbtree.c
+++ b/drivers/base/regmap/regcache-rbtree.c
@@ -413,8 +413,8 @@
 		max = reg + max_dist;
 
 		/* look for an adjacent register to the one we are about to add */
-		for (node = rb_first(&rbtree_ctx->root); node;
-		     node = rb_next(node)) {
+		node = rbtree_ctx->root.rb_node;
+		while (node) {
 			rbnode_tmp = rb_entry(node, struct regcache_rbtree_node,
 					      node);
 
@@ -425,6 +425,11 @@
 				new_base_reg = min(reg, base_reg);
 				new_top_reg = max(reg, top_reg);
 			} else {
+				if (max < base_reg)
+					node = node->rb_left;
+				else
+					node = node->rb_right;
+
 				continue;
 			}