risingwave_common/types/
num256.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::{Display, Formatter, Write};
16use std::hash::Hasher;
17use std::io::Read;
18use std::mem;
19use std::num::ParseIntError;
20use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
21use std::str::FromStr;
22
23use bytes::{BufMut, Bytes};
24use ethnum::{AsI256, i256, u256};
25use num_traits::{
26    CheckedAdd, CheckedDiv, CheckedMul, CheckedNeg, CheckedRem, CheckedSub, Num, One, Zero,
27};
28use risingwave_common_estimate_size::EstimateSize;
29use risingwave_pb::data::ArrayType;
30use serde::{Deserialize, Serialize};
31use to_text::ToText;
32
33use crate::array::ArrayResult;
34use crate::types::to_binary::ToBinary;
35use crate::types::{Buf, DataType, F64, Scalar, ScalarRef, to_text};
36
37/// A 256-bit signed integer.
38#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Default, Hash)]
39pub struct Int256(pub(crate) Box<i256>);
40
41/// A reference to an `Int256` value.
42#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd)]
43pub struct Int256Ref<'a>(pub &'a i256);
44
45impl Display for Int256Ref<'_> {
46    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
47        self.write(f)
48    }
49}
50
51macro_rules! impl_common_for_num256 {
52    ($scalar:ident, $scalar_ref:ident < $gen:tt > , $inner:ty, $array_type:ident) => {
53        impl Scalar for $scalar {
54            type ScalarRefType<$gen> = $scalar_ref<$gen>;
55
56            fn as_scalar_ref(&self) -> Self::ScalarRefType<'_> {
57                $scalar_ref(self.0.as_ref())
58            }
59        }
60
61        impl<$gen> ScalarRef<$gen> for $scalar_ref<$gen> {
62            type ScalarType = $scalar;
63
64            fn to_owned_scalar(&self) -> Self::ScalarType {
65                $scalar((*self.0).into())
66            }
67
68            fn hash_scalar<H: Hasher>(&self, state: &mut H) {
69                use std::hash::Hash as _;
70                self.0.hash(state)
71            }
72        }
73
74        impl FromStr for $scalar {
75            type Err = ParseIntError;
76
77            fn from_str(s: &str) -> Result<Self, Self::Err> {
78                <$inner>::from_str(s).map(Into::into)
79            }
80        }
81
82        impl $scalar {
83            #[inline]
84            pub fn min_value() -> Self {
85                Self::from(<$inner>::MIN)
86            }
87
88            #[inline]
89            pub fn max_value() -> Self {
90                Self::from(<$inner>::MAX)
91            }
92
93            #[inline]
94            pub fn into_inner(self) -> $inner {
95                *self.0
96            }
97
98            #[inline]
99            pub const fn size() -> usize {
100                mem::size_of::<$inner>()
101            }
102
103            #[inline]
104            pub fn array_type() -> ArrayType {
105                ArrayType::$array_type
106            }
107
108            #[inline]
109            pub fn from_ne_bytes(bytes: [u8; mem::size_of::<$inner>()]) -> Self {
110                Self(Box::new(<$inner>::from_ne_bytes(bytes)))
111            }
112
113            #[inline]
114            pub fn from_le_bytes(bytes: [u8; mem::size_of::<$inner>()]) -> Self {
115                Self(Box::new(<$inner>::from_le_bytes(bytes)))
116            }
117
118            #[inline]
119            pub fn from_be_bytes(bytes: [u8; mem::size_of::<$inner>()]) -> Self {
120                Self(Box::new(<$inner>::from_be_bytes(bytes)))
121            }
122
123            pub fn from_protobuf(input: &mut impl Read) -> ArrayResult<Self> {
124                let mut buf = [0u8; mem::size_of::<$inner>()];
125                input.read_exact(&mut buf)?;
126                Ok(Self::from_be_bytes(buf))
127            }
128        }
129
130        impl From<$inner> for $scalar {
131            fn from(value: $inner) -> Self {
132                Self(Box::new(value))
133            }
134        }
135
136        impl $scalar_ref<'_> {
137            #[inline]
138            pub fn to_le_bytes(self) -> [u8; mem::size_of::<$inner>()] {
139                self.0.to_le_bytes()
140            }
141
142            #[inline]
143            pub fn to_be_bytes(self) -> [u8; mem::size_of::<$inner>()] {
144                self.0.to_be_bytes()
145            }
146
147            #[inline]
148            pub fn to_ne_bytes(self) -> [u8; mem::size_of::<$inner>()] {
149                self.0.to_ne_bytes()
150            }
151
152            pub fn to_protobuf<T: std::io::Write>(self, output: &mut T) -> ArrayResult<usize> {
153                output.write(&self.to_be_bytes()).map_err(Into::into)
154            }
155        }
156
157        impl ToText for $scalar_ref<'_> {
158            fn write<W: Write>(&self, f: &mut W) -> std::fmt::Result {
159                write!(f, "{}", self.0)
160            }
161
162            fn write_with_type<W: Write>(&self, _ty: &DataType, f: &mut W) -> std::fmt::Result {
163                self.write(f)
164            }
165        }
166
167        impl ToBinary for $scalar_ref<'_> {
168            fn to_binary_with_type(&self, _ty: &DataType) -> super::to_binary::Result<Bytes> {
169                let mut output = bytes::BytesMut::new();
170                let buffer = self.to_be_bytes();
171                output.put_slice(&buffer);
172                Ok(output.freeze())
173            }
174        }
175
176        impl $scalar {
177            pub fn from_binary(mut input: &[u8]) -> ArrayResult<Self> {
178                let mut buf = [0; Self::size()];
179                input.read_exact(&mut buf)?;
180                Ok(Self::from_be_bytes(buf))
181            }
182        }
183    };
184}
185
186impl_common_for_num256!(Int256, Int256Ref<'a>, i256, Int256);
187
188impl Int256 {
189    // `i256::str_from_hex` and `i256::str_from_prefixed` doesn't support uppercase "0X", so when it
190    // fails it will try to parse the lowercase version of the `src`
191
192    // `from_str_prefixed` function accepts string inputs that start with "0x". If the parsing
193    // fails, it will attempt to parse the input as a decimal value.
194    pub fn from_str_prefixed(src: &str) -> Result<Self, ParseIntError> {
195        u256::from_str_prefixed(src)
196            .or_else(|_| u256::from_str_prefixed(&src.to_lowercase()))
197            .map(|u| u.as_i256().into())
198    }
199
200    // `from_str_hex` function only accepts string inputs that start with "0x".
201    pub fn from_str_hex(src: &str) -> Result<Self, ParseIntError> {
202        u256::from_str_hex(src)
203            .or_else(|_| u256::from_str_hex(&src.to_lowercase()))
204            .map(|u| u.as_i256().into())
205    }
206}
207
208impl Int256Ref<'_> {
209    pub fn memcmp_serialize(
210        &self,
211        serializer: &mut memcomparable::Serializer<impl bytes::BufMut>,
212    ) -> memcomparable::Result<()> {
213        let (hi, lo) = self.0.into_words();
214        (hi, lo as u128).serialize(serializer)
215    }
216}
217
218impl Int256 {
219    pub const MEMCMP_ENCODED_SIZE: usize = 32;
220
221    pub fn memcmp_deserialize(
222        deserializer: &mut memcomparable::Deserializer<impl Buf>,
223    ) -> memcomparable::Result<Self> {
224        let (hi, lo) = <(i128, u128)>::deserialize(deserializer)?;
225        let signed = i256::from_words(hi, lo as i128);
226        Ok(Int256::from(signed))
227    }
228}
229
230macro_rules! impl_convert_from {
231    ($($t:ty),* $(,)?) => {$(
232        impl From<$t> for Int256 {
233            #[inline]
234            fn from(value: $t) -> Self {
235                Self(Box::new(value.as_i256()))
236            }
237        }
238    )*};
239}
240
241impl_convert_from!(i16, i32, i64);
242
243impl<'a> From<Int256Ref<'a>> for F64 {
244    fn from(value: Int256Ref<'a>) -> Self {
245        Self::from(value.0.as_f64())
246    }
247}
248
249// Conversions for mathematical operations.
250impl From<Int256Ref<'_>> for Int256 {
251    fn from(value: Int256Ref<'_>) -> Self {
252        value.to_owned_scalar()
253    }
254}
255
256macro_rules! impl_checked_op {
257    ($trait:ty, $func:ident, $op:tt, $proxied_trait:tt, $proxied_func:ident, $scalar:ident, $scalar_ref:ident < $gen:tt >) => {
258        impl $proxied_trait<Self> for $scalar {
259            type Output = Self;
260
261            fn $proxied_func(self, rhs: Self) -> Self::Output {
262                $scalar::from(self.0.as_ref() $op rhs.0.as_ref())
263            }
264        }
265
266        impl<$gen> $proxied_trait<Self> for $scalar_ref<$gen> {
267            type Output = $scalar;
268
269            fn $proxied_func(self, rhs: Self) -> Self::Output {
270                Int256::from(self.0 $op rhs.0)
271            }
272        }
273
274        impl $trait for $scalar {
275            fn $func(&self, other: &Self) -> Option<Self> {
276                self.0.$func(*other.0).map(Into::into)
277            }
278        }
279    };
280}
281
282impl_checked_op!(CheckedAdd, checked_add, +, Add, add, Int256, Int256Ref<'a>);
283impl_checked_op!(CheckedSub, checked_sub, -, Sub, sub, Int256, Int256Ref<'a>);
284impl_checked_op!(CheckedMul, checked_mul, *, Mul, mul, Int256, Int256Ref<'a>);
285impl_checked_op!(CheckedDiv, checked_div, /, Div, div, Int256, Int256Ref<'a>);
286impl_checked_op!(CheckedRem, checked_rem, %, Rem, rem, Int256, Int256Ref<'a>);
287
288impl Neg for Int256 {
289    type Output = Int256;
290
291    fn neg(self) -> Self::Output {
292        Int256::from(self.0.neg())
293    }
294}
295
296impl CheckedNeg for Int256 {
297    fn checked_neg(&self) -> Option<Self> {
298        self.0.checked_neg().map(Into::into)
299    }
300}
301
302impl Zero for Int256 {
303    fn zero() -> Self {
304        Int256::from(i256::ZERO)
305    }
306
307    fn is_zero(&self) -> bool {
308        !self.0.is_negative() && !self.0.is_positive()
309    }
310}
311
312impl One for Int256 {
313    fn one() -> Self {
314        Self::from(i256::ONE)
315    }
316}
317
318impl Num for Int256 {
319    type FromStrRadixErr = ParseIntError;
320
321    fn from_str_radix(str: &str, radix: u32) -> Result<Self, Self::FromStrRadixErr> {
322        i256::from_str_radix(str, radix).map(Into::into)
323    }
324}
325
326impl EstimateSize for Int256 {
327    fn estimated_heap_size(&self) -> usize {
328        mem::size_of::<i128>() * 2
329    }
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335
336    macro_rules! check_op {
337        ($t:ty, $lhs:expr, $rhs:expr, [$($op:tt),+]) => {
338            $(assert_eq!(
339                Int256::from($lhs as $t) $op Int256::from($rhs as $t),
340                Int256::from(($lhs as $t) $op ($rhs as $t))
341            );)+
342        };
343    }
344
345    macro_rules! check_checked_op {
346        ($t:ty, $lhs:expr, $rhs:expr, [$($f:ident),+]) => {
347            $(assert_eq!(
348                Int256::from($lhs as $t).$f(&Int256::from($rhs as $t)),
349                ($lhs as $t).$f(($rhs as $t)).map(Int256::from)
350            );)+
351        };
352    }
353
354    macro_rules! generate_op_test {
355        ($($t:ty),* $(,)?) => {$(
356            check_op!($t, 1, 1, [+, -, *, /, %]);
357            check_op!($t, -1, 1, [+, -, *, /, %]);
358            check_op!($t, 0, 1, [+, -, *, /, %]);
359            check_op!($t, -12, 34, [+, -, *, /, %]);
360            check_op!($t, 12, -34, [+, -, *, /, %]);
361            check_op!($t, -12, -34, [+, -, *, /, %]);
362            check_op!($t, 12, 34, [+, -, *, /, %]);
363        )*};
364    }
365
366    macro_rules! generate_checked_op_test {
367        ($($t:ty),* $(,)?) => {$(
368            check_checked_op!($t, 1, 1, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
369            check_checked_op!($t, 1, 0, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
370            check_checked_op!($t, -1, 1, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
371            check_checked_op!($t, -1, 0, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
372            check_checked_op!($t, 0, 1, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
373            check_checked_op!($t, -12, 34, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
374            check_checked_op!($t, 12, -34, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
375            check_checked_op!($t, -12, -34, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
376            check_checked_op!($t, 12, 34, [checked_add, checked_sub, checked_mul, checked_div, checked_rem]);
377        )*};
378    }
379
380    #[test]
381    fn test_checked_op() {
382        generate_checked_op_test!(i16, i32, i64);
383    }
384
385    #[test]
386    fn test_op() {
387        generate_op_test!(i16, i32, i64);
388    }
389
390    #[test]
391    fn test_zero() {
392        let zero = Int256::zero();
393        assert_eq!(zero, Int256::from(0));
394        assert!(zero.is_zero());
395    }
396
397    #[test]
398    fn test_neg() {
399        assert_eq!(Int256::from(1).neg(), Int256::from(-1));
400        assert_eq!(-Int256::from(1), Int256::from(-1));
401        assert_eq!(Int256::from(0).neg(), Int256::from(0));
402        assert_eq!(-Int256::from(0), Int256::from(0));
403    }
404
405    #[test]
406    fn test_float64() {
407        let vs: Vec<i64> = vec![-9007199254740990, -100, -1, 0, 1, 100, 9007199254740991];
408
409        for v in vs {
410            let i = Int256::from(v);
411            assert_eq!(F64::from(i.as_scalar_ref()), F64::from(v));
412        }
413    }
414
415    #[test]
416    fn hex_to_int256() {
417        assert_eq!(Int256::from_str_hex("0x0").unwrap(), Int256::from(0));
418        assert_eq!(Int256::from_str_hex("0x1").unwrap(), Int256::from(1));
419        assert_eq!(
420            Int256::from_str_hex(
421                "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
422            )
423            .unwrap(),
424            Int256::from(-1)
425        );
426        assert_eq!(Int256::from_str_hex("0xa").unwrap(), Int256::from(10));
427        assert_eq!(Int256::from_str_hex("0xA").unwrap(), Int256::from(10));
428        assert_eq!(Int256::from_str_hex("0Xff").unwrap(), Int256::from(255));
429
430        assert_eq!(
431            Int256::from_str_hex(
432                "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff01"
433            )
434            .unwrap(),
435            Int256::from(-255)
436        );
437
438        assert_eq!(
439            Int256::from_str_hex("0xf").unwrap(),
440            Int256::from_str("15").unwrap()
441        );
442        assert_eq!(
443            Int256::from_str_hex("0xfffff").unwrap(),
444            Int256::from_str("1048575").unwrap()
445        );
446        assert_eq!(
447            Int256::from_str_hex("0xfffffffff").unwrap(),
448            Int256::from_str("68719476735").unwrap()
449        );
450        assert_eq!(
451            Int256::from_str_hex("0xfffffffffffff").unwrap(),
452            Int256::from_str("4503599627370495").unwrap()
453        );
454        assert_eq!(
455            Int256::from_str_hex("0xfffffffffffffffff").unwrap(),
456            Int256::from_str("295147905179352825855").unwrap()
457        );
458        assert_eq!(
459            Int256::from_str_hex("0xfffffffffffffffffffff").unwrap(),
460            Int256::from_str("19342813113834066795298815").unwrap()
461        );
462        assert_eq!(
463            Int256::from_str_hex("0xfffffffffffffffffffffffff").unwrap(),
464            Int256::from_str("1267650600228229401496703205375").unwrap()
465        );
466        assert_eq!(
467            Int256::from_str_hex("0xfffffffffffffffffffffffffffff").unwrap(),
468            Int256::from_str("83076749736557242056487941267521535").unwrap()
469        );
470        assert_eq!(
471            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffff").unwrap(),
472            Int256::from_str("5444517870735015415413993718908291383295").unwrap()
473        );
474        assert_eq!(
475            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffff").unwrap(),
476            Int256::from_str("356811923176489970264571492362373784095686655").unwrap()
477        );
478        assert_eq!(
479            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffffffff").unwrap(),
480            Int256::from_str("23384026197294446691258957323460528314494920687615").unwrap()
481        );
482        assert_eq!(
483            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffffffffffff").unwrap(),
484            Int256::from_str("1532495540865888858358347027150309183618739122183602175").unwrap()
485        );
486        assert_eq!(
487            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffffffffffffffff").unwrap(),
488            Int256::from_str("100433627766186892221372630771322662657637687111424552206335")
489                .unwrap()
490        );
491        assert_eq!(
492            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffffffffffffffffffff")
493                .unwrap(),
494            Int256::from_str("6582018229284824168619876730229402019930943462534319453394436095")
495                .unwrap()
496        );
497        assert_eq!(
498            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
499                .unwrap(),
500            Int256::from_str(
501                "431359146674410236714672241392314090778194310760649159697657763987455"
502            )
503            .unwrap()
504        );
505        assert_eq!(
506            Int256::from_str_hex("0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
507                .unwrap(),
508            Int256::from_str(
509                "28269553036454149273332760011886696253239742350009903329945699220681916415"
510            )
511            .unwrap()
512        );
513
514        // int256 max
515        assert_eq!(
516            Int256::from_str_hex(
517                "0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
518            )
519            .unwrap(),
520            Int256::max_value(),
521        );
522
523        // int256 min
524        assert_eq!(
525            Int256::from_str_hex(
526                "0x8000000000000000000000000000000000000000000000000000000000000000"
527            )
528            .unwrap(),
529            Int256::min_value(),
530        );
531    }
532
533    #[test]
534    fn test_num256_estimate_size() {
535        let num256 = Int256::min_value();
536        assert_eq!(num256.estimated_size(), 40);
537    }
538}