use probability::distribution::Distribution;
use rustc_hash::FxHashMap;
use crate::buffer::{Buff, BuffMut};
use super::{pre_encode, wrapped::WrappedReedSolomon};
#[derive(Debug)]
pub struct FrameEncoder {
rate_table: FxHashMap<(u8, usize), usize>,
target_loss: u8,
}
impl FrameEncoder {
#[tracing::instrument(level = "trace")]
pub fn new(target_loss: u8) -> Self {
FrameEncoder {
rate_table: FxHashMap::default(),
target_loss,
}
}
#[tracing::instrument(level = "trace", skip(pkts))]
pub fn encode(&mut self, measured_loss: u8, pkts: &[Buff]) -> Vec<Buff> {
let max_length = pkts.iter().map(|v| v.len()).max().unwrap();
let mut padded_pkts: Vec<BuffMut> =
pkts.iter().map(|p| pre_encode(p, max_length + 2)).collect();
let data_shards = pkts.len();
let parity_shards = self.repair_len(measured_loss, pkts.len());
let mut parity_shard_space = vec![vec![0u8; max_length + 2]; parity_shards];
let mut padded_pkts: Vec<&mut [u8]> = padded_pkts.iter_mut().map(|v| v.as_mut()).collect();
for r in parity_shard_space.iter_mut() {
padded_pkts.push(r);
}
if parity_shards > 0 {
let encoder = WrappedReedSolomon::new_cached(data_shards, parity_shards);
encoder
.get_inner()
.encode(&mut padded_pkts)
.expect("can't encode");
}
let mut toret = Vec::with_capacity(data_shards + parity_shards);
toret.extend(padded_pkts.iter().map(|vec| Buff::copy_from_slice(vec)));
toret
}
fn repair_len(&mut self, measured_loss: u8, run_len: usize) -> usize {
let target_loss = self.target_loss;
(*self
.rate_table
.entry((measured_loss, run_len))
.or_insert_with(|| {
for additional_len in 0.. {
let distro = probability::distribution::Binomial::with_failure(
run_len + additional_len,
(measured_loss as f64 / 256.0).max(1e-100).min(1.0 - 1e-100),
);
let result_loss = distro.distribution(run_len as f64);
if result_loss <= target_loss as f64 / 256.0 {
return additional_len.saturating_sub(1usize);
}
}
panic!()
}))
.min(255 - run_len)
.min(run_len * 2)
}
}