5 files changed, 56 insertions, 9 deletions
diff --git a/source/blender/blenlib/BLI_dlrbTree.h b/source/blender/blenlib/BLI_dlrbTree.h
index 3cf849efaef..ad6e350ba2c 100644
--- a/source/blender/blenlib/BLI_dlrbTree.h
+++ b/source/blender/blenlib/BLI_dlrbTree.h
@@ -99,6 +99,12 @@ typedef DLRBT_Node *(*DLRBT_NAlloc_FP)(void *data);
  */
 typedef void (*DLRBT_NUpdate_FP)(void *node, void *data);
 
+/**
+ * Free a node and the wrapped data.
+ * \param node: <DLRBT_Node> the node to free.
+ */
+typedef void (*DLRBT_NFree_FP)(void *node);
+
 /* ********************************************** */
 /* Public API */
 
@@ -122,7 +128,7 @@ void BLI_dlrbTree_init(DLRBT_Tree *tree);
 /**
  * Free the given tree's data but not the tree itself.
  */
-void BLI_dlrbTree_free(DLRBT_Tree *tree);
+void BLI_dlrbTree_free(DLRBT_Tree *tree, DLRBT_NFree_FP free_cb);
 
 /**
  * Make sure the tree's Double-Linked list representation is valid.
diff --git a/source/blender/blenlib/BLI_listbase.h b/source/blender/blenlib/BLI_listbase.h
index f73d1f22502..f970563d2a6 100644
--- a/source/blender/blenlib/BLI_listbase.h
+++ b/source/blender/blenlib/BLI_listbase.h
@@ -58,6 +58,12 @@ ListBase BLI_listbase_from_link(struct Link *some_link);
  */
 void *BLI_findlink(const struct ListBase *listbase, int number) ATTR_WARN_UNUSED_RESULT
     ATTR_NONNULL(1);
+
+/**
+ * Returns the nth element after \a link, numbering from 0.
+ */
+void *BLI_findlinkfrom(struct Link *start, int number) ATTR_WARN_UNUSED_RESULT;
+
 /**
  * Finds the first element of \a listbase which contains the null-terminated
  * string \a id at the specified offset, returning NULL if not found.
diff --git a/source/blender/blenlib/intern/DLRB_tree.c b/source/blender/blenlib/intern/DLRB_tree.c
index 66c394f5eb2..9c22afeb893 100644
--- a/source/blender/blenlib/intern/DLRB_tree.c
+++ b/source/blender/blenlib/intern/DLRB_tree.c
@@ -45,7 +45,7 @@ void BLI_dlrbTree_init(DLRBT_Tree *tree)
 }
 
 /* Helper for traversing tree and freeing sub-nodes */
-static void recursive_tree_free_nodes(DLRBT_Node *node)
+static void recursive_tree_free_nodes(DLRBT_Node *node, DLRBT_NFree_FP free_cb)
 {
   /* sanity check */
   if (node == NULL) {
@@ -53,14 +53,16 @@ static void recursive_tree_free_nodes(DLRBT_Node *node)
   }
 
   /* free child nodes + subtrees */
-  recursive_tree_free_nodes(node->left);
-  recursive_tree_free_nodes(node->right);
+  recursive_tree_free_nodes(node->left, free_cb);
+  recursive_tree_free_nodes(node->right, free_cb);
 
   /* free self */
-  MEM_freeN(node);
+  if (free_cb) {
+    free_cb(node);
+  }
 }
 
-void BLI_dlrbTree_free(DLRBT_Tree *tree)
+void BLI_dlrbTree_free(DLRBT_Tree *tree, DLRBT_NFree_FP free_cb)
 {
   if (tree == NULL) {
     return;
@@ -71,11 +73,19 @@ void BLI_dlrbTree_free(DLRBT_Tree *tree)
    */
   if (tree->first) {
     /* free list */
-    BLI_freelistN((ListBase *)tree);
+    if (free_cb) {
+      LISTBASE_FOREACH_MUTABLE(DLRBT_Node *, node, tree) {
+        free_cb(node);
+      }
+      BLI_listbase_clear((ListBase *)tree);
+    }
+    else {
+      BLI_freelistN((ListBase *)tree);
+    }
   }
   else {
     /* traverse tree, freeing sub-nodes */
-    recursive_tree_free_nodes(tree->root);
+    recursive_tree_free_nodes(tree->root, free_cb);
   }
 
   /* clear pointers */
@@ -584,8 +594,10 @@ DLRBT_Node *BLI_dlrbTree_add(DLRBT_Tree *tree,
       }
       default: /* update the duplicate node as appropriate */
       {
+        /* Return the updated node after calling the callback.  */
+        node = parNode;
         if (update_cb) {
-          update_cb(parNode, data);
+          update_cb(node, data);
         }
         break;
       }
diff --git a/source/blender/blenlib/intern/listbase.c b/source/blender/blenlib/intern/listbase.c
index 513b08a589d..c21b448e505 100644
--- a/source/blender/blenlib/intern/listbase.c
+++ b/source/blender/blenlib/intern/listbase.c
@@ -537,6 +537,21 @@ void *BLI_rfindlink(const ListBase *listbase, int number)
   return link;
 }
 
+void *BLI_findlinkfrom(Link *start, int number)
+{
+  Link *link = NULL;
+
+  if (number >= 0) {
+    link = start;
+    while (link != NULL && number != 0) {
+      number--;
+      link = link->next;
+    }
+  }
+
+  return link;
+}
+
 int BLI_findindex(const ListBase *listbase, const void *vlink)
 {
   Link *link = NULL;
diff --git a/source/blender/blenlib/tests/BLI_listbase_test.cc b/source/blender/blenlib/tests/BLI_listbase_test.cc
index 9e4d7c7dd36..88f597cdf67 100644
--- a/source/blender/blenlib/tests/BLI_listbase_test.cc
+++ b/source/blender/blenlib/tests/BLI_listbase_test.cc
@@ -80,18 +80,26 @@ TEST(listbase, FindLinkOrIndex)
   EXPECT_EQ(BLI_rfindlink(&lb, 0), (void *)nullptr);
   EXPECT_EQ(BLI_rfindlink(&lb, 1), (void *)nullptr);
   EXPECT_EQ(BLI_findindex(&lb, link1), -1);
+  EXPECT_EQ(BLI_findlinkfrom((Link *)lb.first, -1), (void *)nullptr);
+  EXPECT_EQ(BLI_findlinkfrom((Link *)lb.first, 0), (void *)nullptr);
+  EXPECT_EQ(BLI_findlinkfrom((Link *)lb.first, 1), (void *)nullptr);
 
   /* One link */
   BLI_addtail(&lb, link1);
   EXPECT_EQ(BLI_findlink(&lb, 0), link1);
   EXPECT_EQ(BLI_rfindlink(&lb, 0), link1);
   EXPECT_EQ(BLI_findindex(&lb, link1), 0);
+  EXPECT_EQ(BLI_findlinkfrom((Link *)lb.first, 0), link1);
 
   /* Two links */
   BLI_addtail(&lb, link2);
   EXPECT_EQ(BLI_findlink(&lb, 1), link2);
   EXPECT_EQ(BLI_rfindlink(&lb, 0), link2);
   EXPECT_EQ(BLI_findindex(&lb, link2), 1);
+  EXPECT_EQ(BLI_findlinkfrom((Link *)lb.first, 1), link2);
+
+  /* After end of list */
+  EXPECT_EQ(BLI_findlinkfrom((Link *)lb.first, 2), (void *)nullptr);
 
   BLI_freelistN(&lb);
 }