2 files changed, 23 insertions, 14 deletions

diff --git a/include/framework/structure/ExtensionStructure.h b/include/framework/structure/ExtensionStructure.h
index a77088f..0fd737c 100644
--- a/include/framework/structure/ExtensionStructure.h
+++ b/include/framework/structure/ExtensionStructure.h
@@ -170,8 +170,23 @@ public:
   inline reconstruction_results<ShardType> perform_reconstruction(ReconstructionTask task) const { 
     reconstruction_results<ShardType> result;
     result.target_level = task.target;
-    
-    std::vector<const ShardType *> shards;
+
+    /* if there is only one source, then we don't need to actually rebuild */
+    if (task.sources.size() == 1) {
+      auto shid = task.sources[0];
+      if (shid.shard_idx == all_shards_idx && m_levels[shid.level_idx]->get_shard_count() > 1) {
+        /* there's more than one shard, so we need to do the reconstruction */
+      } else {
+        auto raw_shard_ptr = m_levels[shid.level_idx]->get_shard(shid.shard_idx);
+        assert(raw_shard_ptr);
+        result.source_shards.emplace_back(shid.level_idx, raw_shard_ptr);
+        result.new_shard = m_levels[shid.level_idx]->get_shard_ptr(shid.shard_idx).first;
+
+        return result;
+      }
+    }
+        
+    std::vector<const ShardType*> shards;
     for (ShardID shid : task.sources) {
       assert(shid.level_idx < (level_index) m_levels.size());
       assert(shid.shard_idx >= -1);
@@ -233,7 +248,7 @@ public:
         m_levels[shards[i].first]->delete_shard(shard_idx);
       } else {
         fprintf(stderr, "ERROR: failed to delete shard %ld\t%p\n", shards[i].first, shards[i].second);
-        //exit(EXIT_FAILURE);
+        exit(EXIT_FAILURE);
       }
     }
   }
diff --git a/include/framework/structure/InternalLevel.h b/include/framework/structure/InternalLevel.h
index 5659c72..54b3ae2 100644
--- a/include/framework/structure/InternalLevel.h
+++ b/include/framework/structure/InternalLevel.h
@@ -217,17 +217,11 @@ public:
     m_shards.push_back(shard);
   }
 
-  const ShardType *get_shard(ShardID shid) const {
-    if (shid < m_shards.size()) {
-      return m_shards[shid].first.get();
-    }
-
-    return nullptr;
-  }
-
-  const shard_ptr get_shard_ptr(size_t shid) const {
-    if (shid < m_shards.size()) {
-      return m_shards[shid];
+  const shard_ptr get_shard_ptr(ssize_t idx) const {
+    if (idx >= 0 && idx < (ssize_t) m_shards.size()) {
+      return m_shards[idx];
+    } else if (idx == all_shards_idx && m_shards.size() == 1) {
+      return m_shards[0];
     }
 
     return {nullptr, 0};