diff --git a/Cargo.toml b/Cargo.toml index a90a79a..138a710 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ name = "urdf-rs" # Then, CI will publish to crates.io and create a GitHub release. version = "0.9.0" authors = ["Takashi Ogura "] -edition = "2021" +edition = "2024" description = "URDF parser" license = "Apache-2.0" keywords = ["robotics", "robot", "ros", "urdf"] diff --git a/src/deserialize.rs b/src/deserialize.rs index eff6972..6f04302 100644 --- a/src/deserialize.rs +++ b/src/deserialize.rs @@ -1,6 +1,7 @@ use serde::de::Visitor; use serde::{Deserialize, Deserializer, Serialize}; +use std::collections::HashMap; use std::ops::{Deref, DerefMut}; #[derive(Debug, Deserialize, Serialize, Default, Clone)] @@ -487,6 +488,118 @@ pub struct Robot { pub materials: Vec, } +impl Robot { + /// Convert Robot into RobotTree + /// RobotTree is a tree structure of Robot, which makes it easier to traverse the robot structure. + /// you can use `Robot::into_tree` to convert Robot into RobotTree. + /// you can also see the tree structure by using `RobotTree::format_tree` method. + pub fn into_tree(self) -> RobotTree { + assert!( + !self.links.is_empty(), + "Robot::into_tree requires at least one link", + ); + + let mut links_iter = self.links.into_iter(); + let root_link = links_iter + .next() + .expect("checked above for at least one link"); + let mut link_map: HashMap = + links_iter.map(|link| (link.name.clone(), link)).collect(); + + let mut adjacency: HashMap> = HashMap::new(); + for joint in self.joints { + adjacency + .entry(joint.parent.link.clone()) + .or_default() + .push(joint); + } + + let tree = build_robot_tree(root_link, &mut link_map, &mut adjacency); + + if let Some(orphan) = link_map.keys().next() { + panic!( + "Unreachable link `{}` found while building RobotTree", + orphan + ); + } + if let Some(orphan_parent) = adjacency.keys().next() { + panic!( + "Parent link `{}` has outgoing joints but is missing from the tree", + orphan_parent + ); + } + + tree + } +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct RobotTree { + pub link: Link, + pub joint: Vec, + pub children: Vec, +} + +impl std::fmt::Display for RobotTree { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fn walk(node: &RobotTree, prefix: String, out: &mut String) { + assert_eq!(node.joint.len(), node.children.len()); + let child_count = node.children.len(); + for (index, (joint, child)) in node.joint.iter().zip(node.children.iter()).enumerate() { + out.push_str(&prefix); + out.push_str("+-- "); + out.push_str(&child.link.name); + out.push_str(" [joint: "); + out.push_str(&joint.name); + out.push_str("]\n"); + + let mut child_prefix = prefix.clone(); + child_prefix.push_str(if index + 1 == child_count { + " " + } else { + "| " + }); + walk(child, child_prefix, out); + } + } + + let mut out = String::new(); + out.push_str(&self.link.name); + out.push('\n'); + walk(self, String::new(), &mut out); + write!(f, "{}", out) + } +} + +fn build_robot_tree( + link: Link, + links: &mut HashMap, + adjacency: &mut HashMap>, +) -> RobotTree { + let joints = adjacency.remove(&link.name).unwrap_or_default(); + let mut edges = Vec::with_capacity(joints.len()); + let mut children = Vec::with_capacity(joints.len()); + + for joint in joints { + let child_name = joint.child.link.clone(); + let child_link = links.remove(&child_name).unwrap_or_else(|| { + panic!( + "Joint `{}` references missing child link `{}`", + joint.name, child_name + ) + }); + let child_tree = build_robot_tree(child_link, links, adjacency); + edges.push(joint); + children.push(child_tree); + } + + RobotTree { + link, + joint: edges, + children, + } +} + fn de_f64<'de, D>(deserializer: D) -> Result where D: Deserializer<'de>, @@ -538,3 +651,31 @@ impl<'de> Visitor<'de> for OptF64Visitor { Ok(None) } } + +#[cfg(test)] +mod tests { + use crate::read_file; + use std::path::PathBuf; + + #[test] + fn renders_sample_tree() { + let mut sample = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + sample.push("sample.urdf"); + + let robot = read_file(sample).expect("sample urdf should load"); + let tree = robot.into_tree(); + let rendered = format!("{tree}"); + + let expected = "\ +root ++-- shoulder1 [joint: shoulder_yaw] + +-- shoulder2 [joint: shoulder_pitch] + +-- shoulder3 [joint: shoulder_roll] + +-- elbow1 [joint: elbow_pitch] + +-- wrist1 [joint: wrist_yaw] + +-- wrist2 [joint: wrist_pitch] +"; + + assert_eq!(rendered, expected); + } +} diff --git a/src/funcs.rs b/src/funcs.rs index 1a8cee0..5b20e2d 100644 --- a/src/funcs.rs +++ b/src/funcs.rs @@ -126,8 +126,8 @@ pub fn write_to_string(robot: &Robot) -> Result { #[cfg(test)] mod tests { - use crate::{read_from_string, write_to_string}; use crate::{Geometry, JointType, Robot}; + use crate::{read_from_string, write_to_string}; use assert_approx_eq::assert_approx_eq; fn check_robot(robot: &Robot) { @@ -201,20 +201,14 @@ mod tests { _ => panic!("geometry error"), } match &link.visual[1].geometry { - Geometry::Mesh { - ref filename, - scale, - } => { + Geometry::Mesh { filename, scale } => { assert_eq!(filename, "aa.dae"); assert!(scale.is_none()); } _ => panic!("geometry error"), } match &link.visual[2].geometry { - Geometry::Mesh { - ref filename, - scale, - } => { + Geometry::Mesh { filename, scale } => { assert_eq!(filename, "bbb.dae"); let scale = scale.as_ref().unwrap(); assert_approx_eq!(scale[0], 2.0);