feldera-sqllib 0.331.0

//! byte arrays (binary objects in SQL)

use crate::{
    SqlString, some_function1, some_function2, some_function3, some_function4,
    some_polymorphic_function1, some_polymorphic_function2,
};
use base58::{FromBase58, ToBase58};
use base64::prelude::*;
use dbsp::NumEntries;
use feldera_macros::IsNone;
use feldera_types::serde_with_context::{
    DeserializeWithContext, SerializeWithContext, SqlSerdeConfig, serde_config::BinaryFormat,
};
use flate2::read::GzDecoder;
use hex::ToHex;
use md5::{Digest, Md5};
use serde::{
    Deserialize, Deserializer, Serialize, Serializer,
    de::{Error as _, Visitor},
};
use size_of::SizeOf;
use smallvec::{SmallVec, smallvec};
use std::{
    borrow::Cow,
    cmp::{Ordering, min},
    fmt::Debug,
    io::Read,
};
use xxhash_rust::xxh64::xxh64;

/// Values smaller than this size are allocated on the stack
const THRESHOLD: usize = 32; // up to 256 bits
type CompactVec = SmallVec<[u8; THRESHOLD]>;

/// A ByteArray object, representing a SQL value with type
/// `BINARY` or `VARBINARY`.
#[derive(
    Debug,
    Default,
    Clone,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Serialize,
    Deserialize,
    rkyv::Archive,
    rkyv::Serialize,
    rkyv::Deserialize,
    IsNone,
)]
#[archive_attr(derive(Ord, Eq, PartialEq, PartialOrd))]
#[serde(transparent)]
pub struct ByteArray {
    data: CompactVec,
}

impl SizeOf for ByteArray {
    fn size_of_children(&self, context: &mut size_of::Context) {
        self.data.size_of_children(context);
    }
}

impl SerializeWithContext<SqlSerdeConfig> for ByteArray {
    fn serialize_with_context<S>(
        &self,
        serializer: S,
        context: &SqlSerdeConfig,
    ) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match context.binary_format {
            BinaryFormat::Array => self.data.serialize(serializer),
            BinaryFormat::Base64 => serializer.serialize_str(&BASE64_STANDARD.encode(&self.data)),
            BinaryFormat::Base58 => serializer.serialize_str(&self.data.to_base58()),
            BinaryFormat::Bytes => serializer.serialize_bytes(&self.data),
            BinaryFormat::PgHex => {
                serializer.serialize_str(&format!("\\x{}", hex::encode(&self.data)))
            }
            BinaryFormat::CHex => {
                serializer.serialize_str(&format!("0x{}", hex::encode(&self.data)))
            }
        }
    }
}

struct ByteVisitor;

impl Visitor<'_> for ByteVisitor {
    type Value = ByteArray;

    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        formatter.write_str("byte array")
    }

    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        Ok(ByteArray::new(v))
    }
}

impl<'de, AUX> DeserializeWithContext<'de, SqlSerdeConfig, AUX> for ByteArray {
    fn deserialize_with_context<D>(
        deserializer: D,
        config: &'de SqlSerdeConfig,
    ) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        fn parse_hex_string(s: &str, prefix: &str) -> Option<CompactVec> {
            let s = s.strip_prefix(prefix).unwrap_or(s).as_bytes();
            if !s.len().is_multiple_of(2) || !s.iter().all(u8::is_ascii_hexdigit) {
                None
            } else {
                let mut result = CompactVec::with_capacity(s.len() / 2);
                for i in 0..s.len() / 2 {
                    let a = (s[i * 2] as char).to_digit(16).unwrap() as u8;
                    let b = (s[i * 2 + 1] as char).to_digit(16).unwrap() as u8;
                    result.push(a * 16 + b);
                }
                Some(result)
            }
        }

        match config.binary_format {
            BinaryFormat::Array => {
                let data = CompactVec::deserialize(deserializer)?;
                Ok(Self { data })
            }
            BinaryFormat::Base64 => {
                let str: Cow<'de, str> = Deserialize::deserialize(deserializer)?;
                let data = BASE64_STANDARD
                    .decode(&*str)
                    .map_err(|e| D::Error::custom(format!("invalid base64 string: {e}")))?;
                Ok(Self { data: data.into() })
            }
            BinaryFormat::Base58 => {
                let str: Cow<'de, str> = Deserialize::deserialize(deserializer)?;
                let data = str
                    .from_base58()
                    .map_err(|e| D::Error::custom(format!("invalid base58 string: {e:?}")))?;
                Ok(Self { data: data.into() })
            }
            BinaryFormat::Bytes => deserializer.deserialize_bytes(ByteVisitor),
            BinaryFormat::PgHex => Err(D::Error::custom(
                "binary format Postgres Hexadecimal is not supported for input",
            )),
            BinaryFormat::CHex => {
                let str: Cow<'de, str> = Deserialize::deserialize(deserializer)?;
                match parse_hex_string(&str, "0x") {
                    None => Err(D::Error::custom(format!(
                        "Invalid C-style hex string: {str:?}"
                    ))),
                    Some(data) => Ok(Self { data }),
                }
            }
        }
    }
}

#[cfg(test)]
mod test_binary_deserializer {
    use feldera_types::{
        format::json::JsonFlavor,
        serde_with_context::{DeserializeWithContext, SerializeWithContext, SqlSerdeConfig},
    };

    use super::ByteArray;

    #[test]
    fn test_base58() {
        let encoded = "4F85ZySpwyY6FuH7mQYyyr5b8nV9zFRBLj92AJa37w6y";

        let decoded =
            <ByteArray as DeserializeWithContext<SqlSerdeConfig, ()>>::deserialize_with_context(
                serde_json::Value::String(encoded.to_string()),
                &SqlSerdeConfig::from(JsonFlavor::Blockchain),
            )
            .unwrap();

        assert_eq!(
            decoded,
            ByteArray::from(b"012345678901234567890123456789ab".as_slice())
        );

        let mut reencoded = Vec::<u8>::new();

        decoded
            .serialize_with_context(
                &mut serde_json::Serializer::new(&mut reencoded),
                &SqlSerdeConfig::from(JsonFlavor::Blockchain),
            )
            .unwrap();

        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&reencoded).unwrap(),
            serde_json::Value::String(encoded.to_string())
        );
    }
}

#[doc(hidden)]
impl NumEntries for &ByteArray {
    const CONST_NUM_ENTRIES: Option<usize> = None;

    #[doc(hidden)]
    #[inline]
    fn num_entries_shallow(&self) -> usize {
        self.length()
    }

    #[doc(hidden)]
    #[inline]
    fn num_entries_deep(&self) -> usize {
        self.length()
    }
}

impl From<&[u8]> for ByteArray {
    fn from(value: &[u8]) -> Self {
        Self::new(value)
    }
}

impl ByteArray {
    /// Create a ByteArray from a slice of bytes
    pub fn new(d: &[u8]) -> Self {
        Self { data: d.into() }
    }

    pub fn with_size(d: &[u8], size: i32, fixed: bool) -> Self {
        if size < 0 {
            ByteArray::new(d)
        } else {
            let size = size as usize;
            match d.len().cmp(&size) {
                Ordering::Equal => ByteArray::new(d),
                Ordering::Greater => ByteArray::new(&d[..size]),
                Ordering::Less => {
                    if fixed {
                        let mut data: CompactVec = smallvec![0; size];
                        data[..d.len()].copy_from_slice(d);
                        ByteArray { data }
                    } else {
                        ByteArray::new(d)
                    }
                }
            }
        }
    }

    pub fn with_size_truncate_left(d: &[u8], size: i32, fixed: bool) -> Self {
        if size < 0 {
            ByteArray::new(d)
        } else {
            let size = size as usize;
            match d.len().cmp(&size) {
                Ordering::Equal => ByteArray::new(d),
                Ordering::Greater => ByteArray::new(&d[d.len() - size..]),
                Ordering::Less => {
                    if fixed {
                        let mut data: CompactVec = smallvec![0; size];
                        data[size - d.len()..].copy_from_slice(d);
                        ByteArray { data }
                    } else {
                        ByteArray::new(d)
                    }
                }
            }
        }
    }

    pub fn zero(size: usize) -> Self {
        Self {
            data: smallvec![0; size],
        }
    }

    /// Create a ByteArray from a Vector of bytes
    pub fn from_vec(d: Vec<u8>) -> Self {
        Self { data: d.into() }
    }

    /// Length of the byte array in bytes
    pub fn length(&self) -> usize {
        self.data.len()
    }

    #[doc(hidden)]
    /// Combine two byte arrays of the same length using
    /// a pointwise function.  panics if the lengths
    /// are not the same.
    pub fn zip<F>(&self, other: &Self, op: F) -> ByteArray
    where
        F: Fn(&u8, &u8) -> u8,
    {
        let self_len = self.data.len();
        let other_len = other.data.len();
        if self_len != other_len {
            panic!(
                "Cannot operate on BINARY objects of different sizes {} and {}",
                self_len, other_len
            );
        }
        let result: Vec<u8> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(l, r)| op(l, r))
            .collect();
        ByteArray::new(&result)
    }

    #[doc(hidden)]
    /// Bytewise 'and' of two byte arrays of the same length.
    /// Panics if the arrays do not have the same length.
    pub fn and(&self, other: &Self) -> Self {
        self.zip(other, |left, right| left & right)
    }

    #[doc(hidden)]
    /// Bytewise 'or' of two byte arrays of the same length.
    /// Panics if the arrays do not have the same length.
    pub fn or(&self, other: &Self) -> Self {
        self.zip(other, |left, right| left | right)
    }

    #[doc(hidden)]
    /// Bytewise 'xor' of two byte arrays of the same length.
    /// Panics if the arrays do not have the same length.
    pub fn xor(&self, other: &Self) -> Self {
        self.zip(other, |left, right| left ^ right)
    }

    /// Concatenate two byte arrays, produces a new byte array.
    pub fn concat(&self, other: &Self) -> Self {
        let mut r = Vec::<u8>::with_capacity(self.data.len() + other.data.len());
        r.extend(&self.data);
        r.extend(&other.data);
        ByteArray::from_vec(r)
    }

    /// Get a reference to the data in a ByteArray as a byte slice
    pub fn as_slice(&self) -> &[u8] {
        &self.data
    }
}

#[doc(hidden)]
pub fn concat_bytes_bytes(left: ByteArray, right: ByteArray) -> ByteArray {
    left.concat(&right)
}

some_polymorphic_function2!(concat, bytes, ByteArray, bytes, ByteArray, ByteArray);

#[doc(hidden)]
pub fn to_hex_(value: ByteArray) -> SqlString {
    SqlString::from(value.data.encode_hex::<String>())
}

some_function1!(to_hex, ByteArray, SqlString);

#[doc(hidden)]
pub fn octet_length_(value: ByteArray) -> i32 {
    value.length() as i32
}

some_function1!(octet_length, ByteArray, i32);

#[doc(hidden)]
pub fn binary_position__(needle: ByteArray, haystack: ByteArray) -> i32 {
    haystack
        .data
        .windows(needle.data.len())
        .position(|window| *window == *needle.data)
        .map(|v| v + 1)
        .unwrap_or(0) as i32
}

some_function2!(binary_position, ByteArray, ByteArray, i32);

#[doc(hidden)]
pub fn binary_substring2__(source: ByteArray, left: i32) -> ByteArray {
    // SQL indexing starts at 1
    let start = if left < 1 { 0 } else { left - 1 };
    let data = source.data.into_iter().skip(start as usize).collect();

    ByteArray { data }
}

some_function2!(binary_substring2, ByteArray, i32, ByteArray);

#[doc(hidden)]
pub fn binary_substring3___(source: ByteArray, left: i32, count: i32) -> ByteArray {
    // SQL indexing starts at 1
    let start = if left < 1 { 0 } else { left - 1 };
    if count < 0 {
        return ByteArray::default();
    }
    let count = count as usize;

    let data = source
        .data
        .into_iter()
        .skip(start as usize)
        .take(count)
        .collect();

    ByteArray { data }
}

some_function3!(binary_substring3, ByteArray, i32, i32, ByteArray);

#[doc(hidden)]
pub fn left_bytes_i32(source: ByteArray, size: i32) -> ByteArray {
    binary_substring3___(source, 1, size)
}

some_polymorphic_function2!(left, bytes, ByteArray, i32, i32, ByteArray);

#[doc(hidden)]
pub fn right_bytes_i32(source: ByteArray, size: i32) -> ByteArray {
    if size <= 0 {
        return ByteArray::default();
    }
    let size = size as usize;
    let start = if size >= source.length() {
        1
    } else {
        source.length() - size + 1
    };
    binary_substring3___(source, start as i32, size as i32)
}

some_polymorphic_function2!(right, bytes, ByteArray, i32, i32, ByteArray);

#[doc(hidden)]
pub fn binary_overlay3___(source: ByteArray, replacement: ByteArray, position: i32) -> ByteArray {
    let len = replacement.length() as i32;
    binary_overlay4____(source, replacement, position, len)
}

some_function3!(binary_overlay3, ByteArray, ByteArray, i32, ByteArray);

#[doc(hidden)]
pub fn binary_overlay4____(
    source: ByteArray,
    mut replacement: ByteArray,
    position: i32,
    remove: i32,
) -> ByteArray {
    let mut remove = remove;
    if remove < 0 {
        remove = 0;
    }

    if position <= 0 {
        source
    } else if position > source.length() as i32 {
        source.concat(&replacement)
    } else {
        let mut result = binary_substring3___(source.clone(), 0, position - 1);
        result.data.append(&mut replacement.data);
        let mut substr = binary_substring2__(source, position + remove);
        result.data.append(&mut substr.data);
        result
    }
}

some_function4!(binary_overlay4, ByteArray, ByteArray, i32, i32, ByteArray);

#[doc(hidden)]
pub fn gunzip_(source: ByteArray) -> SqlString {
    let mut gz = GzDecoder::new(&source.data[..]);
    let mut s = String::new();

    gz.read_to_string(&mut s)
        .expect("failed to decompress gzipped data");
    SqlString::from(s)
}

some_function1!(gunzip, ByteArray, SqlString);

#[doc(hidden)]
pub fn to_int_(source: ByteArray) -> i32 {
    let mut result = 0;
    for i in 0..min(4, source.length()) {
        result = (result << 8) | (source.data[i] as i32);
    }
    result
}

some_function1!(to_int, ByteArray, i32);

#[doc(hidden)]
pub fn md5_bytes(source: ByteArray) -> SqlString {
    let mut hasher = Md5::new();
    hasher.update(source.data);
    let result = hasher.finalize();
    SqlString::from(format!("{:x}", result))
}

some_polymorphic_function1!(md5, bytes, ByteArray, SqlString);

#[doc(hidden)]
pub fn bin2utf8_(bytes: ByteArray) -> Option<SqlString> {
    std::str::from_utf8(&bytes.data)
        .ok()
        .map(SqlString::from_ref)
}

#[doc(hidden)]
pub fn bin2utf8N(source: Option<ByteArray>) -> Option<SqlString> {
    match source {
        None => None,
        Some(bytes) => bin2utf8_(bytes),
    }
}

#[doc(hidden)]
pub fn xxhash_bytes_i64(source: ByteArray, seed: i64) -> i64 {
    let hash = xxh64(&source.data, seed as u64);
    hash as i64
}

some_polymorphic_function2!(xxhash, bytes, ByteArray, i64, i64, i64);

#[doc(hidden)]
pub fn binary_to_u8(b: ByteArray) -> u8 {
    assert!(b.length() <= 1);
    b.data[0]
}

#[doc(hidden)]
pub fn binary_to_u16(b: ByteArray) -> u16 {
    assert!(b.length() <= 2);
    let mut buf = [0u8; 2];
    buf[2 - b.length()..].copy_from_slice(&b.data);
    u16::from_be_bytes(buf)
}

#[doc(hidden)]
pub fn binary_to_u32(b: ByteArray) -> u32 {
    assert!(b.length() <= 4);
    let mut buf = [0u8; 4];
    buf[4 - b.length()..].copy_from_slice(&b.data);
    u32::from_be_bytes(buf)
}

#[doc(hidden)]
pub fn binary_to_u64(b: ByteArray) -> u64 {
    assert!(b.length() <= 8);
    let mut buf = [0u8; 8];
    buf[8 - b.length()..].copy_from_slice(&b.data);
    u64::from_be_bytes(buf)
}

#[doc(hidden)]
pub fn binary_to_u128(b: ByteArray) -> u128 {
    assert!(b.length() <= 16);
    let mut buf = [0u8; 16];
    buf[16 - b.length()..].copy_from_slice(&b.data);
    u128::from_be_bytes(buf)
}

#[test]
pub fn testBinaryToInteger() {
    let bin = ByteArray::new(&[0x12, 0x34]);
    assert_eq!(0x1234, binary_to_u16(bin.clone()));
    assert_eq!(0x1234, binary_to_u32(bin.clone()));
    assert_eq!(0x1234, binary_to_u64(bin.clone()));
    assert_eq!(0x1234, binary_to_u128(bin));
}

// Used in range aggregates; the compiler guarantees that the ByteArray fits in 64 bits
#[doc(hidden)]
pub fn bytes_to_u64_(b: ByteArray) -> u64 {
    binary_to_u64(b)
}

some_function1!(bytes_to_u64, ByteArray, u64);

// Used in range aggregates; the compiler guarantees that the ByteArray fits in 128 bits
#[doc(hidden)]
pub fn bytes_to_u128_(b: ByteArray) -> u128 {
    binary_to_u128(b)
}

some_function1!(bytes_to_u128, ByteArray, u128);

// Used in range aggregates
#[doc(hidden)]
pub fn u64_to_bytes__(u: u64, precision: i32) -> ByteArray {
    ByteArray::with_size_truncate_left(&u.to_be_bytes(), precision, true)
}

// precision can never be Option<i32>, it's a compile-time constant
#[doc(hidden)]
pub fn u64_to_bytes_N_(u: Option<u64>, precision: i32) -> Option<ByteArray> {
    let u = u?;
    Some(u64_to_bytes__(u, precision))
}

#[doc(hidden)]
pub fn u128_to_bytes__(u: u128, precision: i32) -> ByteArray {
    ByteArray::with_size_truncate_left(&u.to_be_bytes(), precision, true)
}

// precision can never be Option<i32>, it's a compile-time constant
#[doc(hidden)]
pub fn u128_to_bytes_N_(u: Option<u128>, precision: i32) -> Option<ByteArray> {
    let u = u?;
    Some(u128_to_bytes__(u, precision))
}

#[test]
pub fn test_binary_to_range_key() {
    for bin in [ByteArray::new(&[0x12, 0x34]), ByteArray::new(&[0x00, 0x00])] {
        let u = bytes_to_u64_(bin.clone());
        let inv = u64_to_bytes__(u, 2);
        assert_eq!(inv, bin);

        let u = bytes_to_u128_(bin.clone());
        let inv = u128_to_bytes__(u, 2);
        assert_eq!(inv, bin);

        let u = bytes_to_u64N(Some(bin.clone()));
        let inv = u64_to_bytes_N_(u, 2).unwrap();
        assert_eq!(inv, bin);

        let u = bytes_to_u128N(Some(bin.clone()));
        let inv = u128_to_bytes_N_(u, 2).unwrap();
        assert_eq!(inv, bin);
    }
}