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