world-id-primitives 0.11.0

use ruint::aliases::U256;
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};

use crate::FieldElement;

/// Encoded World ID Proof.
///
/// Internally, the first 4 elements are a compressed Groth16 proof
/// [a (G1), b (G2), b (G2), c (G1)]. Proofs also require the root hash of the Merkle tree
/// in the `WorldIDRegistry` as a public input. To simplify transmission, that root is encoded as the last element
/// with the proof.
///
/// The `WorldIDVerifier.sol` contract handles the decoding and verification.
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct ZeroKnowledgeProof {
    /// Array of 5 U256 values: first 4 are compressed Groth16 proof, last is Merkle root.
    inner: [U256; 5],
}

impl ZeroKnowledgeProof {
    /// Outputs the proof as a Solidity-friendly representation.
    #[must_use]
    pub const fn as_ethereum_representation(&self) -> [U256; 5] {
        self.inner
    }

    /// Initializes a proof from an encoded Solidity-friendly representation.
    #[must_use]
    pub const fn from_ethereum_representation(value: [U256; 5]) -> Self {
        Self { inner: value }
    }

    /// Converts the proof to compressed bytes (160 bytes total: 5 × 32 bytes).
    #[must_use]
    pub fn to_compressed_bytes(&self) -> Vec<u8> {
        self.inner
            .iter()
            .flat_map(U256::to_be_bytes::<32>)
            .collect()
    }

    /// Constructs a proof from compressed bytes (must be exactly 160 bytes).
    ///
    /// # Errors
    /// Returns an error if the input is not exactly 160 bytes or if bytes cannot be parsed.
    pub fn from_compressed_bytes(bytes: &[u8]) -> Result<Self, String> {
        if bytes.len() != 160 {
            return Err(format!(
                "Invalid length: expected 160 bytes, got {}",
                bytes.len()
            ));
        }

        let mut inner = [U256::ZERO; 5];
        for (i, chunk) in bytes.chunks_exact(32).enumerate() {
            let mut arr = [0u8; 32];
            arr.copy_from_slice(chunk);
            inner[i] = U256::from_be_bytes(arr);
        }

        Ok(Self { inner })
    }
}

impl Serialize for ZeroKnowledgeProof {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = self.to_compressed_bytes();
        if serializer.is_human_readable() {
            serializer.serialize_str(&hex::encode(bytes))
        } else {
            serializer.serialize_bytes(&bytes)
        }
    }
}

impl<'de> Deserialize<'de> for ZeroKnowledgeProof {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let bytes = if deserializer.is_human_readable() {
            let hex_str = String::deserialize(deserializer)?;
            hex::decode(hex_str).map_err(D::Error::custom)?
        } else {
            Vec::deserialize(deserializer)?
        };

        Self::from_compressed_bytes(&bytes).map_err(D::Error::custom)
    }
}

impl From<ZeroKnowledgeProof> for [U256; 5] {
    fn from(value: ZeroKnowledgeProof) -> Self {
        value.inner
    }
}

/// A WIP-103 Ownership Proof.
///
/// Contains the ZKP and the Merkle root public input that the verifier
/// doesn't initially provide.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct OwnershipProof {
    /// The WHIR R1CS proof from ProveKit.
    pub proof: provekit_common::WhirR1CSProof,
    /// Merkle tree root used for inclusion within the circuit.
    pub merkle_root: FieldElement,
}

#[cfg(test)]
mod tests {
    use ruint::uint;

    use super::*;

    #[test]
    fn test_encoding_round_trip() {
        let proof = ZeroKnowledgeProof::default();
        let compressed_bytes = proof.to_compressed_bytes();

        assert_eq!(compressed_bytes.len(), 160);

        let encoded = serde_json::to_string(&proof).unwrap();
        assert_eq!(
            encoded,
            "\"00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000\""
        );

        let proof_from = ZeroKnowledgeProof::from_compressed_bytes(&compressed_bytes).unwrap();

        assert_eq!(proof.inner, proof_from.inner);
    }

    #[test]
    fn test_json_deserialization() {
        let proof = ZeroKnowledgeProof::default();

        // Test roundtrip serialization
        let json_str = serde_json::to_string(&proof).unwrap();
        let deserialized_proof: ZeroKnowledgeProof = serde_json::from_str(&json_str).unwrap();

        // Verify the roundtrip preserved all values
        assert_eq!(proof.inner, deserialized_proof.inner);
    }

    #[test]
    fn test_from_ethereum_representation() {
        let values = [
            uint!(0x0000000000000000000000000000000000000000000000000000000000000001_U256),
            uint!(0x0000000000000000000000000000000000000000000000000000000000000002_U256),
            uint!(0x0000000000000000000000000000000000000000000000000000000000000003_U256),
            uint!(0x0000000000000000000000000000000000000000000000000000000000000004_U256),
            uint!(0x11d223ce7b91ac212f42cf50f0a3439ae3fcdba4ea32acb7f194d1051ed324c2_U256),
        ];

        let proof = ZeroKnowledgeProof::from_ethereum_representation(values);
        assert_eq!(proof.as_ethereum_representation(), values);

        // Test serialization roundtrip
        let bytes = proof.to_compressed_bytes();
        assert_eq!(bytes.len(), 160);

        let proof_from_bytes = ZeroKnowledgeProof::from_compressed_bytes(&bytes).unwrap();
        assert_eq!(proof.inner, proof_from_bytes.inner);
    }

    #[test]
    fn test_invalid_bytes_length() {
        let too_short = vec![0u8; 159];
        let result = ZeroKnowledgeProof::from_compressed_bytes(&too_short);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));

        let too_long = vec![0u8; 161];
        let result = ZeroKnowledgeProof::from_compressed_bytes(&too_long);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));
    }

    #[test]
    fn test_ownership_proof_json_roundtrip() {
        let whir_proof = provekit_common::WhirR1CSProof {
            narg_string: vec![1, 2, 3],
            hints: vec![4, 5],
            #[cfg(debug_assertions)]
            pattern: vec![],
        };
        let proof = OwnershipProof {
            proof: whir_proof,
            merkle_root: crate::FieldElement::from(999u64),
        };
        let json = serde_json::to_string(&proof).unwrap();
        assert!(json.contains("proof"));
        assert!(json.contains("merkle_root"));
        let decoded: OwnershipProof = serde_json::from_str(&json).unwrap();
        assert_eq!(proof, decoded);
    }

    /// A proof with a wrong (flipped) merkle root must not equal the original.
    ///
    /// This test simulates the data-level check that a verifier would perform:
    /// if the merkle root stored inside an [`OwnershipProof`] has been tampered
    /// with, the proof object is distinguishable from the original.
    #[test]
    fn test_ownership_proof_wrong_merkle_root_is_detected() {
        let whir_proof = provekit_common::WhirR1CSProof {
            narg_string: vec![10, 20, 30],
            hints: vec![],
            #[cfg(debug_assertions)]
            pattern: vec![],
        };
        let proof = OwnershipProof {
            proof: whir_proof.clone(),
            merkle_root: crate::FieldElement::from(12345u64),
        };

        // Construct a proof with a different merkle root (simulating tampering).
        let tampered = OwnershipProof {
            proof: whir_proof,
            merkle_root: crate::FieldElement::from(99999u64),
        };

        assert_ne!(
            proof.merkle_root, tampered.merkle_root,
            "a flipped merkle root must differ from the original"
        );
        assert_ne!(proof, tampered);
    }

    /// A proof with tampered proof bytes must not equal the original.
    ///
    /// This test simulates the data-level check that a verifier would perform:
    /// if the proof bytes inside an [`OwnershipProof`] have been tampered with,
    /// the proof object is distinguishable from the original.
    #[test]
    fn test_ownership_proof_tampered_bytes_is_detected() {
        let original_bytes = vec![0xde, 0xad, 0xbe, 0xef];
        let proof = OwnershipProof {
            proof: provekit_common::WhirR1CSProof {
                narg_string: original_bytes.clone(),
                hints: vec![],
                #[cfg(debug_assertions)]
                pattern: vec![],
            },
            merkle_root: crate::FieldElement::from(42u64),
        };

        // Flip the first byte to simulate proof tampering.
        let mut tampered_bytes = original_bytes;
        tampered_bytes[0] ^= 0xFF;

        let tampered = OwnershipProof {
            proof: provekit_common::WhirR1CSProof {
                narg_string: tampered_bytes,
                hints: vec![],
                #[cfg(debug_assertions)]
                pattern: vec![],
            },
            merkle_root: crate::FieldElement::from(42u64),
        };

        assert_ne!(
            proof.proof.narg_string, tampered.proof.narg_string,
            "tampered proof bytes must differ from the original"
        );
        assert_ne!(proof, tampered);
    }
}