[FFmpeg-devel] [PATCH v2 1/6] avutil/tree: av_tree_find2() to also find the first and last elements comparing equal

[Michael Niedermayer]

This patch also improves the worst case of O(n) (with n equal elements) to O(log n)

This may be used by AVDictionary2 And AVSet for iterating over identical
keys

Signed-off-by: Michael Niedermayer <michael at niedermayer.cc>
---
 libavutil/tree.c | 34 +++++++++++++++++++++++++++++-----
 libavutil/tree.h |  6 ++++++
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/libavutil/tree.c b/libavutil/tree.c
index 7b57b2d39a7..9f58a3f3d49 100644
--- a/libavutil/tree.c
+++ b/libavutil/tree.c
@@ -18,6 +18,7 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */
 
+#include "avassert.h"
 #include "error.h"
 #include "log.h"
 #include "mem.h"
@@ -36,19 +37,36 @@ struct AVTreeNode *av_tree_node_alloc(void)
     return av_mallocz(sizeof(struct AVTreeNode));
 }
 
-void *av_tree_find(const AVTreeNode *t, void *key,
-                   int (*cmp)(const void *key, const void *b), void *next[2])
+static void tree_find_next(const AVTreeNode *t, const void *key,
+                   int (*cmp)(const void *key, const void *b), void *next[4], int nextlen, int direction)
+{
+    if (t) {
+        unsigned int v = cmp(key, t->elem);
+        if (v) {
+            next[direction] = t->elem;
+            av_assert2((v >> 31) == direction);
+            tree_find_next(t->child[!direction], key, cmp, next, nextlen, direction);
+        } else {
+            if (nextlen >= 4)
+                next[2+direction] = t->elem;
+            tree_find_next(t->child[direction], key, cmp, next, nextlen, direction);
+        }
+    }
+}
+
+void *av_tree_find2(const AVTreeNode *t, const void *key,
+                    int (*cmp)(const void *key, const void *b), void *next[4], int nextlen)
 {
     if (t) {
         unsigned int v = cmp(key, t->elem);
         if (v) {
             if (next)
                 next[v >> 31] = t->elem;
-            return av_tree_find(t->child[(v >> 31) ^ 1], key, cmp, next);
+            return av_tree_find2(t->child[(v >> 31) ^ 1], key, cmp, next, nextlen);
         } else {
             if (next) {
-                av_tree_find(t->child[0], key, cmp, next);
-                av_tree_find(t->child[1], key, cmp, next);
+                tree_find_next(t->child[0], key, cmp, next, nextlen, 0);
+                tree_find_next(t->child[1], key, cmp, next, nextlen, 1);
             }
             return t->elem;
         }
@@ -56,6 +74,12 @@ void *av_tree_find(const AVTreeNode *t, void *key,
     return NULL;
 }
 
+void *av_tree_find(const AVTreeNode *t, void *key,
+                   int (*cmp)(const void *key, const void *b), void *next[2])
+{
+    return av_tree_find2(t, key, cmp, next, 2);
+}
+
 void *av_tree_insert(AVTreeNode **tp, void *key,
                      int (*cmp)(const void *key, const void *b), AVTreeNode **next)
 {
diff --git a/libavutil/tree.h b/libavutil/tree.h
index bbb8fbb1262..6bafd47e593 100644
--- a/libavutil/tree.h
+++ b/libavutil/tree.h
@@ -55,6 +55,9 @@ struct AVTreeNode *av_tree_node_alloc(void);
  * @param next If next is not NULL, then next[0] will contain the previous
  *             element and next[1] the next element. If either does not exist,
  *             then the corresponding entry in next is unchanged.
+ *             if nextlen is 4 then next[2] will contain the first element comparing
+ *             equal that is smaller than the returned. Similarly next[3] will
+ *             contain the last element comparing equal that is larger than the returned
  * @param cmp compare function used to compare elements in the tree,
  *            API identical to that of Standard C's qsort
  *            It is guaranteed that the first and only the first argument to cmp()
@@ -63,6 +66,9 @@ struct AVTreeNode *av_tree_node_alloc(void);
  * @return An element with cmp(key, elem) == 0 or NULL if no such element
  *         exists in the tree.
  */
+void *av_tree_find2(const struct AVTreeNode *root, const void *key,
+                   int (*cmp)(const void *key, const void *b), void *next[4], int nextlen);
+
 void *av_tree_find(const struct AVTreeNode *root, void *key,
                    int (*cmp)(const void *key, const void *b), void *next[2]);
 
-- 
2.49.0