Introduce shards.insert

AndersTrier · Aug 4, 2024 · 4c25120 · 4c25120
1 parent 5a6e363
commit 4c25120
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 30 deletions.
diff --git a/src/engine/shards.rs b/src/engine/shards.rs
@@ -32,6 +32,29 @@ impl Shards {
         self.data
             .resize(self.shard_count * self.shard_len_64, [0; 64]);
     }
+
+    pub(crate) fn insert(&mut self, index: usize, shard: &[u8]) {
+        assert_eq!(shard.len() % 2, 0);
+
+        let chunks_64 = shard.len() / 64;
+        let tail_len = shard.len() % 64;
+
+        let (src_chunks, src_tail) = shard.split_at(shard.len() - tail_len);
+
+        let dst = &mut self[index];
+        dst[..chunks_64]
+            .as_flattened_mut()
+            .copy_from_slice(src_chunks);
+
+        // Last chunk is special if shard.len() % 64 != 0.
+        // See src/algorithm.md for an explanation.
+        if tail_len > 0 {
+            let (src_lo, src_hi) = src_tail.split_at(tail_len / 2);
+            let (dst_lo, dst_hi) = dst[chunks_64].split_at_mut(32);
+            dst_lo[..src_lo.len()].copy_from_slice(src_lo);
+            dst_hi[..src_hi.len()].copy_from_slice(src_hi);
+        }
+    }
 }
 
 // ======================================================================

diff --git a/src/rate/decoder_work.rs b/src/rate/decoder_work.rs
@@ -80,16 +80,7 @@ impl DecoderWork {
                 got: original_shard.len(),
             })
         } else {
-            let remaining = self.shard_bytes % 64;
-            let split_index = self.shard_bytes - remaining / 2;
-
-            let (src, src_hi) = original_shard.split_at(split_index);
-
-            let dst = self.shards[pos].as_flattened_mut();
-
-            dst[..split_index].copy_from_slice(src);
-            let offset = dst.len() - 32;
-            dst[offset..offset + src_hi.len()].copy_from_slice(src_hi);
+            self.shards.insert(pos, original_shard);
 
             self.original_received_count += 1;
             self.received.set(pos, true);
@@ -118,16 +109,7 @@ impl DecoderWork {
                 got: recovery_shard.len(),
             })
         } else {
-            let remaining = self.shard_bytes % 64;
-            let split_index = self.shard_bytes - remaining / 2;
-
-            let (src, src_hi) = recovery_shard.split_at(split_index);
-
-            let dst = self.shards[pos].as_flattened_mut();
-
-            dst[..split_index].copy_from_slice(src);
-            let offset = dst.len() - 32;
-            dst[offset..offset + src_hi.len()].copy_from_slice(src_hi);
+            self.shards.insert(pos, recovery_shard);
 
             self.recovery_received_count += 1;
             self.received.set(pos, true);

diff --git a/src/rate/encoder_work.rs b/src/rate/encoder_work.rs
@@ -63,16 +63,8 @@ impl EncoderWork {
                 got: original_shard.len(),
             })
         } else {
-            let remaining = self.shard_bytes % 64;
-            let split_index = self.shard_bytes - remaining / 2;
-
-            let (src, src_hi) = original_shard.split_at(split_index);
-
-            let dst = self.shards[self.original_received_count].as_flattened_mut();
-
-            dst[..split_index].copy_from_slice(src);
-            let offset = dst.len() - 32;
-            dst[offset..offset + src_hi.len()].copy_from_slice(src_hi);
+            self.shards
+                .insert(self.original_received_count, original_shard);
 
             self.original_received_count += 1;
             Ok(())