zrip-core 0.2.0

Shared types and codecs for zrip (internal crate)
Documentation 
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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

#![forbid(unsafe_code)]

#[cfg(feature = "alloc")]
use alloc::vec;
#[cfg(feature = "alloc")]
use alloc::vec::Vec;

use crate::bitstream::writer::BitWriter;

pub struct FseEncodeTable {
    pub symbol_tt: Vec<SymbolTransform>,
    pub state_table: Vec<u16>,
    pub accuracy_log: u8,
}

#[derive(Clone, Copy)]
pub struct SymbolTransform {
    pub delta_nb_bits: u32,
    pub delta_find_state: i32,
}

#[cfg(feature = "alloc")]
impl FseEncodeTable {
    pub fn from_distribution(distribution: &[i16], accuracy_log: u8) -> Self {
        let table_size = 1usize << accuracy_log;
        let mut symbol_tt = vec![
            SymbolTransform {
                delta_nb_bits: 0,
                delta_find_state: 0,
            };
            distribution.len()
        ];
        let mut state_table = vec![0u16; table_size];

        let step = (table_size >> 1) + (table_size >> 3) + 3;
        let mask = table_size - 1;

        let mut high_threshold = table_size - 1;
        let mut position = 0;
        let mut cumul = vec![0u32; distribution.len() + 1];

        for (s, &prob) in distribution.iter().enumerate() {
            if prob == -1 {
                state_table[high_threshold] = s as u16;
                high_threshold -= 1;
                cumul[s + 1] = cumul[s] + 1;
            } else {
                cumul[s + 1] = cumul[s] + prob.max(0) as u32;
            }
        }

        for (s, &prob) in distribution.iter().enumerate() {
            if prob <= 0 {
                continue;
            }
            for _ in 0..prob {
                state_table[position] = s as u16;
                position = (position + step) & mask;
                while position > high_threshold {
                    position = (position + step) & mask;
                }
            }
        }

        let mut next_state_number = vec![0u32; distribution.len()];
        for s in 0..distribution.len() {
            let prob = distribution[s];
            if prob <= 0 {
                let max_nb_bits = accuracy_log;
                let min_state_plus = if prob == -1 { 1u32 } else { 0 };
                symbol_tt[s].delta_nb_bits = ((max_nb_bits as u32 + 1) << 16) - min_state_plus;
                symbol_tt[s].delta_find_state = 0;
                next_state_number[s] = cumul[s];
            } else if prob == 1 {
                let max_bits_out = accuracy_log as u32;
                let min_state_plus = 1u32 << accuracy_log;
                symbol_tt[s].delta_nb_bits = (max_bits_out << 16).wrapping_sub(min_state_plus);
                symbol_tt[s].delta_find_state = cumul[s] as i32 - 1;
                next_state_number[s] = cumul[s];
            } else {
                let prob = prob as u32;
                let max_bits_out = accuracy_log as u32 - high_bit(prob - 1);
                let min_state_plus = prob << max_bits_out;
                symbol_tt[s].delta_nb_bits = (max_bits_out << 16).wrapping_sub(min_state_plus);
                symbol_tt[s].delta_find_state = (cumul[s] as i32) - (prob as i32);
                next_state_number[s] = cumul[s];
            }
        }

        // State table stores values in [table_size, 2*table_size),
        // matching C zstd's FSE_buildCTable convention.
        let mut table_symbol_sorted = vec![0u16; table_size];
        for i in 0..table_size {
            let s = state_table[i] as usize;
            let ns = next_state_number[s];
            table_symbol_sorted[ns as usize] = (table_size + i) as u16;
            next_state_number[s] += 1;
        }

        Self {
            symbol_tt,
            state_table: table_symbol_sorted,
            accuracy_log,
        }
    }
}

pub struct FseEncodeState<'t> {
    table: &'t FseEncodeTable,
    state: u32,
}

#[cfg(feature = "alloc")]
impl<'t> FseEncodeState<'t> {
    pub fn init(table: &'t FseEncodeTable, symbol: u8) -> Self {
        let tt = &table.symbol_tt[symbol as usize];
        let nb_bits_out = tt.delta_nb_bits.wrapping_add(1 << 16) >> 16;
        let value = (nb_bits_out << 16).wrapping_sub(tt.delta_nb_bits);
        let idx = (value >> nb_bits_out) as i32 + tt.delta_find_state;
        let state = table.state_table[idx as usize] as u32;
        Self { table, state }
    }

    pub fn encode_symbol(&mut self, writer: &mut BitWriter, symbol: u8) {
        let tt = &self.table.symbol_tt[symbol as usize];
        let nb_bits_out = (self.state.wrapping_add(tt.delta_nb_bits) >> 16) as u8;
        writer.write_bits(self.state & ((1u32 << nb_bits_out) - 1), nb_bits_out);
        self.state = self.table.state_table
            [((self.state >> nb_bits_out) as i32 + tt.delta_find_state) as usize]
            as u32;
    }

    pub fn flush(&self, writer: &mut BitWriter, accuracy_log: u8) {
        writer.write_bits(self.state & ((1u32 << accuracy_log) - 1), accuracy_log);
    }

    pub fn state(&self) -> u32 {
        self.state
    }
}

fn high_bit(val: u32) -> u32 {
    debug_assert!(val > 0);
    31 - val.leading_zeros()
}