infiniloom_engine/
newtypes.rs

1//! Type-safe wrappers for primitive types
2//!
3//! This module provides newtype wrappers that prevent accidentally mixing
4//! different kinds of values (e.g., token counts vs line counts).
5//!
6//! # Example
7//!
8//! ```rust
9//! use infiniloom_engine::newtypes::{TokenCount, LineNumber, ByteOffset};
10//!
11//! let tokens = TokenCount::new(1000);
12//! let line = LineNumber::new(42);
13//!
14//! // These are different types and can't be accidentally mixed
15//! // tokens + line; // This would be a compile error
16//! ```
17
18use serde::{Deserialize, Serialize};
19use std::fmt;
20use std::ops::{Add, AddAssign, Sub, SubAssign};
21
22/// A count of tokens (for LLM context budgeting)
23#[derive(
24    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Serialize, Deserialize,
25)]
26#[repr(transparent)]
27pub struct TokenCount(u32);
28
29impl TokenCount {
30    /// Create a new token count
31    #[inline]
32    pub const fn new(count: u32) -> Self {
33        Self(count)
34    }
35
36    /// Create a zero token count
37    #[inline]
38    pub const fn zero() -> Self {
39        Self(0)
40    }
41
42    /// Get the inner value
43    #[inline]
44    pub const fn get(self) -> u32 {
45        self.0
46    }
47
48    /// Check if this is zero
49    #[inline]
50    pub const fn is_zero(self) -> bool {
51        self.0 == 0
52    }
53
54    /// Saturating subtraction
55    #[inline]
56    pub const fn saturating_sub(self, rhs: Self) -> Self {
57        Self(self.0.saturating_sub(rhs.0))
58    }
59
60    /// Saturating addition
61    #[inline]
62    pub const fn saturating_add(self, rhs: Self) -> Self {
63        Self(self.0.saturating_add(rhs.0))
64    }
65
66    /// Calculate percentage of another token count
67    #[inline]
68    pub fn percentage_of(self, total: Self) -> f32 {
69        if total.0 == 0 {
70            0.0
71        } else {
72            (self.0 as f32 / total.0 as f32) * 100.0
73        }
74    }
75}
76
77impl Add for TokenCount {
78    type Output = Self;
79
80    #[inline]
81    fn add(self, rhs: Self) -> Self::Output {
82        Self(self.0 + rhs.0)
83    }
84}
85
86impl AddAssign for TokenCount {
87    #[inline]
88    fn add_assign(&mut self, rhs: Self) {
89        self.0 += rhs.0;
90    }
91}
92
93impl Sub for TokenCount {
94    type Output = Self;
95
96    #[inline]
97    fn sub(self, rhs: Self) -> Self::Output {
98        Self(self.0 - rhs.0)
99    }
100}
101
102impl SubAssign for TokenCount {
103    #[inline]
104    fn sub_assign(&mut self, rhs: Self) {
105        self.0 -= rhs.0;
106    }
107}
108
109impl From<u32> for TokenCount {
110    #[inline]
111    fn from(value: u32) -> Self {
112        Self(value)
113    }
114}
115
116impl From<TokenCount> for u32 {
117    #[inline]
118    fn from(value: TokenCount) -> Self {
119        value.0
120    }
121}
122
123impl fmt::Display for TokenCount {
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        write!(f, "{} tokens", self.0)
126    }
127}
128
129impl std::iter::Sum for TokenCount {
130    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
131        iter.fold(Self::zero(), |acc, x| acc + x)
132    }
133}
134
135/// A 1-indexed line number in source code
136#[derive(
137    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Serialize, Deserialize,
138)]
139#[repr(transparent)]
140pub struct LineNumber(u32);
141
142impl LineNumber {
143    /// Create a new line number (1-indexed)
144    #[inline]
145    pub const fn new(line: u32) -> Self {
146        Self(line)
147    }
148
149    /// Get the first line (line 1)
150    #[inline]
151    pub const fn first() -> Self {
152        Self(1)
153    }
154
155    /// Get the inner value
156    #[inline]
157    pub const fn get(self) -> u32 {
158        self.0
159    }
160
161    /// Check if this is a valid line number (> 0)
162    #[inline]
163    pub const fn is_valid(self) -> bool {
164        self.0 > 0
165    }
166
167    /// Convert to 0-indexed offset
168    #[inline]
169    pub const fn to_zero_indexed(self) -> u32 {
170        self.0.saturating_sub(1)
171    }
172
173    /// Create from 0-indexed offset
174    #[inline]
175    pub const fn from_zero_indexed(offset: u32) -> Self {
176        Self(offset + 1)
177    }
178
179    /// Calculate line count between this and another line (inclusive)
180    #[inline]
181    pub const fn lines_to(self, end: Self) -> u32 {
182        if end.0 >= self.0 {
183            end.0 - self.0 + 1
184        } else {
185            1
186        }
187    }
188}
189
190impl From<u32> for LineNumber {
191    #[inline]
192    fn from(value: u32) -> Self {
193        Self(value)
194    }
195}
196
197impl From<LineNumber> for u32 {
198    #[inline]
199    fn from(value: LineNumber) -> Self {
200        value.0
201    }
202}
203
204impl fmt::Display for LineNumber {
205    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
206        write!(f, "L{}", self.0)
207    }
208}
209
210/// A byte offset in a file or string
211#[derive(
212    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Serialize, Deserialize,
213)]
214#[repr(transparent)]
215pub struct ByteOffset(usize);
216
217impl ByteOffset {
218    /// Create a new byte offset
219    #[inline]
220    pub const fn new(offset: usize) -> Self {
221        Self(offset)
222    }
223
224    /// Create a zero offset
225    #[inline]
226    pub const fn zero() -> Self {
227        Self(0)
228    }
229
230    /// Get the inner value
231    #[inline]
232    pub const fn get(self) -> usize {
233        self.0
234    }
235}
236
237impl From<usize> for ByteOffset {
238    #[inline]
239    fn from(value: usize) -> Self {
240        Self(value)
241    }
242}
243
244impl From<ByteOffset> for usize {
245    #[inline]
246    fn from(value: ByteOffset) -> Self {
247        value.0
248    }
249}
250
251impl fmt::Display for ByteOffset {
252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253        write!(f, "@{}", self.0)
254    }
255}
256
257/// A unique identifier for a symbol in the index
258#[derive(
259    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Serialize, Deserialize,
260)]
261#[repr(transparent)]
262pub struct SymbolId(u32);
263
264impl SymbolId {
265    /// Create a new symbol ID
266    #[inline]
267    pub const fn new(id: u32) -> Self {
268        Self(id)
269    }
270
271    /// Create an invalid/unknown symbol ID
272    #[inline]
273    pub const fn unknown() -> Self {
274        Self(0)
275    }
276
277    /// Get the inner value
278    #[inline]
279    pub const fn get(self) -> u32 {
280        self.0
281    }
282
283    /// Check if this is a valid symbol ID
284    #[inline]
285    pub const fn is_valid(self) -> bool {
286        self.0 > 0
287    }
288}
289
290impl From<u32> for SymbolId {
291    #[inline]
292    fn from(value: u32) -> Self {
293        Self(value)
294    }
295}
296
297impl From<SymbolId> for u32 {
298    #[inline]
299    fn from(value: SymbolId) -> Self {
300        value.0
301    }
302}
303
304impl fmt::Display for SymbolId {
305    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
306        write!(f, "#{}", self.0)
307    }
308}
309
310/// File size in bytes
311#[derive(
312    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Serialize, Deserialize,
313)]
314#[repr(transparent)]
315pub struct FileSize(u64);
316
317impl FileSize {
318    /// Create a new file size
319    #[inline]
320    pub const fn new(bytes: u64) -> Self {
321        Self(bytes)
322    }
323
324    /// Create a zero size
325    #[inline]
326    pub const fn zero() -> Self {
327        Self(0)
328    }
329
330    /// Get the inner value in bytes
331    #[inline]
332    pub const fn bytes(self) -> u64 {
333        self.0
334    }
335
336    /// Get size in kilobytes
337    #[inline]
338    pub const fn kilobytes(self) -> u64 {
339        self.0 / 1024
340    }
341
342    /// Get size in megabytes
343    #[inline]
344    pub const fn megabytes(self) -> u64 {
345        self.0 / (1024 * 1024)
346    }
347
348    /// Check if this exceeds a limit
349    #[inline]
350    pub const fn exceeds(self, limit: Self) -> bool {
351        self.0 > limit.0
352    }
353}
354
355impl From<u64> for FileSize {
356    #[inline]
357    fn from(value: u64) -> Self {
358        Self(value)
359    }
360}
361
362impl From<FileSize> for u64 {
363    #[inline]
364    fn from(value: FileSize) -> Self {
365        value.0
366    }
367}
368
369impl fmt::Display for FileSize {
370    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
371        if self.0 >= 1024 * 1024 {
372            write!(f, "{:.1} MB", self.0 as f64 / (1024.0 * 1024.0))
373        } else if self.0 >= 1024 {
374            write!(f, "{:.1} KB", self.0 as f64 / 1024.0)
375        } else {
376            write!(f, "{} bytes", self.0)
377        }
378    }
379}
380
381/// Importance score (0.0 to 1.0)
382#[derive(Debug, Clone, Copy, PartialEq, Default, Serialize, Deserialize)]
383#[repr(transparent)]
384pub struct ImportanceScore(f32);
385
386impl ImportanceScore {
387    /// Create a new importance score, clamping to valid range
388    #[inline]
389    pub fn new(score: f32) -> Self {
390        Self(score.clamp(0.0, 1.0))
391    }
392
393    /// Create a zero importance score
394    #[inline]
395    pub const fn zero() -> Self {
396        Self(0.0)
397    }
398
399    /// Create maximum importance score
400    #[inline]
401    pub const fn max() -> Self {
402        Self(1.0)
403    }
404
405    /// Create default importance score
406    #[inline]
407    pub const fn default_score() -> Self {
408        Self(0.5)
409    }
410
411    /// Get the inner value
412    #[inline]
413    pub const fn get(self) -> f32 {
414        self.0
415    }
416
417    /// Check if this is considered high importance (> 0.7)
418    #[inline]
419    pub fn is_high(self) -> bool {
420        self.0 > 0.7
421    }
422
423    /// Check if this is considered low importance (< 0.3)
424    #[inline]
425    pub fn is_low(self) -> bool {
426        self.0 < 0.3
427    }
428}
429
430impl From<f32> for ImportanceScore {
431    #[inline]
432    fn from(value: f32) -> Self {
433        Self::new(value)
434    }
435}
436
437impl From<ImportanceScore> for f32 {
438    #[inline]
439    fn from(value: ImportanceScore) -> Self {
440        value.0
441    }
442}
443
444impl fmt::Display for ImportanceScore {
445    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
446        write!(f, "{:.2}", self.0)
447    }
448}
449
450impl Eq for ImportanceScore {}
451
452impl PartialOrd for ImportanceScore {
453    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
454        Some(self.cmp(other))
455    }
456}
457
458impl Ord for ImportanceScore {
459    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
460        // Use f32's partial_cmp directly, treating NaN as equal (shouldn't happen with clamped values)
461        self.0
462            .partial_cmp(&other.0)
463            .unwrap_or(std::cmp::Ordering::Equal)
464    }
465}
466
467impl std::hash::Hash for ImportanceScore {
468    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
469        self.0.to_bits().hash(state);
470    }
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476
477    // TokenCount tests
478    #[test]
479    fn test_token_count_operations() {
480        let a = TokenCount::new(100);
481        let b = TokenCount::new(50);
482
483        assert_eq!((a + b).get(), 150);
484        assert_eq!((a - b).get(), 50);
485        assert_eq!(a.saturating_sub(TokenCount::new(200)).get(), 0);
486    }
487
488    #[test]
489    fn test_token_count_percentage() {
490        let part = TokenCount::new(25);
491        let total = TokenCount::new(100);
492
493        assert!((part.percentage_of(total) - 25.0).abs() < 0.01);
494        assert_eq!(part.percentage_of(TokenCount::zero()), 0.0);
495    }
496
497    #[test]
498    fn test_token_count_sum() {
499        let counts = vec![TokenCount::new(10), TokenCount::new(20), TokenCount::new(30)];
500        let sum: TokenCount = counts.into_iter().sum();
501        assert_eq!(sum.get(), 60);
502    }
503
504    // LineNumber tests
505    #[test]
506    fn test_line_number_indexing() {
507        let line = LineNumber::new(10);
508
509        assert_eq!(line.to_zero_indexed(), 9);
510        assert_eq!(LineNumber::from_zero_indexed(9).get(), 10);
511    }
512
513    #[test]
514    fn test_line_number_range() {
515        let start = LineNumber::new(5);
516        let end = LineNumber::new(10);
517
518        assert_eq!(start.lines_to(end), 6);
519        assert_eq!(end.lines_to(start), 1); // Invalid range returns 1
520    }
521
522    // ByteOffset tests
523    #[test]
524    fn test_byte_offset() {
525        let offset = ByteOffset::new(1024);
526        assert_eq!(offset.get(), 1024);
527        assert_eq!(ByteOffset::zero().get(), 0);
528    }
529
530    // SymbolId tests
531    #[test]
532    fn test_symbol_id_validity() {
533        assert!(!SymbolId::unknown().is_valid());
534        assert!(SymbolId::new(1).is_valid());
535        assert!(!SymbolId::new(0).is_valid());
536    }
537
538    // FileSize tests
539    #[test]
540    fn test_file_size_conversions() {
541        let size = FileSize::new(1024 * 1024 + 512 * 1024); // 1.5 MB
542
543        assert_eq!(size.kilobytes(), 1536);
544        assert_eq!(size.megabytes(), 1);
545    }
546
547    #[test]
548    fn test_file_size_display() {
549        assert_eq!(FileSize::new(500).to_string(), "500 bytes");
550        assert_eq!(FileSize::new(2048).to_string(), "2.0 KB");
551        assert_eq!(FileSize::new(1024 * 1024).to_string(), "1.0 MB");
552    }
553
554    // ImportanceScore tests
555    #[test]
556    fn test_importance_score_clamping() {
557        assert_eq!(ImportanceScore::new(-0.5).get(), 0.0);
558        assert_eq!(ImportanceScore::new(1.5).get(), 1.0);
559        assert_eq!(ImportanceScore::new(0.5).get(), 0.5);
560    }
561
562    #[test]
563    fn test_importance_score_classification() {
564        assert!(ImportanceScore::new(0.8).is_high());
565        assert!(!ImportanceScore::new(0.5).is_high());
566        assert!(ImportanceScore::new(0.2).is_low());
567        assert!(!ImportanceScore::new(0.5).is_low());
568    }
569
570    // Display tests
571    #[test]
572    fn test_display_formatting() {
573        assert_eq!(TokenCount::new(100).to_string(), "100 tokens");
574        assert_eq!(LineNumber::new(42).to_string(), "L42");
575        assert_eq!(ByteOffset::new(1000).to_string(), "@1000");
576        assert_eq!(SymbolId::new(5).to_string(), "#5");
577    }
578}