samaharam 0.2.0

//! Verification key registry for heterogeneous proof types.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use crate::crypto::VerificationKey;
use crate::error::Error;
use crate::traits::PairingEngine;

/// Unique identifier for a registered verification key.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct VkId(u64);

impl VkId {
    /// Create a new VkId (primarily for testing).
    pub(crate) fn new(id: u64) -> Self {
        Self(id)
    }

    /// Get the underlying ID value.
    pub fn as_u64(&self) -> u64 {
        self.0
    }
}

/// A registered verification key with metadata.
#[derive(Debug, Clone)]
pub struct RegisteredVk<E: PairingEngine> {
    /// Human-readable name for this circuit
    pub name: String,

    /// The typed verification key
    pub vk: VerificationKey<E>,
}

/// Thread-safe registry for heterogeneous verification keys.
///
/// This registry allows registering VKs from different circuits,
/// enabling heterogeneous proof aggregation.
///
/// # Example
///
/// ```rust,ignore
/// let registry = VkRegistry::<Bn254>::new();
///
/// let transfer_vk = registry.register("transfer", vk_bytes, 5);
/// let deposit_vk = registry.register("deposit", vk_bytes, 3);
///
/// // Proofs can now reference either VK
/// let proof1 = Proof::new(data, inputs, transfer_vk);
/// let proof2 = Proof::new(data, inputs, deposit_vk);
/// ```
pub struct VkRegistry<E: PairingEngine> {
    inner: RwLock<VkRegistryInner<E>>,
}

struct VkRegistryInner<E: PairingEngine> {
    vks: HashMap<VkId, Arc<RegisteredVk<E>>>,
    name_to_id: HashMap<String, VkId>,
    next_id: u64,
}

impl<E: PairingEngine> VkRegistry<E> {
    /// Create a new empty registry.
    pub fn new() -> Self {
        Self {
            inner: RwLock::new(VkRegistryInner {
                vks: HashMap::new(),
                name_to_id: HashMap::new(),
                next_id: 0,
            }),
        }
    }

    /// Register a new verification key.
    ///
    /// # Arguments
    ///
    /// * `name` - Human-readable circuit name
    /// * `vk` - The typed verification key
    ///
    /// # Returns
    ///
    /// A unique ID for this verification key.
    pub fn register(&self, name: impl Into<String>, vk: VerificationKey<E>) -> VkId {
        let name = name.into();
        let mut inner = self.inner.write().unwrap();

        let id = VkId(inner.next_id);
        inner.next_id += 1;

        let registered = Arc::new(RegisteredVk {
            name: name.clone(),
            vk,
        });

        inner.vks.insert(id, registered);
        inner.name_to_id.insert(name, id);

        id
    }

    /// Get a verification key by ID.
    ///
    /// Returns a cheap `Arc` clone of the VK.
    pub fn get(&self, id: VkId) -> Option<Arc<RegisteredVk<E>>> {
        self.inner.read().ok()?.vks.get(&id).cloned()
    }

    /// Get a verification key ID by name.
    pub fn get_by_name(&self, name: &str) -> Option<VkId> {
        self.inner.read().ok()?.name_to_id.get(name).copied()
    }

    /// Check if a VK ID exists in the registry.
    pub fn contains(&self, id: VkId) -> bool {
        self.inner.read().map(|guard| guard.vks.contains_key(&id)).unwrap_or(false)
    }

    /// Get the number of registered verification keys.
    pub fn len(&self) -> usize {
        self.inner.read().map(|guard| guard.vks.len()).unwrap_or(0)
    }

    /// Check if the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Verify a VK exists, returning error if not found.
    pub fn require(&self, id: VkId) -> Result<Arc<RegisteredVk<E>>, Error> {
        self.get(id).ok_or(Error::UnknownVk(id))
    }
}

impl<E: PairingEngine> Default for VkRegistry<E> {
    fn default() -> Self {
        Self::new()
    }
}

// Manual Debug impl to avoid requiring E: Debug on internal state
impl<E: PairingEngine> std::fmt::Debug for VkRegistry<E> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let inner = self.inner.read().unwrap();
        f.debug_struct("VkRegistry").field("num_vks", &inner.vks.len()).finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use group::{Curve, Group};
    use halo2curves::bn256::{G1, G2};
    use rand::rngs::OsRng;

    // Mock engine for testing
    #[derive(Clone, Debug)]
    struct MockEngine;

    impl PairingEngine for MockEngine {
        type Fr = halo2curves::bn256::Fr;
        type G1Affine = halo2curves::bn256::G1Affine;
        type G1 = halo2curves::bn256::G1;
        type G2Affine = halo2curves::bn256::G2Affine;
        type Gt = halo2curves::bn256::Gt;

        fn pairing(_p: &Self::G1Affine, _q: &Self::G2Affine) -> Self::Gt {
            unimplemented!("mock")
        }

        fn multi_pairing<'a>(
            _pairs: impl IntoIterator<Item = (&'a Self::G1Affine, &'a Self::G2Affine)>,
        ) -> Self::Gt {
            unimplemented!("mock")
        }
    }

    /// Create a mock verification key for testing
    fn mock_vk(num_public_inputs: usize) -> VerificationKey<MockEngine> {
        VerificationKey {
            num_public_inputs,
            domain_size: 1024,
            selector_commitments: vec![
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            permutation_commitments: vec![G1::random(OsRng).to_affine()],
            x_g2: G2::random(OsRng).to_affine(),
            g2_generator: G2::generator().to_affine(),
        }
    }

    #[test]
    fn registry_assigns_unique_ids() {
        let registry = VkRegistry::<MockEngine>::new();

        let id1 = registry.register("circuit_a", mock_vk(5));
        let id2 = registry.register("circuit_b", mock_vk(3));

        assert_ne!(id1, id2);
    }

    #[test]
    fn registry_retrieves_by_id() {
        let registry = VkRegistry::<MockEngine>::new();

        let id = registry.register("test", mock_vk(5));
        let retrieved = registry.get(id).unwrap();

        assert_eq!(retrieved.name, "test");
        assert_eq!(retrieved.vk.num_public_inputs, 5);
    }

    #[test]
    fn registry_retrieves_by_name() {
        let registry = VkRegistry::<MockEngine>::new();

        let id = registry.register("transfer", mock_vk(2));
        let found_id = registry.get_by_name("transfer").unwrap();

        assert_eq!(id, found_id);
    }

    #[test]
    fn registry_returns_none_for_unknown() {
        let registry = VkRegistry::<MockEngine>::new();

        assert!(registry.get(VkId(999)).is_none());
        assert!(registry.get_by_name("unknown").is_none());
    }

    #[test]
    fn registry_require_returns_error_for_unknown() {
        let registry = VkRegistry::<MockEngine>::new();

        let result = registry.require(VkId(999));
        assert!(matches!(result, Err(Error::UnknownVk(_))));
    }

    #[test]
    fn registry_len_tracks_registrations() {
        let registry = VkRegistry::<MockEngine>::new();

        assert!(registry.is_empty());
        assert_eq!(registry.len(), 0);

        registry.register("a", mock_vk(1));
        assert!(!registry.is_empty());
        assert_eq!(registry.len(), 1);

        registry.register("b", mock_vk(1));
        assert_eq!(registry.len(), 2);
    }

    #[test]
    fn registry_is_thread_safe() {
        use std::thread;

        let registry = Arc::new(VkRegistry::<MockEngine>::new());
        let registry2 = Arc::clone(&registry);

        let handle = thread::spawn(move || registry2.register("from_thread", mock_vk(1)));

        let id1 = registry.register("from_main", mock_vk(1));
        let id2 = handle.join().unwrap();

        assert_ne!(id1, id2);
        assert_eq!(registry.len(), 2);
    }
}