#8

3 files changed, 85 insertions, 37 deletions

diff --git a/tcejdb/bson.c b/tcejdb/bson.c
index 39c7867..87e4585 100644
--- a/tcejdb/bson.c
+++ b/tcejdb/bson.c
@@ -316,11 +316,11 @@ 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) {
+static void bson_visit_fields_impl(bson_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) {
+    bson_visitor_cmd_t vcmd;
+    while (!(vcmd & BSON_VCMD_TERMINATE) && (t = bson_iterator_next(it)) != BSON_EOO) {
         const char* key = bson_iterator_key(it);
         klen = strlen(key);
         if (curr + klen > BSON_MAX_FPATH_LEN) {
@@ -334,15 +334,19 @@ static void bson_visit_fields_impl(traverse_flags_t flags, char* pstack, int cur
         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);
-            }
+        vcmd = visitor(pstack, curr, key, klen, it, false, op);
+        if (vcmd & BSON_VCMD_TERMINATE) {
+            break;
+        }
+        if (!(vcmd & BSON_VCMD_SKIP_NESTED) &&
+                ((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);
         }
+        vcmd = visitor(pstack, curr, key, klen, it, true, op);
         //POP
         curr -= klen;
         if (curr > 0) {
@@ -351,7 +355,7 @@ static void bson_visit_fields_impl(traverse_flags_t flags, char* pstack, int cur
     }
 }
 
-EJDB_EXPORT void bson_visit_fields(bson_iterator *it, traverse_flags_t flags, BSONVISITOR visitor, void *op) {
+EJDB_EXPORT void bson_visit_fields(bson_iterator *it, bson_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);
 }
@@ -773,9 +777,24 @@ int bson_ensure_space(bson *b, const int bytesNeeded) {
         }
     }
 
-    b->data = bson_realloc(b->data, new_size);
-    if (!b->data)
+    if (b->flags & BSON_FLAG_STACK_ALLOCATED) { //translate stack memory into heap
+        char *odata = b->data;
+        b->data = bson_malloc_func(new_size);
+        if (!b->data) {
+            bson_fatal_msg(!!b->data, "malloc() failed");
+            return BSON_ERROR;
+        }
+        if (odata) {
+            memcpy(b->data, odata, MIN(new_size, b->dataSize));
+        }
+        b->flags &= ~BSON_FLAG_STACK_ALLOCATED; //reset this flag
+    } else {
+        b->data = bson_realloc(b->data, new_size);
+    }
+    if (!b->data) {
         bson_fatal_msg(!!b->data, "realloc() failed");
+        return BSON_ERROR;
+    }
 
     b->dataSize = new_size;
     b->cur += b->data - orig;
@@ -1196,10 +1215,10 @@ EJDB_EXPORT void bson_swap_endian32(void *outp, const void *inp) {
     out[3] = in[0];
 }
 
-EJDB_EXPORT void bson_append_field_from_iterator(bson_iterator *from, bson *into) {
+EJDB_EXPORT int bson_append_field_from_iterator(const bson_iterator *from, bson *into) {
     bson_type t = bson_iterator_type(from);
-    if (t == BSON_EOO) {
-        return;
+    if (t == BSON_EOO || !into || into->finished) {
+        return BSON_ERROR;
     }
     const char* key = bson_iterator_key(from);
     switch (t) {
@@ -1271,7 +1290,7 @@ EJDB_EXPORT void bson_append_field_from_iterator(bson_iterator *from, bson *into
         default:
             break;
     }
-
+    return BSON_OK;
 }
 
 EJDB_EXPORT int bson_merge(bson *b1, bson *b2, bson_bool_t overwrite, bson *out) {
diff --git a/tcejdb/bson.h b/tcejdb/bson.h
index 334da76..fb3de81 100644
--- a/tcejdb/bson.h
+++ b/tcejdb/bson.h
@@ -68,7 +68,8 @@ enum bson_binary_subtype_t {
 };
 
 enum bson_flags_t {
-    BSON_FLAG_QUERY_MODE = 1
+    BSON_FLAG_QUERY_MODE = 1,
+    BSON_FLAG_STACK_ALLOCATED = 1 << 1 /**< If it set BSON data is allocated on stack and realloc should deal with this case */
 };
 
 typedef enum {
@@ -200,9 +201,14 @@ EJDB_EXPORT bson_type bson_find_fieldpath_value2(const char *fpath, int fplen, b
 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);
+} bson_traverse_flags_t;
+typedef enum {
+    BSON_VCMD_OK = 0,
+    BSON_VCMD_TERMINATE = 1,
+    BSON_VCMD_SKIP_NESTED = 1 << 1
+} bson_visitor_cmd_t;
+typedef bson_visitor_cmd_t (*BSONVISITOR)(const char *ipath, int ipathlen, const char *key, int keylen, const bson_iterator *it, bool after, void *op);
+EJDB_EXPORT void bson_visit_fields(bson_iterator *it, bson_traverse_flags_t flags, BSONVISITOR visitor, void *op);
 
 
 EJDB_EXPORT bson_iterator* bson_iterator_create(void);
@@ -1056,7 +1062,7 @@ EJDB_EXPORT void bson_swap_endian64(void *outp, const void *inp);
  * @param from
  * @param into
  */
-EJDB_EXPORT void bson_append_field_from_iterator(bson_iterator *from, bson *into);
+EJDB_EXPORT int bson_append_field_from_iterator(const bson_iterator *from, bson *into);
 
 
 /**
diff --git a/tcejdb/ejdb.c b/tcejdb/ejdb.c
index fb17b5e..c333a63 100644
--- a/tcejdb/ejdb.c
+++ b/tcejdb/ejdb.c
@@ -3,6 +3,7 @@
 
 
 #include <regex.h>
+#include <asm-generic/ioctls.h>
 
 #include "ejdb_private.h"
 #include "ejdbutl.h"
@@ -1410,13 +1411,32 @@ typedef struct {
     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;
+static bson_visitor_cmd_t _bsonstripvisitor(const char *ipath, int ipathlen, const char *key, int keylen,
+        const bson_iterator *it, bool after, void *op) {
     _BSONSTRIPVISITORCTX *ictx = op;
-    assert(ictx);
-
-    return rv;
+    assert(ictx && ictx->sbson && ictx->ifields && ipath && key && it && op);
+    TCMAP *ifields = ictx->ifields;
+    const void *buf;
+    int bufsz;
+    bson_type bt = bson_iterator_type(it);
+    if (bt == BSON_EOO) {
+        return BSON_VCMD_OK;
+    }
+    if (bt != BSON_OBJECT && bt != BSON_ARRAY) {
+        if (after) { //simple primitive case
+            return BSON_VCMD_OK;
+        }
+        //const void *tcmapget(const TCMAP *map, const void *kbuf, int ksiz, int *sp);
+        buf = tcmapget(ifields, ipath, ipathlen, &bufsz);
+        if (buf) {
+            bson_append_field_from_iterator(it, ictx->sbson);
+        }
+        return BSON_VCMD_OK;
+    } else { //more complicated case
+        //TODO
+        
+    }
+    return BSON_VCMD_OK;
 }
 
 /* push bson into rs with only fields listed in ifields */
@@ -1430,8 +1450,10 @@ static void _pushstripbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbufsz)
     bson sbson;
     bson_reset(&sbson);
     sbson.data = tmpbuf;
-    sbson.dataSize = bsbufsz;
-
+    sbson.dataSize = MAX(bsbufsz, JBSBUFFERSZ);
+    if (sbson.data == bstack) {
+        sbson.flags |= BSON_FLAG_STACK_ALLOCATED;
+    }
     _BSONSTRIPVISITORCTX ictx;
     ictx.sbson = &sbson;
     ictx.ifields = ifields;
@@ -1440,16 +1462,17 @@ static void _pushstripbson(TCLIST *rs, TCMAP *ifields, void *bsbuf, int bsbufsz)
     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) {
+    if (bson_finish(&sbson) == BSON_OK) {
         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);
+    //Cleanup
+    char *sdata = sbson.data; //save data ptr to check if it on stack
+    sbson.data = NULL; //this data will be freed at the end of this func
+    bson_destroy(&sbson);
+    if (sdata != bstack) {
+        TCFREE(sdata);
     }
 }