starkom-pcs 5.0.1

use primitive_types::{H256, U128, U256, U512};
use sha2::Digest;
use starkom_ff::Field256;
use std::fmt::Debug;
use std::marker::PhantomData;

/// Describes a hash backend for hashing internal Merkle tree nodes.
pub trait MerkleHasher {
    /// Hashes a binary Merkle tree node.
    fn hash_binary(left: H256, right: H256) -> H256;

    /// Hashes a ternary Merkle tree node.
    fn hash_ternary(children: [H256; 3]) -> H256;
}

/// Describes a hash backend for our proof system.
///
/// In order to warrant sufficient security even under Grover, our proof system works entirely with
/// 256-bit fields. If the arithmetization process uses a smaller field such as Goldilocks, all
/// scalar values must be embedded into a 256-bit field such as Goldilocks^4 before they can be
/// committed.
pub trait Hasher<F: Field256>: MerkleHasher {
    /// Hashes the provided field elements.
    ///
    /// NOTE: this is a raw hash, it doesn't automatically prepend any DSTs or element count; the
    /// caller is responsible for those when necessary.
    fn hash(inputs: impl IntoIterator<Item = F>) -> H256;

    /// Hashes the provided transcript.
    fn hash_transcript(dst: H256, transcript: &[H256]) -> H256;

    /// Generates a ~256-bit Fiat-Shamir challenge from a transcript.
    ///
    /// The provided implementation hashes the transcript with [`Self::hash_transcript`] and
    /// converts the hash to a field element. To avoid distribution bias, the transcript is hashed
    /// twice with slightly different DSTs, and the resulting 512 bits are converted to a 256-bit
    /// field element via modular reduction.
    fn challenge(dst: H256, transcript: &[H256]) -> F {
        let dst_hi = U256::from_big_endian(dst.as_bytes());
        let dst_lo = dst_hi + U256::one();
        let hi = Self::hash_transcript(H256::from_slice(&dst_hi.to_big_endian()), transcript);
        let lo = Self::hash_transcript(H256::from_slice(&dst_lo.to_big_endian()), transcript);
        let mut bytes = [0u8; 64];
        bytes[0..32].copy_from_slice(hi.as_bytes());
        bytes[32..].copy_from_slice(lo.as_bytes());
        let hash = U512::from_big_endian(&bytes);
        let modulus: U512 = F::MODULUS.parse().unwrap();
        let challenge = hash % modulus;
        F::try_from_be_bytes(&challenge.to_big_endian()[32..]).unwrap()
    }
}

mod internal {
    use super::*;

    pub trait LowLevelHash: Debug + Default {
        fn update(&mut self, bytes: &[u8]);
        fn finalize(self) -> H256;
    }

    #[derive(Debug, Default)]
    pub struct LowLevelSha2Hash {
        hasher: sha2::Sha256,
    }

    impl LowLevelHash for LowLevelSha2Hash {
        fn update(&mut self, bytes: &[u8]) {
            self.hasher.update(bytes);
        }

        fn finalize(self) -> H256 {
            H256::from_slice(self.hasher.finalize().as_slice())
        }
    }

    #[derive(Debug, Default)]
    pub struct LowLevelKeccak256Hash {
        hasher: sha3::Keccak256,
    }

    impl LowLevelHash for LowLevelKeccak256Hash {
        fn update(&mut self, bytes: &[u8]) {
            self.hasher.update(bytes);
        }

        fn finalize(self) -> H256 {
            H256::from_slice(self.hasher.finalize().as_slice())
        }
    }

    pub struct HasherImpl<H: LowLevelHash, F: Field256> {
        _data: PhantomData<(H, F)>,
    }

    impl<H: LowLevelHash, F: Field256> Debug for HasherImpl<H, F> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("HasherImpl").finish()
        }
    }

    impl<H: LowLevelHash, F: Field256> Default for HasherImpl<H, F> {
        fn default() -> Self {
            Self { _data: PhantomData }
        }
    }

    impl<H: LowLevelHash, F: Field256> Copy for HasherImpl<H, F> {}

    impl<H: LowLevelHash, F: Field256> Clone for HasherImpl<H, F> {
        fn clone(&self) -> Self {
            *self
        }
    }

    impl<H: LowLevelHash, F: Field256> PartialEq for HasherImpl<H, F> {
        fn eq(&self, _other: &Self) -> bool {
            true
        }
    }

    impl<H: LowLevelHash, F: Field256> Eq for HasherImpl<H, F> {}

    impl<H: LowLevelHash, F: Field256> PartialOrd for HasherImpl<H, F> {
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
            Some(self.cmp(other))
        }
    }

    impl<H: LowLevelHash, F: Field256> Ord for HasherImpl<H, F> {
        fn cmp(&self, _other: &Self) -> std::cmp::Ordering {
            std::cmp::Ordering::Equal
        }
    }

    impl<H: LowLevelHash, F: Field256> MerkleHasher for HasherImpl<H, F> {
        fn hash_binary(left: H256, right: H256) -> H256 {
            let mut hasher = H::default();
            hasher.update(left.as_bytes());
            hasher.update(right.as_bytes());
            hasher.finalize()
        }

        fn hash_ternary(children: [H256; 3]) -> H256 {
            let mut hasher = H::default();
            hasher.update(children[0].as_bytes());
            hasher.update(children[1].as_bytes());
            hasher.update(children[2].as_bytes());
            hasher.finalize()
        }
    }

    impl<H: LowLevelHash, F: Field256> Hasher<F> for HasherImpl<H, F> {
        fn hash(inputs: impl IntoIterator<Item = F>) -> H256 {
            let mut hasher = H::default();
            for input in inputs {
                hasher.update(&input.to_be_bytes());
            }
            hasher.finalize()
        }

        fn hash_transcript(dst: H256, transcript: &[H256]) -> H256 {
            let mut hasher = H::default();
            hasher.update(dst.as_bytes());
            hasher.update(&U128::from(transcript.len() as u64).to_big_endian());
            for element in transcript {
                hasher.update(element.as_bytes());
            }
            hasher.finalize()
        }
    }
}

/// SHA2 hash backend.
pub type Sha2Hash<F> = internal::HasherImpl<internal::LowLevelSha2Hash, F>;

/// Keccak-256 hash backend.
pub type Keccak256Hash<F> = internal::HasherImpl<internal::LowLevelKeccak256Hash, F>;

#[cfg(all(test, feature = "bluesky", feature = "goldilocks"))]
mod tests {
    use super::*;
    use starkom_bluesky::Scalar as BS;
    use starkom_goldilocks::GL4;
    use std::str::FromStr;
    use std::sync::LazyLock;

    static DST: LazyLock<H256> = LazyLock::new(|| {
        let mut hasher = sha3::Sha3_256::new();
        hasher.update(b"starkom/pcs/test");
        H256::from_slice(hasher.finalize().as_slice())
    });

    fn from_u64<F: Field256>(value: u64) -> F {
        F::from(value)
    }

    fn parse<V: FromStr<Err: Debug>>(s: &'static str) -> V {
        s.parse().unwrap()
    }

    fn test_sha2_hash_impl<F: Field256>() {
        assert_eq!(
            Sha2Hash::<F>::hash([from_u64(12)]),
            parse("0xa82872b96246dac512ddf0515f5da862a92ecebebcb92537b6e3e73199694c45")
        );
        assert_eq!(
            Sha2Hash::<F>::hash([from_u64(34), from_u64(56)]),
            parse("0xbdcf24876d0b8979976f54ea123b70112da34a5cb4dc381646a3321f0817a5e8")
        );
        assert_eq!(
            Sha2Hash::<F>::hash([from_u64(78), from_u64(90), from_u64(12)]),
            parse("0xce080de3e477a622d8b9711eb599aec6fe9ddda88c47c1175bff2aaadc43c3b4")
        );
    }

    #[test]
    fn test_sha2_hash() {
        test_sha2_hash_impl::<BS>();
        test_sha2_hash_impl::<GL4>();
    }

    fn test_keccak256_hash_impl<F: Field256>() {
        assert_eq!(
            Keccak256Hash::<F>::hash([from_u64(12)]),
            parse("0xdf6966c971051c3d54ec59162606531493a51404a002842f56009d7e5cf4a8c7")
        );
        assert_eq!(
            Keccak256Hash::<F>::hash([from_u64(34), from_u64(56)]),
            parse("0x72700d0d963d58363ea77095ddabb7ed1a429a1fe618c8ace040205d52391bb9")
        );
        assert_eq!(
            Keccak256Hash::<F>::hash([from_u64(78), from_u64(90), from_u64(12)]),
            parse("0x452eb69ea7065787fb9f51a2894ff6a6ff50ae842aaa68b3ad40ed3580117a8a")
        );
    }

    #[test]
    fn test_keccak256_hash() {
        test_keccak256_hash_impl::<BS>();
        test_keccak256_hash_impl::<GL4>();
    }

    fn test_sha2_hash_binary_impl<F: Field256>() {
        assert_eq!(
            Sha2Hash::<F>::hash_binary(
                parse("0x2c8ea861b3ec34715e3d50053d001e9d9929a774954e93ffa1f1784f94339592"),
                parse("0x57c43629d9178a54c611df8562260309677c1a10a3bd3ca4c9997acbc5f908bb")
            ),
            parse("0xff888724c7e8b1359d4cb4014d8ddcba61b938e5becbfd4f0af454dee20b83be")
        );
        assert_eq!(
            Sha2Hash::<F>::hash_binary(
                parse("0x229045f7089bd1cb6527f1edd7d94a7e91f2dcf607553e2a89fe24d72b2b013c"),
                parse("0x021905a563d6bf385cb3f2eb4bc69d7750812bcb29542ee4bcfcb0fd0e1bc9b8")
            ),
            parse("0x61944633a56b8421a4b1b53555bfa58283989e76fc21b7e6c1634730969f81f4")
        );
    }

    #[test]
    fn test_sha2_hash_binary() {
        test_sha2_hash_binary_impl::<BS>();
        test_sha2_hash_binary_impl::<GL4>();
    }

    fn test_keccak256_hash_binary_impl<F: Field256>() {
        assert_eq!(
            Keccak256Hash::<F>::hash_binary(
                parse("0x2c8ea861b3ec34715e3d50053d001e9d9929a774954e93ffa1f1784f94339592"),
                parse("0x57c43629d9178a54c611df8562260309677c1a10a3bd3ca4c9997acbc5f908bb")
            ),
            parse("0x0887805e67bb6ece5489f58f1863dffec93f52583cd050eca57c33cbd33abb09")
        );
        assert_eq!(
            Keccak256Hash::<F>::hash_binary(
                parse("0x229045f7089bd1cb6527f1edd7d94a7e91f2dcf607553e2a89fe24d72b2b013c"),
                parse("0x021905a563d6bf385cb3f2eb4bc69d7750812bcb29542ee4bcfcb0fd0e1bc9b8")
            ),
            parse("0x3cecf97526033fb79c366cde597b0f66b7774ec26388b71fd591d7275c81dfbe")
        );
    }

    #[test]
    fn test_keccak256_hash_binary() {
        test_keccak256_hash_binary_impl::<BS>();
        test_keccak256_hash_binary_impl::<GL4>();
    }

    fn test_sha2_hash_ternary_impl<F: Field256>() {
        assert_eq!(
            Sha2Hash::<F>::hash_ternary([
                parse("0x7148a9186844710414337e2a454e082f5bc45a823187142bf8519d48153bb3d7"),
                parse("0x6b3cdb310e6ba5e984fb0ff3722a1968d6b22ce0cdad6b1683167ccdb2b3a75a"),
                parse("0x239f4553fa744a1b2a49af348cf9412fe105c3bc0accb23030a40bfe59d7fa2b"),
            ]),
            parse("0xea2a6de526a3b32de298db973d5cbc604d816eeae526beef5a089bd2ae7ca3eb")
        );
        assert_eq!(
            Sha2Hash::<F>::hash_ternary([
                parse("0x73b6d673e9b67cbca809416002ed33ddeb704d58ce0101a6555531e63426383b"),
                parse("0x5b68318ce7baa6d9a0b6e67cee81b1a730dc7f68c34b7f806fd70ec4614f2096"),
                parse("0x5829d1d7b55ee4345afa69b4acf18d722f6962d0e61447a2711b5851b77d1c2b"),
            ]),
            parse("0xcb8497f1c9898641ec978cfe347ff1ccc43cc9a15f9e741f15e18dc0660d50bf")
        );
    }

    #[test]
    fn test_sha2_hash_ternary() {
        test_sha2_hash_ternary_impl::<BS>();
        test_sha2_hash_ternary_impl::<GL4>();
    }

    fn test_keccak256_hash_ternary_impl<F: Field256>() {
        assert_eq!(
            Keccak256Hash::<F>::hash_ternary([
                parse("0x7148a9186844710414337e2a454e082f5bc45a823187142bf8519d48153bb3d7"),
                parse("0x6b3cdb310e6ba5e984fb0ff3722a1968d6b22ce0cdad6b1683167ccdb2b3a75a"),
                parse("0x239f4553fa744a1b2a49af348cf9412fe105c3bc0accb23030a40bfe59d7fa2b"),
            ]),
            parse("0x2153695c5a08dcffb4189075c974db1faddcd9005c493937120b80eeaf3602fb")
        );
        assert_eq!(
            Keccak256Hash::<F>::hash_ternary([
                parse("0x73b6d673e9b67cbca809416002ed33ddeb704d58ce0101a6555531e63426383b"),
                parse("0x5b68318ce7baa6d9a0b6e67cee81b1a730dc7f68c34b7f806fd70ec4614f2096"),
                parse("0x5829d1d7b55ee4345afa69b4acf18d722f6962d0e61447a2711b5851b77d1c2b"),
            ]),
            parse("0x45b7f5b14de8d4ccf714a72bc733123c271484bdfe449b55d892f9ad1a26ea3a")
        );
    }

    #[test]
    fn test_keccak256_hash_ternary() {
        test_keccak256_hash_ternary_impl::<BS>();
        test_keccak256_hash_ternary_impl::<GL4>();
    }

    #[test]
    fn test_sha2_challenge_bluesky() {
        assert_eq!(
            Sha2Hash::<BS>::challenge(
                *DST,
                &[parse(
                    "0x5b584bf4398b7ef509abeb33ba8521c96a4a497ffba046a492cc43ac34174c16"
                )]
            ),
            parse("0x60555de477c0592d48a0f76f80c88c39a0c127d7209e2ba893de300778792682")
        );
        assert_eq!(
            Sha2Hash::<BS>::challenge(
                *DST,
                &[
                    parse("0x2a1d8bd2a5dc960774c53b77e3e8d677b225bb45ec5d9caaf544e4f229a8afcd"),
                    parse("0x39de8bea57c1ab4082270791ce637189f07d169852f8bba5d05784368b505a12"),
                ]
            ),
            parse("0x2abcb323efcc551be2b473c75ccc1c3f563530df6a279dc4f053282deafccabc")
        );
        assert_eq!(
            Sha2Hash::<BS>::challenge(
                *DST,
                &[
                    parse("0x2cbe7a924ef4a68b49dc6eac0fcf7c8504cc9ecfcb1628cf2b9686d8597e088e"),
                    parse("0x4f6c036f1eaa65e8b761c7fc972156ca4a8340d47d26cd091c9a63655b415896"),
                    parse("0x32c783c083c0fafa4dba39e2176fc3a14791b83d4f1a51372aa483279a93d687"),
                ]
            ),
            parse("0x1e4a2f56b812c4842720698bb2fb8708eec162e1febe52e90aceac66bb57ae3e")
        );
    }

    #[test]
    fn test_sha2_challenge_goldilocks() {
        assert_eq!(
            Sha2Hash::<GL4>::challenge(
                *DST,
                &[parse(
                    "0x5b584bf4398b7ef509abeb33ba8521c96a4a497ffba046a492cc43ac34174c16"
                )]
            ),
            parse("0x3db58df62452afc50e241f9ae42230b788cf37d4aa0b13a43831946bed72e170")
        );
        assert_eq!(
            Sha2Hash::<GL4>::challenge(
                *DST,
                &[
                    parse("0x2a1d8bd2a5dc960774c53b77e3e8d677b225bb45ec5d9caaf544e4f229a8afcd"),
                    parse("0x39de8bea57c1ab4082270791ce637189f07d169852f8bba5d05784368b505a12"),
                ]
            ),
            parse("0x10fecd76d44d659f9b0502705a45dd0c54236211eb6f88680df390d92e5d5d06")
        );
        assert_eq!(
            Sha2Hash::<GL4>::challenge(
                *DST,
                &[
                    parse("0x2cbe7a924ef4a68b49dc6eac0fcf7c8504cc9ecfcb1628cf2b9686d8597e088e"),
                    parse("0x4f6c036f1eaa65e8b761c7fc972156ca4a8340d47d26cd091c9a63655b415896"),
                    parse("0x32c783c083c0fafa4dba39e2176fc3a14791b83d4f1a51372aa483279a93d687"),
                ]
            ),
            parse("0x19f2ce3d8ba0c988a6ed215951f72abb675b50da360d4b0c9868360105b69e19")
        );
    }

    #[test]
    fn test_keccak256_challenge_bluesky() {
        assert_eq!(
            Keccak256Hash::<BS>::challenge(
                *DST,
                &[parse(
                    "0x5b584bf4398b7ef509abeb33ba8521c96a4a497ffba046a492cc43ac34174c16"
                )]
            ),
            parse("0x66647f27ced95c30909f94ab8649419a5457026e3b7d956902401b521672b1bd")
        );
        assert_eq!(
            Keccak256Hash::<BS>::challenge(
                *DST,
                &[
                    parse("0x2a1d8bd2a5dc960774c53b77e3e8d677b225bb45ec5d9caaf544e4f229a8afcd"),
                    parse("0x39de8bea57c1ab4082270791ce637189f07d169852f8bba5d05784368b505a12"),
                ]
            ),
            parse("0x448a77521c4edee2ead1cbaedb7d2b9cdf2f72b7f54abbb6824241ce3f273355")
        );
        assert_eq!(
            Keccak256Hash::<BS>::challenge(
                *DST,
                &[
                    parse("0x2cbe7a924ef4a68b49dc6eac0fcf7c8504cc9ecfcb1628cf2b9686d8597e088e"),
                    parse("0x4f6c036f1eaa65e8b761c7fc972156ca4a8340d47d26cd091c9a63655b415896"),
                    parse("0x32c783c083c0fafa4dba39e2176fc3a14791b83d4f1a51372aa483279a93d687"),
                ]
            ),
            parse("0x48e3fc25084ef8f0183b3ed2d827ad8b701a7129c25ba03f471fada5ab59568a")
        );
    }

    #[test]
    fn test_keccak256_challenge_goldilocks() {
        assert_eq!(
            Keccak256Hash::<GL4>::challenge(
                *DST,
                &[parse(
                    "0x5b584bf4398b7ef509abeb33ba8521c96a4a497ffba046a492cc43ac34174c16"
                )]
            ),
            parse("0xe80eb8ea838e3bf1e5734f4ac0efe313d2faee4b74593eaa52e1772040b22e3f")
        );
        assert_eq!(
            Keccak256Hash::<GL4>::challenge(
                *DST,
                &[
                    parse("0x2a1d8bd2a5dc960774c53b77e3e8d677b225bb45ec5d9caaf544e4f229a8afcd"),
                    parse("0x39de8bea57c1ab4082270791ce637189f07d169852f8bba5d05784368b505a12"),
                ]
            ),
            parse("0xc7327dd09bc9df340622777b370ed74e9cb57e4882bc7fc9a23127bcae403499")
        );
        assert_eq!(
            Keccak256Hash::<GL4>::challenge(
                *DST,
                &[
                    parse("0x2cbe7a924ef4a68b49dc6eac0fcf7c8504cc9ecfcb1628cf2b9686d8597e088e"),
                    parse("0x4f6c036f1eaa65e8b761c7fc972156ca4a8340d47d26cd091c9a63655b415896"),
                    parse("0x32c783c083c0fafa4dba39e2176fc3a14791b83d4f1a51372aa483279a93d687"),
                ]
            ),
            parse("0x4123d9b2786928a9e646df736814516967778af7c0c0e28de6719e6c24c0bde8")
        );
    }
}