arcium-primitives 0.4.2

use blake3;
use hybrid_array::Array;

use crate::{
    algebra::field::FieldExtension,
    constants::CollisionResistanceBytes,
    types::SessionId,
};

pub type Digest = Array<u8, CollisionResistanceBytes>;

/// A generic trait for hashing with a tag.
pub trait HashWith: AsRef<[u8]> + From<Digest> {
    /// Hash the current value with the given tag, returning a new value.
    fn hash_with(&self, tag: &[u8]) -> Self {
        hash(&[self.as_ref(), tag]).into()
    }
}

impl<T: AsRef<[u8]> + From<Digest>> HashWith for T {}

/// Utility function to hash a list of byte slices into a fixed-size digest. Uses BLAKE3.
pub fn hash(slices: &[&[u8]]) -> Digest {
    let mut hasher = blake3::Hasher::new();
    for slice in slices {
        hasher.update(slice);
    }
    Into::<[u8; 32]>::into(hasher.finalize()).into()
}

/// Utility function to hash a list of byte slices into the provided output buffer. Uses BLAKE3.
pub fn hash_into<T: AsRef<[u8]>, I: IntoIterator<Item = T>>(slices: I, out: &mut [u8]) {
    let mut hasher = blake3::Hasher::new();
    for slice in slices {
        hasher.update(slice.as_ref());
    }
    hasher.finalize_xof().fill(out.as_mut());
}

/// Hashes the given session ID and seed into a field element of type F.
pub fn hash_to_field<T: AsRef<[u8]>, F: FieldExtension>(session_id: &SessionId, seed: &T) -> F {
    let mut hasher = blake3::Hasher::new();
    let mut output = Array::<u8, F::UniformBytes>::default();

    hasher.update(session_id.as_ref());
    hasher.update(seed.as_ref());
    hasher.finalize_xof().fill(&mut output);

    F::from_uniform_bytes(&output)
}

/// Utility function to flatten a list of byte slices into a single contiguous vector.
pub fn flatten_slices<T: AsRef<[u8]>>(slices: &[T]) -> Vec<u8> {
    let total_len = slices.iter().map(|slice| slice.as_ref().len()).sum();

    let mut flattened = Vec::with_capacity(total_len);
    slices.iter().for_each(|slice| {
        flattened.extend_from_slice(slice.as_ref());
    });

    flattened
}

/// Utility function to flatten a list of length-prepended byte slices into a single contiguous
/// vector.
pub fn flatten_slices_with_length_prefixes<T: AsRef<[u8]>>(slices: &[T]) -> Vec<u8> {
    let mut flattened = Vec::new();
    slices.iter().for_each(|slice| {
        let slice_ref = slice.as_ref();
        // Cast length to u64 to remove platform-dependence and ensure a fixed 8-byte prefix
        let len_prefix = (slice_ref.len() as u64).to_le_bytes();
        flattened.extend_from_slice(&len_prefix);
        flattened.extend_from_slice(slice_ref);
    });

    flattened
}

#[cfg(test)]
mod tests {
    use crate::hashing::{flatten_slices, flatten_slices_with_length_prefixes, hash_into};

    #[test]
    fn test_hash_into_different_results() {
        let (mut seed0, mut seed1, mut seed2, mut seed3) = ([0; 16], [0; 16], [0; 16], [0; 16]);
        hash_into([b"0", b"1"], &mut seed0);
        hash_into([b"0", b"12".as_slice()], &mut seed1);
        hash_into([b"01", b"12"], &mut seed2);
        hash_into([b"01", b"1".as_slice()], &mut seed3);

        assert_ne!(seed0, seed1);
        assert_ne!(seed0, seed2);
        assert_ne!(seed0, seed3);
        assert_ne!(seed1, seed2);
        assert_ne!(seed1, seed3);
        assert_ne!(seed2, seed3);
    }

    // Verify that length-prefixed flattening prevents collisions that plain concatenation allows.
    // flatten_slices(["AB", "CD"]) == flatten_slices(["A", "BCD"]) == flatten_slices(["ABCD"])
    // but the length-prefixed variants must all differ.
    #[test]
    fn test_length_prefixes_prevent_concatenation_collisions() {
        let splits: Vec<Vec<&[u8]>> = vec![
            vec![b"AB", b"CD"],
            vec![b"A", b"BCD"],
            vec![b"ABCD"],
            vec![b"ABC", b"D"],
        ];

        // Plain concatenation produces the same bytes for all splits
        let plain: Vec<_> = splits.iter().map(|s| flatten_slices(s)).collect();
        for p in &plain {
            assert_eq!(p, &plain[0], "plain concatenation should be identical");
        }

        // Length-prefixed flattening must produce distinct outputs
        let prefixed: Vec<_> = splits
            .iter()
            .map(|s| flatten_slices_with_length_prefixes(s))
            .collect();
        for i in 0..prefixed.len() {
            for j in (i + 1)..prefixed.len() {
                assert_ne!(
                    prefixed[i], prefixed[j],
                    "length-prefixed outputs for splits {i} and {j} should differ"
                );
            }
        }
    }

    #[test]
    fn test_length_prefixes_empty_slices() {
        // Empty slice vs no slices
        let a = flatten_slices_with_length_prefixes::<&[u8]>(&[]);
        let b = flatten_slices_with_length_prefixes(&[b"".as_slice()]);
        let c = flatten_slices_with_length_prefixes(&[b"".as_slice(), b"".as_slice()]);
        assert!(a.is_empty());
        assert_ne!(a, b, "zero slices vs one empty slice must differ");
        assert_ne!(b, c, "one empty slice vs two empty slices must differ");
    }

    #[test]
    fn test_length_prefixes_roundtrip_structure() {
        let slices: &[&[u8]] = &[b"hello", b"", b"world"];
        let prefixed = flatten_slices_with_length_prefixes(slices);

        // Manually decode and verify structure
        let mut cursor = 0;
        for original in slices {
            let len_bytes = &prefixed[cursor..cursor + 8];
            let len = u64::from_le_bytes(len_bytes.try_into().unwrap());
            assert_eq!(len, original.len() as u64);
            cursor += 8;

            let data = &prefixed[cursor..cursor + len as usize];
            assert_eq!(data, *original);
            cursor += len as usize;
        }
        assert_eq!(cursor, prefixed.len(), "no trailing bytes");
    }
}