Use wrapping arithmetic for Frame of Reference (#178)

fixes #54

Author: lwwmanning

vortex-array/src/scalar/primitive.rs

@@ -211,9 +211,9 @@ macro_rules! pscalar {
                         ..
                     }) => match pscalar {
                         PScalar::$ptype(v) => Ok(*v),
-                        _ => Err(vortex_err!(MismatchedTypes: "$T", pscalar.ptype())),
+                        _ => Err(vortex_err!(MismatchedTypes: any::type_name::<Self>(), pscalar.ptype())),
                     },
-                    _ => Err(vortex_err!("can't extract $T from scalar: {}", value)),
+                    _ => Err(vortex_err!("can't extract {} from scalar: {}", any::type_name::<Self>(), value)),
                 }
             }
         }

vortex-fastlanes/src/for/compress.rs

@@ -1,5 +1,5 @@
 use itertools::Itertools;
-use num_traits::PrimInt;
+use num_traits::{PrimInt, WrappingAdd, WrappingSub};
 
 use vortex::array::constant::ConstantArray;
 use vortex::array::downcast::DowncastArrayBuiltin;
@@ -70,7 +70,7 @@ impl EncodingCompression for FoREncoding {
     }
 }
 
-fn compress_primitive<T: NativePType + PrimInt>(
+fn compress_primitive<T: NativePType + WrappingSub + PrimInt>(
     parray: &PrimitiveArray,
     shift: u8,
 ) -> PrimitiveArray {
@@ -86,13 +86,13 @@ fn compress_primitive<T: NativePType + PrimInt>(
             .typed_data::<T>()
             .iter()
             .map(|&v| v >> shift as usize)
-            .map(|v| v - shifted_min)
+            .map(|v| v.wrapping_sub(&shifted_min))
             .collect_vec()
     } else {
         parray
             .typed_data::<T>()
             .iter()
-            .map(|&v| v - min)
+            .map(|&v| v.wrapping_sub(&min))
             .collect_vec()
     };
 
@@ -112,16 +112,23 @@ pub fn decompress(array: &FoRArray) -> VortexResult<PrimitiveArray> {
     }))
 }
 
-fn decompress_primitive<T: NativePType + PrimInt>(values: &[T], reference: T, shift: u8) -> Vec<T> {
+fn decompress_primitive<T: NativePType + WrappingAdd + PrimInt>(
+    values: &[T],
+    reference: T,
+    shift: u8,
+) -> Vec<T> {
     if shift > 0 {
         let shifted_reference = reference << shift as usize;
         values
             .iter()
             .map(|&v| v << shift as usize)
-            .map(|v| v + shifted_reference)
+            .map(|v| v.wrapping_add(&shifted_reference))
             .collect_vec()
     } else {
-        values.iter().map(|&v| v + reference).collect_vec()
+        values
+            .iter()
+            .map(|&v| v.wrapping_add(&reference))
+            .collect_vec()
     }
 }
 
@@ -144,6 +151,7 @@ mod test {
     use std::sync::Arc;
 
     use vortex::array::{Encoding, EncodingRef};
+    use vortex::compute::scalar_at::ScalarAtFn;
 
     use crate::BitPackedEncoding;
 
@@ -183,4 +191,30 @@ mod test {
         let decompressed = flatten_primitive(compressed.as_ref()).unwrap();
         assert_eq!(decompressed.typed_data::<u32>(), array.typed_data::<u32>());
     }
+
+    #[test]
+    fn test_overflow() {
+        let ctx = compress_ctx();
+
+        // Create a range offset by a million
+        let array = PrimitiveArray::from((i8::MIN..i8::MAX).collect_vec());
+        let compressed = FoREncoding {}.compress(&array, None, ctx).unwrap();
+        let compressed = compressed.as_for();
+        assert_eq!(i8::MIN, compressed.reference().try_into().unwrap());
+
+        let encoded = flatten_primitive(compressed.encoded()).unwrap();
+        let bitcast: &[u8] = unsafe { std::mem::transmute(encoded.typed_data::<i8>()) };
+        let unsigned: Vec<u8> = (0..u8::MAX).collect_vec();
+        assert_eq!(bitcast, unsigned.as_slice());
+
+        let decompressed = flatten_primitive(compressed).unwrap();
+        assert_eq!(decompressed.typed_data::<i8>(), array.typed_data::<i8>());
+        array
+            .typed_data::<i8>()
+            .iter()
+            .enumerate()
+            .for_each(|(i, v)| {
+                assert_eq!(*v, compressed.scalar_at(i).unwrap().try_into().unwrap());
+            });
+    }
 }

vortex-fastlanes/src/for/compute.rs

@@ -48,8 +48,9 @@ impl ScalarAtFn for FoRArray {
             (Scalar::Primitive(p), Scalar::Primitive(r)) => match p.value() {
                 None => Ok(encoded_scalar),
                 Some(pv) => match_each_integer_ptype!(pv.ptype(), |$P| {
+                    use num_traits::WrappingAdd;
                     Ok(PrimitiveScalar::try_new::<$P>(
-                        Some((p.typed_value::<$P>().unwrap() << self.shift()) + r.typed_value::<$P>().unwrap()),
+                        Some((p.typed_value::<$P>().unwrap() << self.shift()).wrapping_add(r.typed_value::<$P>().unwrap())),
                         p.dtype().nullability()
                     ).unwrap().into())
                 }),