world-id-primitives 0.11.0

use std::{
    ops::{Deref, DerefMut},
    str::FromStr,
};

use ark_babyjubjub::{EdwardsAffine, Fq};
use ark_ff::AdditiveGroup;
use arrayvec::ArrayVec;
use eddsa_babyjubjub::EdDSAPublicKey;
use ruint::aliases::U256;
use serde::{Deserialize, Serialize};

use crate::PrimitiveError;

/// The maximum number of authenticator public keys that can be registered at any time
/// per World ID Account.
///
/// This constrained is introduced to maintain proof performance reasonable even
/// in devices with limited resources.
pub const MAX_AUTHENTICATOR_KEYS: usize = 7;

/// Errors for decoding sparse authenticator pubkey slots.
#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
pub enum SparseAuthenticatorPubkeysError {
    /// The input contains a non-empty slot index outside supported bounds.
    #[error(
        "invalid authenticator pubkey slot {slot_index}; max supported slot is {max_supported_slot}"
    )]
    SlotOutOfBounds {
        /// Slot index returned by the source.
        slot_index: usize,
        /// Highest supported slot index.
        max_supported_slot: usize,
    },
    /// A slot contained bytes that are not a valid compressed public key.
    #[error("invalid authenticator public key at slot {slot_index}: {reason}")]
    InvalidCompressedPubkey {
        /// Slot index of the invalid entry.
        slot_index: usize,
        /// Parse error details.
        reason: String,
    },
}

/// A set of **off-chain** authenticator public keys for a World ID Account.
///
/// Each World ID Account has a number of public keys for each authorized authenticator;
/// a commitment to the entire set of public keys is stored in the `WorldIDRegistry` contract.
///
/// Removed authenticator slots are represented as `None` to preserve stable slot indices.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthenticatorPublicKeySet(ArrayVec<Option<EdDSAPublicKey>, MAX_AUTHENTICATOR_KEYS>);

impl Default for AuthenticatorPublicKeySet {
    fn default() -> Self {
        Self(ArrayVec::new())
    }
}

impl AuthenticatorPublicKeySet {
    /// Creates a new authenticator public key set with the provided active public keys.
    ///
    /// The provided keys are inserted in order.
    /// Sparse holes can be represented later via [`Self::try_clear_at_index`].
    ///
    /// # Errors
    /// Returns an error if the number of public keys exceeds [`MAX_AUTHENTICATOR_KEYS`].
    pub fn new(pubkeys: Vec<EdDSAPublicKey>) -> Result<Self, PrimitiveError> {
        if pubkeys.len() > MAX_AUTHENTICATOR_KEYS {
            return Err(PrimitiveError::OutOfBounds);
        }

        Ok(Self(
            pubkeys
                .into_iter()
                .map(Some)
                .collect::<ArrayVec<_, MAX_AUTHENTICATOR_KEYS>>(),
        ))
    }

    /// Decodes sparse authenticator pubkey slots while preserving slot positions.
    ///
    /// Input may contain `None` entries for removed authenticators. Trailing empty slots are
    /// trimmed, interior holes are preserved as `None`, and used slots beyond
    /// [`MAX_AUTHENTICATOR_KEYS`] are rejected.
    ///
    /// # Errors
    /// Returns [`SparseAuthenticatorPubkeysError`] if a used slot is out of bounds or any
    /// compressed key bytes are invalid.
    pub fn from_sparse_encoded_pubkeys(
        pubkeys: Vec<Option<U256>>,
    ) -> Result<Self, SparseAuthenticatorPubkeysError> {
        let last_present_idx = pubkeys.iter().rposition(Option::is_some);
        if let Some(idx) = last_present_idx
            && idx >= MAX_AUTHENTICATOR_KEYS
        {
            return Err(SparseAuthenticatorPubkeysError::SlotOutOfBounds {
                slot_index: idx,
                max_supported_slot: MAX_AUTHENTICATOR_KEYS - 1,
            });
        }

        let normalized_len = last_present_idx.map_or(0, |idx| idx + 1);
        let decoded_pubkeys = pubkeys
            .into_iter()
            .take(normalized_len)
            .enumerate()
            .map(|(idx, pubkey)| match pubkey {
                Some(pubkey) => EdDSAPublicKey::from_compressed_bytes(pubkey.to_le_bytes())
                    .map(Some)
                    .map_err(
                        |e| SparseAuthenticatorPubkeysError::InvalidCompressedPubkey {
                            slot_index: idx,
                            reason: e.to_string(),
                        },
                    ),
                None => Ok(None),
            })
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Self(
            decoded_pubkeys
                .into_iter()
                .collect::<ArrayVec<_, MAX_AUTHENTICATOR_KEYS>>(),
        ))
    }

    /// Converts the set of public keys to a fixed-length array of Affine points.
    ///
    /// This is usually used to serialize to circuit inputs, which require defaulted Affine points
    /// for empty (`None`) slots.
    pub fn as_affine_array(&self) -> [EdwardsAffine; MAX_AUTHENTICATOR_KEYS] {
        let mut array = [EdwardsAffine::default(); MAX_AUTHENTICATOR_KEYS];
        for (i, pubkey) in self.0.iter().enumerate() {
            array[i] = pubkey
                .as_ref()
                .map_or_else(EdwardsAffine::default, |pubkey| pubkey.pk);
        }
        array
    }

    /// Returns the number of public keys in the set.
    #[must_use]
    pub const fn len(&self) -> usize {
        self.0.len()
    }

    /// Returns `true` if the set is empty.
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Returns the public key at the given index.
    ///
    /// It will return `None` if the index is out of bounds, even if it's less than
    /// `MAX_AUTHENTICATOR_KEYS` but the key is not initialized.
    #[must_use]
    pub fn get(&self, index: usize) -> Option<&EdDSAPublicKey> {
        self.0.get(index).and_then(Option::as_ref)
    }

    /// Sets a new public key at the given index if it's within bounds of the initialized set.
    ///
    /// # Errors
    /// Returns an error if the index is out of bounds.
    pub fn try_set_at_index(
        &mut self,
        index: usize,
        pubkey: EdDSAPublicKey,
    ) -> Result<(), PrimitiveError> {
        if index >= self.len() || index >= MAX_AUTHENTICATOR_KEYS {
            return Err(PrimitiveError::OutOfBounds);
        }
        self.0[index] = Some(pubkey);
        Ok(())
    }

    /// Clears the public key at the given index while preserving slot position.
    ///
    /// # Errors
    /// Returns an error if the index is out of bounds.
    pub fn try_clear_at_index(&mut self, index: usize) -> Result<(), PrimitiveError> {
        if index >= self.len() || index >= MAX_AUTHENTICATOR_KEYS {
            return Err(PrimitiveError::OutOfBounds);
        }
        self.0[index] = None;
        Ok(())
    }

    /// Pushes a new public key onto the set.
    ///
    /// # Errors
    /// Returns an error if the set is full.
    pub fn try_push(&mut self, pubkey: EdDSAPublicKey) -> Result<(), PrimitiveError> {
        self.0
            .try_push(Some(pubkey))
            .map_err(|_| PrimitiveError::OutOfBounds)
    }

    /// Inserts a public key into the first empty slot, or appends it if no empty slot exists.
    ///
    /// Returns the slot index that now contains the inserted key.
    ///
    /// # Errors
    /// Returns an error if the set is already at capacity.
    pub fn insert_or_reuse(
        &mut self,
        new_authenticator_pubkey: EdDSAPublicKey,
    ) -> Result<usize, PrimitiveError> {
        if let Some(index) = self.iter().position(Option::is_none) {
            self.try_set_at_index(index, new_authenticator_pubkey)?;
            Ok(index)
        } else {
            self.try_push(new_authenticator_pubkey)?;
            Ok(self.len() - 1)
        }
    }

    /// Computes the Poseidon2 leaf hash commitment for this key set as stored in the `WorldIDRegistry`.
    ///
    /// # Panics
    /// Panics if the domain separator constant cannot be converted into an `Fq`.
    #[must_use]
    pub fn leaf_hash(&self) -> Fq {
        let mut input = [Fq::ZERO; 16];

        input[0] =
            Fq::from_str("105702839725298824521994315").expect("domain separator fits in field");

        let pk_array = self.as_affine_array();
        for i in 0..MAX_AUTHENTICATOR_KEYS {
            input[i * 2 + 1] = pk_array[i].x;
            input[i * 2 + 2] = pk_array[i].y;
        }

        poseidon2::bn254::t16::permutation(&input)[1]
    }
}

impl Deref for AuthenticatorPublicKeySet {
    type Target = ArrayVec<Option<EdDSAPublicKey>, MAX_AUTHENTICATOR_KEYS>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for AuthenticatorPublicKeySet {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{MAX_AUTHENTICATOR_KEYS, Signer};
    use ark_babyjubjub::EdwardsAffine;
    use ark_serialize::CanonicalSerialize as _;

    fn create_test_pubkey() -> EdDSAPublicKey {
        EdDSAPublicKey {
            pk: EdwardsAffine::default(),
        }
    }

    fn test_pubkey(seed_byte: u8) -> EdDSAPublicKey {
        Signer::from_seed_bytes(&[seed_byte; 32])
            .unwrap()
            .offchain_signer_pubkey()
    }

    fn encoded_test_pubkey(seed_byte: u8) -> U256 {
        let mut compressed = Vec::new();
        test_pubkey(seed_byte)
            .pk
            .serialize_compressed(&mut compressed)
            .unwrap();
        U256::from_le_slice(&compressed)
    }

    #[test]
    fn test_try_set_at_index_within_bounds() {
        let mut key_set = AuthenticatorPublicKeySet::default();
        let pubkey = create_test_pubkey();

        let result = key_set.try_set_at_index(0, pubkey.clone());
        assert!(result.is_err(), "Should not panic, should return error");

        key_set.try_push(pubkey.clone()).unwrap();
        key_set.try_set_at_index(0, pubkey).unwrap();
    }

    #[test]
    fn test_try_set_at_index_at_length() {
        let mut key_set = AuthenticatorPublicKeySet::default();
        let pubkey = create_test_pubkey();

        key_set.try_push(pubkey.clone()).unwrap();

        let result = key_set.try_set_at_index(1, pubkey);
        assert!(result.is_err(), "Should not panic when index equals length");
    }

    #[test]
    fn test_try_set_at_index_out_of_bounds() {
        let mut key_set = AuthenticatorPublicKeySet::default();
        let pubkey = create_test_pubkey();

        let result = key_set.try_set_at_index(MAX_AUTHENTICATOR_KEYS, pubkey.clone());
        assert!(
            result.is_err(),
            "Should not panic when index >= MAX_AUTHENTICATOR_KEYS"
        );

        let result = key_set.try_set_at_index(MAX_AUTHENTICATOR_KEYS + 1, pubkey);
        assert!(
            result.is_err(),
            "Should not panic when index > MAX_AUTHENTICATOR_KEYS"
        );
    }

    #[test]
    fn test_try_push_within_capacity() {
        let mut key_set = AuthenticatorPublicKeySet::default();
        let pubkey = create_test_pubkey();

        for i in 0..MAX_AUTHENTICATOR_KEYS {
            let result = key_set.try_push(pubkey.clone());
            assert!(
                result.is_ok(),
                "Should not panic when pushing element {} of {}",
                i + 1,
                MAX_AUTHENTICATOR_KEYS
            );
        }

        assert_eq!(key_set.len(), MAX_AUTHENTICATOR_KEYS);

        let result = key_set.try_push(pubkey);
        assert!(result.is_err());
    }

    #[test]
    fn test_as_affine_array_empty_set() {
        let key_set = AuthenticatorPublicKeySet::default();
        let array = key_set.as_affine_array();

        assert_eq!(array.len(), MAX_AUTHENTICATOR_KEYS);
        for affine in &array {
            assert_eq!(*affine, EdwardsAffine::default());
        }
    }

    #[test]
    fn test_as_affine_array_partial_set() {
        let mut key_set = AuthenticatorPublicKeySet::default();
        let pubkey = create_test_pubkey();

        for _ in 0..3 {
            key_set.try_push(pubkey.clone()).unwrap();
        }

        let array = key_set.as_affine_array();

        assert_eq!(array.len(), MAX_AUTHENTICATOR_KEYS);

        for item in array.iter().take(3) {
            assert_eq!(item, &pubkey.pk);
        }

        for item in array.iter().take(MAX_AUTHENTICATOR_KEYS).skip(3) {
            assert_eq!(item, &EdwardsAffine::default());
        }
    }

    #[test]
    fn test_decode_sparse_pubkeys_trims_trailing_empty_slots() {
        let mut encoded_pubkeys = vec![Some(encoded_test_pubkey(1)), Some(encoded_test_pubkey(2))];
        encoded_pubkeys.extend(vec![None; MAX_AUTHENTICATOR_KEYS + 5]);

        let key_set =
            AuthenticatorPublicKeySet::from_sparse_encoded_pubkeys(encoded_pubkeys).unwrap();

        assert_eq!(key_set.len(), 2);
        assert_eq!(key_set[0].as_ref().unwrap().pk, test_pubkey(1).pk);
        assert_eq!(key_set[1].as_ref().unwrap().pk, test_pubkey(2).pk);
    }

    #[test]
    fn test_decode_sparse_pubkeys_preserves_interior_holes() {
        let key_set = AuthenticatorPublicKeySet::from_sparse_encoded_pubkeys(vec![
            Some(encoded_test_pubkey(1)),
            None,
            Some(encoded_test_pubkey(2)),
        ])
        .unwrap();

        assert_eq!(key_set.len(), 3);
        assert_eq!(key_set[0].as_ref().unwrap().pk, test_pubkey(1).pk);
        assert_eq!(key_set[1], None);
        assert_eq!(key_set[2].as_ref().unwrap().pk, test_pubkey(2).pk);
    }

    #[test]
    fn test_decode_sparse_pubkeys_rejects_used_slot_beyond_max() {
        let mut encoded_pubkeys = vec![None; MAX_AUTHENTICATOR_KEYS + 1];
        encoded_pubkeys[MAX_AUTHENTICATOR_KEYS] = Some(encoded_test_pubkey(1));

        let error =
            AuthenticatorPublicKeySet::from_sparse_encoded_pubkeys(encoded_pubkeys).unwrap_err();
        assert!(matches!(
            error,
            SparseAuthenticatorPubkeysError::SlotOutOfBounds {
                slot_index,
                max_supported_slot
            } if slot_index == MAX_AUTHENTICATOR_KEYS && max_supported_slot == MAX_AUTHENTICATOR_KEYS - 1
        ));
    }
}