#8

3 files changed, 90 insertions, 10 deletions

diff --git a/tcejdb/bson.c b/tcejdb/bson.c
index f575a9e..39c7867 100644
--- a/tcejdb/bson.c
+++ b/tcejdb/bson.c
@@ -316,6 +316,46 @@ EJDB_EXPORT bson_type bson_find_from_buffer(bson_iterator *it, const char *buffe
     return bson_iterator_type(it);
 }
 
+static void bson_visit_fields_impl(traverse_flags_t flags, char* pstack, int curr, bson_iterator *it, BSONVISITOR visitor, void *op) {
+    int klen = 0;
+    bson_type t;
+    bool vret = true;
+    while (vret && (t = bson_iterator_next(it)) != BSON_EOO) {
+        const char* key = bson_iterator_key(it);
+        klen = strlen(key);
+        if (curr + klen > BSON_MAX_FPATH_LEN) {
+            continue;
+        }
+        //PUSH
+        if (curr > 0) { //add leading dot
+            memset(pstack + curr, '.', 1);
+            curr++;
+        }
+        memcpy(pstack + curr, key, klen);
+        curr += klen;
+        //Call visitor
+        bool vret = visitor(pstack, curr, key, klen, it, op);
+        if (vret) {
+            if ((t == BSON_OBJECT && (flags & BSON_TRAVERSE_OBJECTS_EXCLUDED) == 0) ||
+                    (t == BSON_ARRAY && (flags & BSON_TRAVERSE_ARRAYS_EXCLUDED) == 0)) {
+                bson_iterator sit;
+                bson_iterator_subiterator(it, &sit);
+                bson_visit_fields_impl(flags, pstack, curr, it, visitor, op);
+            }
+        }
+        //POP
+        curr -= klen;
+        if (curr > 0) {
+            curr--; //remove leading dot
+        }
+    }
+}
+
+EJDB_EXPORT void bson_visit_fields(bson_iterator *it, traverse_flags_t flags, BSONVISITOR visitor, void *op) {
+    char pstack[BSON_MAX_FPATH_LEN + 1];
+    bson_visit_fields_impl(flags, pstack, 0, it, visitor, op);
+}
+
 static bson_type bson_find_fieldpath_value_impl(char* pstack, int curr, const char *fpath, int fplen, bson_iterator *it) {
     int i;
     int klen = 0;
@@ -360,12 +400,12 @@ EJDB_EXPORT bson_type bson_find_fieldpath_value(const char *fpath, bson_iterator
 }
 
 EJDB_EXPORT bson_type bson_find_fieldpath_value2(const char *fpath, int fplen, bson_iterator *it) {
-    char pstackstack[BSON_MAX_FPATH_LEN];
+    char pstackstack[BSON_MAX_FPATH_LEN + 1];
     char *pstack;
-    if (fplen < BSON_MAX_FPATH_LEN) {
+    if (fplen <= BSON_MAX_FPATH_LEN) {
         pstack = pstackstack;
     } else {
-        pstack = MYMALLOC((fplen + 1) * sizeof(char));
+        pstack = MYMALLOC((fplen + 1) * sizeof (char));
         if (!pstack) {
             return BSON_EOO;
         }
@@ -762,7 +802,9 @@ EJDB_EXPORT int bson_finish(bson *b) {
 
 EJDB_EXPORT void bson_destroy(bson *b) {
     if (b) {
-        bson_free(b->data);
+        if (b->data) {
+            bson_free(b->data);
+        }
         b->err = 0;
         b->data = 0;
         b->cur = 0;
@@ -1301,7 +1343,7 @@ EJDB_EXPORT int bson_compare_fpaths(const void *bsdata1, const void *bsdata2, co
         int l2 = bson_iterator_bin_len(&it2);
         return memcmp(bson_iterator_bin_data(&it1), bson_iterator_bin_data(&it2), MIN(l1, l2));
     } else if (t1 == BSON_OID && t2 == BSON_OID) {
-        return memcmp(bson_iterator_oid(&it1), bson_iterator_oid(&it2), sizeof(bson_oid_t));
+        return memcmp(bson_iterator_oid(&it1), bson_iterator_oid(&it2), sizeof (bson_oid_t));
     }
     return 0;
 }
diff --git a/tcejdb/bson.h b/tcejdb/bson.h
index 8d39a21..334da76 100644
--- a/tcejdb/bson.h
+++ b/tcejdb/bson.h
@@ -191,6 +191,19 @@ EJDB_EXPORT bson_type bson_find_from_buffer(bson_iterator *it, const char *buffe
 EJDB_EXPORT bson_type bson_find_fieldpath_value(const char *fieldpath, bson_iterator *it);
 EJDB_EXPORT bson_type bson_find_fieldpath_value2(const char *fpath, int fplen, bson_iterator *it);
 
+/**
+ * BSON object visitor
+ * @param it bson iterator to traverse
+ * @param visitor Visitor function
+ * @param op Opaque data for visitor
+ */
+typedef enum  {
+    BSON_TRAVERSE_ARRAYS_EXCLUDED = 1,
+    BSON_TRAVERSE_OBJECTS_EXCLUDED = 1 << 1
+} traverse_flags_t;
+typedef bool (*BSONVISITOR)(const char *ipath, int ipathlen, const char *key, int keylen, const bson_iterator *curr, void *op);
+EJDB_EXPORT void bson_visit_fields(bson_iterator *it, traverse_flags_t flags, BSONVISITOR visitor, void *op);
+
 
 EJDB_EXPORT bson_iterator* bson_iterator_create(void);
 EJDB_EXPORT void bson_iterator_dispose(bson_iterator*);
diff --git a/tcejdb/ejdb.c b/tcejdb/ejdb.c
index 2151d71..fb17b5e 100644
--- a/tcejdb/ejdb.c
+++ b/tcejdb/ejdb.c
@@ -34,7 +34,7 @@
 /* Default size of stack allocated buffer for string conversions eg. tcicaseformat() */
 #define JBSTRINOPBUFFERSZ 512
 
-/* Default size of tmp bson buffer on stack for field stripping in _pushstrippedbson() */
+/* Default size of tmp bson buffer on stack for field stripping in _pushstripbson() */
 #define JBSBUFFERSZ 8192
 
 /* string processing/conversion flags */
@@ -96,7 +96,7 @@ static bool _qryallcondsmatch(bool onlycount, int anum, EJCOLL *jcoll, const EJQ
         const void *pkbuf, int pkbufsz, void **bsbuf, int *bsbufsz);
 static void _qrydup(const EJQ *src, EJQ *target, uint32_t qflags);
 static void _qrydel(EJQ *q, bool freequery);
-static void _pushstrippedbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbufsz);
+static void _pushstripbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbufsz);
 static TCLIST* _qrysearch(EJCOLL *jcoll, const EJQ *q, uint32_t *count, int qflags, TCXSTR *log);
 EJDB_INLINE void _nufetch(_EJDBNUM *nu, const char *sval, bson_type bt);
 EJDB_INLINE int _nucmp(_EJDBNUM *nu, const char *sval, bson_type bt);
@@ -1405,9 +1405,22 @@ static void _qrydup(const EJQ *src, EJQ *target, uint32_t qflags) {
     }
 }
 
+typedef struct {
+    bson *sbson;
+    TCMAP *ifields;
+} _BSONSTRIPVISITORCTX;
+
+static bool _bsonstripvisitor(const char *ipath, int ipathlen, const char *key, int keylen,
+        const bson_iterator *it, void *op) {
+    bool rv = true;
+    _BSONSTRIPVISITORCTX *ictx = op;
+    assert(ictx);
+
+    return rv;
+}
 
 /* push bson into rs with only fields listed in ifields */
-static void _pushstrippedbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbufsz) {
+static void _pushstripbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbufsz) {
     if (!ifields || TCMAPRNUM(ifields) <= 0) {
         TCLISTPUSH(rs, bsbuf, bsbufsz);
         return;
@@ -1419,9 +1432,21 @@ static void _pushstrippedbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbuf
     sbson.data = tmpbuf;
     sbson.dataSize = bsbufsz;
 
+    _BSONSTRIPVISITORCTX ictx;
+    ictx.sbson = &sbson;
+    ictx.ifields = ifields;
+
     //Now copy filtered bson fields
-    
+    bson_iterator it;
+    bson_iterator_from_buffer(&it, bsbuf);
+    bson_visit_fields(&it, BSON_TRAVERSE_ARRAYS_EXCLUDED, _bsonstripvisitor, &ictx);
+    bson_finish(&sbson);
 
+    if (!sbson.err) {
+        TCLISTPUSH(rs, bson_data(&sbson), bson_size(&sbson));
+    }
+    sbson.data = NULL; //this data will be freed at the end of this func
+    bson_destroy(&sbson); //destroy
 
     if (tmpbuf != bstack) {
         TCFREE(tmpbuf);
@@ -1543,7 +1568,7 @@ static TCLIST* _qrysearch(EJCOLL *jcoll, const EJQ *q, uint32_t *outcount, int q
     ++count; \
     if (!onlycount && (all || count > skip)) { \
         if (ifields) {\
-            _pushstrippedbson(res, ifields, (_bsbuf), (_bsbufsz)); \
+            _pushstripbson(res, ifields, (_bsbuf), (_bsbufsz)); \
         } else { \
             TCLISTPUSH(res, (_bsbuf), (_bsbufsz)); \
         } \