1 files changed, 20 insertions, 8 deletions

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;
       }