openigtlink_rust/protocol/types/
qtdata.rs

1//! QTDATA (QuaternionTrackingData) message type implementation
2//!
3//! The QTDATA message type is used to transfer 3D positions and orientations
4//! of surgical tools, markers, etc. using quaternions for orientation.
5
6use crate::error::{IgtlError, Result};
7use crate::protocol::message::Message;
8use bytes::{Buf, BufMut};
9
10/// Instrument type
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12#[repr(u8)]
13pub enum InstrumentType {
14    Tracker = 1,
15    Instrument6D = 2,
16    Instrument3D = 3,
17    Instrument5D = 4,
18}
19
20impl InstrumentType {
21    fn from_u8(value: u8) -> Result<Self> {
22        match value {
23            1 => Ok(InstrumentType::Tracker),
24            2 => Ok(InstrumentType::Instrument6D),
25            3 => Ok(InstrumentType::Instrument3D),
26            4 => Ok(InstrumentType::Instrument5D),
27            _ => Err(IgtlError::InvalidHeader(format!(
28                "Invalid instrument type: {}",
29                value
30            ))),
31        }
32    }
33}
34
35/// Tracking data element with name, type, position and quaternion
36#[derive(Debug, Clone, PartialEq)]
37pub struct TrackingElement {
38    /// Name/ID of the instrument/tracker (max 20 chars)
39    pub name: String,
40    /// Type of instrument
41    pub instrument_type: InstrumentType,
42    /// Position (x, y, z) in millimeters
43    pub position: [f32; 3],
44    /// Orientation as quaternion (qx, qy, qz, w)
45    pub quaternion: [f32; 4],
46}
47
48impl TrackingElement {
49    /// Create a new tracking element
50    pub fn new(
51        name: impl Into<String>,
52        instrument_type: InstrumentType,
53        position: [f32; 3],
54        quaternion: [f32; 4],
55    ) -> Self {
56        TrackingElement {
57            name: name.into(),
58            instrument_type,
59            position,
60            quaternion,
61        }
62    }
63}
64
65/// QTDATA message containing multiple tracking elements
66///
67/// # OpenIGTLink Specification
68/// - Message type: "QTDATA"
69/// - Each element: NAME (`char[20]`) + TYPE (uint8) + Reserved (uint8) + POSITION (`float32[3]`) + QUATERNION (`float32[4]`)
70/// - Element size: 20 + 1 + 1 + 12 + 16 = 50 bytes
71#[derive(Debug, Clone, PartialEq)]
72pub struct QtDataMessage {
73    /// List of tracking elements
74    pub elements: Vec<TrackingElement>,
75}
76
77impl QtDataMessage {
78    /// Create a new QTDATA message with elements
79    pub fn new(elements: Vec<TrackingElement>) -> Self {
80        QtDataMessage { elements }
81    }
82
83    /// Create an empty QTDATA message
84    pub fn empty() -> Self {
85        QtDataMessage {
86            elements: Vec::new(),
87        }
88    }
89
90    /// Add a tracking element
91    pub fn add_element(&mut self, element: TrackingElement) {
92        self.elements.push(element);
93    }
94
95    /// Get number of tracking elements
96    pub fn len(&self) -> usize {
97        self.elements.len()
98    }
99
100    /// Check if message has no elements
101    pub fn is_empty(&self) -> bool {
102        self.elements.is_empty()
103    }
104}
105
106impl Message for QtDataMessage {
107    fn message_type() -> &'static str {
108        "QTDATA"
109    }
110
111    fn encode_content(&self) -> Result<Vec<u8>> {
112        let mut buf = Vec::with_capacity(self.elements.len() * 50);
113
114        for element in &self.elements {
115            // Encode NAME (`char[20]`)
116            let mut name_bytes = [0u8; 20];
117            let name_str = element.name.as_bytes();
118            let copy_len = name_str.len().min(19); // Reserve 1 byte for null terminator
119            name_bytes[..copy_len].copy_from_slice(&name_str[..copy_len]);
120            buf.extend_from_slice(&name_bytes);
121
122            // Encode TYPE (uint8)
123            buf.put_u8(element.instrument_type as u8);
124
125            // Encode Reserved (uint8)
126            buf.put_u8(0);
127
128            // Encode POSITION (`float32[3]`)
129            for &coord in &element.position {
130                buf.put_f32(coord);
131            }
132
133            // Encode QUATERNION (`float32[4]`)
134            for &comp in &element.quaternion {
135                buf.put_f32(comp);
136            }
137        }
138
139        Ok(buf)
140    }
141
142    fn decode_content(mut data: &[u8]) -> Result<Self> {
143        let mut elements = Vec::new();
144
145        while data.len() >= 50 {
146            // Decode NAME (`char[20]`)
147            let name_bytes = &data[..20];
148            data.advance(20);
149
150            // Find null terminator or use full length
151            let name_len = name_bytes.iter().position(|&b| b == 0).unwrap_or(20);
152            let name = String::from_utf8(name_bytes[..name_len].to_vec())?;
153
154            // Decode TYPE (uint8)
155            let instrument_type = InstrumentType::from_u8(data.get_u8())?;
156
157            // Decode Reserved (uint8)
158            let _reserved = data.get_u8();
159
160            // Decode POSITION (`float32[3]`)
161            let position = [data.get_f32(), data.get_f32(), data.get_f32()];
162
163            // Decode QUATERNION (`float32[4]`)
164            let quaternion = [
165                data.get_f32(),
166                data.get_f32(),
167                data.get_f32(),
168                data.get_f32(),
169            ];
170
171            elements.push(TrackingElement {
172                name,
173                instrument_type,
174                position,
175                quaternion,
176            });
177        }
178
179        if !data.is_empty() {
180            return Err(IgtlError::InvalidSize {
181                expected: 0,
182                actual: data.len(),
183            });
184        }
185
186        Ok(QtDataMessage { elements })
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn test_message_type() {
196        assert_eq!(QtDataMessage::message_type(), "QTDATA");
197    }
198
199    #[test]
200    fn test_instrument_type() {
201        assert_eq!(InstrumentType::Tracker as u8, 1);
202        assert_eq!(InstrumentType::Instrument6D as u8, 2);
203        assert_eq!(InstrumentType::Instrument3D as u8, 3);
204        assert_eq!(InstrumentType::Instrument5D as u8, 4);
205    }
206
207    #[test]
208    fn test_empty() {
209        let msg = QtDataMessage::empty();
210        assert!(msg.is_empty());
211        assert_eq!(msg.len(), 0);
212    }
213
214    #[test]
215    fn test_new() {
216        let elem = TrackingElement::new(
217            "Tool1",
218            InstrumentType::Instrument6D,
219            [1.0, 2.0, 3.0],
220            [0.0, 0.0, 0.0, 1.0],
221        );
222        let msg = QtDataMessage::new(vec![elem]);
223        assert_eq!(msg.len(), 1);
224    }
225
226    #[test]
227    fn test_add_element() {
228        let mut msg = QtDataMessage::empty();
229        msg.add_element(TrackingElement::new(
230            "Tool1",
231            InstrumentType::Tracker,
232            [0.0, 0.0, 0.0],
233            [0.0, 0.0, 0.0, 1.0],
234        ));
235        assert_eq!(msg.len(), 1);
236    }
237
238    #[test]
239    fn test_encode_single_element() {
240        let elem = TrackingElement::new(
241            "Tool",
242            InstrumentType::Instrument6D,
243            [10.0, 20.0, 30.0],
244            [0.1, 0.2, 0.3, 0.9],
245        );
246        let msg = QtDataMessage::new(vec![elem]);
247        let encoded = msg.encode_content().unwrap();
248
249        assert_eq!(encoded.len(), 50);
250        // Check TYPE field
251        assert_eq!(encoded[20], 2); // Instrument6D
252                                    // Check Reserved field
253        assert_eq!(encoded[21], 0);
254    }
255
256    #[test]
257    fn test_roundtrip_single() {
258        let original = QtDataMessage::new(vec![TrackingElement::new(
259            "Tracker1",
260            InstrumentType::Tracker,
261            [100.5, 200.5, 300.5],
262            [0.1, 0.2, 0.3, 0.9],
263        )]);
264
265        let encoded = original.encode_content().unwrap();
266        let decoded = QtDataMessage::decode_content(&encoded).unwrap();
267
268        assert_eq!(decoded.elements.len(), 1);
269        assert_eq!(decoded.elements[0].name, "Tracker1");
270        assert_eq!(decoded.elements[0].instrument_type, InstrumentType::Tracker);
271        assert_eq!(decoded.elements[0].position, [100.5, 200.5, 300.5]);
272        assert_eq!(decoded.elements[0].quaternion, [0.1, 0.2, 0.3, 0.9]);
273    }
274
275    #[test]
276    fn test_roundtrip_multiple() {
277        let original = QtDataMessage::new(vec![
278            TrackingElement::new(
279                "Tool1",
280                InstrumentType::Instrument6D,
281                [1.0, 2.0, 3.0],
282                [0.0, 0.0, 0.0, 1.0],
283            ),
284            TrackingElement::new(
285                "Tool2",
286                InstrumentType::Instrument3D,
287                [4.0, 5.0, 6.0],
288                [0.1, 0.2, 0.3, 0.9],
289            ),
290        ]);
291
292        let encoded = original.encode_content().unwrap();
293        let decoded = QtDataMessage::decode_content(&encoded).unwrap();
294
295        assert_eq!(decoded.elements.len(), 2);
296        assert_eq!(decoded.elements[0].name, "Tool1");
297        assert_eq!(decoded.elements[1].name, "Tool2");
298    }
299
300    #[test]
301    fn test_name_truncation() {
302        let long_name = "ThisIsAVeryLongNameThatExceedsTwentyCharacters";
303        let elem = TrackingElement::new(
304            long_name,
305            InstrumentType::Tracker,
306            [0.0, 0.0, 0.0],
307            [0.0, 0.0, 0.0, 1.0],
308        );
309        let msg = QtDataMessage::new(vec![elem]);
310
311        let encoded = msg.encode_content().unwrap();
312        let decoded = QtDataMessage::decode_content(&encoded).unwrap();
313
314        // Name should be truncated to 19 chars (20th is null terminator)
315        assert!(decoded.elements[0].name.len() <= 19);
316    }
317
318    #[test]
319    fn test_empty_message() {
320        let msg = QtDataMessage::empty();
321        let encoded = msg.encode_content().unwrap();
322        let decoded = QtDataMessage::decode_content(&encoded).unwrap();
323
324        assert_eq!(decoded.elements.len(), 0);
325        assert_eq!(encoded.len(), 0);
326    }
327
328    #[test]
329    fn test_decode_invalid_size() {
330        let data = vec![0u8; 49]; // One byte short of a complete element
331        let result = QtDataMessage::decode_content(&data);
332        assert!(result.is_err());
333    }
334}