1use crate::dsa::rpo_falcon512::FALCON_ENCODING_BITS;
  2
  3use super::{
  4    super::{Rpo256, LOG_N, N, PK_LEN},
  5    ByteReader, ByteWriter, Deserializable, DeserializationError, FalconFelt, Felt, Polynomial,
  6    Serializable, Signature, Word,
  7};
  8use alloc::string::ToString;
  9use core::ops::Deref;
 10use num::Zero;
 11
 12// PUBLIC KEY
 13// ================================================================================================
 14
 15/// A public key for verifying signatures.
 16///
 17/// The public key is a [Word] (i.e., 4 field elements) that is the hash of the coefficients of
 18/// the polynomial representing the raw bytes of the expanded public key. The hash is computed
 19/// using Rpo256.
 20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 21pub struct PublicKey(Word);
 22
 23impl PublicKey {
 24    /// Returns a new [PublicKey] which is a commitment to the provided expanded public key.
 25    pub fn new(pub_key: Word) -> Self {
 26        Self(pub_key)
 27    }
 28
 29    /// Verifies the provided signature against provided message and this public key.
 30    pub fn verify(&self, message: Word, signature: &Signature) -> bool {
 31        signature.verify(message, self.0)
 32    }
 33}
 34
 35impl From<PubKeyPoly> for PublicKey {
 36    fn from(pk_poly: PubKeyPoly) -> Self {
 37        let pk_felts: Polynomial<Felt> = pk_poly.0.into();
 38        let pk_digest = Rpo256::hash_elements(&pk_felts.coefficients).into();
 39        Self(pk_digest)
 40    }
 41}
 42
 43impl From<PublicKey> for Word {
 44    fn from(key: PublicKey) -> Self {
 45        key.0
 46    }
 47}
 48
 49// PUBLIC KEY POLYNOMIAL
 50// ================================================================================================
 51
 52#[derive(Debug, Clone, PartialEq, Eq)]
 53pub struct PubKeyPoly(pub Polynomial<FalconFelt>);
 54
 55impl Deref for PubKeyPoly {
 56    type Target = Polynomial<FalconFelt>;
 57
 58    fn deref(&self) -> &Self::Target {
 59        &self.0
 60    }
 61}
 62
 63impl From<Polynomial<FalconFelt>> for PubKeyPoly {
 64    fn from(pk_poly: Polynomial<FalconFelt>) -> Self {
 65        Self(pk_poly)
 66    }
 67}
 68
 69impl Serializable for &PubKeyPoly {
 70    fn write_into<W: ByteWriter>(&self, target: &mut W) {
 71        let mut buf = [0_u8; PK_LEN];
 72        buf[0] = LOG_N;
 73
 74        let mut acc = 0_u32;
 75        let mut acc_len: u32 = 0;
 76
 77        let mut input_pos = 1;
 78        for c in self.0.coefficients.iter() {
 79            let c = c.value();
 80            acc = (acc << FALCON_ENCODING_BITS) | c as u32;
 81            acc_len += FALCON_ENCODING_BITS;
 82            while acc_len >= 8 {
 83                acc_len -= 8;
 84                buf[input_pos] = (acc >> acc_len) as u8;
 85                input_pos += 1;
 86            }
 87        }
 88        if acc_len > 0 {
 89            buf[input_pos] = (acc >> (8 - acc_len)) as u8;
 90        }
 91
 92        target.write(buf);
 93    }
 94}
 95
 96impl Deserializable for PubKeyPoly {
 97    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
 98        let buf = source.read_array::<PK_LEN>()?;
 99
100        if buf[0] != LOG_N {
101            return Err(DeserializationError::InvalidValue(format!(
102                "Failed to decode public key: expected the first byte to be {LOG_N} but was {}",
103                buf[0]
104            )));
105        }
106
107        let mut acc = 0_u32;
108        let mut acc_len = 0;
109
110        let mut output = [FalconFelt::zero(); N];
111        let mut output_idx = 0;
112
113        for &byte in buf.iter().skip(1) {
114            acc = (acc << 8) | (byte as u32);
115            acc_len += 8;
116
117            if acc_len >= FALCON_ENCODING_BITS {
118                acc_len -= FALCON_ENCODING_BITS;
119                let w = (acc >> acc_len) & 0x3FFF;
120                let element = w.try_into().map_err(|err| {
121                    DeserializationError::InvalidValue(format!(
122                        "Failed to decode public key: {err}"
123                    ))
124                })?;
125                output[output_idx] = element;
126                output_idx += 1;
127            }
128        }
129
130        if (acc & ((1u32 << acc_len) - 1)) == 0 {
131            Ok(Polynomial::new(output.to_vec()).into())
132        } else {
133            Err(DeserializationError::InvalidValue(
134                "Failed to decode public key: input not fully consumed".to_string(),
135            ))
136        }
137    }
138}