octree.rs

extern crate rand;
extern crate ndarray;

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use ndarray::{ArrayView3, array, s, Axis, Ix3, Array3};
use nalgebra::{zero, Vector3};
use rand::Rng;

struct Octree {
    side_length: i32,
    offset: Vector3<i32>,
    has_value: bool,
    children: Option<Box<[Octree; 8]>>,
    value: i32
}

fn ceil_to_pow2(x: i32) -> i32 {
    let mut v = x - 1;
    v |= v >> 1;
    v |= v >> 2;
    v |= v >> 4;
    v |= v >> 8;
    v |= v >> 16;
    v += 1;
    return v;
}

fn create_octree(data: ArrayView3<i32>, offset: Vector3<i32>) -> Octree {
    let x_size = data.len_of(Axis { 0: 0 }) as i32;
    let z_size = data.len_of(Axis { 0: 1 }) as i32;
    let y_size = data.len_of(Axis { 0: 2}) as i32;

    // for now
    assert_eq!(x_size, ceil_to_pow2(x_size));
    assert_eq!(z_size, ceil_to_pow2(z_size));
    assert_eq!(y_size, ceil_to_pow2(y_size));
    assert_eq!(x_size, y_size);
    assert_eq!(y_size, z_size);

    let side_length = x_size;

    let first_val = data[Ix3(0, 0, 0)];
    let mut split = false;

    for val in data {
        if *val != first_val {
            split = true;
            break;
        }
    }

    if split {
        fn create_child(old_data: ArrayView3<i32>, offset: Vector3<i32>, side_length: i32) -> Octree  {
            let new_data = old_data.slice(s!(offset[0]..(offset[0] + side_length), offset[1]..(offset[1] + side_length), offset[2]..(offset[2] + side_length)));
            create_octree(new_data, offset)
        };

        let new_side_length = side_length/2;
        Octree {
            side_length,
            offset,
            has_value: false,
            children: Some(Box::new([
                create_child(data, Vector3::new(0, 0, 0), new_side_length),
                create_child(data, Vector3::new(new_side_length, 0, 0), new_side_length),
                create_child(data, Vector3::new(0, new_side_length, 0), new_side_length),
                create_child(data, Vector3::new(0, 0, new_side_length), new_side_length),
                create_child(data, Vector3::new(new_side_length, new_side_length, 0), new_side_length),
                create_child(data, Vector3::new(new_side_length, 0, new_side_length), new_side_length),
                create_child(data, Vector3::new(0, new_side_length, new_side_length), new_side_length),
                create_child(data, Vector3::new(new_side_length, new_side_length, new_side_length), new_side_length),
            ])),
            value: 0
        }
    }
    else {
        Octree {
            side_length,
            offset,
            has_value: true,
            children: None,
            value: first_val
        }
    }
}

fn criterion_benchmark(c: &mut Criterion) {
    let mut rng = rand::thread_rng();
    const SIDE_LENGTH: usize = 64;
    let mut data : Array3<i32> = Array3::zeros((SIDE_LENGTH, SIDE_LENGTH, SIDE_LENGTH));
    for d1 in data.iter_mut() {
        *d1 = rng.gen();
    }
    let data_view = data.view();
    c.bench_function(concat!("octree ", stringify!(IDE_LENGTH)), |b| b.iter(|| create_octree(black_box(data_view), zero())));
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);