1 files changed, 27 insertions, 14 deletions

diff --git a/drivers/base/regmap/regcache-rbtree.c b/drivers/base/regmap/regcache-rbtree.c
index aa0875f6f1b7..3a000145ac6e 100644
--- a/drivers/base/regmap/regcache-rbtree.c
+++ b/drivers/base/regmap/regcache-rbtree.c
@@ -304,6 +304,27 @@ static int regcache_rbtree_insert_to_block(struct regmap *map,
         return 0;
 }
 
+static struct regcache_rbtree_node *
+regcache_rbtree_node_alloc(struct regmap *map, unsigned int reg)
+{
+        struct regcache_rbtree_node *rbnode;
+
+        rbnode = kzalloc(sizeof(*rbnode), GFP_KERNEL);
+        if (!rbnode)
+                return NULL;
+
+        rbnode->blklen = sizeof(*rbnode);
+        rbnode->base_reg = reg;
+        rbnode->block = kmalloc(rbnode->blklen * map->cache_word_size,
+                                GFP_KERNEL);
+        if (!rbnode->block) {
+                kfree(rbnode);
+                return NULL;
+        }
+
+        return rbnode;
+}
+
 static int regcache_rbtree_write(struct regmap *map, unsigned int reg,
                                  unsigned int value)
 {
@@ -354,23 +375,15 @@ static int regcache_rbtree_write(struct regmap *map, unsigned int reg,
                                 return 0;
                         }
                 }
-                /* we did not manage to find a place to insert it in an existing
-                 * block so create a new rbnode with a single register in its block.
-                 * This block will get populated further if any other adjacent
-                 * registers get modified in the future.
+
+                /* We did not manage to find a place to insert it in
+                 * an existing block so create a new rbnode.
                  */
-                rbnode = kzalloc(sizeof *rbnode, GFP_KERNEL);
+                rbnode = regcache_rbtree_node_alloc(map, reg);
                 if (!rbnode)
                         return -ENOMEM;
-                rbnode->blklen = sizeof(*rbnode);
-                rbnode->base_reg = reg;
-                rbnode->block = kmalloc(rbnode->blklen * map->cache_word_size,
-                                        GFP_KERNEL);
-                if (!rbnode->block) {
-                        kfree(rbnode);
-                        return -ENOMEM;
-                }
-                regcache_rbtree_set_register(map, rbnode, 0, value);
+                regcache_rbtree_set_register(map, rbnode,
+                                             reg - rbnode->base_reg, value);
                 regcache_rbtree_insert(map, &rbtree_ctx->root, rbnode);
                 rbtree_ctx->cached_rbnode = rbnode;
         }