openbim_step/
partition.rs

1//! Record-aligned data partitioning.
2//!
3//! Arbitrary byte splitting is unsafe: a target offset can land inside a
4//! quoted string, comment, aggregate, or record. This module tokenizes first,
5//! identifies complete `#id=...;` records, and only emits boundaries between
6//! those records.
7
8use crate::lexer::{Lexer, Token};
9use crate::{Span, StepError};
10
11/// A half-open byte range containing one or more complete data records.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct Partition {
14    /// Inclusive offset aligned to a `#id=` record start.
15    pub start: usize,
16    /// Exclusive offset aligned to the next record start or final semicolon.
17    pub end: usize,
18}
19
20impl Partition {
21    /// Converts this partition to a source span.
22    #[must_use]
23    pub const fn span(self) -> Span {
24        Span::new(self.start, self.end)
25    }
26}
27
28/// Locates complete data records.
29///
30/// For a complete exchange only records inside `DATA; ... ENDSEC;` are
31/// returned. For a record-only partition (with no `DATA` marker), all top-level
32/// `#id=...;` records are returned. Returned spans are relative to `input`.
33/// # Errors
34///
35/// Returns a lexical diagnostic when the input contains malformed tokens.
36pub fn data_record_spans(input: &[u8]) -> Result<Vec<Span>, StepError> {
37    let tokens = Lexer::new(input).collect::<Result<Vec<_>, _>>()?;
38    let has_data = tokens.windows(2).any(|window| {
39        matches!(&window[0].value, Token::Name(name) if name.eq_ignore_ascii_case(b"DATA"))
40            && window[1].value == Token::Semicolon
41    });
42    let mut in_data = !has_data;
43    let mut active_start = None;
44    let mut records = Vec::new();
45    let mut index = 0;
46
47    while index < tokens.len() {
48        let token = &tokens[index];
49        if let Token::Name(name) = &token.value {
50            if has_data && name.eq_ignore_ascii_case(b"DATA") && active_start.is_none() {
51                in_data = true;
52                index += 1;
53                continue;
54            }
55            if has_data && name.eq_ignore_ascii_case(b"ENDSEC") && active_start.is_none() {
56                in_data = false;
57                index += 1;
58                continue;
59            }
60        }
61
62        if in_data && active_start.is_none() {
63            if matches!(token.value, Token::Id(_))
64                && tokens
65                    .get(index + 1)
66                    .is_some_and(|next| next.value == Token::Equals)
67            {
68                active_start = Some(token.span.start);
69            }
70        } else if let (Some(start), Token::Semicolon) = (active_start, &token.value) {
71            records.push(Span::new(start, token.span.end));
72            active_start = None;
73        }
74        index += 1;
75    }
76
77    if let Some(start) = active_start {
78        return Err(StepError::syntax(
79            Span::new(start, input.len()),
80            "unterminated data record",
81        ));
82    }
83    Ok(records)
84}
85
86/// Splits data records into at most `partition_count` balanced groups.
87///
88/// Every partition starts at a record start; every non-final end is the next
89/// partition's record start. Inter-record whitespace and comments are assigned
90/// to the preceding partition. Empty partitions are never returned.
91/// # Errors
92///
93/// Returns a diagnostic when `partition_count` is zero or tokenization fails.
94pub fn partition_data_records(
95    input: &[u8],
96    partition_count: usize,
97) -> Result<Vec<Partition>, StepError> {
98    if partition_count == 0 {
99        return Err(StepError::invalid_argument(
100            "partition_count must be greater than zero",
101        ));
102    }
103    let records = data_record_spans(input)?;
104    if records.is_empty() {
105        return Ok(Vec::new());
106    }
107    let count = partition_count.min(records.len());
108    let mut starts = Vec::with_capacity(count + 1);
109    let base_size = records.len() / count;
110    let remainder = records.len() % count;
111    for partition in 0..count {
112        // The first `remainder` groups receive one extra record.
113        let record_index = partition * base_size + partition.min(remainder);
114        starts.push(records[record_index].start);
115    }
116    starts.push(records[records.len() - 1].end);
117
118    Ok(starts
119        .windows(2)
120        .map(|window| Partition {
121            start: window[0],
122            end: window[1],
123        })
124        .collect())
125}
openbim_step/partition.rs

openbim_step/
partition.rs
