From 32beb06382182abb357fd2d8c678a875e79b036b Mon Sep 17 00:00:00 2001
From: Paul Masurel <paul.masurel@datadoghq.com>
Date: Wed, 3 Dec 2025 13:02:10 +0100
Subject: [PATCH] plastic surgery

---
 src/indexer/merger.rs             |  6 +++---
 src/indexer/segment_serializer.rs |  2 +-
 src/spatial/bkd.rs                | 25 +++++++++++++------------
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/src/indexer/merger.rs b/src/indexer/merger.rs
index 152a31cf7..1e5caddea 100644
--- a/src/indexer/merger.rs
+++ b/src/indexer/merger.rs
@@ -533,9 +533,9 @@ impl IndexMerger {
         /// Unfortunately, there are no special trick to merge segments.
         /// We need to rebuild a BKD-tree based off the list of triangles.
         ///
-        /// Because the data can be large, we do this by writing the sequence of triangles to disk,
-        /// and mmapping it as mutable slice, and calling the same code as what is done for the
-        /// segment serialization.
+        /// Because the data can be large, we do this by writing the sequence of triangles to
+        /// disk, and mmapping it as mutable slice, and calling the same code as what
+        /// is done for the segment serialization.
         ///
         /// The OS is in charge of deciding how to handle its page cache.
         /// This is the same as what would have happened with swapping,
diff --git a/src/indexer/segment_serializer.rs b/src/indexer/segment_serializer.rs
index 473dbd1e8..34b832e7e 100644
--- a/src/indexer/segment_serializer.rs
+++ b/src/indexer/segment_serializer.rs
@@ -70,7 +70,7 @@ impl SegmentSerializer {
         &mut self.fast_field_write
     }
 
-    /// HUSH
+    /// Accessor to the `SpatialSerializer`
     pub fn extract_spatial_serializer(&mut self) -> Option<SpatialSerializer> {
         self.spatial_serializer.take()
     }
diff --git a/src/spatial/bkd.rs b/src/spatial/bkd.rs
index 5cda76182..f4d0829cd 100644
--- a/src/spatial/bkd.rs
+++ b/src/spatial/bkd.rs
@@ -296,6 +296,8 @@ fn write_branch_nodes(
     }
 }
 
+const VERSION: u16 = 1u16;
+
 /// Builds and serializes a block kd-tree for spatial indexing of triangles.
 ///
 /// Takes a collection of triangles and constructs a complete block kd-tree, writing both the
@@ -314,16 +316,14 @@ pub fn write_block_kd_tree(
     triangles: &mut [Triangle],
     write: &mut CountingWriter<WritePtr>,
 ) -> io::Result<()> {
-    assert_eq!(
-        triangles.as_ptr() as usize % std::mem::align_of::<Triangle>(),
-        0
-    );
-    write.write_all(&1u16.to_le_bytes())?;
+    write.write_all(&VERSION.to_le_bytes())?;
+
     let tree = write_leaf_pages(triangles, write)?;
     let current = write.written_bytes();
-    let aligned = (current + 31) & !31;
+    let aligned = current.next_multiple_of(32);
     let padding = aligned - current;
     write.write_all(&vec![0u8; padding as usize])?;
+
     write_leaf_nodes(&tree, write)?;
     let branch_position = write.written_bytes();
     let mut branch_offset: i32 = 0;
@@ -336,15 +336,16 @@ pub fn write_block_kd_tree(
     Ok(())
 }
 
-fn decompress_leaf(data: &[u8]) -> io::Result<Vec<Triangle>> {
-    let count = u16::from_le_bytes([data[0], data[1]]) as usize;
-    let mut offset = 2;
-    let mut triangles = Vec::with_capacity(count);
-    offset += decompress::<u32, _>(&data[offset..], count, |_, doc_id| {
+fn decompress_leaf(mut data: &[u8]) -> io::Result<Vec<Triangle>> {
+    use common::BinarySerializable;
+    let triangle_count: usize = u16::deserialize(&mut data)? as usize;
+    let mut offset: usize = 0;
+    let mut triangles: Vec<Triangle> = Vec::with_capacity(triangle_count);
+    offset += decompress::<u32, _>(&data[offset..], triangle_count, |_, doc_id| {
         triangles.push(Triangle::skeleton(doc_id))
     })?;
     for i in 0..7 {
-        offset += decompress::<i32, _>(&data[offset..], count, |j, word| {
+        offset += decompress::<i32, _>(&data[offset..], triangle_count, |j, word| {
             triangles[j].words[i] = word
         })?;
     }