vortex_scalar/bigint/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod bigcast;
5
6use std::fmt::Display;
7use std::ops::{Add, BitOr, Div, Mul, Rem, Shl, Shr, Sub};
8
9pub use bigcast::*;
10use num_traits::{CheckedAdd, CheckedSub, ConstZero, One, WrappingAdd, WrappingSub, Zero};
11use vortex_error::VortexExpect;
12
13/// Signed 256-bit integer type.
14///
15/// This is one of the physical representations of `DecimalScalar` values and can be safely converted
16/// back and forth with Arrow's [`i256`][arrow_buffer::i256].
17#[repr(transparent)]
18#[allow(non_camel_case_types)]
19#[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Hash, PartialOrd, Ord)]
20pub struct i256(arrow_buffer::i256);
21
22#[allow(clippy::same_name_method)]
23impl i256 {
24    /// The zero value for `i256`.
25    pub const ZERO: Self = Self(arrow_buffer::i256::ZERO);
26    /// The one value for `i256`.
27    pub const ONE: Self = Self(arrow_buffer::i256::ONE);
28    /// The maximum value for `i256`.
29    pub const MAX: Self = Self(arrow_buffer::i256::MAX);
30    /// The minimum value for `i256`.
31    pub const MIN: Self = Self(arrow_buffer::i256::MIN);
32
33    /// Construct a new `i256` from an unsigned `lower` bits and a signed `upper` bits.
34    pub const fn from_parts(lower: u128, upper: i128) -> Self {
35        Self(arrow_buffer::i256::from_parts(lower, upper))
36    }
37
38    /// Create an `i256` value from a signed 128-bit value.
39    pub const fn from_i128(i: i128) -> Self {
40        Self(arrow_buffer::i256::from_i128(i))
41    }
42
43    /// Attempts to convert this i256 to an i128.
44    ///
45    /// Returns None if the value is too large to fit in an i128.
46    pub fn maybe_i128(self) -> Option<i128> {
47        self.0.to_i128()
48    }
49
50    /// Create an integer value from its little-endian byte array representation.
51    pub const fn from_le_bytes(bytes: [u8; 32]) -> Self {
52        Self(arrow_buffer::i256::from_le_bytes(bytes))
53    }
54
55    /// Split the 256-bit signed integer value into an unsigned lower bits and a signed upper bits.
56    ///
57    /// This version gives us ownership of the value.
58    pub const fn into_parts(self) -> (u128, i128) {
59        self.0.to_parts()
60    }
61
62    /// Split the 256-bit signed integer value into an unsigned lower bits and a signed upper bits.
63    pub const fn to_parts(&self) -> (u128, i128) {
64        self.0.to_parts()
65    }
66
67    /// Raises self to the power of `exp`, wrapping around on overflow.
68    pub fn wrapping_pow(&self, exp: u32) -> Self {
69        Self(self.0.wrapping_pow(exp))
70    }
71
72    /// Wrapping (modular) addition. Computes `self + other`, wrapping around at the boundary.
73    pub fn wrapping_add(&self, other: Self) -> Self {
74        Self(self.0.wrapping_add(other.0))
75    }
76
77    /// Return the memory representation of this integer as a byte array in little-endian byte order.
78    #[inline]
79    pub const fn to_le_bytes(&self) -> [u8; 32] {
80        self.0.to_le_bytes()
81    }
82
83    /// Return the memory representation of this integer as a byte array in big-endian byte order.
84    #[inline]
85    pub const fn to_be_bytes(&self) -> [u8; 32] {
86        self.0.to_be_bytes()
87    }
88}
89
90impl From<i256> for arrow_buffer::i256 {
91    fn from(i: i256) -> Self {
92        i.0
93    }
94}
95
96impl From<arrow_buffer::i256> for i256 {
97    fn from(i: arrow_buffer::i256) -> Self {
98        Self(i)
99    }
100}
101
102impl Display for i256 {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        write!(f, "{}", self.0)
105    }
106}
107
108impl Add for i256 {
109    type Output = Self;
110
111    fn add(self, rhs: Self) -> Self::Output {
112        Self(self.0.add(rhs.0))
113    }
114}
115
116impl Sub for i256 {
117    type Output = Self;
118
119    fn sub(self, rhs: Self) -> Self::Output {
120        Self(self.0.sub(rhs.0))
121    }
122}
123
124impl Mul<Self> for i256 {
125    type Output = Self;
126
127    fn mul(self, rhs: Self) -> Self::Output {
128        Self(self.0.mul(rhs.0))
129    }
130}
131
132impl Div<Self> for i256 {
133    type Output = Self;
134
135    fn div(self, rhs: Self) -> Self::Output {
136        Self(self.0.div(rhs.0))
137    }
138}
139
140impl Rem<Self> for i256 {
141    type Output = Self;
142
143    fn rem(self, rhs: Self) -> Self::Output {
144        Self(self.0.rem(rhs.0))
145    }
146}
147
148impl Zero for i256 {
149    fn zero() -> Self {
150        Self::default()
151    }
152
153    fn is_zero(&self) -> bool {
154        *self == Self::zero()
155    }
156}
157
158impl ConstZero for i256 {
159    const ZERO: Self = Self(arrow_buffer::i256::ZERO);
160}
161
162impl One for i256 {
163    fn one() -> Self {
164        Self(arrow_buffer::i256::ONE)
165    }
166}
167
168impl CheckedAdd for i256 {
169    fn checked_add(&self, v: &Self) -> Option<Self> {
170        self.0.checked_add(v.0).map(Self)
171    }
172}
173
174impl WrappingAdd for i256 {
175    fn wrapping_add(&self, v: &Self) -> Self {
176        Self(self.0.wrapping_add(v.0))
177    }
178}
179
180impl CheckedSub for i256 {
181    fn checked_sub(&self, v: &Self) -> Option<Self> {
182        self.0.checked_sub(v.0).map(Self)
183    }
184}
185
186impl WrappingSub for i256 {
187    fn wrapping_sub(&self, v: &Self) -> Self {
188        Self(self.0.wrapping_sub(v.0))
189    }
190}
191
192impl Shr<Self> for i256 {
193    type Output = Self;
194
195    fn shr(self, rhs: Self) -> Self::Output {
196        use num_traits::ToPrimitive;
197
198        Self(
199            self.0.shr(
200                rhs.0
201                    .to_u8()
202                    .vortex_expect("Can't shift more than 256 bits"),
203            ),
204        )
205    }
206}
207
208impl Shl<usize> for i256 {
209    type Output = Self;
210
211    fn shl(self, rhs: usize) -> Self::Output {
212        use num_traits::ToPrimitive;
213        Self(
214            self.0
215                .shl(rhs.to_u8().vortex_expect("Can't shift more than 256 bits")),
216        )
217    }
218}
219
220impl BitOr<Self> for i256 {
221    type Output = Self;
222
223    fn bitor(self, rhs: Self) -> Self::Output {
224        Self(self.0.bitor(rhs.0))
225    }
226}
227
228impl num_traits::ToPrimitive for i256 {
229    fn to_i64(&self) -> Option<i64> {
230        self.maybe_i128().and_then(|v| v.to_i64())
231    }
232
233    fn to_i128(&self) -> Option<i128> {
234        self.maybe_i128()
235    }
236
237    fn to_u64(&self) -> Option<u64> {
238        self.maybe_i128().and_then(|v| v.to_u64())
239    }
240
241    fn to_u128(&self) -> Option<u128> {
242        self.maybe_i128().and_then(|v| v.to_u128())
243    }
244}
245
246#[cfg(test)]
247#[allow(clippy::many_single_char_names)]
248mod tests {
249    use num_traits::ToPrimitive;
250
251    use super::*;
252
253    #[test]
254    fn test_i256_constants() {
255        assert_eq!(i256::ZERO, i256::from_i128(0));
256        assert_eq!(i256::ONE, i256::from_i128(1));
257        assert!(i256::MIN < i256::ZERO);
258        assert!(i256::MAX > i256::ZERO);
259        assert!(i256::MIN < i256::MAX);
260    }
261
262    #[test]
263    fn test_i256_from_i128() {
264        let value = i256::from_i128(123456789);
265        assert_eq!(value.maybe_i128(), Some(123456789));
266
267        let negative = i256::from_i128(-987654321);
268        assert_eq!(negative.maybe_i128(), Some(-987654321));
269
270        let max_i128 = i256::from_i128(i128::MAX);
271        assert_eq!(max_i128.maybe_i128(), Some(i128::MAX));
272
273        let min_i128 = i256::from_i128(i128::MIN);
274        assert_eq!(min_i128.maybe_i128(), Some(i128::MIN));
275    }
276
277    #[test]
278    fn test_i256_from_parts() {
279        let value = i256::from_parts(1000, 2000);
280        let (lower, upper) = value.into_parts();
281        assert_eq!(lower, 1000);
282        assert_eq!(upper, 2000);
283
284        // Test to_parts (non-consuming)
285        let (lower2, upper2) = value.to_parts();
286        assert_eq!(lower2, 1000);
287        assert_eq!(upper2, 2000);
288    }
289
290    #[test]
291    fn test_i256_byte_conversions() {
292        let original = i256::from_i128(123456789012345);
293
294        // Test little-endian
295        let le_bytes = original.to_le_bytes();
296        let recovered_le = i256::from_le_bytes(le_bytes);
297        assert_eq!(original, recovered_le);
298
299        // Test big-endian
300        let be_bytes = original.to_be_bytes();
301        assert_ne!(le_bytes, be_bytes); // Should be different unless value is symmetric
302
303        // Test zero
304        let zero_le = i256::ZERO.to_le_bytes();
305        assert_eq!(zero_le, [0u8; 32]);
306    }
307
308    #[test]
309    fn test_i256_display() {
310        let value = i256::from_i128(42);
311        assert_eq!(format!("{value}"), "42");
312
313        let negative = i256::from_i128(-42);
314        assert_eq!(format!("{negative}"), "-42");
315    }
316
317    #[test]
318    fn test_i256_arithmetic_add() {
319        let a = i256::from_i128(100);
320        let b = i256::from_i128(200);
321        let sum = a + b;
322        assert_eq!(sum.maybe_i128(), Some(300));
323
324        // Test negative addition
325        let c = i256::from_i128(-50);
326        let sum2 = a + c;
327        assert_eq!(sum2.maybe_i128(), Some(50));
328    }
329
330    #[test]
331    fn test_i256_arithmetic_sub() {
332        let a = i256::from_i128(500);
333        let b = i256::from_i128(200);
334        let diff = a - b;
335        assert_eq!(diff.maybe_i128(), Some(300));
336
337        // Test negative result
338        let diff2 = b - a;
339        assert_eq!(diff2.maybe_i128(), Some(-300));
340    }
341
342    #[test]
343    fn test_i256_arithmetic_mul() {
344        let a = i256::from_i128(100);
345        let b = i256::from_i128(200);
346        let product = a * b;
347        assert_eq!(product.maybe_i128(), Some(20000));
348
349        // Test negative multiplication
350        let c = i256::from_i128(-5);
351        let product2 = a * c;
352        assert_eq!(product2.maybe_i128(), Some(-500));
353    }
354
355    #[test]
356    fn test_i256_arithmetic_div() {
357        let a = i256::from_i128(1000);
358        let b = i256::from_i128(25);
359        let quotient = a / b;
360        assert_eq!(quotient.maybe_i128(), Some(40));
361
362        // Test negative division
363        let c = i256::from_i128(-1000);
364        let quotient2 = c / b;
365        assert_eq!(quotient2.maybe_i128(), Some(-40));
366    }
367
368    #[test]
369    fn test_i256_arithmetic_rem() {
370        let a = i256::from_i128(103);
371        let b = i256::from_i128(10);
372        let remainder = a % b;
373        assert_eq!(remainder.maybe_i128(), Some(3));
374
375        // Test negative remainder
376        let c = i256::from_i128(-103);
377        let remainder2 = c % b;
378        assert_eq!(remainder2.maybe_i128(), Some(-3));
379    }
380
381    #[test]
382    fn test_i256_wrapping_pow() {
383        let base = i256::from_i128(2);
384        let result = base.wrapping_pow(10);
385        assert_eq!(result.maybe_i128(), Some(1024));
386
387        let base2 = i256::from_i128(10);
388        let result2 = base2.wrapping_pow(3);
389        assert_eq!(result2.maybe_i128(), Some(1000));
390
391        // Test with 0 exponent
392        let result3 = base.wrapping_pow(0);
393        assert_eq!(result3.maybe_i128(), Some(1));
394    }
395
396    #[test]
397    fn test_i256_wrapping_add() {
398        let a = i256::from_i128(100);
399        let b = i256::from_i128(200);
400        let result = a.wrapping_add(b);
401        assert_eq!(result.maybe_i128(), Some(300));
402
403        // Test the method version
404        let result2 = a.wrapping_add(b);
405        assert_eq!(result2.maybe_i128(), Some(300));
406    }
407
408    #[test]
409    fn test_i256_zero_trait() {
410        assert!(i256::zero().is_zero());
411        assert!(!i256::from_i128(1).is_zero());
412        assert!(!i256::from_i128(-1).is_zero());
413
414        // Test ConstZero
415        assert_eq!(i256::ZERO, <i256 as ConstZero>::ZERO);
416    }
417
418    #[test]
419    fn test_i256_one_trait() {
420        assert_eq!(i256::one(), i256::from_i128(1));
421        assert!(!i256::one().is_zero());
422    }
423
424    #[test]
425    fn test_i256_checked_add() {
426        let a = i256::from_i128(100);
427        let b = i256::from_i128(200);
428        let result = a.checked_add(&b);
429        assert_eq!(result, Some(i256::from_i128(300)));
430
431        // Note: Testing overflow would require values larger than i128
432    }
433
434    #[test]
435    fn test_i256_wrapping_add_trait() {
436        let a = i256::from_i128(100);
437        let b = i256::from_i128(200);
438        let result = <i256 as WrappingAdd>::wrapping_add(&a, &b);
439        assert_eq!(result.maybe_i128(), Some(300));
440    }
441
442    #[test]
443    fn test_i256_checked_sub() {
444        let a = i256::from_i128(500);
445        let b = i256::from_i128(200);
446        let result = a.checked_sub(&b);
447        assert_eq!(result, Some(i256::from_i128(300)));
448
449        // Test negative result
450        let result2 = b.checked_sub(&a);
451        assert_eq!(result2, Some(i256::from_i128(-300)));
452    }
453
454    #[test]
455    fn test_i256_wrapping_sub_trait() {
456        let a = i256::from_i128(500);
457        let b = i256::from_i128(200);
458        let result = <i256 as WrappingSub>::wrapping_sub(&a, &b);
459        assert_eq!(result.maybe_i128(), Some(300));
460    }
461
462    #[test]
463    fn test_i256_shift_right() {
464        let value = i256::from_i128(128);
465        let shift_amount = i256::from_i128(1);
466        let result = value >> shift_amount;
467        assert_eq!(result.maybe_i128(), Some(64));
468
469        let shift_amount2 = i256::from_i128(2);
470        let result2 = value >> shift_amount2;
471        assert_eq!(result2.maybe_i128(), Some(32));
472
473        // Shift by 0
474        let shift_zero = i256::from_i128(0);
475        let result3 = value >> shift_zero;
476        assert_eq!(result3.maybe_i128(), Some(128));
477    }
478
479    #[test]
480    fn test_i256_shift_left() {
481        let value = i256::from_i128(32);
482        let result = value << 1;
483        assert_eq!(result.maybe_i128(), Some(64));
484
485        let result2 = value << 2;
486        assert_eq!(result2.maybe_i128(), Some(128));
487
488        // Shift by 0
489        let result3 = value << 0;
490        assert_eq!(result3.maybe_i128(), Some(32));
491    }
492
493    #[test]
494    fn test_i256_bitor() {
495        let a = i256::from_i128(0b1010);
496        let b = i256::from_i128(0b1100);
497        let result = a | b;
498        assert_eq!(result.maybe_i128(), Some(0b1110));
499
500        // Test with zero
501        let result2 = a | i256::ZERO;
502        assert_eq!(result2.maybe_i128(), Some(0b1010));
503    }
504
505    #[test]
506    fn test_i256_to_primitive() {
507        let value = i256::from_i128(1000);
508
509        // Test to_i64
510        assert_eq!(value.to_i64(), Some(1000i64));
511
512        // Test to_i128
513        assert_eq!(value.to_i128(), Some(1000i128));
514
515        // Test to_u64
516        assert_eq!(value.to_u64(), Some(1000u64));
517
518        // Test to_u128
519        assert_eq!(value.to_u128(), Some(1000u128));
520
521        // Test negative value
522        let negative = i256::from_i128(-500);
523        assert_eq!(negative.to_i64(), Some(-500i64));
524        assert_eq!(negative.to_i128(), Some(-500i128));
525        assert_eq!(negative.to_u64(), None); // Can't convert negative to unsigned
526        assert_eq!(negative.to_u128(), None);
527    }
528
529    #[test]
530    fn test_i256_arrow_buffer_conversion() {
531        let arrow_value = arrow_buffer::i256::from_i128(42);
532        let our_value: i256 = arrow_value.into();
533        assert_eq!(our_value.maybe_i128(), Some(42));
534
535        // Convert back
536        let arrow_again: arrow_buffer::i256 = our_value.into();
537        assert_eq!(arrow_again, arrow_value);
538    }
539
540    #[test]
541    fn test_i256_default() {
542        let default_value = i256::default();
543        assert_eq!(default_value, i256::ZERO);
544        assert!(default_value.is_zero());
545    }
546
547    #[test]
548    fn test_i256_ordering() {
549        let a = i256::from_i128(100);
550        let b = i256::from_i128(200);
551        let c = i256::from_i128(-50);
552
553        assert!(a < b);
554        assert!(b > a);
555        assert!(c < a);
556        assert!(c < b);
557        assert_eq!(a, a);
558        assert_ne!(a, b);
559    }
560
561    #[test]
562    fn test_i256_hash() {
563        use std::collections::hash_map::DefaultHasher;
564        use std::hash::{Hash, Hasher};
565
566        let value1 = i256::from_i128(42);
567        let value2 = i256::from_i128(42);
568        let value3 = i256::from_i128(43);
569
570        let mut hasher1 = DefaultHasher::new();
571        value1.hash(&mut hasher1);
572        let hash1 = hasher1.finish();
573
574        let mut hasher2 = DefaultHasher::new();
575        value2.hash(&mut hasher2);
576        let hash2 = hasher2.finish();
577
578        let mut hasher3 = DefaultHasher::new();
579        value3.hash(&mut hasher3);
580        let hash3 = hasher3.finish();
581
582        assert_eq!(hash1, hash2); // Same values should have same hash
583        assert_ne!(hash1, hash3); // Different values should (likely) have different hash
584    }
585
586    #[test]
587    fn test_i256_large_value_loses_precision() {
588        // Create a value that doesn't fit in i128
589        let large_value = i256::from_parts(u128::MAX, 1);
590        assert_eq!(large_value.maybe_i128(), None);
591
592        // The parts should be preserved
593        let (lower, upper) = large_value.to_parts();
594        assert_eq!(lower, u128::MAX);
595        assert_eq!(upper, 1);
596    }
597}