Skip to content

Commit

Permalink
begin stateful poly
Browse files Browse the repository at this point in the history
  • Loading branch information
supinie committed Feb 8, 2024
1 parent 35bd46c commit 29643f5
Showing 1 changed file with 272 additions and 36 deletions.
308 changes: 272 additions & 36 deletions src/polynomials.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,26 +5,161 @@ use crate::{
};
use core::num::TryFromIntError;

struct StatefulPoly<State: NormalisedState> {
struct StatefulPoly<S: State> {
pub(crate) coeffs: [i16; N],
state: State,
state: S,
}

struct Normalised;
struct Unnormalised;

trait NormalisedState {}
impl NormalisedState for Normalised {}
impl NormalisedState for Unnormalised {}
trait State {}
impl State for Normalised {}
impl State for Unnormalised {}

// impl Default for StatefulPoly<Normalised> {
// fn default() -> Self {
// Self {
// coeffs: [0; N],
// state: Unnormalised,
// }
// }
// }
impl Default for StatefulPoly<Unnormalised> {
fn default() -> Self {
Self {
coeffs: [0; N],
state: Unnormalised,
}
}
}

impl<S: State> StatefulPoly<S> {
// const function equivelent of default (default is needed for ArrayVec)
// Example:
// let poly = Poly::new();
const fn new() -> StatefulPoly<Unnormalised> {
StatefulPoly {
coeffs: [0; N],
state: Unnormalised,
}
}

const fn from(array: &[i16; N]) -> StatefulPoly<Unnormalised> {
StatefulPoly {
coeffs: *array,
state: Unnormalised,
}
}

// Sets self to self + x
// Example:
// poly1.add(&poly2);
fn add(&mut self, x: &Self) {
for i in 0..N {
self.coeffs[i] += x.coeffs[i];
}
// // self.state = Unnormalised
}

// Sets self to self - x
// Example:
// poly1.sub(&poly2);
pub(crate) fn sub(&mut self, x: &Self) {
for i in 0..N {
self.coeffs[i] -= x.coeffs[i];
}
// // self.state = Unnormalised;
}

// Barrett reduces all coefficients of given polynomial
// Example:
// poly.barrett_reduce();
pub(crate) fn barrett_reduce(&mut self) {
for coeff in &mut self.coeffs {
*coeff = barrett_reduce(*coeff);
}
// // self.state = Unnormalised;
}

// Converts all coefficients of the given polynomial to Mongomery form
// Example:
// poly.mont_form();
pub(crate) fn mont_form(&mut self) {
for coeff in &mut self.coeffs {
*coeff = mont_form(*coeff);
}
// // self.state = Unnormalised;
}

// Unpacks a buffer of bytes into a polynomial
// poly will NOT be normalised, but 0 <= coeffs < 4096
// Example:
// poly.unpack(buf);
pub(crate) fn unpack(&mut self, buf: &[u8]) {
for i in 0..N / 2 {
self.coeffs[2 * i] = i16::from(buf[3 * i]) | ((i16::from(buf[3 * i + 1]) << 8) & 0xfff);
self.coeffs[2 * i + 1] =
i16::from(buf[3 * i + 1] >> 4) | ((i16::from(buf[3 * i + 2]) << 4) & 0xfff);
}
// self.state = Unnormalised;
}

// Converts a message buffer into a polynomial
// msg should be of length SYMBYTES (32)
// poly will not be normalised
// Example:
// poly.read_msg(msg_buf);
pub(crate) fn read_msg(&mut self, msg: &[u8]) -> Result<(), TryFromIntError> {
for i in 0..SYMBYTES {
for j in 0..8 {
let mask = ((i16::from(msg[i]) >> j) & 1).wrapping_neg();
self.coeffs[8 * i + j] = mask & i16::try_from((Q + 1) / 2)?;
}
}
// self.state = Unnormalised;
Ok(())
}

// Decompresses buffer into a polynomial
// is dependent on the security level
// buf should be of length poly_compressed_bytes
// output poly is normalised
// Example:
// poly.decompress(buf, k);
pub(crate) fn decompress(
&mut self,
buf: &[u8],
sec_level: &SecurityLevel,
) -> Result<(), TryFromIntError> {
let mut k = 0usize;

match sec_level {
SecurityLevel::FiveOneTwo { .. } | SecurityLevel::SevenSixEight { .. } => {
for (i, &byte) in buf.iter().take(N / 2).enumerate() {
self.coeffs[2 * i] = i16::try_from((usize::from(byte & 15) * Q + 8) >> 4)?;
self.coeffs[2 * i + 1] = i16::try_from((usize::from(byte >> 4) * Q + 8) >> 4)?;
}
// self.state = Normalised;
Ok(())
}
SecurityLevel::TenTwoFour { .. } => {
let mut t = [0u8; 8];
for i in 0..N / 8 {
t[0] = buf[k];
t[1] = (buf[k] >> 5) | (buf[k + 1] << 3);
t[2] = buf[k + 1] >> 2;
t[3] = (buf[k + 1] >> 7) | (buf[k + 2] << 1);
t[4] = (buf[k + 2] >> 4) | (buf[k + 3] << 4);
t[5] = buf[k + 3] >> 1;
t[6] = (buf[k + 3] >> 6) | (buf[k + 4] << 2);
t[7] = buf[k + 4] >> 3;
k += 5;

for (j, t_elem) in t.iter().enumerate() {
self.coeffs[8 * i + j] = i16::try_from(
((u32::from(*t_elem) & 31) * u32::try_from(Q)? + 16) >> 5,
)?;
}
}
// self.state = Normalised;
Ok(())
}
}
}
}

impl StatefulPoly<Unnormalised> {
// Normalise coefficients of given polynomial
Expand All @@ -34,44 +169,145 @@ impl StatefulPoly<Unnormalised> {
for coeff in &mut self.coeffs {
*coeff = conditional_sub_q(barrett_reduce(*coeff));
}
// self.state = Normalised;
}
}

impl StatefulPoly<Normalised> {
// Sets self to self + x
// Pointwise multiplication of two polynomials,
// assumes inputs are of montgomery form.
// Example:
// poly1.add(&poly2);
fn add(&mut self, x: &Self) {
for i in 0..N {
self.coeffs[i] += x.coeffs[i];
// poly1.pointwise_mul(&poly2);
pub(crate) fn pointwise_mul(&mut self, x: &Self) {
for ((chunk, x_chunk), &zeta) in self
.coeffs
.chunks_mut(4)
.zip(x.coeffs.chunks(4))
.zip(ZETAS.iter().skip(64))
{
let mut temp = [0i16; 4];

for (i, coeff) in temp.iter_mut().enumerate() {
if i % 2 == 0 {
let sign: i16 = if i == 2 { -1 } else { 1 };
*coeff = montgomery_reduce(i32::from(chunk[i + 1]) * i32::from(x_chunk[i + 1]));
*coeff = sign * montgomery_reduce(i32::from(*coeff) * i32::from(zeta));
*coeff += montgomery_reduce(i32::from(chunk[i]) * i32::from(x_chunk[i]));
} else {
*coeff = montgomery_reduce(i32::from(chunk[i - 1]) * i32::from(x_chunk[i]));
*coeff += montgomery_reduce(i32::from(chunk[i]) * i32::from(x_chunk[i - 1]));
}
}
chunk.copy_from_slice(&temp);
}
// self.state = Unnormalised;
}
}

// impl<S> HttpResponse<S>
// where S: SendingState
// {
// fn spam_spam_spam(&mut self);
// }
// Packs given poly into a 384-byte (POLYBYTES size) buffer
// must be normalised
// Example:
// poly.pack(buf);
#[allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_possible_wrap
)]
pub(crate) fn pack(&self, buf: &mut [u8]) {
for i in 0..N / 2 {
let mut t0 = self.coeffs[2 * i];
t0 += (t0 >> 15) & Q as i16;
let mut t1 = self.coeffs[2 * i + 1];
t1 += (t1 >> 15) & Q as i16;

impl<State: NormalisedState> StatefulPoly<State>
//where State: NormalisedState
{
// const function equivelent of default (default is needed for ArrayVec)
buf[3 * i] = t0 as u8;
buf[3 * i + 1] = ((t0 >> 8) | (t1 << 4)) as u8;
buf[3 * i + 2] = (t1 >> 4) as u8;
}
}

// Convert a given polynomial into a SYMBYTES (32-byte) message
// poly should be normalised
// Example:
// let poly = Poly::new();
const fn new() -> StatefulPoly<Unnormalised> {
StatefulPoly {
coeffs: [0; N],
state: Unnormalised,
// poly.write_msg(msg_buf);
pub(crate) fn write_msg(&self, buf: &mut [u8]) -> Result<(), TryFromIntError> {
#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
let q_16 = i16::try_from(Q)?;
for i in 0..N / 8 {
buf[i] = 0;
for j in 0..8 {
let mut x = self.coeffs[8 * i + j];
x += (x >> 15) & q_16;
x = (((x << 1) + q_16 / 2) / q_16) & 1;
buf[i] |= u8::try_from(x << j)?;
}
}
Ok(())
}

// fn from(array: [i16; N]) -> Self {
// Self { coeffs: array, state: Normalised }
// }
// Compress polynomial to a buffer
// buf must have space for poly_compressed_bytes
// poly should be normalised
// Example:
// self.compress(buf);
pub(crate) fn compress(
&self,
buf: &mut [u8],
sec_level: &SecurityLevel,
) -> Result<(), TryFromIntError> {
let mut k = 0usize;
let mut t = [0u8; 8];

match sec_level {
SecurityLevel::FiveOneTwo { .. } | SecurityLevel::SevenSixEight { .. } => {
for i in 0..N / 8 {
for j in 0..8 {
let mut u = self.coeffs[8 * i + j];
u += (u >> 15) & i16::try_from(Q)?;
t[j] = u8::try_from(
((((u16::try_from(u)?) << 4) + u16::try_from(Q)? / 2)
/ u16::try_from(Q)?)
& 15,
)?;
}

buf[k..k + 4].copy_from_slice(&[
t[0] | (t[1] << 4),
t[2] | (t[3] << 4),
t[4] | (t[5] << 4),
t[6] | (t[7] << 4),
]);
k += 4;
}
Ok(())
}
SecurityLevel::TenTwoFour { .. } => {
for i in 0..N / 8 {
for j in 0..8 {
let mut u = self.coeffs[8 * i + j];
u += (u >> 15) & i16::try_from(Q)?;
t[j] = u8::try_from(
((((u32::try_from(u)?) << 5) + u32::try_from(Q)? / 2)
/ u32::try_from(Q)?)
& 31,
)?;
}

buf[k..k + 5].copy_from_slice(&[
t[0] | (t[1] << 5),
(t[1] >> 3) | (t[2] << 2) | (t[3] << 7),
(t[3] >> 1) | (t[4] << 4),
(t[4] >> 4) | (t[5] << 1) | (t[6] << 6),
(t[6] >> 2) | (t[7] << 3),
]);
k += 5;
}
Ok(())
}
}
}
}


#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Poly {
pub(crate) coeffs: [i16; N],
Expand Down

0 comments on commit 29643f5

Please sign in to comment.