use crate::{
core::Buffer,
errors::{BufferError, TransformError},
geometry::{Localized, Quaternion, Transform, Vector3},
time::{TimePoint, Timestamp},
};
use alloc::{
collections::{BTreeSet, VecDeque},
string::String,
};
use hashbrown::HashMap;
use core::time::Duration;
#[derive(Debug)]
pub struct Registry<T = Timestamp>
where
T: TimePoint,
{
data: HashMap<String, Buffer<T>>,
max_age: Option<Duration>,
}
impl<T> Registry<T>
where
T: TimePoint,
{
#[must_use]
pub fn new() -> Self {
Self {
data: HashMap::new(),
max_age: None,
}
}
#[must_use]
pub fn with_max_age(max_age: Duration) -> Self {
Self {
data: HashMap::new(),
max_age: Some(max_age),
}
}
pub fn add_transform(
&mut self,
t: Transform<T>,
) -> Result<(), BufferError> {
Self::process_add_transform(t, &mut self.data, self.max_age)
}
pub fn get_transform(
&self,
from: &str,
to: &str,
timestamp: T,
) -> Result<Transform<T>, TransformError> {
Self::process_get_transform(from, to, timestamp, &self.data)
}
pub fn get_transform_for<U>(
&self,
value: &U,
target_frame: &str,
) -> Result<Transform<T>, TransformError>
where
U: Localized<T>,
{
self.get_transform(target_frame, value.frame(), value.timestamp())
}
pub fn get_transform_at(
&self,
target_frame: &str,
target_time: T,
source_frame: &str,
source_time: T,
fixed_frame: &str,
) -> Result<Transform<T>, TransformError> {
Self::process_get_transform_at(
target_frame,
target_time,
source_frame,
source_time,
fixed_frame,
&self.data,
)
}
pub fn delete_transforms_before(
&mut self,
timestamp: T,
) {
for buffer in self.data.values_mut() {
buffer.delete_before(timestamp);
}
self.data.retain(|_, buffer| !buffer.is_empty());
}
pub fn remove_frame(
&mut self,
child: &str,
) -> bool {
self.data.remove(child).is_some()
}
fn process_add_transform(
t: Transform<T>,
data: &mut HashMap<String, Buffer<T>>,
max_age: Option<Duration>,
) -> Result<(), BufferError> {
if !data.contains_key(&t.child) && Self::creates_cycle(&t.child, &t.parent, data) {
return Err(BufferError::CycleDetected);
}
if let Some(buffer) = data.get_mut(&t.child) {
return buffer.insert(t);
}
let mut buffer = match max_age {
Some(max_age) => Buffer::with_max_age(max_age),
None => Buffer::new(),
};
let child = t.child.clone();
buffer.insert(t)?;
data.insert(child, buffer);
Ok(())
}
fn creates_cycle(
child: &str,
parent: &str,
data: &HashMap<String, Buffer<T>>,
) -> bool {
let mut visited = BTreeSet::new();
let mut current = parent;
while let Some(buffer) = data.get(current) {
if !visited.insert(current) {
return true;
}
match buffer.parent() {
Some(next) => {
if next == child {
return true;
}
current = next;
}
None => return false,
}
}
false
}
fn process_get_transform(
from: &str,
to: &str,
timestamp: T,
data: &HashMap<String, Buffer<T>>,
) -> Result<Transform<T>, TransformError> {
if from == to {
return Ok(Transform {
translation: Vector3::zero(),
rotation: Quaternion::identity(),
timestamp,
parent: from.into(),
child: to.into(),
});
}
let reached = |chain: &VecDeque<Transform<T>>, target: &str| {
chain.back().is_some_and(|tf| tf.parent == target)
};
let from_chain = Self::get_transform_chain(from, to, timestamp, data);
let result = match from_chain {
Ok(from_chain) if reached(&from_chain, to) => {
Self::combine_transforms(from_chain, VecDeque::new())
}
from_chain => match (
from_chain,
Self::get_transform_chain(to, from, timestamp, data),
) {
(_, Ok(mut to_chain)) if reached(&to_chain, from) => {
Self::reverse_and_invert_transforms(&mut to_chain)?;
Self::combine_transforms(VecDeque::new(), to_chain)
}
(Ok(mut from_chain), Ok(mut to_chain)) => {
Self::truncate_at_common_parent(&mut from_chain, &mut to_chain);
let connected = match (from_chain.back(), to_chain.back()) {
(Some(from_top), Some(to_top)) => from_top.parent == to_top.parent,
_ => false,
};
if connected {
Self::reverse_and_invert_transforms(&mut to_chain)?;
Self::combine_transforms(from_chain, to_chain)
} else {
Err(TransformError::NotFound(from.into(), to.into()))
}
}
(Ok(from_chain), Err(_)) => Self::combine_transforms(from_chain, VecDeque::new()),
(Err(_), Ok(mut to_chain)) => {
Self::reverse_and_invert_transforms(&mut to_chain)?;
Self::combine_transforms(VecDeque::new(), to_chain)
}
(Err(_), Err(_)) => Err(TransformError::NotFound(from.into(), to.into())),
},
}?;
if result.parent != from || result.child != to {
return Err(TransformError::NotFound(from.into(), to.into()));
}
let mut result = result;
result.timestamp = timestamp;
Ok(result)
}
fn process_get_transform_at(
target_frame: &str,
target_time: T,
source_frame: &str,
source_time: T,
fixed_frame: &str,
data: &HashMap<String, Buffer<T>>,
) -> Result<Transform<T>, TransformError> {
if source_frame == fixed_frame && target_frame == fixed_frame {
return Ok(Transform {
translation: Vector3::zero(),
rotation: Quaternion::identity(),
timestamp: target_time,
parent: target_frame.into(),
child: source_frame.into(),
});
}
if source_frame == fixed_frame {
let mut result =
Self::process_get_transform(fixed_frame, target_frame, target_time, data)?
.inverse()?;
result.timestamp = target_time;
return Ok(result);
}
if target_frame == fixed_frame {
let mut result =
Self::process_get_transform(fixed_frame, source_frame, source_time, data)?;
result.timestamp = target_time;
return Ok(result);
}
let mut source_to_fixed =
Self::process_get_transform(fixed_frame, source_frame, source_time, data)?;
let mut target_to_fixed =
Self::process_get_transform(fixed_frame, target_frame, target_time, data)?;
source_to_fixed.timestamp = T::static_timestamp();
target_to_fixed.timestamp = T::static_timestamp();
let mut result = (target_to_fixed.inverse()? * source_to_fixed)?;
result.timestamp = target_time;
Ok(result)
}
fn get_transform_chain(
from: &str,
to: &str,
timestamp: T,
data: &HashMap<String, Buffer<T>>,
) -> Result<VecDeque<Transform<T>>, TransformError> {
let mut transforms = VecDeque::new();
let mut current_frame: String = from.into();
let mut remaining = data.len();
while let Some(frame_buffer) = data.get(¤t_frame) {
if remaining == 0 {
return Err(TransformError::NotFound(from.into(), to.into()));
}
remaining -= 1;
match frame_buffer.get(×tamp) {
Ok(tf) => {
current_frame.clone_from(&tf.parent);
transforms.push_back(tf);
}
Err(_) => break,
}
if current_frame == to {
break;
}
}
if transforms.is_empty() {
Err(TransformError::NotFound(from.into(), to.into()))
} else {
Ok(transforms)
}
}
fn truncate_at_common_parent(
from_chain: &mut VecDeque<Transform<T>>,
to_chain: &mut VecDeque<Transform<T>>,
) {
let mut start_idx = 0;
for (i, j) in from_chain.iter().rev().zip(to_chain.iter().rev()) {
if i == j {
start_idx += 1;
} else {
break;
}
}
from_chain.truncate(from_chain.len() - start_idx);
to_chain.truncate(to_chain.len() - start_idx);
}
fn combine_transforms(
mut from_chain: VecDeque<Transform<T>>,
mut to_chain: VecDeque<Transform<T>>,
) -> Result<Transform<T>, TransformError> {
from_chain.append(&mut to_chain);
let mut iter = from_chain.into_iter();
let Some(mut final_transform) = iter.next() else {
return Err(TransformError::TransformTreeEmpty);
};
for transform in iter {
final_transform = (transform * final_transform)?;
}
final_transform.inverse()
}
fn reverse_and_invert_transforms(
chain: &mut VecDeque<Transform<T>>
) -> Result<(), TransformError> {
let reversed_and_inverted = chain
.iter()
.rev()
.map(Transform::inverse)
.collect::<Result<VecDeque<Transform<T>>, TransformError>>()?;
*chain = reversed_and_inverted;
Ok(())
}
}
impl<T> Default for Registry<T>
where
T: TimePoint,
{
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;