use crate::{
core::Buffer,
errors::TransformError,
geometry::{Localized, Quaternion, Transform, Vector3},
time::{TimePoint, Timestamp},
};
use alloc::{
collections::{BTreeSet, VecDeque},
string::String,
};
use hashbrown::{hash_map::Entry, HashMap};
mod error;
#[cfg(feature = "std")]
use core::time::Duration;
pub struct Registry<T = Timestamp>
where
T: TimePoint,
{
pub data: HashMap<String, Buffer<T>>,
#[cfg(feature = "std")]
max_age: Duration,
}
impl<T> Registry<T>
where
T: TimePoint,
{
#[cfg(feature = "std")]
#[allow(clippy::new_without_default)]
#[must_use = "The Registry should be used to manage transforms."]
pub fn new(max_age: Duration) -> Self {
Self {
data: HashMap::new(),
max_age,
}
}
#[allow(clippy::new_without_default)]
#[cfg(not(feature = "std"))]
#[must_use = "The Registry should be used to manage transforms."]
pub fn new() -> Self {
Self {
data: HashMap::new(),
}
}
pub fn add_transform(
&mut self,
t: Transform<T>,
) {
#[cfg(not(feature = "std"))]
Self::process_add_transform(t, &mut self.data);
#[cfg(feature = "std")]
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>,
{
if value.frame() == target_frame {
return Ok(Transform {
translation: Vector3::new(0.0, 0.0, 0.0),
rotation: Quaternion::identity(),
timestamp: value.timestamp(),
parent: target_frame.into(),
child: target_frame.into(),
});
}
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);
}
}
#[cfg(not(feature = "std"))]
fn process_add_transform(
t: Transform<T>,
data: &mut HashMap<String, Buffer<T>>,
) {
match data.entry(t.child.clone()) {
Entry::Occupied(mut entry) => {
entry.get_mut().insert(t);
}
Entry::Vacant(entry) => {
let buffer = Buffer::new();
let buffer = entry.insert(buffer);
buffer.insert(t);
}
}
}
#[cfg(feature = "std")]
fn process_add_transform(
t: Transform<T>,
data: &mut HashMap<String, Buffer<T>>,
max_age: Duration,
) {
match data.entry(t.child.clone()) {
Entry::Occupied(mut entry) => {
entry.get_mut().insert(t);
}
Entry::Vacant(entry) => {
let buffer = Buffer::new(max_age);
let buffer = entry.insert(buffer);
buffer.insert(t);
}
}
}
fn process_get_transform(
from: &str,
to: &str,
timestamp: T,
data: &HashMap<String, Buffer<T>>,
) -> Result<Transform<T>, TransformError> {
let from_chain = Self::get_transform_chain(from, to, timestamp, data);
let to_chain = Self::get_transform_chain(to, from, timestamp, data);
match (from_chain, to_chain) {
(Ok(mut from_chain), Ok(mut to_chain)) => {
Self::truncate_at_common_parent(&mut from_chain, &mut to_chain);
Self::reverse_and_invert_transforms(&mut to_chain)?;
Self::combine_transforms(from_chain, to_chain)
}
(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())),
}
}
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> {
let mut source_to_fixed = if source_frame == fixed_frame {
Transform {
translation: crate::geometry::Vector3::new(0.0, 0.0, 0.0),
rotation: crate::geometry::Quaternion::identity(),
timestamp: source_time,
parent: fixed_frame.into(),
child: source_frame.into(),
}
} else {
Self::process_get_transform(fixed_frame, source_frame, source_time, data)?
};
let mut target_to_fixed = if target_frame == fixed_frame {
Transform {
translation: crate::geometry::Vector3::new(0.0, 0.0, 0.0),
rotation: crate::geometry::Quaternion::identity(),
timestamp: target_time,
parent: fixed_frame.into(),
child: target_frame.into(),
}
} else {
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 visited = BTreeSet::new();
let mut current_frame: String = from.into();
while let Some(frame_buffer) = data.get(¤t_frame) {
if !visited.insert(current_frame.clone()) {
return Err(TransformError::NotFound(from.into(), to.into()));
}
match frame_buffer.get(×tamp) {
Ok(tf) => {
transforms.push_back(tf.clone());
current_frame.clone_from(&tf.parent);
}
Err(_) => 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(())
}
}
#[cfg(test)]
mod tests;