Fix demo condition justification

Signed-off-by: begeekmyfriend <begeekmyfriend@gmail.com>
begeekmyfriend · Jan 23, 2018 · 42e8303 · 42e8303
1 parent 727d763
commit 42e8303
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/lib/bplustree.c b/lib/bplustree.c
@@ -1068,13 +1068,22 @@ struct bplus_tree *bplus_tree_init(char *filename, int block_size)
 {
         int i;
         struct bplus_node node;
-        assert((block_size & (block_size - 1)) == 0);
 
         if (strlen(filename) >= 1024) {
                 fprintf(stderr, "Index file name too long!\n");
                 return NULL;
         }
 
+        if (block_size & (block_size - 1) != 0) {
+                fprintf(stderr, "Block size must be pow of 2!\n");
+                return NULL;
+        }
+
+        if (block_size < sizeof(node)) {
+                fprintf(stderr, "block size is too small for one node!\n");
+                return NULL;
+        }
+
         _block_size = block_size;
         _max_order = (block_size - sizeof(node)) / (sizeof(key_t) + sizeof(off_t));
         _max_entries = (block_size - sizeof(node)) / (sizeof(key_t) + sizeof(long));

diff --git a/tests/bplustree_demo.c b/tests/bplustree_demo.c
@@ -22,8 +22,8 @@ static int bplus_tree_setting(struct bplus_tree_config *config)
         int i, size, ret = 0, again = 1;
 
         printf("\n-- B+tree setting...\n");
-        printf("Set data index file name (e.g. /tmp/data.index): ");
         while (again) {
+                printf("Set data index file name (e.g. /tmp/data.index): ");
                 switch (i = getchar()) {
                 case EOF:
                         printf("\n");
@@ -47,8 +47,8 @@ static int bplus_tree_setting(struct bplus_tree_config *config)
         }
 
         again = 1;
-        printf("Set index file block size (bytes, power of 2, e.g. 4096): ");
         while (again) {
+                printf("Set index file block size (bytes, power of 2, e.g. 4096): ");
                 switch (i = getchar()) {
                 case EOF:
                         printf("\n");
@@ -64,11 +64,20 @@ static int bplus_tree_setting(struct bplus_tree_config *config)
                         if (!ret || getchar() != '\n') {
                                 stdin_flush();
                                 again = 1;
+                        } else if (size <= 0 || (size & (size - 1)) != 0) {
+                		fprintf(stderr, "Block size must be positive and pow of 2!\n");
+                                again = 1;
                         } else if (size <= 0 || (size & (size - 1)) != 0) {
                                 again = 1;
                         } else {
-                                config->block_size = size;
-                                again = 0;
+                                int order = (size - sizeof(struct bplus_node)) / (sizeof(key_t) + sizeof(off_t));
+                                if ((size_t) size < sizeof(struct bplus_node) || order <= 2) {
+                                        fprintf(stderr, "block size is too small for one node!\n");
+                                        again = 1;
+                                } else {
+                                        config->block_size = size;
+                                        again = 0;
+                                }
                         }
                         break;
                 }