Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
use std::fmt::Debug;
use super::{Predictor, predict::WeightedPredictorState};
use crate::{
error::Result,
frame::modular::{
Tree,
predict::PredictionData,
tree::{PredictionResult, TreeNode, compute_properties},
},
image::Image,
util::NewWithCapacity,
};
/// Flattened tree node for optimized traversal.
/// Stores parent + info about both children to evaluate 3 nodes per iteration.
#[derive(Debug, Clone, Copy)]
pub(super) enum FlatTreeNode {
Split {
properties: [u8; 3],
splitvals: [i32; 3],
child_id: u32,
},
Leaf {
predictor: Predictor,
multiplier: u32,
context: u32,
offset: i32,
},
}
#[inline(always)]
#[allow(clippy::too_many_arguments, unsafe_code)]
pub(super) fn predict_flat(
flat_tree: &[FlatTreeNode],
prediction_data: PredictionData,
wp_state: Option<&mut WeightedPredictorState>,
pos: (usize, usize),
references: &Image<i32>,
property_buffer: &mut [i32; 256],
) -> PredictionResult {
let wp_pred = compute_properties(
prediction_data,
wp_state,
pos.0,
pos.1,
references,
property_buffer,
);
let mut pos = 0;
loop {
// Removing this bound check doesn't seem to have a significant effect.
let node = flat_tree[pos];
match node {
FlatTreeNode::Split {
properties,
splitvals,
child_id,
} => {
// This bound check is elided by virtue of `property_buffer` having 256 elements.
let props = properties.map(|x| property_buffer[x as usize]);
let p0 = props[0] <= splitvals[0];
let p1 = props[1] <= splitvals[1];
let p2 = props[2] <= splitvals[2];
pos = child_id as usize + if p0 { 2 | p2 as usize } else { p1 as usize };
}
FlatTreeNode::Leaf {
predictor,
multiplier,
context,
offset,
} => {
let pred = predictor.predict_one(prediction_data, wp_pred);
return PredictionResult {
guess: pred + offset as i64,
multiplier,
context,
};
}
};
}
}
impl Tree {
/// Build flat tree using BFS traversal.
/// Each flat node stores parent + both children info to reduce branches.
pub(super) fn build_flat_tree(nodes: &[TreeNode]) -> Result<Vec<FlatTreeNode>> {
use std::collections::VecDeque;
if nodes.is_empty() {
return Ok(vec![]);
}
let mut flat_nodes = Vec::new_with_capacity(nodes.len())?;
let mut queue: VecDeque<usize> = VecDeque::new();
queue.push_back(0); // Start with root
while let Some(cur_idx) = queue.pop_front() {
match nodes[cur_idx] {
TreeNode::Leaf {
predictor,
offset,
multiplier,
id,
} => {
flat_nodes.push(FlatTreeNode::Leaf {
predictor,
offset,
multiplier,
context: id,
});
}
TreeNode::Split {
property,
val,
left,
right,
} => {
// childID points to first of 4 grandchildren in output
let child_id = (flat_nodes.len() + queue.len() + 1) as u32;
let mut splitvals = [val, 0, 0];
let mut properties = [property, 0, 0];
// Process left (i=0) and right (i=1) children
for (i, &child_idx) in [left as usize, right as usize].iter().enumerate() {
match &nodes[child_idx] {
TreeNode::Leaf { .. } => {
// Child is leaf: enqueue leaf twice
queue.push_back(child_idx);
queue.push_back(child_idx);
}
TreeNode::Split {
property: cp,
val: cv,
left: cl,
right: cr,
} => {
// Child is split: store property/splitval and enqueue grandchildren
properties[i + 1] = *cp;
splitvals[i + 1] = *cv;
queue.push_back(*cl as usize);
queue.push_back(*cr as usize);
}
}
}
flat_nodes.push(FlatTreeNode::Split {
properties,
splitvals,
child_id,
});
}
}
}
Ok(flat_nodes)
}
}