use crate::DecodingError;
use bitvec::prelude::*;
use h3o::{CellIndex, Resolution};
use std::io::{self, Write};
const END_OF_TREE_MARKER: u8 = 0;
const NODE_BITSIZE: usize = 7;
const MAX_PATH_LEN: usize = 15 + 1;
const BITRATE: usize = 6;
pub fn encode<W: Write>(
writer: &mut W,
cells: impl IntoIterator<Item = CellIndex>,
) -> Result<(), io::Error> {
let iter = cells.into_iter();
let (lower_bound, upper_bound) = iter.size_hint();
let estimated_size = upper_bound.unwrap_or(lower_bound) * BITRATE;
let mut tree = <BitVec<u8, Lsb0>>::with_capacity(estimated_size);
let mut path: Vec<(Step, usize)> = Vec::with_capacity(MAX_PATH_LEN);
path.push((Step::BaseCell(124), 0));
for index in iter {
for (level, cell) in steps_from_cell_index(index).enumerate() {
match path.get(level) {
Some(&(current_cell, _)) => {
if cell != current_cell {
path.truncate(level + 1);
add_cell(&mut tree, &mut path, Some((cell, level)));
}
}
None => add_cell(&mut tree, &mut path, Some((cell, level))),
}
}
add_cell(&mut tree, &mut path, None);
}
tree.extend_from_raw_slice(&[tag(END_OF_TREE_MARKER)]);
writer.write_all(tree.as_raw_slice())
}
fn add_cell(
tree: &mut BitVec<u8, Lsb0>,
path: &mut Vec<(Step, usize)>,
new_entry: Option<(Step, usize)>,
) {
match new_entry {
Some((cell @ Step::BaseCell(id), _)) => {
path[0] = (cell, tree.len());
tree.extend_from_raw_slice(&[tag(id + 1)]);
}
Some((direction @ Step::Direction(id), level)) => {
if let Some(&mut (ref mut current_cell, index)) =
path.get_mut(level)
{
*current_cell = direction;
tree.set(index + 1 + usize::from(id), true);
} else {
path.push((direction, tree.len()));
tree.extend_from_raw_slice(&[tag(1 << id)]);
}
}
None => tree.push(false),
}
}
const fn tag(value: u8) -> u8 {
value << 1 | 1
}
pub fn decode(bytes: &[u8]) -> Iter<'_> {
Iter::new(bytes)
}
pub struct Iter<'a> {
tree: &'a BitSlice<u8, Lsb0>,
position: usize,
path: Vec<Node>,
}
impl<'a> Iter<'a> {
fn new(bytes: &'a [u8]) -> Self {
Self {
tree: BitSlice::from_slice(bytes),
position: 0,
path: Vec::with_capacity(MAX_PATH_LEN),
}
}
fn next_object(&mut self) -> Result<Object, DecodingError> {
let is_node = self
.tree
.get(self.position)
.ok_or_else(|| DecodingError::missing_tag(self.position))?;
self.position += 1;
if !is_node {
return Ok(Object::Leaf);
}
let bits = self
.tree
.get(self.position..self.position + NODE_BITSIZE)
.ok_or_else(DecodingError::not_enough_data)?;
let bits = bits.load_le::<u8>();
self.position += NODE_BITSIZE;
if bits == END_OF_TREE_MARKER {
return Ok(Object::EndOfTree);
}
Ok(if self.path.is_empty() {
Node::Root(bits - 1).into()
} else {
Node::from(Children(bits)).into()
})
}
fn remove_index_from_path(&mut self) {
for i in (0..self.path.len()).rev() {
#[allow(clippy::match_on_vec_items)] match self.path[i] {
Node::Root(_) => {
self.path.clear();
break;
}
Node::Intermediate(ref mut node) => {
node.remove();
if !node.is_empty() {
break;
}
self.path.pop();
}
}
}
}
}
impl<'a> Iterator for Iter<'a> {
type Item = Result<CellIndex, DecodingError>;
fn next(&mut self) -> Option<Self::Item> {
let mut steps = Vec::with_capacity(MAX_PATH_LEN);
loop {
match self.next_object() {
Ok(Object::EndOfTree) => return None,
Ok(Object::Node(node)) => self.path.push(node),
Ok(Object::Leaf) => {
steps.clear();
for cell in self.path.iter().copied().map(From::from) {
steps.push(cell);
}
let index = cell_index_from_steps(steps.as_slice());
self.remove_index_from_path();
return Some(index);
}
Err(err) => return Some(Err(err)),
}
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Step {
BaseCell(u8),
Direction(u8),
}
fn steps_from_cell_index(index: CellIndex) -> impl Iterator<Item = Step> {
std::iter::once(Step::BaseCell(index.base_cell().into())).chain(
Resolution::range(Resolution::One, index.resolution()).map(
move |resolution| {
Step::Direction(
index.direction_at(resolution).expect("direction").into(),
)
},
),
)
}
fn cell_index_from_steps(steps: &[Step]) -> Result<CellIndex, DecodingError> {
let mut index = 0x8001fffffffffff;
let (base_cell, directions) = match steps {
&[head, ref tail @ ..] => {
if tail.len() > Resolution::Fifteen.into() {
return Err(DecodingError::bad_index(
"too many steps in path",
None,
));
}
(head, tail)
}
_ => return Err(DecodingError::bad_index("empty path", None)),
};
#[allow(clippy::cast_possible_truncation)]
let resolution = directions.len() as u8;
index |= u64::from(resolution) << 52;
if let Step::BaseCell(cell) = base_cell {
index |= u64::from(cell) << 45;
} else {
unreachable!("missing base cell");
}
for (i, step) in directions.iter().enumerate() {
if let Step::Direction(direction) = *step {
#[allow(clippy::cast_possible_truncation)]
let resolution = (i + 1) as u8;
let offset = (15 - resolution) * 3;
index =
(index & !(0b111 << offset)) | (u64::from(direction) << offset);
} else {
unreachable!("more than one base cell");
}
}
CellIndex::try_from(index).map_err(|err| {
DecodingError::bad_index("invalid cell index", Some(err))
})
}
#[derive(Clone, Copy, Debug)]
struct Children(u8);
impl Children {
const fn peek(self) -> u8 {
self.first_child_index()
}
fn remove(&mut self) {
self.0 &= !(1 << self.first_child_index());
}
const fn is_empty(self) -> bool {
self.0 == 0
}
#[allow(clippy::cast_possible_truncation)] const fn first_child_index(self) -> u8 {
self.0.trailing_zeros() as u8
}
}
impl From<Children> for Node {
fn from(value: Children) -> Self {
Self::Intermediate(value)
}
}
#[derive(Clone, Copy, Debug)]
enum Node {
Root(u8),
Intermediate(Children),
}
impl From<Node> for Object {
fn from(value: Node) -> Self {
Self::Node(value)
}
}
impl From<Node> for Step {
fn from(value: Node) -> Self {
match value {
Node::Root(cell) => Self::BaseCell(cell),
Node::Intermediate(children) => Self::Direction(children.peek()),
}
}
}
#[derive(Clone, Copy, Debug)]
enum Object {
Node(Node),
Leaf,
EndOfTree,
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
macro_rules! cells {
($($x: expr),*) => {{
vec![
$(CellIndex::try_from($x).expect("valid cell"),)*
]
}}
}
fn run_encode(input: &[u64]) -> BitVec<u8, Lsb0> {
let mut buffer = Cursor::new(vec![]);
let cells = input
.iter()
.map(|value| CellIndex::try_from(*value).expect("valid cell"));
encode(&mut buffer, cells).expect("encode");
BitVec::from_vec(buffer.into_inner())
}
#[test]
fn encode_base_cell() {
#[rustfmt::skip]
let expected = bitvec![
1, 0,1,1,0,1,0,0, 0, 1, 0,0,0,0,0,0,0, ];
let result = &run_encode(&[0x802bfffffffffff])[..expected.len()];
assert_eq!(result, expected);
}
#[test]
fn encode_one_cell() {
#[rustfmt::skip]
let expected = bitvec![
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,0,0,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 0,0,0,0,0,0,0, ];
let result = &run_encode(&[0x8c2bae305336bff])[..expected.len()];
assert_eq!(result, expected);
}
#[test]
fn encode_siblings() {
#[rustfmt::skip]
let expected = bitvec![
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,0,0,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 1, 0,0,0,0,0,0,0, ];
let result = &run_encode(&[0x8c2bae305336bff, 0x8c2bae305336dff])
[..expected.len()];
assert_eq!(result, expected);
}
#[test]
fn encode_leaf_midway() {
#[rustfmt::skip]
let expected = bitvec![
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,1,0,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 0, 1, 0,0,0,0,0,0,0, ];
let result = &run_encode(&[
0x8c2bae305336bff,
0x8c2bae305336dff,
0x862bae30fffffff,
])[..expected.len()];
assert_eq!(result, expected);
}
#[test]
fn encode_two_branches() {
#[rustfmt::skip]
let expected = bitvec![
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,1,1,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 0,0,0,0,0,0,0, ];
let result = &run_encode(&[
0x8c2bae305336bff,
0x8c2bae305336dff,
0x862bae30fffffff,
0x8c2bae315336bff,
])[..expected.len()];
assert_eq!(result, expected);
}
#[test]
fn encode_two_roots() {
#[rustfmt::skip]
let expected = bitvec![
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,1,1,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 1,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,1,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 0,0,0,0,0,0,0, ];
let result = &run_encode(&[
0x8c2bae305336bff,
0x8c2bae305336dff,
0x862bae30fffffff,
0x8c2bae315336bff,
0x8c2dae315336bff,
])[..expected.len()];
assert_eq!(result, expected);
}
#[test]
fn decode_base_cell() {
let expected_cells = cells![0x802bfffffffffff];
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>().expect("valid input");
assert_eq!(cells, expected_cells);
}
#[test]
fn decode_one_cell() {
let expected_cells = cells![0x8c2bae305336bff];
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,0,0,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>().expect("valid input");
assert_eq!(cells, expected_cells);
}
#[test]
fn decode_siblings() {
let expected_cells = cells![0x8c2bae305336bff, 0x8c2bae305336dff];
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,0,0,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>().expect("valid input");
assert_eq!(cells, expected_cells);
}
#[test]
fn decode_leaf_midway() {
let expected_cells =
cells![0x8c2bae305336bff, 0x8c2bae305336dff, 0x862bae30fffffff];
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,1,0,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>().expect("valid input");
assert_eq!(cells, expected_cells);
}
#[test]
fn decode_two_branches() {
let expected_cells = cells![
0x8c2bae305336bff,
0x8c2bae305336dff,
0x862bae30fffffff,
0x8c2bae315336bff
];
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,1,1,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>().expect("valid input");
assert_eq!(cells, expected_cells);
}
#[test]
fn decode_two_roots() {
let expected_cells = cells![
0x8c2bae305336bff,
0x8c2bae305336dff,
0x862bae30fffffff,
0x8c2bae315336bff,
0x8c2dae315336bff
];
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 1,1,1,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,1, 0, 0, 0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 1,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 1, 0,0,0,0,0,0,1, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,1,0,0,0,0, 1, 0,0,0,0,0,1,0, 1, 0,1,0,0,0,0,0, 1, 0,0,0,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>().expect("valid input");
assert_eq!(cells, expected_cells);
}
#[test]
fn decode_invalid_base_cell() {
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 1,1,1,1,1,1,1, 0, 1, 0,0,0,0,0,0,0, ].as_raw_slice()).collect::<Result<Vec<_>, _>>();
assert!(cells.is_err());
}
#[test]
fn decode_corrupted() {
#[rustfmt::skip]
let cells = decode(bitvec![u8, Lsb0;
1, 0,1,1,0,1,0,0, 1, 0,0,0,0,0,0,1, 1, 0,0,0,0,0,1,0, 0, 0,0,0,1 ].as_raw_slice()).collect::<Result<Vec<_>, _>>();
assert!(cells.is_err());
}
#[test]
fn decode_empty() {
let cells = decode(vec![].as_slice()).collect::<Result<Vec<_>, _>>();
assert!(cells.is_err());
}
#[test]
fn rountrip_full_node() {
let cells = cells![
0x8b184584a21efff,
0x8b184584a246fff,
0x8b184584a2a8fff,
0x8b184584a2cbfff,
0x8b184584a329fff,
0x8b184584a366fff,
0x8b184584a389fff
];
let mut buffer = Cursor::new(vec![]);
encode(&mut buffer, cells.clone()).expect("encode");
let bytes = buffer.into_inner();
let result = decode(&bytes)
.collect::<Result<Vec<_>, _>>()
.expect("valid input");
assert_eq!(result, cells);
}
#[test]
fn rountrip_base_cell_0() {
let cells = cells![0x8b0153a16521fff];
let mut buffer = Cursor::new(vec![]);
encode(&mut buffer, cells.clone()).expect("encode");
let bytes = buffer.into_inner();
let result = decode(&bytes)
.collect::<Result<Vec<_>, _>>()
.expect("valid input");
assert_eq!(result, cells);
}
#[test]
fn decode_missing_tag() {
let bytes = [0x81, 0x23];
let error = decode(bytes.as_slice())
.collect::<Result<Vec<_>, _>>()
.expect_err("invalid input");
assert!(matches!(error, DecodingError::MissingTag(16)));
}
#[test]
fn decode_truncated_input() {
let bytes = [0x91, 0x9d, 0xe6];
let error = decode(bytes.as_slice())
.collect::<Result<Vec<_>, _>>()
.expect_err("invalid input");
assert!(matches!(error, DecodingError::NotEnoughData));
}
#[test]
fn decode_cell_too_long() {
let bytes = [
0x27, 0x9b, 0x94, 0x94, 0x94, 0x94, 0x94, 0x94, 0x9b, 0x9b, 0x9b,
0x9b, 0x9b, 0x7a, 0x91, 0x23, 0xe3, 0x01, 0xce,
];
let error = decode(bytes.as_slice())
.collect::<Result<Vec<_>, _>>()
.expect_err("invalid input");
assert!(matches!(error, DecodingError::InvalidCellIndex { .. }));
}
}