begin to write import iterator
[platform/upstream/libpinyin.git] / src / storage / flexible_ngram.h
1 /* 
2  *  libpinyin
3  *  Library to deal with pinyin.
4  *  
5  *  Copyright (C) 2011 Peng Wu <alexepico@gmail.com>
6  *  
7  *  This program is free software; you can redistribute it and/or modify
8  *  it under the terms of the GNU General Public License as published by
9  *  the Free Software Foundation; either version 2 of the License, or
10  *  (at your option) any later version.
11  * 
12  *  This program is distributed in the hope that it will be useful,
13  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
14  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15  *  GNU General Public License for more details.
16  *  
17  *  You should have received a copy of the GNU General Public License
18  *  along with this program; if not, write to the Free Software
19  *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
20  */
21
22
23
24 #ifndef FLEXIBLE_NGRAM_H
25 #define FLEXIBLE_NGRAM_H
26
27 #include <db.h>
28 #include <errno.h>
29
30 /* Note: the signature of the template parameters.
31  * struct MagicHeader, ArrayHeader, ArrayItem.
32  */
33
34 namespace pinyin{
35
36 typedef GArray * FlexibleBigramPhraseArray;
37
38 /**
39  * FlexibleSingleGram:
40  * @ArrayHeader: the struct ArrayHeader.
41  * @ArrayItem: the struct ArrayItem.
42  *
43  * The flexible single gram is mainly used for training purpose.
44  *
45  */
46
47 template<typename ArrayHeader, typename ArrayItem>
48 class FlexibleSingleGram{
49     template<typename MH, typename AH,
50              typename AI>
51     friend class FlexibleBigram;
52 private:
53     MemoryChunk m_chunk;
54     FlexibleSingleGram(void * buffer, size_t length){
55         m_chunk.set_chunk(buffer, length, NULL);
56     }
57 public:
58     /**
59      * ArrayItemWithToken:
60      *
61      * Define the struct ArrayItemWithToken type.
62      *
63      */
64     typedef struct{
65         phrase_token_t m_token;
66         ArrayItem m_item;
67     } ArrayItemWithToken;
68
69 private:
70     static bool token_less_than(const ArrayItemWithToken & lhs,
71                                 const ArrayItemWithToken & rhs){
72         return lhs.m_token < rhs.m_token;
73     }
74
75 public:
76     /**
77      * FlexibleSingleGram::FlexibleSingleGram:
78      *
79      * The constructor of the FlexibleSingleGram.
80      *
81      */
82     FlexibleSingleGram(){
83         m_chunk.set_size(sizeof(ArrayHeader));
84         memset(m_chunk.begin(), 0, sizeof(ArrayHeader));
85     }
86
87     /**
88      * FlexibleSingleGram::retrieve_all:
89      * @array: the array to store all items in this single gram.
90      * @returns: whether the retrieve operation is successful.
91      *
92      * Retrieve all items in this single gram.
93      *
94      */
95     bool retrieve_all(/* out */ FlexibleBigramPhraseArray array){
96         const ArrayItemWithToken * begin = (const ArrayItemWithToken *)
97             ((const char *)(m_chunk.begin()) + sizeof(ArrayHeader));
98         const ArrayItemWithToken * end = (const ArrayItemWithToken *)
99             m_chunk.end();
100
101         ArrayItemWithToken item;
102         for ( const ArrayItemWithToken * cur_item = begin;
103               cur_item != end;
104               ++cur_item){
105             /* Note: optimize this with g_array_append_vals? */
106             item.m_token = cur_item->m_token;
107             item.m_item = cur_item->m_item;
108             g_array_append_val(array, item);
109         }
110
111         return true;
112     }
113
114     /**
115      * FlexibleSingleGram::search:
116      * @range: the token range.
117      * @array: the array to store the array items with token in the range.
118      * @returns: whether the search operation is successful.
119      *
120      * Search the array items with token in the range.
121      *
122      * Note: The array result may contain many items.
123      *
124      */
125     bool search(/* in */ PhraseIndexRange * range,
126                 /* out */ FlexibleBigramPhraseArray array){
127         const ArrayItemWithToken * begin = (const ArrayItemWithToken *)
128             ((const char *)(m_chunk.begin()) + sizeof(ArrayHeader));
129         const ArrayItemWithToken * end = (const ArrayItemWithToken *)
130             m_chunk.end();
131
132         ArrayItemWithToken compare_item;
133         compare_item.m_token = range->m_range_begin;
134         const ArrayItemWithToken * cur_item = std_lite::lower_bound
135             (begin, end, compare_item, token_less_than);
136
137         ArrayItemWithToken item;
138         for ( ; cur_item != end; ++cur_item){
139             if ( cur_item->m_token >= range->m_range_end )
140                 break;
141             item.m_token = cur_item->m_token;
142             item.m_item = cur_item->m_item;
143             g_array_append_val(array, item);
144         }
145
146         return true;
147     }
148
149     /**
150      * FlexibleSingleGram::insert_array_item:
151      * @token: the phrase token to be inserted.
152      * @item: the array item of this token.
153      * @returns: whether the insert operation is successful.
154      *
155      * Insert the array item of the token.
156      *
157      */
158     bool insert_array_item(/* in */ phrase_token_t token,
159                            /* in */ const ArrayItem & item){
160         ArrayItemWithToken * begin = (ArrayItemWithToken *)
161             ((const char *)(m_chunk.begin()) + sizeof(ArrayHeader));
162         ArrayItemWithToken * end = (ArrayItemWithToken *)
163             m_chunk.end();
164
165         ArrayItemWithToken compare_item;
166         compare_item.m_token = token;
167         ArrayItemWithToken * cur_item = std_lite::lower_bound
168             (begin, end, compare_item, token_less_than);
169
170         ArrayItemWithToken insert_item;
171         insert_item.m_token = token;
172         insert_item.m_item = item;
173
174         for ( ; cur_item != end; ++cur_item ){
175             if ( cur_item->m_token > token ){
176                 size_t offset = sizeof(ArrayHeader) +
177                     sizeof(ArrayItemWithToken) * (cur_item - begin);
178                 m_chunk.insert_content(offset, &insert_item,
179                                        sizeof(ArrayItemWithToken));
180                 return true;
181             }
182             if ( cur_item->m_token == token ){
183                 return false;
184             }
185         }
186         m_chunk.insert_content(m_chunk.size(), &insert_item,
187                                sizeof(ArrayItemWithToken));
188         return true;
189     }
190
191     /**
192      * FlexibleSingleGram::remove_array_item:
193      * @token: the phrase token to be removed.
194      * @item: the content of the removed array item.
195      * @returns: whether the remove operation is successful.
196      *
197      * Remove the array item of the token.
198      *
199      */
200     bool remove_array_item(/* in */ phrase_token_t token,
201                            /* out */ ArrayItem & item)
202     {
203         /* clear retval */
204         memset(&item, 0, sizeof(ArrayItem));
205
206         const ArrayItemWithToken * begin = (const ArrayItemWithToken *)
207             ((const char *)(m_chunk.begin()) + sizeof(ArrayHeader));
208         const ArrayItemWithToken * end = (const ArrayItemWithToken *)
209             m_chunk.end();
210
211         ArrayItemWithToken compare_item;
212         compare_item.m_token = token;
213         const ArrayItemWithToken * cur_item = std_lite::lower_bound
214             (begin, end, compare_item, token_less_than);
215
216         for ( ; cur_item != end; ++cur_item){
217             if ( cur_item->m_token > token )
218                 return false;
219             if ( cur_item->m_token == token ){
220                 memcpy(&item, &(cur_item->m_item), sizeof(ArrayItem));
221                 size_t offset = sizeof(ArrayHeader) +
222                     sizeof(ArrayItemWithToken) * (cur_item - begin);
223                 m_chunk.remove_content(offset, sizeof(ArrayItemWithToken));
224                 return true;
225             }
226         }
227         return false;
228     }
229
230     /**
231      * FlexibleSingleGram::get_array_item:
232      * @token: the phrase token.
233      * @item: the array item of the token.
234      * @returns: whether the get operation is successful.
235      *
236      * Get the array item of the token.
237      *
238      */
239     bool get_array_item(/* in */ phrase_token_t token,
240                         /* out */ ArrayItem & item)
241     {
242         /* clear retval */
243         memset(&item, 0, sizeof(ArrayItem));
244
245         const ArrayItemWithToken * begin = (const ArrayItemWithToken *)
246             ((const char *)(m_chunk.begin()) + sizeof(ArrayHeader));
247         const ArrayItemWithToken * end = (const ArrayItemWithToken *)
248             m_chunk.end();
249
250         ArrayItemWithToken compare_item;
251         compare_item.m_token = token;
252         const ArrayItemWithToken * cur_item = std_lite::lower_bound
253             (begin, end, compare_item, token_less_than);
254
255         for ( ; cur_item != end; ++cur_item){
256             if ( cur_item->m_token > token )
257                 return false;
258             if ( cur_item->m_token == token ){
259                 memcpy(&item, &(cur_item->m_item), sizeof(ArrayItem));
260                 return true;
261             }
262         }
263         return false;
264     }
265
266     /**
267      * FlexibleSingleGram::set_array_item:
268      * @token: the phrase token.
269      * @item: the array item of the token.
270      * @returns: whether the set operation is successful.
271      *
272      * Set the array item of the token.
273      *
274      */
275     bool set_array_item(/* in */ phrase_token_t token,
276                         /* in */ const ArrayItem & item){
277         ArrayItemWithToken * begin = (ArrayItemWithToken *)
278             ((const char *)(m_chunk.begin()) + sizeof(ArrayHeader));
279         ArrayItemWithToken * end = (ArrayItemWithToken *)
280             m_chunk.end();
281
282         ArrayItemWithToken compare_item;
283         compare_item.m_token = token;
284         ArrayItemWithToken * cur_item = std_lite::lower_bound
285             (begin, end, compare_item, token_less_than);
286
287         for ( ; cur_item != end; ++cur_item ){
288             if ( cur_item->m_token > token ){
289                 return false;
290             }
291             if ( cur_item->m_token == token ){
292                 memcpy(&(cur_item->m_item), &item, sizeof(ArrayItem));
293                 return true;
294             }
295         }
296         return false;
297     }
298
299     /**
300      * FlexibleSingleGram::get_array_header:
301      * @header: the array header of this single gram.
302      * @returns: whether the get operation is successful.
303      *
304      * Get the array header of this single gram.
305      *
306      */
307     bool get_array_header(/* out */ ArrayHeader & header){
308         /* clear retval */
309         memset(&header, 0, sizeof(ArrayHeader));
310         char * buf_begin = (char *)m_chunk.begin();
311         memcpy(&header, buf_begin, sizeof(ArrayHeader));
312         return true;
313     }
314
315     /**
316      * FlexibleSingleGram::set_array_header:
317      * @header: the array header of this single gram.
318      * @returns: whether the set operation is successful.
319      *
320      * Set the array header of this single gram.
321      *
322      */
323     bool set_array_header(/* in */ const ArrayHeader & header){
324         char * buf_begin = (char *)m_chunk.begin();
325         memcpy(buf_begin, &header, sizeof(ArrayHeader));
326         return true;
327     }
328 };
329
330 /**
331  * FlexibleBigram:
332  * @MagicHeader: the struct type of the magic header.
333  * @ArrayHeader: the struct type of the array header.
334  * @ArrayItem: the struct type of the array item.
335  *
336  * The flexible bi-gram is mainly used for training purpose.
337  *
338  */
339 template<typename MagicHeader, typename ArrayHeader,
340          typename ArrayItem>
341 class FlexibleBigram{
342     /* Note: some flexible bi-gram file format check should be here. */
343 private:
344     DB * m_db;
345
346     phrase_token_t m_magic_header_index[2];
347
348     char m_magic_number[4];
349
350     void reset(){
351         if ( m_db ){
352             m_db->sync(m_db, 0);
353             m_db->close(m_db, 0);
354             m_db = NULL;
355         }
356     }
357
358 public:
359     /**
360      * FlexibleBigram::FlexibleBigram:
361      * @magic_number: the 4 bytes magic number of the flexible bi-gram.
362      *
363      * The constructor of the FlexibleBigram.
364      *
365      */
366     FlexibleBigram(const char * magic_number){
367         m_db = NULL;
368         m_magic_header_index[0] = null_token;
369         m_magic_header_index[1] = null_token;
370
371         memcpy(m_magic_number, magic_number, sizeof(m_magic_number));
372     }
373
374     /**
375      * FlexibleBigram::~FlexibleBigram:
376      *
377      * The destructor of the FlexibleBigram.
378      *
379      */
380     ~FlexibleBigram(){
381         reset();
382     }
383
384     /**
385      * FlexibleBigram::attach:
386      * @dbfile: the path name of the flexible bi-gram.
387      * @flags: the attach flags for the Berkeley DB.
388      * @returns: whether the attach operation is successful.
389      *
390      * Attach Berkeley DB on filesystem for training purpose.
391      *
392      */
393     bool attach(const char * dbfile, guint32 flags){
394         reset();
395         u_int32_t db_flags = 0;
396
397         if ( flags & ATTACH_READONLY )
398             db_flags |= DB_RDONLY;
399         if ( flags & ATTACH_READWRITE )
400             assert( !(flags & ATTACH_READONLY ) );
401
402         if ( !dbfile )
403             return false;
404         int ret = db_create(&m_db, NULL, 0);
405         if ( ret != 0 )
406             assert(false);
407
408         ret = m_db->open(m_db, NULL, dbfile, NULL, DB_HASH, db_flags, 0644);
409         if ( ret != 0 && (flags & ATTACH_CREATE) ) {
410             db_flags |= DB_CREATE;
411             /* Create database file here, and write the signature. */
412             ret = m_db->open(m_db, NULL, dbfile, NULL, DB_HASH, db_flags, 0644);
413             if ( ret != 0 )
414                 return false;
415
416             DBT db_key;
417             memset(&db_key, 0, sizeof(DBT));
418             db_key.data = m_magic_header_index;
419             db_key.size = sizeof(m_magic_header_index);
420             DBT db_data;
421             memset(&db_data, 0, sizeof(DBT));
422             db_data.data = m_magic_number;
423             db_data.size = sizeof(m_magic_number);
424             db_data.flags = DB_DBT_PARTIAL;
425             db_data.doff = 0;
426             db_data.dlen = sizeof(m_magic_number);
427
428             ret = m_db->put(m_db, NULL, &db_key, &db_data, 0);
429             return ret == 0;
430         }
431
432         /* check the signature. */
433         DBT db_key;
434         memset(&db_key, 0, sizeof(DBT));
435         db_key.data = m_magic_header_index;
436         db_key.size = sizeof(m_magic_header_index);
437         DBT db_data;
438         memset(&db_data, 0, sizeof(DBT));
439         db_data.flags = DB_DBT_PARTIAL;
440         db_data.doff = 0;
441         db_data.dlen = sizeof(m_magic_number);
442         ret = m_db->get(m_db, NULL, &db_key, &db_data, 0);
443         if ( ret != 0 )
444             return false;
445         if ( sizeof(m_magic_number) != db_data.size )
446             return false;
447         if ( memcmp(db_data.data, m_magic_number,
448                     sizeof(m_magic_number)) == 0 )
449             return true;
450         return false;
451     }
452
453     /**
454      * FlexibleBigram::load:
455      * @index: the previous token in the flexible bi-gram.
456      * @single_gram: the single gram of the previous token.
457      * @returns: whether the load operation is successful.
458      *
459      * Load the single gram of the previous token.
460      *
461      */
462     bool load(phrase_token_t index,
463               FlexibleSingleGram<ArrayHeader, ArrayItem> * & single_gram){
464         if ( !m_db )
465             return false;
466
467         DBT db_key;
468         memset(&db_key, 0, sizeof(DBT));
469         db_key.data = &index;
470         db_key.size = sizeof(phrase_token_t);
471
472         single_gram = NULL;
473
474         DBT db_data;
475         memset(&db_data, 0, sizeof(DBT));
476         int ret = m_db->get(m_db, NULL, &db_key, &db_data, 0);
477         if ( ret != 0)
478             return false;
479
480         single_gram = new FlexibleSingleGram<ArrayHeader, ArrayItem>
481             (db_data.data, db_data.size);
482
483         return true;
484     }
485
486     /**
487      * FlexibleBigram::store:
488      * @index: the previous token in the flexible bi-gram.
489      * @single_gram: the single gram of the previous token.
490      * @returns: whether the store operation is successful.
491      *
492      * Store the single gram of the previous token.
493      *
494      */
495     bool store(phrase_token_t index,
496                FlexibleSingleGram<ArrayHeader, ArrayItem> * single_gram){
497         if ( !m_db )
498             return false;
499
500         DBT db_key;
501         memset(&db_key, 0, sizeof(DBT));
502         db_key.data = &index;
503         db_key.size = sizeof(phrase_token_t);
504         DBT db_data;
505         memset(&db_data, 0, sizeof(DBT));
506         db_data.data = single_gram->m_chunk.begin();
507         db_data.size = single_gram->m_chunk.size();
508
509         int ret = m_db->put(m_db, NULL, &db_key, &db_data, 0);
510         return ret == 0;
511     }
512
513     /**
514      * FlexibleBigram::remove:
515      * @index: the previous token in the flexible bi-gram.
516      * @returns: whether the remove operation is successful.
517      *
518      * Remove the single gram of the previous token.
519      *
520      */
521     bool remove(phrase_token_t index){
522         if ( !m_db )
523             return false;
524
525         DBT db_key;
526         memset(&db_key, 0, sizeof(DBT));
527         db_key.data = &index;
528         db_key.size = sizeof(phrase_token_t);
529
530         int ret = m_db->del(m_db, NULL, &db_key, 0);
531         return ret == 0;
532     }
533
534     /**
535      * FlexibleBigram::get_all_items:
536      * @items: the GArray to store all previous tokens.
537      * @returns: whether the get operation is successful.
538      *
539      * Get the array of all previous tokens for parameter estimation.
540      *
541      */
542     bool get_all_items(GArray * items){
543         g_array_set_size(items, 0);
544         if ( !m_db )
545             return false;
546
547         DBC * cursorp;
548         DBT key, data;
549         int ret;
550
551         /* Get a cursor */
552         m_db->cursor(m_db, NULL, &cursorp, 0);
553
554         /* Initialize our DBTs. */
555         memset(&key, 0, sizeof(DBT));
556         memset(&data, 0, sizeof(DBT));
557
558         /* Iterate over the database, retrieving each record in turn. */
559         while ((ret =  cursorp->c_get(cursorp, &key, &data, DB_NEXT)) == 0 ){
560             if (key.size != sizeof(phrase_token_t)){
561                 /* skip magic header. */
562                 continue;
563             }
564             phrase_token_t * token = (phrase_token_t *) key.data;
565             g_array_append_val(items, *token);
566         }
567
568         if ( ret != DB_NOTFOUND ){
569             fprintf(stderr, "training db error, exit!");
570             exit(EIO);
571         }
572
573         /* Cursors must be closed */
574         if (cursorp != NULL)
575             cursorp->c_close(cursorp);
576         return true;
577     }
578
579     /**
580      * FlexibleBigram::get_magic_header:
581      * @header: the magic header.
582      * @returns: whether the get operation is successful.
583      *
584      * Get the magic header of the flexible bi-gram.
585      *
586      */
587     bool get_magic_header(MagicHeader & header){
588         /* clear retval */
589         memset(&header, 0, sizeof(MagicHeader));
590
591         if ( !m_db )
592             return false;
593
594         DBT db_key;
595         memset(&db_key, 0, sizeof(DBT));
596         db_key.data = m_magic_header_index;
597         db_key.size = sizeof(m_magic_header_index);
598         DBT db_data;
599         memset(&db_data, 0, sizeof(DBT));
600         db_data.flags = DB_DBT_PARTIAL;
601         db_data.doff = sizeof(m_magic_number);
602         db_data.dlen = sizeof(MagicHeader);
603         
604         int ret = m_db->get(m_db, NULL, &db_key, &db_data, 0);
605         if ( ret != 0 )
606             return false;
607
608         if ( sizeof(MagicHeader) != db_data.size )
609             return false;
610
611         memcpy(&header, db_data.data, sizeof(MagicHeader));
612         return true;
613     }
614
615     /**
616      * FlexibleBigram::set_magic_header:
617      * @header: the magic header.
618      * @returns: whether the set operation is successful.
619      *
620      * Set the magic header of the flexible bi-gram.
621      *
622      */
623     bool set_magic_header(const MagicHeader & header){
624         if ( !m_db )
625             return false;
626
627         DBT db_key;
628         memset(&db_key, 0, sizeof(DBT));
629         db_key.data = m_magic_header_index;
630         db_key.size = sizeof(m_magic_header_index);
631         DBT db_data;
632         memset(&db_data, 0, sizeof(DBT));
633         db_data.data = (void *) &header;
634         db_data.size = sizeof(MagicHeader);
635         db_data.flags = DB_DBT_PARTIAL;
636         db_data.doff = sizeof(m_magic_number);
637         db_data.dlen = sizeof(MagicHeader);
638
639         int ret = m_db->put(m_db, NULL, &db_key, &db_data, 0);
640         return ret == 0;
641     }
642
643     /**
644      * FlexibleBigram::get_array_header:
645      * @index: the previous token in the flexible bi-gram.
646      * @header: the array header in the single gram of the previous token.
647      * @returns: whether the get operation is successful.
648      *
649      * Get the array header in the single gram of the previous token.
650      *
651      */
652     bool get_array_header(phrase_token_t index, ArrayHeader & header){
653         /* clear retval */
654         memset(&header, 0, sizeof(ArrayHeader));
655
656         if ( !m_db )
657             return false;
658
659         DBT db_key;
660         memset(&db_key, 0, sizeof(DBT));
661         db_key.data = &index;
662         db_key.size = sizeof(phrase_token_t);
663
664         DBT db_data;
665         memset(&db_data, 0, sizeof(DBT));
666         db_data.flags = DB_DBT_PARTIAL;
667         db_data.doff = 0;
668         db_data.dlen = sizeof(ArrayHeader);
669         int ret = m_db->get(m_db, NULL, &db_key, &db_data, 0);
670         if ( ret != 0 )
671             return false;
672
673         assert(db_data.size == sizeof(ArrayHeader));
674         memcpy(&header, db_data.data, sizeof(ArrayHeader));
675         return true;
676     }
677
678     /**
679      * FlexibleBigram::set_array_header:
680      * @index: the previous token of the flexible bi-gram.
681      * @header: the array header in the single gram of the previous token.
682      * @returns: whether the set operation is successful.
683      *
684      * Set the array header in the single gram of the previous token.
685      *
686      */
687     bool set_array_header(phrase_token_t index, const ArrayHeader & header){
688         if ( !m_db )
689             return false;
690
691         DBT db_key;
692         memset(&db_key, 0, sizeof(DBT));
693         db_key.data = &index;
694         db_key.size = sizeof(phrase_token_t);
695         DBT db_data;
696         memset(&db_data, 0, sizeof(DBT));
697         db_data.data = (void *)&header;
698         db_data.size = sizeof(ArrayHeader);
699         db_data.flags = DB_DBT_PARTIAL;
700         db_data.doff = 0;
701         db_data.dlen = sizeof(ArrayHeader);
702
703         int ret = m_db->put(m_db, NULL, &db_key, &db_data, 0);
704         return ret == 0;
705     }
706
707 };
708
709 };
710
711 #endif