gguf_rs_lib/format/
alignment.rs

1//! Data alignment utilities for GGUF files
2
3use crate::error::{GGUFError, Result};
4use crate::format::constants::GGUF_DEFAULT_ALIGNMENT;
5
6#[cfg(not(feature = "std"))]
7extern crate alloc;
8#[cfg(not(feature = "std"))]
9use alloc::{format, vec, vec::Vec};
10#[cfg(not(feature = "std"))]
11use core::{fmt, mem};
12
13/// Calculate the padding needed to align to the specified boundary
14pub fn calculate_padding(current_position: usize, alignment: usize) -> usize {
15    if alignment == 0 || alignment == 1 {
16        return 0;
17    }
18
19    let remainder = current_position % alignment;
20    if remainder == 0 {
21        0
22    } else {
23        alignment - remainder
24    }
25}
26
27/// Align a position to the specified boundary.
28///
29/// Returns `usize::MAX` if the aligned position cannot be represented. Use
30/// [`checked_align_to`] when overflow must be distinguished from that value.
31pub fn align_to(position: usize, alignment: usize) -> usize {
32    checked_align_to(position, alignment).unwrap_or(usize::MAX)
33}
34
35/// Align a position without overflowing.
36pub fn checked_align_to(position: usize, alignment: usize) -> Option<usize> {
37    position.checked_add(calculate_padding(position, alignment))
38}
39
40/// Align a position to the default GGUF alignment boundary
41pub fn align_to_default(position: usize) -> usize {
42    align_to(position, GGUF_DEFAULT_ALIGNMENT)
43}
44
45/// Calculate the padding needed for default alignment
46pub fn calculate_default_padding(position: usize) -> usize {
47    calculate_padding(position, GGUF_DEFAULT_ALIGNMENT)
48}
49
50/// Check if a position is aligned to the specified boundary
51pub fn is_aligned(position: usize, alignment: usize) -> bool {
52    if alignment == 0 || alignment == 1 {
53        return true;
54    }
55    position.is_multiple_of(alignment)
56}
57
58/// Check if a position is aligned to the default boundary
59pub fn is_aligned_default(position: usize) -> bool {
60    is_aligned(position, GGUF_DEFAULT_ALIGNMENT)
61}
62
63/// Utility to create padding bytes (zeros)
64pub fn create_padding(size: usize) -> Vec<u8> {
65    vec![0u8; size]
66}
67
68/// Fallibly create a zero-filled padding buffer.
69pub fn try_create_padding(size: usize) -> Result<Vec<u8>> {
70    let mut padding = Vec::new();
71    padding.try_reserve_exact(size).map_err(|_| {
72        GGUFError::InvalidTensorData(format!(
73            "Unable to allocate alignment padding of {} bytes",
74            size
75        ))
76    })?;
77    padding.resize(size, 0);
78    Ok(padding)
79}
80
81/// Calculate the aligned size for a given unaligned size.
82///
83/// Returns `usize::MAX` on overflow; use [`checked_aligned_size`] for a fallible result.
84pub fn aligned_size(size: usize, alignment: usize) -> usize {
85    checked_aligned_size(size, alignment).unwrap_or(usize::MAX)
86}
87
88/// Calculate an aligned size without overflowing.
89pub fn checked_aligned_size(size: usize, alignment: usize) -> Option<usize> {
90    checked_align_to(size, alignment)
91}
92
93/// Calculate the aligned size using default alignment
94pub fn aligned_size_default(size: usize) -> usize {
95    aligned_size(size, GGUF_DEFAULT_ALIGNMENT)
96}
97
98/// Check if an alignment value is valid (power of 2)
99pub fn is_valid_alignment(alignment: usize) -> bool {
100    alignment > 0 && (alignment & (alignment - 1)) == 0
101}
102
103/// Get the next power of 2 greater than or equal to the given value.
104///
105/// Returns `usize::MAX` when the mathematical result is not representable.
106pub fn next_power_of_2(mut n: usize) -> usize {
107    if n == 0 {
108        return 1;
109    }
110
111    if n & (n - 1) == 0 {
112        return n; // Already a power of 2
113    }
114
115    // Find the next power of 2
116    n -= 1;
117    n |= n >> 1;
118    n |= n >> 2;
119    n |= n >> 4;
120    n |= n >> 8;
121    n |= n >> 16;
122    #[cfg(feature = "std")]
123    {
124        if std::mem::size_of::<usize>() > 4 {
125            n |= n >> 32;
126        }
127    }
128    #[cfg(not(feature = "std"))]
129    {
130        if mem::size_of::<usize>() > 4 {
131            n |= n >> 32;
132        }
133    }
134    n.saturating_add(1)
135}
136
137/// Alignment information for a data section
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
139pub struct AlignmentInfo {
140    /// Current position in the file/stream
141    pub position: usize,
142    /// Required alignment boundary
143    pub alignment: usize,
144    /// Calculated padding needed
145    pub padding: usize,
146    /// Final aligned position
147    pub aligned_position: usize,
148}
149
150impl AlignmentInfo {
151    /// Create alignment information for a given position and alignment.
152    ///
153    /// `aligned_position` saturates at `usize::MAX`; arithmetic-sensitive callers
154    /// should use [`checked_align_to`] directly.
155    pub fn new(position: usize, alignment: usize) -> Self {
156        let padding = calculate_padding(position, alignment);
157        let aligned_position = position.saturating_add(padding);
158
159        Self { position, alignment, padding, aligned_position }
160    }
161
162    /// Create alignment information using default alignment
163    pub fn new_default(position: usize) -> Self {
164        Self::new(position, GGUF_DEFAULT_ALIGNMENT)
165    }
166
167    /// Check if padding is needed
168    pub fn needs_padding(&self) -> bool {
169        self.padding > 0
170    }
171
172    /// Get the padding as a vector of zero bytes
173    pub fn padding_bytes(&self) -> Vec<u8> {
174        create_padding(self.padding)
175    }
176
177    /// Fallibly get the padding as a vector of zero bytes.
178    pub fn try_padding_bytes(&self) -> Result<Vec<u8>> {
179        try_create_padding(self.padding)
180    }
181}
182
183#[cfg(feature = "std")]
184impl std::fmt::Display for AlignmentInfo {
185    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        write!(
187            f,
188            "AlignmentInfo {{ pos: {}, align: {}, padding: {}, aligned_pos: {} }}",
189            self.position, self.alignment, self.padding, self.aligned_position
190        )
191    }
192}
193
194#[cfg(not(feature = "std"))]
195impl fmt::Display for AlignmentInfo {
196    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
197        write!(
198            f,
199            "AlignmentInfo {{ pos: {}, align: {}, padding: {}, aligned_pos: {} }}",
200            self.position, self.alignment, self.padding, self.aligned_position
201        )
202    }
203}
204
205/// Helper struct for tracking alignment during file writing
206#[derive(Debug, Clone)]
207pub struct AlignmentTracker {
208    /// Current position in the output stream
209    pub position: usize,
210    /// Default alignment to use
211    pub default_alignment: usize,
212}
213
214impl AlignmentTracker {
215    /// Create a new alignment tracker
216    pub fn new(default_alignment: usize) -> Self {
217        Self { position: 0, default_alignment }
218    }
219
220    /// Create a new alignment tracker with GGUF default alignment
221    pub fn new_default() -> Self {
222        Self::new(GGUF_DEFAULT_ALIGNMENT)
223    }
224
225    /// Advance the position by the given amount, saturating at `usize::MAX`.
226    pub fn advance(&mut self, bytes: usize) {
227        self.position = self.position.saturating_add(bytes);
228    }
229
230    /// Advance the position, returning `None` if it would overflow.
231    pub fn checked_advance(&mut self, bytes: usize) -> Option<()> {
232        self.position = self.position.checked_add(bytes)?;
233        Some(())
234    }
235
236    /// Calculate padding needed for default alignment
237    pub fn calculate_padding(&self) -> usize {
238        calculate_padding(self.position, self.default_alignment)
239    }
240
241    /// Calculate padding needed for specific alignment
242    pub fn calculate_padding_for(&self, alignment: usize) -> usize {
243        calculate_padding(self.position, alignment)
244    }
245
246    /// Align to default boundary and return padding info
247    pub fn align_default(&mut self) -> AlignmentInfo {
248        let info = AlignmentInfo::new(self.position, self.default_alignment);
249        self.position = info.aligned_position;
250        info
251    }
252
253    /// Align to specific boundary and return padding info
254    pub fn align_to(&mut self, alignment: usize) -> AlignmentInfo {
255        let info = AlignmentInfo::new(self.position, alignment);
256        self.position = info.aligned_position;
257        info
258    }
259
260    /// Get current alignment info without advancing
261    pub fn current_alignment_info(&self) -> AlignmentInfo {
262        AlignmentInfo::new(self.position, self.default_alignment)
263    }
264
265    /// Check if currently aligned to default boundary
266    pub fn is_aligned(&self) -> bool {
267        is_aligned(self.position, self.default_alignment)
268    }
269
270    /// Check if currently aligned to specific boundary
271    pub fn is_aligned_to(&self, alignment: usize) -> bool {
272        is_aligned(self.position, alignment)
273    }
274}
275
276#[cfg(all(test, feature = "std"))]
277mod tests {
278    use super::*;
279
280    #[test]
281    fn test_calculate_padding() {
282        assert_eq!(calculate_padding(0, 32), 0);
283        assert_eq!(calculate_padding(1, 32), 31);
284        assert_eq!(calculate_padding(16, 32), 16);
285        assert_eq!(calculate_padding(32, 32), 0);
286        assert_eq!(calculate_padding(33, 32), 31);
287
288        // Edge cases
289        assert_eq!(calculate_padding(10, 1), 0);
290        assert_eq!(calculate_padding(10, 0), 0);
291    }
292
293    #[test]
294    fn test_align_to() {
295        assert_eq!(align_to(0, 32), 0);
296        assert_eq!(align_to(1, 32), 32);
297        assert_eq!(align_to(16, 32), 32);
298        assert_eq!(align_to(32, 32), 32);
299        assert_eq!(align_to(33, 32), 64);
300    }
301
302    #[test]
303    fn test_is_aligned() {
304        assert!(is_aligned(0, 32));
305        assert!(is_aligned(32, 32));
306        assert!(is_aligned(64, 32));
307        assert!(!is_aligned(1, 32));
308        assert!(!is_aligned(33, 32));
309
310        // Edge cases
311        assert!(is_aligned(42, 1));
312        assert!(is_aligned(42, 0));
313    }
314
315    #[test]
316    fn test_aligned_size() {
317        assert_eq!(aligned_size(0, 32), 0);
318        assert_eq!(aligned_size(1, 32), 32);
319        assert_eq!(aligned_size(16, 32), 32);
320        assert_eq!(aligned_size(32, 32), 32);
321        assert_eq!(aligned_size(33, 32), 64);
322    }
323
324    #[test]
325    fn test_is_valid_alignment() {
326        assert!(is_valid_alignment(1));
327        assert!(is_valid_alignment(2));
328        assert!(is_valid_alignment(4));
329        assert!(is_valid_alignment(8));
330        assert!(is_valid_alignment(16));
331        assert!(is_valid_alignment(32));
332
333        assert!(!is_valid_alignment(0));
334        assert!(!is_valid_alignment(3));
335        assert!(!is_valid_alignment(5));
336        assert!(!is_valid_alignment(6));
337        assert!(!is_valid_alignment(12));
338    }
339
340    #[test]
341    fn test_next_power_of_2() {
342        assert_eq!(next_power_of_2(0), 1);
343        assert_eq!(next_power_of_2(1), 1);
344        assert_eq!(next_power_of_2(2), 2);
345        assert_eq!(next_power_of_2(3), 4);
346        assert_eq!(next_power_of_2(4), 4);
347        assert_eq!(next_power_of_2(5), 8);
348        assert_eq!(next_power_of_2(15), 16);
349        assert_eq!(next_power_of_2(16), 16);
350        assert_eq!(next_power_of_2(17), 32);
351    }
352
353    #[test]
354    fn test_alignment_info() {
355        let info = AlignmentInfo::new(17, 32);
356        assert_eq!(info.position, 17);
357        assert_eq!(info.alignment, 32);
358        assert_eq!(info.padding, 15);
359        assert_eq!(info.aligned_position, 32);
360        assert!(info.needs_padding());
361
362        let info_aligned = AlignmentInfo::new(32, 32);
363        assert_eq!(info_aligned.padding, 0);
364        assert!(!info_aligned.needs_padding());
365    }
366
367    #[test]
368    fn test_alignment_tracker() {
369        let mut tracker = AlignmentTracker::new_default();
370        assert_eq!(tracker.position, 0);
371        assert!(tracker.is_aligned());
372
373        tracker.advance(17);
374        assert_eq!(tracker.position, 17);
375        assert!(!tracker.is_aligned());
376        assert_eq!(tracker.calculate_padding(), 15);
377
378        let info = tracker.align_default();
379        assert_eq!(info.padding, 15);
380        assert_eq!(tracker.position, 32);
381        assert!(tracker.is_aligned());
382    }
383
384    #[test]
385    fn test_default_alignment_functions() {
386        assert_eq!(align_to_default(17), 32);
387        assert_eq!(calculate_default_padding(17), 15);
388        assert!(is_aligned_default(32));
389        assert!(!is_aligned_default(17));
390        assert_eq!(aligned_size_default(17), 32);
391    }
392
393    #[test]
394    fn test_create_padding() {
395        let padding = create_padding(5);
396        assert_eq!(padding.len(), 5);
397        assert_eq!(padding, vec![0u8; 5]);
398
399        let empty_padding = create_padding(0);
400        assert!(empty_padding.is_empty());
401
402        assert_eq!(try_create_padding(5).unwrap(), vec![0u8; 5]);
403        assert!(try_create_padding(usize::MAX).is_err());
404        assert_eq!(AlignmentInfo::new(17, 32).try_padding_bytes().unwrap().len(), 15);
405    }
406
407    #[test]
408    fn test_alignment_info_display() {
409        let info = AlignmentInfo::new(17, 32);
410        let display_str = format!("{}", info);
411        assert!(display_str.contains("17"));
412        assert!(display_str.contains("32"));
413        assert!(display_str.contains("15"));
414    }
415
416    #[test]
417    fn test_alignment_overflow_is_explicit() {
418        assert_eq!(checked_align_to(usize::MAX, 2), None);
419        assert_eq!(align_to(usize::MAX, 2), usize::MAX);
420        assert_eq!(checked_aligned_size(usize::MAX, 8), None);
421        assert_eq!(next_power_of_2(usize::MAX), usize::MAX);
422
423        let mut tracker = AlignmentTracker::new(8);
424        tracker.position = usize::MAX;
425        assert_eq!(tracker.checked_advance(1), None);
426        assert_eq!(tracker.position, usize::MAX);
427    }
428}
gguf_rs_lib/format/alignment.rs

gguf_rs_lib/format/
alignment.rs
