Adjusting light tree implementation in prep for configurable arity.

src/accel/light_tree.rs

@@ -9,20 +9,51 @@ use shading::surface_closure::SurfaceClosure;
 use super::LightAccel;
 use super::objects_split::sah_split;
 
+// const LEVEL_COLLAPSE: usize = 2;  // Number of levels of the binary tree to collapse together (1 = no collapsing)
+// const ARITY: usize = 1 << LEVEL_COLLAPSE;  // Arity of the final tree
+
 
 #[derive(Copy, Clone, Debug)]
 pub struct LightTree<'a> {
-    nodes: &'a [Node],
+    root: Option<&'a Node<'a>>,
-    bounds: &'a [BBox],
     depth: usize,
 }
 
 #[derive(Copy, Clone, Debug)]
-struct Node {
+enum Node<'a> {
-    is_leaf: bool,
+    Inner {
-    bounds_range: (usize, usize),
+        children: &'a [Node<'a>],
-    energy: f32,
+        bounds: &'a [BBox],
-    child_index: usize,
+        energy: f32,
+    },
+    Leaf {
+        light_index: usize,
+        bounds: &'a [BBox],
+        energy: f32,
+    },
+}
+
+impl<'a> Node<'a> {
+    fn bounds(&self) -> &'a [BBox] {
+        match self {
+            &Node::Inner { ref bounds, .. } => bounds,
+            &Node::Leaf { ref bounds, .. } => bounds,
+        }
+    }
+
+    fn energy(&self) -> f32 {
+        match self {
+            &Node::Inner { energy, .. } => energy,
+            &Node::Leaf { energy, .. } => energy,
+        }
+    }
+
+    fn light_index(&self) -> usize {
+        match self {
+            &Node::Inner { .. } => panic!(),
+            &Node::Leaf { light_index, .. } => light_index,
+        }
+    }
 }
 
 impl<'a> LightTree<'a> {
@@ -34,13 +65,60 @@ impl<'a> LightTree<'a> {
     where
         F: 'b + Fn(&T) -> (&'b [BBox], f32),
     {
-        let mut builder = LightTreeBuilder::new();
+        if objects.len() == 0 {
-        builder.recursive_build(0, 0, objects, &info_getter);
+            LightTree {
+                root: None,
+                depth: 0,
+            }
+        } else {
+            let mut builder = LightTreeBuilder::new();
+            builder.recursive_build(0, 0, objects, &info_getter);
 
-        LightTree {
+            let mut root = unsafe { arena.alloc_uninitialized::<Node>() };
-            nodes: arena.copy_slice(&builder.nodes),
+            LightTree::construct_from_builder(arena, &builder, builder.root_node_index(), root);
-            bounds: arena.copy_slice(&builder.bounds),
+
-            depth: 0,
+            LightTree {
+                root: Some(root),
+                depth: builder.depth,
+            }
+        }
+    }
+
+    fn construct_from_builder(
+        arena: &'a MemArena,
+        base: &LightTreeBuilder,
+        node_index: usize,
+        node_mem: &mut Node<'a>,
+    ) {
+        if base.nodes[node_index].is_leaf {
+            // Leaf
+            let bounds_range = base.nodes[node_index].bounds_range;
+            let bounds = arena.copy_slice(&base.bounds[bounds_range.0..bounds_range.1]);
+
+            *node_mem = Node::Leaf {
+                light_index: base.nodes[node_index].child_index,
+                bounds: bounds,
+                energy: base.nodes[node_index].energy,
+            };
+        } else {
+            // Inner
+            let bounds_range = base.nodes[node_index].bounds_range;
+            let bounds = arena.copy_slice(&base.bounds[bounds_range.0..bounds_range.1]);
+
+            let mut children = unsafe { arena.alloc_array_uninitialized::<Node>(2) };
+            LightTree::construct_from_builder(arena, base, node_index + 1, &mut children[0]);
+            LightTree::construct_from_builder(
+                arena,
+                base,
+                base.nodes[node_index].child_index,
+                &mut children[1],
+            );
+
+            *node_mem = Node::Inner {
+                children: children,
+                bounds: bounds,
+                energy: base.nodes[node_index].energy,
+            };
         }
     }
 }
@@ -56,18 +134,13 @@ impl<'a> LightAccel for LightTree<'a> {
         time: f32,
         n: f32,
     ) -> Option<(usize, f32, f32)> {
-        if self.nodes.len() == 0 {
+        if self.root.is_none() {
             return None;
         }
 
-        let mut node_index = 0;
-        let mut tot_prob = 1.0;
-        let mut n = n;
-
         // Calculates the selection probability for a node
         let node_prob = |node_ref: &Node| {
-            let bounds = &self.bounds[node_ref.bounds_range.0..node_ref.bounds_range.1];
+            let bbox = lerp_slice(node_ref.bounds(), time);
-            let bbox = lerp_slice(bounds, time);
             let d = bbox.center() - pos;
             let dist2 = d.length2();
             let r = bbox.diagonal() * 0.5;
@@ -96,42 +169,49 @@ impl<'a> LightAccel for LightTree<'a> {
                 approx_contrib * fstep
             };
 
-            node_ref.energy * inv_surface_area * approx_contrib
+            node_ref.energy() * inv_surface_area * approx_contrib
         };
 
         // Traverse down the tree, keeping track of the relative probabilities
-        while !self.nodes[node_index].is_leaf {
+        let mut node = self.root.unwrap();
-            // Calculate the relative probabilities of the two children
+        let mut tot_prob = 1.0;
-            let (p1, p2) = {
+        let mut n = n;
-                let p1 = node_prob(&self.nodes[node_index + 1]);
+        loop {
-                let p2 = node_prob(&self.nodes[self.nodes[node_index].child_index]);
+            if let Node::Inner { children, bounds, .. } = *node {
-                let total = p1 + p2;
+                // Calculate the relative probabilities of the two children
+                let (p1, p2) = {
+                    let p1 = node_prob(&children[0]);
+                    let p2 = node_prob(&children[1]);
+                    let total = p1 + p2;
 
-                if total <= 0.0 {
+                    if total <= 0.0 {
-                    (0.5, 0.5)
+                        (0.5, 0.5)
+                    } else {
+                        (p1 / total, p2 / total)
+                    }
+                };
+
+                if n <= p1 {
+                    tot_prob *= p1;
+                    node = &children[0];
+                    n /= p1;
                 } else {
-                    (p1 / total, p2 / total)
+                    tot_prob *= p2;
+                    node = &children[1];
+                    n = (n - p1) / p2;
                 }
-            };
-
-            if n <= p1 {
-                tot_prob *= p1;
-                node_index = node_index + 1;
-                n /= p1;
             } else {
-                tot_prob *= p2;
+                break;
-                node_index = self.nodes[node_index].child_index;
-                n = (n - p1) / p2;
             }
         }
 
         // Found our light!
-        Some((self.nodes[node_index].child_index, tot_prob, n))
+        Some((node.light_index(), tot_prob, n))
     }
 
     fn approximate_energy(&self) -> f32 {
-        if self.nodes.len() > 0 {
+        if let Some(ref node) = self.root {
-            self.nodes[0].energy
+            node.energy()
         } else {
             0.0
         }
@@ -140,11 +220,19 @@ impl<'a> LightAccel for LightTree<'a> {
 
 
 struct LightTreeBuilder {
-    nodes: Vec<Node>,
+    nodes: Vec<BuilderNode>,
     bounds: Vec<BBox>,
     depth: usize,
 }
 
+#[derive(Copy, Clone, Debug)]
+struct BuilderNode {
+    is_leaf: bool,
+    bounds_range: (usize, usize),
+    energy: f32,
+    child_index: usize,
+}
+
 impl LightTreeBuilder {
     fn new() -> LightTreeBuilder {
         LightTreeBuilder {
@@ -154,6 +242,10 @@ impl LightTreeBuilder {
         }
     }
 
+    pub fn root_node_index(&self) -> usize {
+        0
+    }
+
     fn recursive_build<'a, T, F>(
         &mut self,
         offset: usize,
@@ -173,7 +265,7 @@ impl LightTreeBuilder {
             let bi = self.bounds.len();
             let (obj_bounds, energy) = info_getter(&objects[0]);
             self.bounds.extend(obj_bounds);
-            self.nodes.push(Node {
+            self.nodes.push(BuilderNode {
                 is_leaf: true,
                 bounds_range: (bi, self.bounds.len()),
                 energy: energy,
@@ -187,7 +279,7 @@ impl LightTreeBuilder {
             return (me_index, (bi, self.bounds.len()));
         } else {
             // Not a leaf node
-            self.nodes.push(Node {
+            self.nodes.push(BuilderNode {
                 is_leaf: false,
                 bounds_range: (0, 0),
                 energy: 0.0,
@@ -221,7 +313,7 @@ impl LightTreeBuilder {
 
             // Set node
             let energy = self.nodes[me_index + 1].energy + self.nodes[c2_index].energy;
-            self.nodes[me_index] = Node {
+            self.nodes[me_index] = BuilderNode {
                 is_leaf: false,
                 bounds_range: (bi, self.bounds.len()),
                 energy: energy,