refactor: codec abstraction and benchmarks

.github/workflows/bench.yml

@@ -0,0 +1,40 @@
+name: test
+
+on:
+  pull_request:
+    types: [ labeled ]
+    branches: [ "main" ]
+
+env:
+  CARGO_TERM_COLOR: always
+  RUST_BACKTRACE: 1
+
+jobs:
+  bench:
+    if: ${{ github.event.label.name == 'bench' }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [windows-latest, macos-latest]
+
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/cache@v3
+      with:
+        path: |
+          ~/.cargo/bin/
+          ~/.cargo/registry/index/
+          ~/.cargo/registry/cache/
+          ~/.cargo/git/db/
+          target/
+        key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }}
+
+    - name: criterion compare
+      uses: boa-dev/[email protected]
+      with:
+        branchName: ${{ github.base_ref }}
+        token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/test.yml

@@ -33,9 +33,15 @@ jobs:
           target/
         key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }}
 
+    - name: check
+      run: cargo check
+
     - name: build
       run: cargo build
 
+    - name: build_tools
+      run: cargo build --bin ply_to_gcloud
+
     - name: lint
       run: cargo clippy
 

Cargo.toml

@@ -14,6 +14,14 @@ exclude = [".devcontainer", ".github", "docs", "dist", "build", "assets", "credi
 default-run = "viewer"
 
 
+[features]
+default = ["io_flexbuffers"]
+
+io_bincode2 = []
+io_capnp = []
+io_flexbuffers = []
+
+
 # TODO: separate dev-dependencies
 [dependencies]
 bevy-inspector-egui = "0.21"
@@ -49,6 +57,10 @@ features = [
 ]
 
 
+[dev-dependencies]
+criterion = { version = "0.5", features = ["html_reports"] }
+
+
 [profile.dev.package."*"]
 opt-level = 3
 
@@ -77,3 +89,8 @@ path = "viewer/viewer.rs"
 [[bin]]
 name = "ply_to_gcloud"
 path = "tools/ply_to_gcloud.rs"
+
+
+[[bench]]
+name = "io"
+harness = false

benches/io.rs

@@ -0,0 +1,47 @@
+use criterion::{
+    BenchmarkId,
+    criterion_group,
+    criterion_main,
+    Criterion,
+    Throughput,
+};
+
+use bevy_gaussian_splatting::{
+    Gaussian,
+    GaussianCloud,
+    io::codec::GaussianCloudCodec,
+    random_gaussians,
+};
+
+
+const GAUSSIAN_COUNTS: [usize; 4] = [
+    1000,
+    10000,
+    84_348,
+    1_244_819,
+    // 6_131_954,
+];
+
+fn gaussian_cloud_decode_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("encode gaussian clouds");
+    for count in GAUSSIAN_COUNTS.iter() {
+        group.throughput(Throughput::Bytes(*count as u64 * std::mem::size_of::<Gaussian>() as u64));
+        group.bench_with_input(
+            BenchmarkId::new("decode", count),
+            &count,
+            |b, &count| {
+                let gaussians = random_gaussians(*count);
+                let bytes = gaussians.encode();
+
+                b.iter(|| GaussianCloud::decode(bytes.as_slice()));
+            },
+        );
+    }
+}
+
+criterion_group!{
+    name = io_benches;
+    config = Criterion::default().sample_size(10);
+    targets = gaussian_cloud_decode_benchmark
+}
+criterion_main!(io_benches);

src/gaussian.rs

@@ -1,39 +1,26 @@
-use rand::{seq::SliceRandom, prelude::Distribution, Rng};
-use std::{
-    io::{
-        BufReader,
-        Cursor,
-        ErrorKind,
-    },
-    marker::Copy,
+use rand::{
+    seq::SliceRandom,
+    prelude::Distribution,
+    Rng,
 };
+use std::marker::Copy;
 
 use bevy::{
     prelude::*,
-    asset::{
-        AssetLoader,
-        AsyncReadExt,
-        LoadContext,
-        io::Reader,
-    },
     reflect::TypeUuid,
     render::render_resource::ShaderType,
-    utils::BoxedFuture,
 };
 use bytemuck::{
     Pod,
     Zeroable,
 };
-use flexbuffers::Reader as FlexReader;
 use serde::{
     Deserialize,
     Serialize,
     Serializer,
     ser::SerializeTuple,
 };
 
-use crate::ply::parse_ply;
-
 
 const fn num_sh_coefficients(degree: usize) -> usize {
     if degree == 0 {
@@ -230,53 +217,6 @@ impl Default for GaussianCloudSettings {
     }
 }
 
-
-#[derive(Default)]
-pub struct GaussianCloudLoader;
-
-impl AssetLoader for GaussianCloudLoader {
-    type Asset = GaussianCloud;
-    type Settings = ();
-    type Error = std::io::Error;
-
-    fn load<'a>(
-        &'a self,
-        reader: &'a mut Reader,
-        _settings: &'a Self::Settings,
-        load_context: &'a mut LoadContext,
-    ) -> BoxedFuture<'a, Result<Self::Asset, Self::Error>> {
-
-        Box::pin(async move {
-            let mut bytes = Vec::new();
-            reader.read_to_end(&mut bytes).await?;
-
-            match load_context.path().extension() {
-                Some(ext) if ext == "ply" => {
-                    let cursor = Cursor::new(bytes);
-                    let mut f = BufReader::new(cursor);
-
-                    let ply_cloud = parse_ply(&mut f)?;
-                    let cloud = GaussianCloud(ply_cloud);
-
-                    Ok(cloud)
-                },
-                Some(ext) if ext == "gcloud" => {
-                    let reader = FlexReader::get_root(bytes.as_slice()).expect("failed to read flexbuffer");
-                    let cloud = GaussianCloud::deserialize(reader).expect("deserialization failed");
-
-                    Ok(cloud)
-                },
-                _ => Err(std::io::Error::new(ErrorKind::Other, "only .ply and .gcloud supported")),
-            }
-        })
-    }
-
-    fn extensions(&self) -> &[&str] {
-        &["ply", "gcloud"]
-    }
-}
-
-
 impl Distribution<Gaussian> for rand::distributions::Standard {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Gaussian {
         Gaussian {

src/io/codec.rs

@@ -0,0 +1,6 @@
+
+// TODO: support streamed codecs
+pub trait GaussianCloudCodec {
+    fn encode(&self) -> Vec<u8>;
+    fn decode(data: &[u8]) -> Self;
+}

src/io/gcloud/capnp.rs

@@ -0,0 +1,2 @@
+
+

src/io/gcloud/flexbuffers.rs

@@ -0,0 +1,30 @@
+use flexbuffers::{
+    FlexbufferSerializer,
+    Reader,
+};
+use serde::{
+    Deserialize,
+    Serialize,
+};
+
+use crate::{
+    GaussianCloud,
+    io::codec::GaussianCloudCodec,
+};
+
+
+impl GaussianCloudCodec for GaussianCloud {
+    fn encode(&self) -> Vec<u8> {
+        let mut serializer = FlexbufferSerializer::new();
+        self.serialize(&mut serializer).expect("failed to serialize cloud");
+
+        serializer.view().to_vec()
+    }
+
+    fn decode(data: &[u8]) -> Self {
+        let reader = Reader::get_root(data).expect("failed to read flexbuffer");
+        let cloud = GaussianCloud::deserialize(reader).expect("deserialization failed");
+
+        cloud
+    }
+}

src/io/gcloud/mod.rs

@@ -0,0 +1,8 @@
+#[cfg(feature = "io_bincode2")]
+pub mod bincode2;
+
+#[cfg(feature = "io_capnp")]
+pub mod capnp;
+
+#[cfg(feature = "io_flexbuffers")]
+pub mod flexbuffers;

src/io/loader.rs

@@ -0,0 +1,68 @@
+use std::io::{
+    BufReader,
+    Cursor,
+    ErrorKind,
+};
+
+use bevy::{
+    asset::{
+        AssetLoader,
+        AsyncReadExt,
+        LoadContext,
+        io::Reader,
+    },
+    utils::BoxedFuture,
+};
+
+use crate::{
+    GaussianCloud,
+    io::{
+        codec::GaussianCloudCodec,
+        ply::parse_ply,
+    },
+};
+
+
+#[derive(Default)]
+pub struct GaussianCloudLoader;
+
+impl AssetLoader for GaussianCloudLoader {
+    type Asset = GaussianCloud;
+    type Settings = ();
+    type Error = std::io::Error;
+
+    fn load<'a>(
+        &'a self,
+        reader: &'a mut Reader,
+        _settings: &'a Self::Settings,
+        load_context: &'a mut LoadContext,
+    ) -> BoxedFuture<'a, Result<Self::Asset, Self::Error>> {
+
+        Box::pin(async move {
+            let mut bytes = Vec::new();
+            reader.read_to_end(&mut bytes).await?;
+
+            match load_context.path().extension() {
+                Some(ext) if ext == "ply" => {
+                    let cursor = Cursor::new(bytes);
+                    let mut f = BufReader::new(cursor);
+
+                    let ply_cloud = parse_ply(&mut f)?;
+                    let cloud = GaussianCloud(ply_cloud);
+
+                    Ok(cloud)
+                },
+                Some(ext) if ext == "gcloud" => {
+                    let cloud = GaussianCloud::decode(bytes.as_slice());
+
+                    Ok(cloud)
+                },
+                _ => Err(std::io::Error::new(ErrorKind::Other, "only .ply and .gcloud supported")),
+            }
+        })
+    }
+
+    fn extensions(&self) -> &[&str] {
+        &["ply", "gcloud"]
+    }
+}

src/io/mod.rs

@@ -0,0 +1,4 @@
+pub mod codec;
+pub mod gcloud;
+pub mod loader;
+pub mod ply;

src/lib.rs

@@ -3,16 +3,17 @@ use bevy::prelude::*;
 pub use gaussian::{
     Gaussian,
     GaussianCloud,
-    GaussianCloudLoader,
     GaussianCloudSettings,
     SphericalHarmonicCoefficients,
     random_gaussians,
 };
 
+use io::loader::GaussianCloudLoader;
+
 use render::RenderPipelinePlugin;
 
 pub mod gaussian;
-pub mod ply;
+pub mod io;
 pub mod render;
 pub mod utils;