Skip to content

Commit

Permalink
refactor decompress
Browse files Browse the repository at this point in the history
  • Loading branch information
@supinie
supinie committed Feb 14, 2024
1 parent 32bcb0d commit 355479f
Showing 1 changed file with 67 additions and 63 deletions.
130 changes: 67 additions & 63 deletions src/polynomials.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,52 +195,51 @@ impl StatefulPoly<Normalised> {
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)
for (coeff_chunk, buf_chunk) in self.coeffs.chunks_exact(8)
.zip(buf.chunks_exact_mut(4)) {
for (coeff, t_elem) in coeff_chunk.iter().zip(t.iter_mut()) {
let mut temp = *coeff;
temp += (temp >> 15) & i16::try_from(Q)?;
*t_elem = u8::try_from(
((((u16::try_from(temp)?) << 4) + u16::try_from(Q)? / 2)
/ u16::try_from(Q)?)
& 15,
)?;
}

buf[k..k + 4].copy_from_slice(&[
buf_chunk.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)
for (coeff_chunk, buf_chunk) in self.coeffs.chunks_exact(8)
.zip(buf.chunks_exact_mut(5)) {
for (coeff, t_elem) in coeff_chunk.iter().zip(t.iter_mut()) {
let mut temp = *coeff;
temp += (temp >> 15) & i16::try_from(Q)?;
*t_elem = u8::try_from(
((((u32::try_from(temp)?) << 5) + u32::try_from(Q)? / 2)
/ u32::try_from(Q)?)
& 31,
)?;
}

buf[k..k + 5].copy_from_slice(&[
buf_chunk.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(())
}
Expand Down Expand Up @@ -290,52 +289,57 @@ fn read_msg_to_poly(msg: &[u8]) -> Result<StatefulPoly<Unnormalised>, TryFromInt
})
}

// // 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(())
// }
// }
// }
// 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);
fn decompress(buf: &[u8], sec_level: &SecurityLevel) -> Result<StatefulPoly<Normalised>, TryFromIntError> {
match sec_level {
SecurityLevel::FiveOneTwo { .. } | SecurityLevel::SevenSixEight { .. } => {
let coeffs_arr: [i16; N] = buf.iter()
.flat_map(|&byte| (0..2).map(move |i| {
if i == 0 {
(usize::from(byte & 15) * Q + 8) >> 4
} else {
(usize::from(byte >> 4) * Q + 8) >> 4
}
}))
.map(i16::try_from)
.collect::<Result<ArrayVec<[i16; N]>, TryFromIntError>>()?
.into_inner();

Ok(StatefulPoly {
coeffs: coeffs_arr,
state: Normalised,
})
}
SecurityLevel::TenTwoFour { .. } => {
let mut coeffs_arr = [0i16; N];
for (coeffs_chunk, buf_chunk) in coeffs_arr.chunks_exact_mut(8)
.zip(buf.chunks_exact(5)) {
let temp: [u8; 8] = [
buf_chunk[0],
(buf_chunk[0] >> 5) | (buf_chunk[1] << 3),
buf_chunk[1] >> 2,
(buf_chunk[1] >> 7) | (buf_chunk[2] << 1),
(buf_chunk[2] >> 4) | (buf_chunk[3] << 4),
buf_chunk[3] >> 1,
(buf_chunk[3] >> 6) | (buf_chunk[4] << 2),
buf_chunk[4] >> 3
];
for (coeff, t_elem) in coeffs_chunk.iter_mut().zip(temp.iter()) {
*coeff = i16::try_from(((u32::from(*t_elem) & 31) * u32::try_from(Q)? + 16) >> 5)?;
}
}
Ok(StatefulPoly {
coeffs: coeffs_arr,
state: Normalised
})
}
}
}


#[derive(Copy, Clone, Debug, PartialEq, Eq)]
Expand Down

0 comments on commit 355479f

Please sign in to comment.