use std::cmp::min;
use std::io::{self, Read};
use std::mem::size_of;
#[cfg(not(target_arch = "aarch64"))]
use byte_slice_cast::AsSliceOf;
use byte_slice_cast::AsByteSlice;
const NUM_FRS_PER_BLOCK: usize = 4;
const IN_BITS_FR: usize = 254;
const OUT_BITS_FR: usize = 256;
const NUM_BYTES_IN_BLOCK: usize = NUM_FRS_PER_BLOCK * IN_BITS_FR / 8;
const NUM_BYTES_OUT_BLOCK: usize = NUM_FRS_PER_BLOCK * OUT_BITS_FR / 8;
const NUM_U128S_PER_BLOCK: usize = NUM_BYTES_OUT_BLOCK / size_of::<u128>();
const MASK_SKIP_HIGH_2: u128 = 0b0011_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111_1111;
#[repr(align(16))]
struct AlignedBuffer([u8; NUM_BYTES_IN_BLOCK + 1]);
pub struct Fr32Reader<R> {
source: R,
in_buffer: AlignedBuffer,
out_buffer: [u128; NUM_U128S_PER_BLOCK],
out_offset: usize,
available_frs: usize,
done: bool,
}
macro_rules! process_fr {
(
$in_buffer:expr,
$out0:expr,
$out1:expr,
$bit_offset:expr
) => {{
$out0 = $in_buffer[0] >> 128 - $bit_offset;
$out0 |= $in_buffer[1] << $bit_offset;
$out1 = $in_buffer[1] >> 128 - $bit_offset;
$out1 |= $in_buffer[2] << $bit_offset;
$out1 &= MASK_SKIP_HIGH_2; }};
}
impl<R: Read> Fr32Reader<R> {
pub fn new(source: R) -> Self {
Fr32Reader {
source,
in_buffer: AlignedBuffer([0; NUM_BYTES_IN_BLOCK + 1]),
out_buffer: [0; NUM_U128S_PER_BLOCK],
out_offset: 0,
available_frs: 0,
done: false,
}
}
fn process_block(&mut self) {
let in_buffer: &[u128] = {
#[cfg(target_arch = "aarch64")]
unsafe {
&mut (*(&self.in_buffer.0 as *const [u8] as *mut [u128]))
}
#[cfg(not(target_arch = "aarch64"))]
self.in_buffer.0.as_slice_of::<u128>().unwrap()
};
let out = &mut self.out_buffer;
{
out[0] = in_buffer[0];
out[1] = in_buffer[1] & MASK_SKIP_HIGH_2;
}
process_fr!(&in_buffer[1..], out[2], out[3], 2);
process_fr!(&in_buffer[3..], out[4], out[5], 4);
process_fr!(&in_buffer[5..], out[6], out[7], 6);
self.out_offset = 0;
}
fn fill_in_buffer(&mut self) -> io::Result<usize> {
let mut bytes_read = 0;
let mut buf = &mut self.in_buffer.0[..NUM_BYTES_IN_BLOCK];
while !buf.is_empty() {
match self.source.read(buf) {
Ok(0) => {
break;
}
Ok(n) => {
buf = &mut buf[n..];
bytes_read += n;
}
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
Err(e) => return Err(e),
}
}
for val in &mut self.in_buffer.0[bytes_read..NUM_BYTES_IN_BLOCK] {
*val = 0;
}
Ok(bytes_read)
}
}
#[inline]
const fn div_ceil(x: usize, y: usize) -> usize {
1 + ((x - 1) / y)
}
impl<R: Read> Read for Fr32Reader<R> {
fn read(&mut self, target: &mut [u8]) -> io::Result<usize> {
if self.done || target.is_empty() {
return Ok(0);
}
let mut bytes_read = 0;
let bytes_to_read = target.len();
while bytes_read < bytes_to_read {
if self.available_frs == 0 {
let bytes_read = self.fill_in_buffer()?;
if bytes_read == 0 {
self.done = true;
break;
}
self.process_block();
self.available_frs = div_ceil(bytes_read * 8, IN_BITS_FR);
}
{
let available_bytes = self.available_frs * (OUT_BITS_FR / 8);
let target_start = bytes_read;
let target_end = min(target_start + available_bytes, bytes_to_read);
let len = target_end - target_start;
let out_start = self.out_offset;
let out_end = out_start + len;
target[target_start..target_end]
.copy_from_slice(&self.out_buffer.as_byte_slice()[out_start..out_end]);
bytes_read += len;
self.out_offset += len;
self.available_frs -= div_ceil(len * 8, OUT_BITS_FR);
}
}
Ok(bytes_read)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
use bitvec::{order::Lsb0 as LittleEndian, vec::BitVec};
use itertools::Itertools;
use pretty_assertions::assert_eq;
use rand::random;
use crate::bytes_into_fr;
const DATA_BITS: u64 = 254;
const TARGET_BITS: u64 = 256;
#[test]
fn test_simple_short() {
let data = vec![3u8; 30];
let mut reader = Fr32Reader::new(Cursor::new(&data));
let mut padded = Vec::new();
reader
.read_to_end(&mut padded)
.expect("in-memory read failed");
assert_eq!(padded.len(), 32);
assert_eq!(&data[..], &padded[..30]);
assert_eq!(padded.into_boxed_slice(), bit_vec_padding(data));
}
#[test]
fn test_simple_single() {
let data = vec![255u8; 32];
let mut padded = Vec::new();
let mut reader = Fr32Reader::new(Cursor::new(&data));
reader
.read_to_end(&mut padded)
.expect("in-memory read failed");
assert_eq!(&padded[0..31], &data[0..31]);
assert_eq!(padded[31], 0b0011_1111);
assert_eq!(padded[32], 0b0000_0011);
assert_eq!(padded.len(), 64);
let bv = bit_vec_padding(data);
assert_eq!(bv.len(), 64);
assert_eq!(padded.into_boxed_slice(), bv);
}
#[test]
fn test_simple_127() {
let data = vec![255u8; 127];
let mut padded = Vec::new();
let mut reader = Fr32Reader::new(Cursor::new(&data));
reader
.read_to_end(&mut padded)
.expect("in-memory read failed");
assert_eq!(&padded[0..31], &data[0..31]);
assert_eq!(padded[31], 0b0011_1111);
assert_eq!(padded[32], 0b1111_1111);
assert_eq!(padded.len(), 128);
assert_eq!(padded.into_boxed_slice(), bit_vec_padding(data));
}
#[test]
fn test_chained_byte_source() {
let random_bytes: Vec<u8> = (0..127).map(|_| random::<u8>()).collect();
let output_x = {
let input_x = Cursor::new(random_bytes.clone());
let mut reader = Fr32Reader::new(input_x);
let mut buf_x = Vec::new();
reader.read_to_end(&mut buf_x).expect("could not seek");
buf_x
};
for n in 1..127 {
let random_bytes = random_bytes.clone();
let output_y = {
let input_y =
Cursor::new(random_bytes.iter().take(n).cloned().collect::<Vec<u8>>()).chain(
Cursor::new(random_bytes.iter().skip(n).cloned().collect::<Vec<u8>>()),
);
let mut reader = Fr32Reader::new(input_y);
let mut buf_y = Vec::new();
reader.read_to_end(&mut buf_y).expect("could not seek");
buf_y
};
assert_eq!(&output_x, &output_y, "should have written same bytes");
assert_eq!(
output_x.clone().into_boxed_slice(),
bit_vec_padding(random_bytes)
);
}
}
#[test]
fn test_full() {
let data = vec![255u8; 127];
let mut buf = Vec::new();
let mut reader = Fr32Reader::new(Cursor::new(&data));
reader.read_to_end(&mut buf).expect("in-memory read failed");
assert_eq!(buf.clone().into_boxed_slice(), bit_vec_padding(data));
validate_fr32(&buf);
}
#[test]
#[ignore]
fn test_long() {
use rand::{thread_rng, RngCore};
let mut rng = thread_rng();
for i in 1..100 {
for j in 1..50 {
let mut data = vec![0u8; i * j];
rng.fill_bytes(&mut data);
let mut buf = Vec::new();
let mut reader = Fr32Reader::new(Cursor::new(&data));
reader.read_to_end(&mut buf).expect("in-memory read failed");
assert_eq!(
buf.into_boxed_slice(),
bit_vec_padding(data),
"{} - {}",
i,
j
);
}
}
}
fn bit_vec_padding(raw_data: Vec<u8>) -> Box<[u8]> {
let mut padded_data: BitVec<LittleEndian, u8> = BitVec::new();
let raw_data: BitVec<LittleEndian, u8> = BitVec::from(raw_data);
for data_unit in raw_data.into_iter().chunks(DATA_BITS as usize).into_iter() {
padded_data.extend(data_unit);
if padded_data.len() % 8 != 0 {
for _ in 0..(TARGET_BITS - DATA_BITS) {
padded_data.push(false);
}
}
}
while padded_data.len() % (TARGET_BITS as usize) != 0 {
padded_data.push(false);
}
padded_data.into_boxed_slice()
}
fn validate_fr32(bytes: &[u8]) {
let chunks = (bytes.len() as f64 / 32_f64).ceil() as usize;
for (i, chunk) in bytes.chunks(32).enumerate() {
let _ = bytes_into_fr(chunk).unwrap_or_else(|_| {
panic!(
"chunk {}/{} cannot be converted to valid Fr: {:?}",
i + 1,
chunks,
chunk
)
});
}
}
#[test]
fn test_exotic() {
let mut source = vec![
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
25, 26, 27, 28, 29, 30, 31, 0xff, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 0xff, 9, 9,
];
source.extend(vec![9, 0xff]);
let mut buf = Vec::new();
let mut reader = Fr32Reader::new(Cursor::new(&source));
reader.read_to_end(&mut buf).expect("in-memory read failed");
for (i, &byte) in buf.iter().enumerate().take(31) {
assert_eq!(byte, i as u8 + 1);
}
assert_eq!(buf[31], 63); assert_eq!(buf[32], (1 << 2) | 0b11); for (i, &byte) in buf.iter().enumerate().skip(33).take(30) {
assert_eq!(byte, (i as u8 - 31) << 2);
}
assert_eq!(buf[63], (0x0f << 2)); assert_eq!(buf[64], 0x0f | 9 << 4); assert_eq!(buf[65], 9 << 4); assert_eq!(buf[66], 9 << 4); assert_eq!(buf[67], 0xf0); assert_eq!(buf[68], 0x0f);
assert_eq!(buf.into_boxed_slice(), bit_vec_padding(source));
}
}