1 files changed, 7 insertions, 12 deletions

diff --git a/include/framework/reconstruction/BSMPolicy.h b/include/framework/reconstruction/BSMPolicy.h
index 42ae506..65554fb 100644
--- a/include/framework/reconstruction/BSMPolicy.h
+++ b/include/framework/reconstruction/BSMPolicy.h
@@ -36,7 +36,7 @@ public:
     ReconstructionVector reconstructions;
     auto levels = version->get_structure()->get_level_vector();
 
-    level_index target_level = find_reconstruction_target(levels);
+    level_index target_level = find_reconstruction_target(levels, version->get_structure()->get_record_count());
     assert(target_level != -1);
     level_index source_level = 0;
 
@@ -48,22 +48,17 @@ public:
     ReconstructionTask task;
     task.target = target_level;
 
-    if (target_level == 1 &&
-        (levels.size() == 1 || levels[1]->get_record_count() == 0)) {
-      /* if the first level is empty, then we just append the buffer to it */
-      task.type = ReconstructionType::Append;
-    } else {
-      /* otherwise, we'll need to do a merge of at least two shards */
-      task.type = ReconstructionType::Merge;
-    }
-
     size_t reccnt = 0;
     if (target_level < (ssize_t)levels.size() && levels[target_level]->get_record_count() > 0) {
       task.sources.push_back({target_level, all_shards_idx});
+      task.type = ReconstructionType::Merge;
+    } else {
+      task.type = ReconstructionType::Append;
     }
 
     for (level_index i = target_level - 1; i >= source_level; i--) {
       assert(i < (ssize_t)levels.size());
+      assert(levels[i]->get_record_count() > 0 || i == 0);
       task.sources.push_back({i, all_shards_idx});
       reccnt += levels[i]->get_record_count();
     }
@@ -76,11 +71,11 @@ public:
   }
 
 private:
-  level_index find_reconstruction_target(LevelVector &levels) const {
+  level_index find_reconstruction_target(LevelVector &levels, size_t reccnt) const {
     level_index target_level = invalid_level_idx;
 
     for (level_index i = 1; i < (level_index)levels.size(); i++) {
-      if (levels[i]->get_record_count() == 0) {
+      if (levels[i]->get_record_count() < capacity(i, reccnt)) {
         target_level = i;
         break;
       }