inflate.rs

//! deflate の展開 (RFC 1951)
//!
//! ZIP のエントリはほとんどが deflate で縮めてある。ここはそれを元に戻すところである。
//!
//! # 表引きで符号を読む
//!
//! ハフマン符号は 1 ビットずつ木をたどっても読めるが、記号ごとに平均 8 回まわることになる。
//! 先に 9 ビット分の表を作っておくと、短い符号は 1 回引くだけで記号と符号長が分かる。
//! 9 ビットを超える符号だけ、1 ビットずつたどる遅いほうへ落とす。
//!
//! # 入口も出口もためこまない
//!
//! 巨大なシートは、縮めたままでも数百 MB、展開すると数 GB になる。どちらもメモリには置けない。
//! [`inflate_stream`] は縮めたバイトを少しずつ受け取り、展開できたぶんを少しずつ渡す。
//! 手元に残すのは、読みかけの入力と、後ろ向きの参照に要る 32KB だけである。
//! 小さなパートは [`inflate`] で一度に受け取ればよい。

use crate::error::Error;


/// 一度に引けるビット数。これを超える符号は遅いほうで読む
const FAST_BITS: u32 = 9;

/// 後ろ向きの参照が届く距離。deflate の決まりで 32768 バイト
const HISTORY: usize = 32768;

/// これだけたまったら呼び出し側へ渡す
const FLUSH_LIMIT: usize = 256 * 1024;

/// 読み終えた入力をこれだけためたら捨てる
const COMPACT: usize = 256 * 1024;

/// 長さ符号 257..285 の基準値
const LENGTH_BASE: [u16; 29] = [
    3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131,
    163, 195, 227, 258,
];

/// 長さ符号 257..285 の追加ビット数
const LENGTH_EXTRA: [u32; 29] = [
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0,
];

/// 距離符号 0..29 の基準値
const DISTANCE_BASE: [u16; 30] = [
    1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537,
    2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577,
];

/// 距離符号 0..29 の追加ビット数
const DISTANCE_EXTRA: [u32; 30] = [
    0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13,
    13,
];

/// 符号長そのものを送る順番
const ORDER: [usize; 19] = [16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15];


/// 縮めたバイト列を元に戻して、全部まとめて返す
pub fn inflate(data: &[u8]) -> Result<Vec<u8>, Error> {
    inflate_sized(data, 0)
}


/// 展開後の大きさが分かっているときは、先に場所を取っておく
pub fn inflate_sized(data: &[u8], size: usize) -> Result<Vec<u8>, Error> {
    let mut output = Output {
        window: Vec::with_capacity(size),
        sent: 0,
        sink: None,
    };

    let mut once = Some(data.to_vec());
    let mut supply = move || Ok(once.take());
    run(&mut supply, &mut output)?;
    Ok(output.window)
}


/// そろっているバイト列を展開しながら、出てきたぶんを少しずつ sink へ渡す
pub fn inflate_with(
    data: &[u8],
    sink: &mut dyn FnMut(&[u8]) -> Result<(), Error>,
) -> Result<(), Error> {
    let mut once = Some(data.to_vec());
    let mut supply = move || Ok(once.take());
    inflate_stream(&mut supply, sink)
}


/// 縮めたバイトを少しずつ受け取り、展開したぶんを少しずつ sink へ渡す
///
/// supply は続きを返す。もう続きがなければ None を返す。
/// sink は出力の続きを順に受け取る。同じバイトが二度渡ることはない。
pub fn inflate_stream(
    supply: &mut dyn FnMut() -> Result<Option<Vec<u8>>, Error>,
    sink: &mut dyn FnMut(&[u8]) -> Result<(), Error>,
) -> Result<(), Error> {
    let mut output = Output {
        window: Vec::with_capacity(FLUSH_LIMIT + HISTORY),
        sent: 0,
        sink: Some(sink),
    };

    run(supply, &mut output)?;
    output.flush(true)
}


/// 出力を貯めるところ
struct Output<'a> {
    /// 出したバイト。sink があるときは、後ろ向きの参照に要るぶんだけを残す
    window: Vec<u8>,
    /// window の先頭から、すでに sink へ渡したバイト数
    sent: usize,
    /// 渡す先。ない場合は window に全部ためる
    sink: Option<&'a mut dyn FnMut(&[u8]) -> Result<(), Error>>,
}


impl Output<'_> {
    /// 1 バイト出す
    fn push(&mut self, byte: u8) {
        self.window.push(byte);
    }

    /// distance だけ戻ったところから length バイト写す
    fn copy(&mut self, distance: usize, length: usize) -> Result<(), Error> {
        if distance > self.window.len() {
            return Err(Error::new("圧縮データが届かないところを参照しています"));
        }

        let start = self.window.len() - distance;

        if distance >= length {
            self.window.extend_from_within(start..start + length);
            return Ok(());
        }

        // 重なっているので 1 バイトずつ。写したそばから写す元になる
        for index in 0..length {
            let byte = self.window[start + index];
            self.window.push(byte);
        }

        Ok(())
    }

    /// たまったぶんを渡す。force でないときは、たまるまで何もしない
    fn flush(&mut self, force: bool) -> Result<(), Error> {
        let Some(sink) = self.sink.as_mut() else {
            return Ok(());
        };

        if !force && self.window.len() < FLUSH_LIMIT + HISTORY {
            return Ok(());
        }

        if self.window.len() > self.sent {
            sink(&self.window[self.sent..])?;
            self.sent = self.window.len();
        }

        if self.window.len() > HISTORY {
            let start = self.window.len() - HISTORY;
            self.window.copy_within(start.., 0);
            self.window.truncate(HISTORY);
            self.sent = HISTORY;
        }

        Ok(())
    }
}


/// ビットを下位から順に取り出す
///
/// data は縮めたバイトの読みかけの区間である。読み終えた前のほうは捨て、
/// 足りなくなったら supply から続きを受け取る。
struct Bits<'a> {
    data: Vec<u8>,
    /// data のうち、hold へ取り込んだところまでの位置
    position: usize,
    /// 取り込んだビット。下位が先に読むほう
    hold: u64,
    /// hold に入っている有効なビット数
    count: u32,
    /// 続きをくれるもの
    supply: &'a mut dyn FnMut() -> Result<Option<Vec<u8>>, Error>,
    /// supply が尽きたか
    ended: bool,
}


impl Bits<'_> {
    /// 続きを受け取る。もうなければ false
    fn pull(&mut self) -> Result<bool, Error> {
        if self.ended {
            return Ok(false);
        }

        match (self.supply)()? {
            Some(piece) => {
                self.data.extend_from_slice(&piece);
                Ok(true)
            }

            None => {
                self.ended = true;
                Ok(false)
            }
        }
    }

    /// 読み終えた前のほうを捨てる
    fn compact(&mut self) {
        let keep = self.byte_position();

        if keep < COMPACT {
            return;
        }

        self.data.copy_within(keep.., 0);
        self.data.truncate(self.data.len() - keep);
        self.position -= keep;
    }

    /// hold を埋められるだけ埋める
    fn refill(&mut self) -> Result<(), Error> {
        while self.count <= 56 {
            if self.position >= self.data.len() {
                self.compact();

                if !self.pull()? {
                    return Ok(());
                }

                continue;
            }

            self.hold |= (self.data[self.position] as u64) << self.count;
            self.position += 1;
            self.count += 8;
        }

        Ok(())
    }

    /// n ビット取り出す
    fn take(&mut self, n: u32) -> Result<u32, Error> {
        if self.count < n {
            self.refill()?;
        }

        if self.count < n {
            return Err(Error::new("圧縮データが途中で終わっています"));
        }

        let value = (self.hold & ((1u64 << n) - 1)) as u32;
        self.hold >>= n;
        self.count -= n;
        Ok(value)
    }

    /// バイトの境目までビットを捨てる
    fn align(&mut self) {
        let extra = self.count % 8;
        self.hold >>= extra;
        self.count -= extra;
    }

    /// 境目にそろえたあとの、data 上の読み出し位置
    fn byte_position(&self) -> usize {
        self.position - (self.count / 8) as usize
    }

    /// 読み出し位置から need バイト用意する
    fn ready(&mut self, need: usize) -> Result<(), Error> {
        loop {
            if self.data.len() - self.byte_position() >= need {
                return Ok(());
            }

            self.compact();

            if !self.pull()? {
                return Err(Error::new("圧縮データが途中で終わっています"));
            }
        }
    }

    /// 境目にそろえたあと、length バイトぶん進める
    fn skip(&mut self, length: usize) {
        let start = self.byte_position();
        self.position = start + length;
        self.hold = 0;
        self.count = 0;
    }

    /// 記号を 1 つ読む
    fn decode(&mut self, huffman: &Huffman) -> Result<u16, Error> {
        if self.count < FAST_BITS {
            self.refill()?;
        }

        let entry = huffman.fast[(self.hold & ((1 << FAST_BITS) - 1)) as usize];
        let length = (entry >> 12) as u32;

        if length != 0 && length <= self.count {
            self.hold >>= length;
            self.count -= length;
            return Ok(entry & 0x0fff);
        }

        self.decode_long(huffman)
    }

    /// 9 ビットを超える符号を 1 ビットずつたどる
    fn decode_long(&mut self, huffman: &Huffman) -> Result<u16, Error> {
        let mut code: i32 = 0;
        let mut first: i32 = 0;
        let mut index: i32 = 0;

        for length in 1..16 {
            code |= self.take(1)? as i32;
            let count = huffman.count[length] as i32;

            if code - first < count {
                return Ok(huffman.symbol[(index + code - first) as usize]);
            }

            index += count;
            first = (first + count) << 1;
            code <<= 1;
        }

        Err(Error::new("圧縮データの符号が壊れています"))
    }
}


/// 符号長の並びから作った引き表
struct Huffman {
    /// 符号長ごとの記号数
    count: [u16; 16],
    /// 符号の順に並べた記号
    symbol: Vec<u16>,
    /// FAST_BITS ビットの表。上位 4 ビットが符号長、下位 12 ビットが記号。0 は載っていない印
    fast: Vec<u16>,
}


/// 符号長の並びから引き表を作る
///
/// 記号が 1 つだけの符号は、埋まりきらない不完全な符号になる。
/// 中身が空の区画や、一致がまったくない区画で実際に出てくるので、single を渡して受ける。
/// 受けるのは記号が 1 つ以下のときだけで、それ以外の不完全な符号は壊れたデータとして弾く。
/// 符号長そのものを送る符号 (ORDER のほう) では、これも許さない。
fn build(length: &[u8], single: bool) -> Result<Huffman, Error> {
    let mut count = [0u16; 16];

    for value in length {
        count[*value as usize] += 1;
    }

    count[0] = 0;
    let mut left: i32 = 1;

    for bits in 1..16 {
        left <<= 1;
        left -= count[bits] as i32;

        if left < 0 {
            return Err(Error::new("圧縮データの符号長が多すぎます"));
        }
    }

    let mut longest = 0;

    for bits in 1..16 {
        if count[bits] != 0 {
            longest = bits;
        }
    }

    if left > 0 && !(single && longest <= 1) {
        return Err(Error::new("圧縮データの符号長が足りません"));
    }

    let mut offset = [0u16; 16];

    for bits in 1..15 {
        offset[bits + 1] = offset[bits] + count[bits];
    }

    let total: usize = count.iter().map(|value| *value as usize).sum();
    let mut symbol = vec![0u16; total.max(1)];

    for (value, bits) in length.iter().enumerate() {
        if *bits != 0 {
            symbol[offset[*bits as usize] as usize] = value as u16;
            offset[*bits as usize] += 1;
        }
    }

    let mut next = [0u32; 16];
    let mut code: u32 = 0;

    for bits in 1..16 {
        code = (code + count[bits - 1] as u32) << 1;
        next[bits] = code;
    }

    let mut fast = vec![0u16; 1 << FAST_BITS];

    for (value, bits) in length.iter().enumerate() {
        let bits = *bits as u32;

        if bits == 0 {
            continue;
        }

        let assigned = next[bits as usize];
        next[bits as usize] += 1;

        if bits > FAST_BITS {
            continue;
        }

        // 流れの先頭ビットが符号の最上位なので、表を引くには符号を逆に並べる
        let mut place = reverse(assigned, bits) as usize;
        let step = 1usize << bits;
        let entry = ((bits as u16) << 12) | value as u16;

        while place < (1 << FAST_BITS) {
            fast[place] = entry;
            place += step;
        }
    }

    Ok(Huffman {count, symbol, fast})
}


/// 下位 bits ビットを逆に並べる
fn reverse(value: u32, bits: u32) -> u32 {
    let mut result = 0;

    for index in 0..bits {
        result |= ((value >> index) & 1) << (bits - 1 - index);
    }

    result
}


/// 区画を順に読んで展開する
fn run(
    supply: &mut dyn FnMut() -> Result<Option<Vec<u8>>, Error>,
    output: &mut Output,
) -> Result<(), Error> {
    let mut bits = Bits {
        data: Vec::with_capacity(COMPACT * 2),
        position: 0,
        hold: 0,
        count: 0,
        supply,
        ended: false,
    };

    loop {
        let last = bits.take(1)?;
        let kind = bits.take(2)?;

        match kind {
            0 => stored(&mut bits, output)?,
            1 => {
                let (literal, distance) = fixed()?;
                block(&mut bits, output, &literal, &distance)?;
            }

            2 => {
                let (literal, distance) = dynamic(&mut bits)?;
                block(&mut bits, output, &literal, &distance)?;
            }

            _ => return Err(Error::new("圧縮データに知らない区画の種類があります")),
        }

        output.flush(false)?;

        if last == 1 {
            return Ok(());
        }
    }
}


/// 縮めていない区画をそのまま写す
fn stored(bits: &mut Bits, output: &mut Output) -> Result<(), Error> {
    bits.align();
    bits.ready(4)?;
    let start = bits.byte_position();
    let length = u16::from_le_bytes([bits.data[start], bits.data[start + 1]]) as usize;
    let check = u16::from_le_bytes([bits.data[start + 2], bits.data[start + 3]]) as usize;

    if length != !check & 0xffff {
        return Err(Error::new("縮めていない区画の長さが合いません"));
    }

    bits.ready(4 + length)?;
    let start = bits.byte_position();
    output.window.extend_from_slice(&bits.data[start + 4..start + 4 + length]);
    bits.skip(4 + length);
    Ok(())
}


/// 決め打ちの符号表
fn fixed() -> Result<(Huffman, Huffman), Error> {
    let mut literal = [0u8; 288];

    for value in 0..288 {
        literal[value] = match value {
            0..=143 => 8,
            144..=255 => 9,
            256..=279 => 7,
            _ => 8,
        };
    }

    // 距離符号は 32 個ぶん 5 ビットで置く。30 と 31 は現れないが、符号としては埋まっている
    let distance = [5u8; 32];
    Ok((build(&literal, false)?, build(&distance, false)?))
}


/// 区画の頭に載っている符号表を読む
fn dynamic(bits: &mut Bits) -> Result<(Huffman, Huffman), Error> {
    let literal_count = bits.take(5)? as usize + 257;
    let distance_count = bits.take(5)? as usize + 1;
    let order_count = bits.take(4)? as usize + 4;

    if literal_count > 286 || distance_count > 30 {
        return Err(Error::new("圧縮データの符号表が大きすぎます"));
    }

    let mut order_length = [0u8; 19];

    for index in 0..order_count {
        order_length[ORDER[index]] = bits.take(3)? as u8;
    }

    let order = build(&order_length, false)?;
    let mut length = vec![0u8; literal_count + distance_count];
    let mut index = 0;

    while index < length.len() {
        let symbol = bits.decode(&order)?;

        match symbol {
            0..=15 => {
                length[index] = symbol as u8;
                index += 1;
            }

            16 => {
                if index == 0 {
                    return Err(Error::new("圧縮データの符号長に写す元がありません"));
                }

                let previous = length[index - 1];
                let repeat = bits.take(2)? as usize + 3;
                index = fill(&mut length, index, repeat, previous)?;
            }

            17 => {
                let repeat = bits.take(3)? as usize + 3;
                index = fill(&mut length, index, repeat, 0)?;
            }

            18 => {
                let repeat = bits.take(7)? as usize + 11;
                index = fill(&mut length, index, repeat, 0)?;
            }

            _ => return Err(Error::new("圧縮データの符号長が壊れています")),
        }
    }

    let literal = build(&length[..literal_count], true)?;
    let distance = build(&length[literal_count..], true)?;
    Ok((literal, distance))
}


/// 符号長を repeat 個ぶん埋める
fn fill(length: &mut [u8], index: usize, repeat: usize, value: u8) -> Result<usize, Error> {
    if index + repeat > length.len() {
        return Err(Error::new("圧縮データの符号長が多すぎます"));
    }

    for place in index..index + repeat {
        length[place] = value;
    }

    Ok(index + repeat)
}


/// 符号表が決まったあとの本体。記号を読んでは出していく
fn block(
    bits: &mut Bits,
    output: &mut Output,
    literal: &Huffman,
    distance: &Huffman,
) -> Result<(), Error> {
    loop {
        if output.window.len() >= FLUSH_LIMIT + HISTORY {
            output.flush(false)?;
        }

        let symbol = bits.decode(literal)?;

        if symbol < 256 {
            output.push(symbol as u8);
            continue;
        }

        if symbol == 256 {
            return Ok(());
        }

        let index = symbol as usize - 257;

        if index >= LENGTH_BASE.len() {
            return Err(Error::new("圧縮データの長さ符号が範囲を超えています"));
        }

        let length = LENGTH_BASE[index] as usize + bits.take(LENGTH_EXTRA[index])? as usize;
        let code = bits.decode(distance)? as usize;

        if code >= DISTANCE_BASE.len() {
            return Err(Error::new("圧縮データの距離符号が範囲を超えています"));
        }

        let reach = DISTANCE_BASE[code] as usize + bits.take(DISTANCE_EXTRA[code])? as usize;
        output.copy(reach, length)?;
    }
}