Skip to main content

mig_assembly/
tokenize.rs

1//! EDIFACT tokenization helper — collects segments from input into a Vec.
2//!
3//! The `edifact-parser` crate uses an event-driven (SAX-style) API.
4//! This module provides a convenience function that collects all parsed
5//! segments into an owned `Vec<OwnedSegment>` for two-pass processing.
6//!
7//! `OwnedSegment` itself lives in `mig-types::segment` — re-exported here
8//! for backward compatibility.
9
10use std::sync::Arc;
11
12use edifact_primitives::{Control, EdifactDelimiters, RawSegment};
13
14// Re-export OwnedSegment from mig-types so existing `use crate::tokenize::OwnedSegment` paths work.
15pub use mig_types::segment::OwnedSegment;
16
17/// A single EDIFACT message (UNH...UNT) with its interchange envelope.
18#[derive(Debug, Clone)]
19pub struct MessageChunk {
20    /// Interchange envelope segments (UNA, UNB) — shared across all messages via `Arc`.
21    pub envelope: Arc<Vec<OwnedSegment>>,
22    /// The UNH segment itself.
23    pub unh: OwnedSegment,
24    /// Segments between UNH and UNT (exclusive of both).
25    pub body: Vec<OwnedSegment>,
26    /// The UNT segment itself.
27    pub unt: OwnedSegment,
28}
29
30impl MessageChunk {
31    /// Reconstruct the full segment list for this message (envelope + UNH + body + UNT).
32    ///
33    /// Note: only use this when the MIG schema includes envelope segments (UNA, UNB).
34    /// For MIG schemas that start at UNH, use `message_segments()` instead.
35    pub fn all_segments(&self) -> Vec<OwnedSegment> {
36        let mut segs = Vec::with_capacity(self.envelope.len() + 2 + self.body.len());
37        segs.extend_from_slice(&self.envelope);
38        segs.push(self.unh.clone());
39        segs.extend(self.body.iter().cloned());
40        segs.push(self.unt.clone());
41        segs
42    }
43
44    /// Message-only segments (UNH + body + UNT), excluding interchange envelope.
45    ///
46    /// Use this for assembly when the MIG schema starts at UNH (most message types).
47    /// UTILMD MIG schemas include UNA/UNB, so those still need `all_segments()`.
48    pub fn message_segments(&self) -> Vec<OwnedSegment> {
49        let mut segs = Vec::with_capacity(2 + self.body.len());
50        segs.push(self.unh.clone());
51        segs.extend(self.body.iter().cloned());
52        segs.push(self.unt.clone());
53        segs
54    }
55
56    /// Return the appropriate segments for assembly based on whether the MIG
57    /// includes envelope segments (UNA/UNB) or starts at UNH.
58    pub fn segments_for_mig(&self, mig: &mig_types::schema::mig::MigSchema) -> Vec<OwnedSegment> {
59        if mig.includes_envelope() {
60            self.all_segments()
61        } else {
62            self.message_segments()
63        }
64    }
65}
66
67/// A complete EDIFACT interchange split into per-message chunks.
68#[derive(Debug, Clone)]
69pub struct InterchangeChunks {
70    /// Interchange envelope segments (UNA, UNB) — shared across all messages.
71    pub envelope: Vec<OwnedSegment>,
72    /// One entry per UNH/UNT pair.
73    pub messages: Vec<MessageChunk>,
74    /// The UNZ segment (interchange trailer), if present.
75    pub unz: Option<OwnedSegment>,
76}
77
78/// Remove release character escaping from a component value.
79///
80/// The EDIFACT tokenizer preserves escape sequences in borrowed slices (e.g., `?+`
81/// stays as `?+`). When converting to owned `OwnedSegment` strings, we strip the
82/// release characters so values are stored clean (e.g., `+`). The renderer then
83/// re-escapes when writing EDIFACT output.
84fn unescape_edifact(value: &str, release: u8) -> String {
85    EdifactDelimiters {
86        release,
87        ..EdifactDelimiters::default()
88    }
89    .unescape(value)
90}
91
92/// Handler that collects all segments into owned copies.
93struct SegmentCollector {
94    segments: Vec<OwnedSegment>,
95    /// Release character for unescaping component values.
96    release: u8,
97}
98
99impl edifact_parser::EdifactHandler for SegmentCollector {
100    fn on_segment(&mut self, segment: &RawSegment<'_>) -> Control {
101        let release = self.release;
102        self.segments.push(OwnedSegment {
103            id: segment.id.to_string(),
104            elements: segment
105                .elements
106                .iter()
107                .map(|e| e.iter().map(|c| unescape_edifact(c, release)).collect())
108                .collect(),
109            segment_number: segment.position.segment_number,
110        });
111        Control::Continue
112    }
113
114    fn on_delimiters(&mut self, delimiters: &EdifactDelimiters, _explicit_una: bool) {
115        self.release = delimiters.release;
116    }
117
118    fn on_interchange_start(&mut self, _unb: &RawSegment<'_>) -> Control {
119        Control::Continue
120    }
121
122    fn on_message_start(&mut self, _unh: &RawSegment<'_>) -> Control {
123        Control::Continue
124    }
125
126    fn on_message_end(&mut self, _unt: &RawSegment<'_>) {}
127
128    fn on_interchange_end(&mut self, _unz: &RawSegment<'_>) {}
129}
130
131/// Parse an EDIFACT message into a list of owned segments.
132///
133/// This is "pass 1" of the two-pass assembler. It uses the streaming
134/// `edifact-parser` to tokenize the input and collects all segments
135/// into owned data structures suitable for random access.
136pub fn parse_to_segments(input: &[u8]) -> Result<Vec<OwnedSegment>, crate::AssemblyError> {
137    let mut collector = SegmentCollector {
138        segments: Vec::new(),
139        release: EdifactDelimiters::default().release,
140    };
141    edifact_parser::EdifactStreamParser::parse(input, &mut collector)
142        .map_err(|e| crate::AssemblyError::ParseError(e.to_string()))?;
143    Ok(collector.segments)
144}
145
146/// Split a flat segment list into per-message chunks at UNH/UNT boundaries.
147///
148/// Each message gets a copy of the interchange envelope (UNB and any segments
149/// before the first UNH) so it can be independently assembled.
150///
151/// # Errors
152///
153/// Returns an error if no UNH/UNT pairs are found.
154pub fn split_messages(
155    segments: Vec<OwnedSegment>,
156) -> Result<InterchangeChunks, crate::AssemblyError> {
157    let mut envelope: Vec<OwnedSegment> = Vec::with_capacity(4);
158    // Collect (unh, body, unt) tuples first, then wrap with shared envelope Arc.
159    let mut raw_messages: Vec<(OwnedSegment, Vec<OwnedSegment>, OwnedSegment)> = Vec::new();
160    let mut unz: Option<OwnedSegment> = None;
161
162    // State machine
163    let mut current_unh: Option<OwnedSegment> = None;
164    let mut current_body: Vec<OwnedSegment> = Vec::with_capacity(32);
165    let mut seen_first_unh = false;
166
167    for seg in segments {
168        let id_upper = seg.id.to_uppercase();
169        match id_upper.as_str() {
170            "UNH" => {
171                seen_first_unh = true;
172                current_unh = Some(seg);
173                current_body.clear();
174            }
175            "UNT" => {
176                if let Some(unh) = current_unh.take() {
177                    raw_messages.push((unh, std::mem::take(&mut current_body), seg));
178                }
179            }
180            "UNZ" => {
181                unz = Some(seg);
182            }
183            _ => {
184                if seen_first_unh {
185                    current_body.push(seg);
186                } else {
187                    envelope.push(seg);
188                }
189            }
190        }
191    }
192
193    if raw_messages.is_empty() {
194        return Err(crate::AssemblyError::ParseError(
195            "No UNH/UNT message pairs found in interchange".to_string(),
196        ));
197    }
198
199    // Share the envelope via Arc across all messages to avoid N clones.
200    let envelope_arc = Arc::new(envelope);
201    let messages = raw_messages
202        .into_iter()
203        .map(|(unh, body, unt)| MessageChunk {
204            envelope: Arc::clone(&envelope_arc),
205            unh,
206            body,
207            unt,
208        })
209        .collect();
210
211    Ok(InterchangeChunks {
212        envelope: (*envelope_arc).clone(),
213        messages,
214        unz,
215    })
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn test_parse_to_segments_minimal() {
224        let input = b"UNA:+.? 'UNB+UNOC:3+SENDER+RECEIVER+210101:1200+REF001'UNH+MSG001+UTILMD:D:11A:UN:S2.1'BGM+E03+DOC001'UNT+3+MSG001'UNZ+1+REF001'";
225        let segments = parse_to_segments(input).unwrap();
226
227        // Should have UNB, UNH, BGM, UNT, UNZ
228        assert_eq!(segments.len(), 5);
229        assert!(segments[0].is("UNB"));
230        assert!(segments[1].is("UNH"));
231        assert!(segments[2].is("BGM"));
232        assert!(segments[3].is("UNT"));
233        assert!(segments[4].is("UNZ"));
234    }
235
236    #[test]
237    fn test_parse_to_segments_element_access() {
238        let input = b"UNA:+.? 'UNB+UNOC:3'UNH+001+UTILMD:D:11A'BGM+E03+DOC001'UNT+2+001'UNZ+1'";
239        let segments = parse_to_segments(input).unwrap();
240
241        let bgm = &segments[2];
242        assert_eq!(bgm.id, "BGM");
243        assert_eq!(bgm.get_element(0), "E03");
244        assert_eq!(bgm.get_element(1), "DOC001");
245        assert_eq!(bgm.get_element(99), "");
246    }
247
248    #[test]
249    fn test_parse_to_segments_composite_access() {
250        let input = b"UNA:+.? 'UNH+001+UTILMD:D:11A:UN:S2.1'UNT+1+001'";
251        let segments = parse_to_segments(input).unwrap();
252
253        let unh = &segments[0]; // UNH
254        assert_eq!(unh.get_component(1, 0), "UTILMD");
255        assert_eq!(unh.get_component(1, 1), "D");
256        assert_eq!(unh.get_component(1, 4), "S2.1");
257    }
258
259    #[test]
260    fn test_message_chunk_struct_exists() {
261        let chunk = MessageChunk {
262            envelope: Arc::new(vec![]),
263            unh: OwnedSegment {
264                id: "UNH".to_string(),
265                elements: vec![],
266                segment_number: 0,
267            },
268            body: vec![],
269            unt: OwnedSegment {
270                id: "UNT".to_string(),
271                elements: vec![],
272                segment_number: 1,
273            },
274        };
275        assert_eq!(chunk.unh.id, "UNH");
276        assert_eq!(chunk.unt.id, "UNT");
277        assert!(chunk.envelope.is_empty());
278        assert!(chunk.body.is_empty());
279    }
280
281    #[test]
282    fn test_interchange_chunks_struct_exists() {
283        let chunks = InterchangeChunks {
284            envelope: vec![],
285            messages: vec![],
286            unz: None,
287        };
288        assert!(chunks.messages.is_empty());
289        assert!(chunks.unz.is_none());
290    }
291
292    #[test]
293    fn test_split_messages_single_message() {
294        let input = b"UNA:+.? 'UNB+UNOC:3+SENDER+RECEIVER+210101:1200+REF001'UNH+MSG001+UTILMD:D:11A:UN:S2.1'BGM+E03+DOC001'UNT+3+MSG001'UNZ+1+REF001'";
295        let segments = parse_to_segments(input).unwrap();
296        let chunks = split_messages(segments).unwrap();
297
298        assert_eq!(chunks.messages.len(), 1);
299        assert_eq!(chunks.envelope.len(), 1); // UNB only (UNA not emitted by parser)
300        assert!(chunks.unz.is_some());
301
302        let msg = &chunks.messages[0];
303        assert!(msg.unh.is("UNH"));
304        assert!(msg.unt.is("UNT"));
305        assert_eq!(msg.body.len(), 1); // BGM only
306        assert!(msg.body[0].is("BGM"));
307
308        // all_segments() should reconstruct: UNB, UNH, BGM, UNT
309        let all = msg.all_segments();
310        assert_eq!(all.len(), 4);
311        assert!(all[0].is("UNB"));
312        assert!(all[1].is("UNH"));
313        assert!(all[2].is("BGM"));
314        assert!(all[3].is("UNT"));
315    }
316
317    #[test]
318    fn test_split_messages_two_messages() {
319        let input = b"UNA:+.? 'UNB+UNOC:3+SENDER+RECEIVER+210101:1200+REF001'UNH+001+UTILMD:D:11A:UN:S2.1'BGM+E01+DOC001'UNT+2+001'UNH+002+UTILMD:D:11A:UN:S2.1'BGM+E03+DOC002'DTM+137:20250101:102'UNT+3+002'UNZ+2+REF001'";
320        let segments = parse_to_segments(input).unwrap();
321        let chunks = split_messages(segments).unwrap();
322
323        assert_eq!(chunks.messages.len(), 2);
324
325        // First message: UNH, BGM, UNT
326        let msg1 = &chunks.messages[0];
327        assert_eq!(msg1.unh.get_element(0), "001");
328        assert_eq!(msg1.body.len(), 1);
329        assert!(msg1.body[0].is("BGM"));
330
331        // Second message: UNH, BGM, DTM, UNT
332        let msg2 = &chunks.messages[1];
333        assert_eq!(msg2.unh.get_element(0), "002");
334        assert_eq!(msg2.body.len(), 2);
335        assert!(msg2.body[0].is("BGM"));
336        assert!(msg2.body[1].is("DTM"));
337
338        // Both messages share the same envelope
339        assert_eq!(msg1.envelope.len(), msg2.envelope.len());
340        assert!(msg1.envelope[0].is("UNB"));
341    }
342
343    #[test]
344    fn test_split_messages_envelope_preserved_per_message() {
345        // Each message's all_segments() should start with envelope
346        let input = b"UNA:+.? 'UNB+UNOC:3+SEND+RECV+210101:1200+REF'UNH+001+UTILMD:D:11A:UN:S2.1'UNT+1+001'UNH+002+UTILMD:D:11A:UN:S2.1'UNT+1+002'UNZ+2+REF'";
347        let segments = parse_to_segments(input).unwrap();
348        let chunks = split_messages(segments).unwrap();
349
350        for msg in &chunks.messages {
351            let all = msg.all_segments();
352            assert!(all[0].is("UNB"), "First segment should be UNB");
353            assert!(all[1].is("UNH"), "Second segment should be UNH");
354            assert!(all.last().unwrap().is("UNT"), "Last segment should be UNT");
355        }
356    }
357
358    #[test]
359    fn test_split_messages_no_messages_errors() {
360        let input = b"UNA:+.? 'UNB+UNOC:3+S+R+210101:1200+REF'UNZ+0+REF'";
361        let segments = parse_to_segments(input).unwrap();
362        let result = split_messages(segments);
363        assert!(result.is_err());
364    }
365
366    #[test]
367    fn test_owned_segment_is_case_insensitive() {
368        let input = b"UNA:+.? 'UNB+UNOC:3'UNZ+0'";
369        let segments = parse_to_segments(input).unwrap();
370        assert!(segments[0].is("unb"));
371        assert!(segments[0].is("UNB"));
372        assert!(segments[0].is("Unb"));
373    }
374}