parallel bulk insert

2 files changed, 82 insertions, 14 deletions

diff --git a/mtree-test/src/bin/encode.rs b/mtree-test/src/bin/encode.rs
index 5a166c8..80f3499 100644
--- a/mtree-test/src/bin/encode.rs
+++ b/mtree-test/src/bin/encode.rs
@@ -1,6 +1,10 @@
 use framework::{BitstreamFilter, bitstream_filter_main};
 use glam::{I16Vec2, i16vec2};
 use mtree_test::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
+use rayon::{
+    iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator},
+    join,
+};
 use std::{collections::VecDeque, io::Result, time::Instant};
 
 fn main() -> Result<()> {
@@ -11,6 +15,9 @@ struct Enc {
     res: I16Vec2,
     last: LastFrames,
     tree: MTree,
+    frame_num: u64,
+
+    use_bulk_insert: bool,
 }
 impl BitstreamFilter for Enc {
     const INPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
@@ -27,6 +34,8 @@ impl BitstreamFilter for Enc {
                 off: i16vec2(0, 0),
                 frame: u64::MAX,
             }),
+            frame_num: 0,
+            use_bulk_insert: true,
         }
     }
     fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
@@ -51,7 +60,7 @@ impl BitstreamFilter for Enc {
                         boff,
                     )
                 };
-                if rdist < (BLOCK_SIZE * BLOCK_SIZE * 20) as u32 {
+                if rdist < BLOCK_SIZE as u32 * BLOCK_SIZE as u32 * 50 {
                     out.push(1);
                     out.extend(r.frame.to_le_bytes());
                     out.extend(r.off.x.to_le_bytes());
@@ -76,23 +85,39 @@ impl BitstreamFilter for Enc {
         self.last.frames.push_back(frame);
         let frame_index = self.last.frame_offset + self.last.frames.len() as u64 - 1;
 
-        let t = Instant::now();
-        for xo in 0..self.res.x - BLOCK_SIZE {
-            for yo in 0..self.res.y - BLOCK_SIZE {
-                self.tree.insert(
-                    self.res,
-                    &self.last,
-                    AbsRef {
-                        frame: frame_index,
-                        off: i16vec2(xo, yo),
-                    },
-                )
+        if self.frame_num % 30 == 0 {
+            let t = Instant::now();
+            if self.use_bulk_insert {
+                let mut rs = Vec::new();
+                for xo in 0..self.res.x - BLOCK_SIZE {
+                    for yo in 0..self.res.y - BLOCK_SIZE {
+                        rs.push(AbsRef {
+                            frame: frame_index,
+                            off: i16vec2(xo, yo),
+                        });
+                    }
+                }
+                self.tree.bulk_insert(self.res, &self.last, rs);
+            } else {
+                for xo in 0..self.res.x - BLOCK_SIZE {
+                    for yo in 0..self.res.y - BLOCK_SIZE {
+                        self.tree.insert(
+                            self.res,
+                            &self.last,
+                            AbsRef {
+                                frame: frame_index,
+                                off: i16vec2(xo, yo),
+                            },
+                        )
+                    }
+                }
             }
+            eprintln!("insert\t{:?}", t.elapsed());
         }
-        eprintln!("insert\t{:?}", t.elapsed());
 
         eprintln!("depth={} refs={num_refs}", self.tree.depth());
 
+        self.frame_num += 1;
         out
     }
 }
@@ -114,6 +139,49 @@ impl MTree {
             MTree::Leaf(r) => *r,
         }
     }
+    pub fn bulk_insert(&mut self, res: I16Vec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
+        match self {
+            MTree::Branch(b) => {
+                let dirs = rs
+                    .par_iter()
+                    .map(|r| {
+                        distance_absref(res, last, b[0].center(), *r)
+                            < distance_absref(res, last, b[1].center(), *r)
+                    })
+                    .collect::<Vec<bool>>();
+
+                let mut rs0 = Vec::new();
+                let mut rs1 = Vec::new();
+                for (r, d) in rs.into_iter().zip(dirs) {
+                    if d {
+                        rs0.push(r);
+                    } else {
+                        rs1.push(r);
+                    }
+                }
+                let [b0, b1] = b.as_mut();
+                join(
+                    || b0.bulk_insert(res, last, rs0),
+                    || b1.bulk_insert(res, last, rs1),
+                );
+            }
+            MTree::Leaf(l) => {
+                if rs.is_empty() {
+                    return;
+                }
+                if l.frame == u64::MAX {
+                    *l = rs.pop().unwrap();
+                    if rs.is_empty() {
+                        return;
+                    }
+                }
+                let mut b =
+                    MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(rs.pop().unwrap())]));
+                b.bulk_insert(res, last, rs);
+                *self = b;
+            }
+        }
+    }
     pub fn insert(&mut self, res: I16Vec2, last: &LastFrames, r: AbsRef) {
         match self {
             MTree::Branch(b) => {
diff --git a/mtree-test/src/lib.rs b/mtree-test/src/lib.rs
index bf0117d..fd282db 100644
--- a/mtree-test/src/lib.rs
+++ b/mtree-test/src/lib.rs
@@ -1,7 +1,7 @@
 use glam::{I16Vec2, i16vec2};
 use std::{collections::VecDeque, ops::Range};
 
-pub const BLOCK_SIZE: i16 = 4;
+pub const BLOCK_SIZE: i16 = 8;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct AbsRef {