From 2f0a3a3c8e897468e77bd409fd8918ed85d070be Mon Sep 17 00:00:00 2001
From: Gregory Lirent <opensource@lirent.ru>
Date: Wed, 8 Jun 2022 09:58:16 +0300
Subject: [PATCH] Update set (extra) foreach implementation

---
 include/extra/set.h | 11 +++++++----
 src/set/extra.c     | 26 ++++++++++++++++++++++----
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/include/extra/set.h b/include/extra/set.h
index 7f37db5..8849e09 100644
--- a/include/extra/set.h
+++ b/include/extra/set.h
@@ -6,11 +6,14 @@
 #ifndef LIBCDSB_EXTRA_SET_H
 #define LIBCDSB_EXTRA_SET_H
 
-#define vset_foreach libcdsb_vset_foreach
+typedef int (*vset_foreach_callback)(const void* value, vtype type, void* data);
 
-extern _Bool libcdsb_vset_find  (vtype_value* x, vtype_set* s, const void* value, vtype type, _Bool cut);
-extern _Bool libcdsb_vset_insert(vtype_set*   x, const void* value, vtype type);
 
-extern int libcdsb_vset_foreach(const vtype_set* x, void* data, int (*callback)(const void* value, vtype type, void* data)) LIBCDSB_nt__ LIBCDSB_nn13__;
+#define vset_foreach(x, data, callback) libcdsb_vset_foreach(x, data, callback, 0)
+
+extern _Bool libcdsb_vset_find  (vtype_value* x, vtype_set* s, const void* value, vtype type, _Bool cut) LIBCDSB_nt__ LIBCDSB_nn23__;
+extern _Bool libcdsb_vset_insert(vtype_set*   x, const void* value, vtype type) LIBCDSB_nt__ LIBCDSB_nn12__;
+
+extern int libcdsb_vset_foreach(vtype_set* x, void* data, vset_foreach_callback, _Bool flush) LIBCDSB_nt__ LIBCDSB_nn13__;
 
 #endif /* LIBCDSB_EXTRA_SET_H */
diff --git a/src/set/extra.c b/src/set/extra.c
index da5bb36..6c841f6 100644
--- a/src/set/extra.c
+++ b/src/set/extra.c
@@ -70,7 +70,7 @@ _Bool libcdsb_vset_insert(set_t* x, const void* v, vtype t) {
 }
 
 
-int libcdsb_vset_foreach(const set_t* x, void* data, int (*callback)(const void* value, vtype type, void* data)) {
+int libcdsb_vset_foreach(set_t* x, void* data, vset_foreach_callback callback, _Bool flush) {
     stack_t   z = { .prev = 0, .value = x->root };
     int       r = 0;
     rbnode_t* c;
@@ -78,14 +78,32 @@ int libcdsb_vset_foreach(const set_t* x, void* data, int (*callback)(const void*
     if (rbnode_is_empty(x->root)) return 0;
 
     while ((c = stack_pop(&z))) {
-        if ((r = callback(vnode_peek(&c->value, x->type), x->type, data))) {
-            stack_flush(&z);
+        if ((r = callback(vnode_peek(&c->value, x->type), x->type, data)))
             break;
-        }
+
 
         if (!rbnode_is_empty(c->right)) stack_push(&z, c->right);
         if (!rbnode_is_empty(c->left))  stack_push(&z, c->left);
+
+        if (flush) {
+            vnode_free(&c->value, x->type);
+            free(c);
+        }
     }
 
+    if (flush) {
+        while (c) {
+            if (!rbnode_is_empty(c->right)) stack_push(&z, c->right);
+            if (!rbnode_is_empty(c->left))  stack_push(&z, c->left);
+
+            vnode_free(&c->value, x->type);
+            free(c);
+
+            c = stack_pop(&z);
+        }
+
+        memset(x, 0, sizeof(*x));
+    } else stack_flush(&z);
+
     return r;
 }