use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
use crate::{integer, string};
pub trait Chooser: std::fmt::Debug + Sync {
fn name(&self) -> &'static str;
fn narrow_strings(
&self,
values: &[&[u8]],
offered: &[string::Kind],
depth: u8,
) -> Vec<string::Kind>;
fn narrow_integers(
&self,
values: &[i64],
offered: &[integer::Kind],
depth: u8,
) -> Vec<integer::Kind>;
fn considers_integer(&self, kind: integer::Kind, depth: u8) -> bool {
let _ = (kind, depth);
true
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Exhaustive;
pub const EXHAUSTIVE: Exhaustive = Exhaustive;
impl Chooser for Exhaustive {
fn name(&self) -> &'static str {
"exhaustive"
}
fn narrow_strings(
&self,
_values: &[&[u8]],
offered: &[string::Kind],
_depth: u8,
) -> Vec<string::Kind> {
offered.to_vec()
}
fn narrow_integers(
&self,
_values: &[i64],
offered: &[integer::Kind],
_depth: u8,
) -> Vec<integer::Kind> {
offered.to_vec()
}
}
#[derive(Debug, Clone, Copy)]
pub struct Sampled {
window: usize,
regions: usize,
}
const WINDOW: usize = 1024;
const REGIONS: usize = 8;
const FLOOR: usize = 256 * 1024;
impl Default for Sampled {
fn default() -> Self {
Self { window: WINDOW, regions: REGIONS }
}
}
impl Sampled {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn over(window: usize, regions: usize) -> Self {
Self { window: window.max(1), regions: regions.max(1) }
}
#[must_use]
pub fn size(self) -> usize {
self.window * self.regions
}
fn worth_it(self, count: usize, bytes: usize) -> bool {
count > self.size() && bytes >= FLOOR
}
}
impl Chooser for Sampled {
fn name(&self) -> &'static str {
"sampled"
}
fn narrow_strings(
&self,
values: &[&[u8]],
offered: &[string::Kind],
depth: u8,
) -> Vec<string::Kind> {
let bytes = values.iter().map(|value| value.len()).sum();
if offered.len() < 2 || !self.worth_it(values.len(), bytes) {
return offered.to_vec();
}
let sample = sample(values, self.window, self.regions);
let mut best: Option<(string::Kind, usize)> = None;
for &kind in offered {
let Ok(Some(size)) = string::size_as(kind, &sample, depth) else {
continue;
};
if best.is_none_or(|(_, smallest)| size < smallest) {
best = Some((kind, size));
}
}
best.map_or_else(|| offered.to_vec(), |(kind, _)| vec![kind])
}
fn narrow_integers(
&self,
values: &[i64],
offered: &[integer::Kind],
depth: u8,
) -> Vec<integer::Kind> {
if offered.len() < 2 || !self.worth_it(values.len(), values.len() * 8) {
return offered.to_vec();
}
let sample = sample(values, self.window, self.regions);
let mut best: Option<(integer::Kind, usize)> = None;
for &kind in offered {
let Ok(Some(size)) = integer::size_as(kind, &sample, depth) else {
continue;
};
if best.is_none_or(|(_, smallest)| size < smallest) {
best = Some((kind, size));
}
}
best.map_or_else(|| offered.to_vec(), |(kind, _)| vec![kind])
}
}
#[derive(Debug, Clone)]
pub struct Settled {
strings: Vec<string::Kind>,
integers: Vec<integer::Kind>,
}
impl Settled {
#[must_use]
pub fn new(strings: Vec<string::Kind>, integers: Vec<integer::Kind>) -> Self {
Self { strings, integers }
}
#[must_use]
pub fn strings(&self) -> &[string::Kind] {
&self.strings
}
}
impl Chooser for Settled {
fn name(&self) -> &'static str {
"settled"
}
fn narrow_strings(
&self,
_values: &[&[u8]],
offered: &[string::Kind],
depth: u8,
) -> Vec<string::Kind> {
match self.strings.get(depth as usize) {
Some(kind) if offered.contains(kind) => vec![*kind],
_ => offered.to_vec(),
}
}
fn narrow_integers(
&self,
_values: &[i64],
offered: &[integer::Kind],
depth: u8,
) -> Vec<integer::Kind> {
match self.integers.get(depth as usize) {
Some(kind) if offered.contains(kind) => vec![*kind],
_ => offered.to_vec(),
}
}
}
#[derive(Debug)]
pub struct Replay<'a> {
kinds: &'a [integer::Kind],
next: AtomicUsize,
lost: AtomicBool,
first: AtomicU8,
expected: Option<u8>,
fallback: &'a dyn Chooser,
}
impl<'a> Replay<'a> {
#[must_use]
pub fn new(kinds: &'a [integer::Kind], fallback: &'a dyn Chooser) -> Self {
Self {
kinds,
next: AtomicUsize::new(0),
lost: AtomicBool::new(false),
first: AtomicU8::new(0),
expected: None,
fallback,
}
}
#[must_use]
pub fn expecting(mut self, offered: &[integer::Kind]) -> Self {
self.expected = Some(bits(offered));
self
}
#[must_use]
pub fn first_offered(&self) -> Vec<integer::Kind> {
let first = self.first.load(Ordering::Relaxed);
integer::Kind::ALL.into_iter().filter(|kind| first & (1 << *kind as u8) != 0).collect()
}
#[must_use]
pub fn held(&self) -> bool {
!self.lost.load(Ordering::Relaxed) && self.next.load(Ordering::Relaxed) == self.kinds.len()
}
}
impl Chooser for Replay<'_> {
fn name(&self) -> &'static str {
"replay"
}
fn narrow_strings(
&self,
values: &[&[u8]],
offered: &[string::Kind],
depth: u8,
) -> Vec<string::Kind> {
self.fallback.narrow_strings(values, offered, depth)
}
fn narrow_integers(
&self,
values: &[i64],
offered: &[integer::Kind],
depth: u8,
) -> Vec<integer::Kind> {
if depth == 0 {
self.first.store(bits(offered), Ordering::Relaxed);
if self.expected.is_some_and(|expected| expected != bits(offered)) {
self.lost.store(true, Ordering::Relaxed);
}
}
if !self.lost.load(Ordering::Relaxed) {
let at = self.next.fetch_add(1, Ordering::Relaxed);
match self.kinds.get(at) {
Some(kind) if offered.contains(kind) => return vec![*kind],
_ => self.lost.store(true, Ordering::Relaxed),
}
}
self.fallback.narrow_integers(values, offered, depth)
}
fn considers_integer(&self, kind: integer::Kind, depth: u8) -> bool {
self.kinds.contains(&kind) || self.fallback.considers_integer(kind, depth)
}
}
fn bits(kinds: &[integer::Kind]) -> u8 {
kinds.iter().fold(0, |set, kind| set | 1 << *kind as u8)
}
pub(crate) fn sample<T: Copy>(values: &[T], window: usize, regions: usize) -> Vec<T> {
let wanted = window * regions;
if values.len() <= wanted {
return values.to_vec();
}
let last = values.len() - window;
let mut out = Vec::with_capacity(wanted);
for region in 0..regions {
let from = if regions == 1 { 0 } else { region * last / (regions - 1) };
out.extend_from_slice(&values[from..from + window]);
}
out
}
#[cfg(test)]
mod tests {
use super::{Chooser, EXHAUSTIVE, Replay, Sampled, sample};
use crate::{integer, string};
fn shaped_columns() -> Vec<Vec<i64>> {
let mut state = 0x9e37_79b9_7f4a_7c15_u64;
let mut noise = || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
(state % 1_000_000) as i64
};
vec![
(0..2048).map(|row| 1_600_000_000_000_000 + row * 1_000_000 + row % 7).collect(),
(0..2048).map(|row| row / 300).collect(),
(0..2048).map(|row| if row % 97 == 0 { row } else { 42 }).collect(),
(0..2048).map(|row| 5 + row * 1_000_000).collect(),
(0..2048).map(|_| noise()).collect(),
(0..37).map(|row| row * row).collect(),
]
}
#[test]
fn a_chunk_replayed_through_its_own_shape_comes_out_the_same() {
for values in shaped_columns() {
let searched = integer::encode_with(&values, &EXHAUSTIVE).unwrap();
let kinds = integer::shape(&searched).unwrap();
let replay = Replay::new(&kinds, &EXHAUSTIVE);
let replayed = integer::encode_with(&values, &replay).unwrap();
assert_eq!(replayed, searched, "{}", integer::describe(&searched).unwrap());
assert!(replay.held(), "{}", integer::describe(&searched).unwrap());
}
}
#[test]
fn a_shape_searched_under_another_offer_searches_again() {
let columns = shaped_columns();
let (noise, sparse) = (&columns[4], &columns[2]);
let first = Replay::new(&[], &EXHAUSTIVE);
let searched = integer::encode_with(noise, &first).unwrap();
let kinds = integer::shape(&searched).unwrap();
let offered = first.first_offered();
assert_eq!(offered, integer::offered(noise));
let blind = Replay::new(&kinds, &EXHAUSTIVE);
let packed = integer::encode_with(sparse, &blind).unwrap();
assert!(blind.held(), "packing fits any column, which is the trouble");
let checked = Replay::new(&kinds, &EXHAUSTIVE).expecting(&offered);
let written = integer::encode_with(sparse, &checked).unwrap();
assert!(!checked.held());
assert_eq!(written, integer::encode_with(sparse, &EXHAUSTIVE).unwrap());
assert!(written.len() * 4 < packed.len(), "{} against {}", written.len(), packed.len());
}
#[test]
fn a_shape_that_does_not_fit_still_writes_values_that_read_back() {
let columns = shaped_columns();
for from in &columns {
let kinds = integer::shape(&integer::encode_with(from, &EXHAUSTIVE).unwrap()).unwrap();
for values in &columns {
let replay = Replay::new(&kinds, &EXHAUSTIVE);
let bytes = integer::encode_with(values, &replay).unwrap();
assert_eq!(&integer::decode(&bytes).unwrap(), values);
}
}
}
#[test]
fn a_sample_covers_the_whole_input_and_not_one_end_of_it() {
let values: Vec<i64> = (0..8000).collect();
let taken = sample(&values, 10, 4);
assert_eq!(taken.len(), 40);
assert_eq!(taken[0], 0);
assert_eq!(taken[10], 2663);
assert_eq!(taken[20], 5326);
assert_eq!(taken[30], 7990);
assert_eq!(taken[39], 7999);
}
#[test]
fn an_input_no_bigger_than_the_sample_is_the_sample() {
let values: Vec<i64> = (0..30).collect();
assert_eq!(sample(&values, 10, 4), values);
}
#[test]
fn the_last_window_does_not_run_off_the_end() {
let values: Vec<i64> = (0..100).collect();
let taken = sample(&values, 40, 2);
assert_eq!(taken.len(), 80);
assert_eq!(*taken.last().expect("the sample is not empty"), 99);
}
#[test]
fn the_exhaustive_chooser_hands_back_exactly_what_it_was_offered() {
let offered = [string::Kind::Plain, string::Kind::Fsst, string::Kind::Dict];
assert_eq!(EXHAUSTIVE.narrow_strings(&[b"a".as_slice()], &offered, 0), offered);
let offered = [integer::Kind::Packed, integer::Kind::Delta];
assert_eq!(EXHAUSTIVE.narrow_integers(&[1, 2], &offered, 0), offered);
}
#[test]
fn a_chunk_no_bigger_than_the_sample_is_not_narrowed_at_all() {
let sampled = Sampled::over(4, 2);
let values: Vec<i64> = (0..8).collect();
let offered = [integer::Kind::Packed, integer::Kind::Delta];
assert_eq!(sampled.narrow_integers(&values, &offered, 0), offered);
}
#[test]
fn a_sampled_chooser_returns_one_of_what_it_was_offered() {
let sampled = Sampled::over(16, 2);
let values: Vec<i64> = (0..40_000).map(|index| index / 200).collect();
let offered = [integer::Kind::Packed, integer::Kind::Rle, integer::Kind::Dict];
let narrowed = sampled.narrow_integers(&values, &offered, 0);
assert_eq!(narrowed.len(), 1);
assert!(offered.contains(&narrowed[0]), "{narrowed:?}");
}
#[test]
fn a_chunk_with_plenty_of_values_and_hardly_any_bytes_is_not_sampled() {
let sampled = Sampled::over(16, 2);
let empty = Vec::new();
let values: Vec<&[u8]> = vec![empty.as_slice(); 40_000];
let offered = [string::Kind::Plain, string::Kind::Fsst, string::Kind::Dict];
assert_eq!(sampled.narrow_strings(&values, &offered, 0), offered);
}
}