add get/set array header to flexible bi-gram

2 files changed, 63 insertions, 3 deletions

diff --git a/src/storage/flexible_ngram.h b/src/storage/flexible_ngram.h
index 10bad8f..3cce47d 100644
--- a/src/storage/flexible_ngram.h
+++ b/src/storage/flexible_ngram.h
@@ -231,14 +231,15 @@ public:
     /* load/store one array. */
     bool load(phrase_token_t index,
               FlexibleSingleGram<ArrayHeader, ArrayItem> * & single_gram){
+        if ( !m_db )
+            return false;
+
         DBT db_key;
         memset(&db_key, 0, sizeof(DBT));
         db_key.data = &index;
         db_key.size = sizeof(phrase_token_t);
 
         single_gram = NULL;
-        if ( !m_db )
-            return false;
 
         DBT db_data;
         memset(&db_data, 0, sizeof(DBT));
@@ -342,6 +343,50 @@ public:
         int ret = m_db->put(m_db, NULL, &db_key, &db_data, 0);
         return ret == 0;
     }
+
+    bool get_array_header(phrase_token_t index, ArrayHeader & header){
+        if ( !m_db )
+            return false;
+
+        DBT db_key;
+        memset(&db_key, 0, sizeof(DBT));
+        db_key.data = &index;
+        db_key.size = sizeof(phrase_token_t);
+
+        DBT db_data;
+        memset(&db_data, 0, sizeof(DBT));
+        db_data.flags = DB_DBT_PARTIAL;
+        db_data.doff = 0;
+        db_data.dlen = sizeof(ArrayHeader);
+        int ret = m_db->get(m_db, NULL, &db_key, &db_data, 0);
+        if ( ret != 0 )
+            return false;
+
+        assert(db_data.size == sizeof(ArrayHeader));
+        memcpy(&header, db_data.data, sizeof(ArrayHeader));
+        return true;
+    }
+
+    bool set_array_header(phrase_token_t index, const ArrayHeader & header){
+        if ( !m_db )
+            return false;
+
+        DBT db_key;
+        memset(&db_key, 0, sizeof(DBT));
+        db_key.data = &index;
+        db_key.size = sizeof(phrase_token_t);
+        DBT db_data;
+        memset(&db_data, 0, sizeof(DBT));
+        db_data.data = (void *)&header;
+        db_data.size = sizeof(ArrayHeader);
+        db_data.flags = DB_DBT_PARTIAL;
+        db_data.doff = 0;
+        db_data.dlen = sizeof(ArrayHeader);
+
+        int ret = m_db->put(m_db, NULL, &db_key, &db_data, 0);
+        return ret == 0;
+    }
+
 };
 
 #endif
diff --git a/tests/storage/test_flexible_ngram.cpp b/tests/storage/test_flexible_ngram.cpp
index f915a1f..16eb225 100644
--- a/tests/storage/test_flexible_ngram.cpp
+++ b/tests/storage/test_flexible_ngram.cpp
@@ -38,7 +38,6 @@ int main(int argc, char * argv[]) {
     single_gram.set_array_header(32);
     single_gram.get_array_header(freq);
     printf("new array header:%d\n", freq);
-
     bigram.store(2, &single_gram);
 
     for (int m = 1; m <= 2; ++m ){
@@ -69,4 +68,20 @@ int main(int argc, char * argv[]) {
     bigram.get_magic_header(freq);
     assert(total_freq == freq);
     printf("magic header:%d\n", freq);
+
+    printf("-----------------------array header---------------------\n");
+    for ( size_t i = 1; i <= 2; ++i){
+        bigram.get_array_header(i, freq);
+        printf("single gram: %d, freq:%d\n", i, freq);
+    }
+
+    bigram.set_array_header(1, 1);
+
+    printf("-----------------------array header---------------------\n");
+    for ( size_t i = 1; i <= 2; ++i){
+        bigram.get_array_header(i, freq);
+        printf("single gram: %d, freq:%d\n", i, freq);
+    }
+
+    return 0;
 }