26 Commits

Author SHA1 Message Date
af5fe898bf add --no-repeated-offsets to help 2022-09-25 16:24:24 +02:00
331857a711 add option to disable repeated offsets 2022-09-25 16:23:11 +02:00
12e6f95fe8 add remaining encoding config options + presets for x86 and z80 2022-09-24 22:00:50 +02:00
23872b3222 implement encoding options 2022-09-24 20:52:39 +02:00
ced6cc8c32 some more risc-v optimizations 2022-09-24 08:45:14 +02:00
8c9e4311b9 first (poorly optimized) risc-v unpacker 2022-09-23 22:40:47 +02:00
31c31bdcfb clean up command line interface 2022-09-21 22:45:06 +02:00
8f33ae0b1e add reverse compression option 2022-09-21 21:37:30 +02:00
f5fc9bd005 implement optional parity contexts 2022-09-20 23:24:19 +02:00
cc41feb5cd alternative way to write state/prob update 2022-09-19 18:33:02 +02:00
5c7aee046a optimize decode_bit some more -> 166b 2022-09-18 23:11:26 +02:00
612084a5bf decode_length returns negative value -> 172b 2022-09-18 22:36:31 +02:00
ad731c2e75 Merge pull request #4 from Ferdi265/master
unpack_armv6m: update comment headers and remove unneeded pushed register
2022-09-18 18:51:37 +02:00
Ferdinand Bachmann
52f9778c0f unpack_armv6m: update comment headers and remove unneeded pushed register 2022-09-18 18:49:20 +02:00
49a611e8ba some more optimizations -> 176 bytes 2022-09-18 17:17:37 +02:00
2f820316e3 change prob_index update to save two instructions -> 184b 2022-09-18 16:27:21 +02:00
5bc3f88564 invert was_match -> 188 bytes 2022-09-18 15:58:31 +02:00
434769b591 simple dev setup for asm unpackers 2022-09-18 15:40:23 +02:00
629c5fce7d optimize c_unpacker state update a bit, add -b flag to --help 2022-09-09 19:10:31 +02:00
a205473ad6 slight optimization 2022-09-09 08:32:03 +02:00
4903ac3786 c unpacker now works 2022-09-09 00:47:33 +02:00
f817dc9254 first try of c decompressor, not working yet 2022-09-08 23:42:03 +02:00
d93aec186c add compressed_size function 2022-06-19 23:08:47 +02:00
3902425922 add bitstream variant, could be useful on 8bit platforms 2021-12-28 23:59:56 +01:00
2e7983fc65 use u8 context entries 2021-12-27 22:35:53 +01:00
f7f891e154 fix creating broken files if they end with a match 2021-12-27 22:34:14 +01:00
21 changed files with 1071 additions and 136 deletions

16
Cargo.lock generated
View File

@@ -62,6 +62,12 @@ version = "1.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
[[package]]
name = "lexopt"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "478ee9e62aaeaf5b140bd4138753d1f109765488581444218d3ddda43234f3e8"
[[package]] [[package]]
name = "libc" name = "libc"
version = "0.2.108" version = "0.2.108"
@@ -89,12 +95,6 @@ dependencies = [
"winapi", "winapi",
] ]
[[package]]
name = "pico-args"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "db8bcd96cb740d03149cbad5518db9fd87126a10ab519c011893b1754134c468"
[[package]] [[package]]
name = "sacabase" name = "sacabase"
version = "2.0.0" version = "2.0.0"
@@ -117,12 +117,12 @@ dependencies = [
[[package]] [[package]]
name = "upkr" name = "upkr"
version = "0.1.0" version = "0.2.0-pre1"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"cdivsufsort", "cdivsufsort",
"lexopt",
"pbr", "pbr",
"pico-args",
] ]
[[package]] [[package]]

View File

@@ -1,12 +1,12 @@
[package] [package]
name = "upkr" name = "upkr"
version = "0.1.0" version = "0.2.0-pre2"
edition = "2021" edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies] [dependencies]
cdivsufsort = "2" cdivsufsort = "2"
pico-args = "0.4" lexopt = "0.2.1"
anyhow = "1" anyhow = "1"
pbr = "1" pbr = "1"

1
asm_unpackers/.gitignore vendored Normal file
View File

@@ -0,0 +1 @@
/build/

51
asm_unpackers/Makefile Normal file
View File

@@ -0,0 +1,51 @@
build/unpack_riscv64: ../c_unpacker/main.c unpack_riscv.S
mkdir -p build
riscv64-linux-gnu-gcc -g -static -o $@ $^
test_riscv64: build/unpack_riscv64
qemu-riscv64 $< test_data.upk /tmp/out.bin
cmp test_data.bin /tmp/out.bin
build/unpack_riscv64.o: unpack_riscv.S
mkdir -p build
riscv64-linux-gnu-gcc -c -o $@ $?
build/unpack_riscv64.bin: build/unpack_riscv64.o
riscv64-linux-gnu-objcopy -O binary --only-section=.text $? $@
disas-riscv64: build/unpack_riscv64.o
riscv64-linux-gnu-objdump -d $?
build/unpack_riscv32.o: unpack_riscv.S
mkdir -p build
riscv64-linux-gnu-gcc -march=rv32imc -mabi=ilp32 -c -o $@ $?
build/unpack_riscv32.bin: build/unpack_riscv32.o
riscv64-linux-gnu-objcopy -O binary --only-section=.text $? $@
disas-riscv32: build/unpack_riscv32.o
riscv64-linux-gnu-objdump -d $?
build/unpack_armv6m: ../c_unpacker/main.c unpack_armv6m.S
mkdir -p build
arm-linux-gnueabihf-gcc -g -static -o $@ $^
test_armv6m: build/unpack_armv6m
qemu-arm $< test_data.upk /tmp/out.bin
cmp test_data.bin /tmp/out.bin
build/unpack_armv6m.bin: unpack_armv6m.S
mkdir -p build
arm-none-eabi-gcc -march=armv6-m -c -o build/unpack_armv6m.o $?
arm-none-eabi-objcopy -O binary --only-section=.text build/unpack_armv6m.o $@
build/unpack_c: ../c_unpacker/main.c ../c_unpacker/unpack.c
mkdir -p build
gcc -g -o $@ $^
test_c: build/unpack_c
$< test_data.upk /tmp/out.bin
cmp test_data.bin /tmp/out.bin
sizes: build/unpack_armv6m.bin build/unpack_riscv64.bin build/unpack_riscv32.bin
ls -l build/*.bin

View File

@@ -0,0 +1,99 @@
typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned long u32;
u8* upkr_data_ptr;
u8 upkr_probs[1 + 255 + 1 + 2*32 + 2*32];
#ifdef UPKR_BITSTREAM
u16 upkr_state;
u8 upkr_current_byte;
int upkr_bits_left;
#else
u32 upkr_state;
#endif
int upkr_decode_bit(int context_index) {
#ifdef UPKR_BITSTREAM
while(upkr_state < 32768) {
if(upkr_bits_left == 0) {
upkr_current_byte = *upkr_data_ptr++;
upkr_bits_left = 8;
}
upkr_state = (upkr_state << 1) + (upkr_current_byte & 1);
upkr_current_byte >>= 1;
--upkr_bits_left;
}
#else
while(upkr_state < 4096) {
upkr_state = (upkr_state << 8) | *upkr_data_ptr++;
}
#endif
int prob = upkr_probs[context_index];
int bit = (upkr_state & 255) < prob ? 1 : 0;
int tmp = prob;
if(!bit) {
tmp = 256 - tmp;
}
upkr_state = tmp * (upkr_state >> 8) + (upkr_state & 255);
tmp += (256 - tmp + 8) >> 4;
if(!bit) {
upkr_state -= prob;
tmp = 256 - tmp;
}
upkr_probs[context_index] = tmp;
return bit;
}
int upkr_decode_length(int context_index) {
int length = 0;
int bit_pos = 0;
while(upkr_decode_bit(context_index)) {
length |= upkr_decode_bit(context_index + 1) << bit_pos++;
context_index += 2;
}
return length | (1 << bit_pos);
}
void* upkr_unpack(void* destination, void* compressed_data) {
upkr_data_ptr = (u8*)compressed_data;
upkr_state = 0;
#ifdef UPKR_BITSTREAM
upkr_bits_left = 0;
#endif
for(int i = 0; i < sizeof(upkr_probs); ++i)
upkr_probs[i] = 128;
u8* write_ptr = (u8*)destination;
int prev_was_match = 0;
int offset = 0;
for(;;) {
if(upkr_decode_bit(0)) {
if(prev_was_match || upkr_decode_bit(256)) {
offset = upkr_decode_length(257) - 1;
if(offset == 0) {
break;
}
}
int length = upkr_decode_length(257 + 64);
while(length--) {
*write_ptr = write_ptr[-offset];
++write_ptr;
}
prev_was_match = 1;
} else {
int byte = 1;
while(byte < 256) {
int bit = upkr_decode_bit(byte);
byte = (byte << 1) + bit;
}
*write_ptr++ = byte;
prev_was_match = 0;
}
}
return write_ptr;
}

BIN
asm_unpackers/test_data.upk Normal file

Binary file not shown.

View File

@@ -0,0 +1,162 @@
// armv6-m upkr unpacker by yrlf
// some optimizations by exoticorn
.syntax unified
.thumb
.section .text
#define ALIGNUP(n, align) (((n) + (align) - 1) & ~((align) - 1))
#define PROB_LEN (1 + 255 + 1 + 2*32 + 2*32)
#define FRAME_SIZE ALIGNUP(PROB_LEN, 4)
// auto upkr_unpack(uint8_t * out, uint8_t * in) -> tuple<uint8_t *, uint8_t *>
.global upkr_unpack
.type upkr_unpack, %function
// r0 .. out_ptr (returned)
// r1 .. in_ptr (returned)
// r2 .. state
// r3 .. offset
// r4 .. prev_was_literal / decode_length ret
// r5 .. subroutine arg (preserved)
// r6 .. decode_bit ret
// r7 .. probs ptr
upkr_unpack:
push { r4, r5, r6, r7, lr }
sub sp, sp, #FRAME_SIZE
mov r7, sp
movs r2, #255
adds r2, r2, #(PROB_LEN - 255)
movs r3, #128
.Lclear:
subs r2, r2, #1
strb r3, [r7, r2]
bne .Lclear
.Lloop:
movs r5, #0
bl upkr_decode_bit
beq .Ldata
.Lmatch:
// r6 = 1
lsls r5, r6, #8
cmp r4, #0
beq 1f
bl upkr_decode_bit
beq 2f
1:
bl upkr_decode_length
adds r3, r4, #1
beq .Lend
2:
adds r5, r5, #64
bl upkr_decode_length
.Lcopy_loop:
ldrb r5, [r0, r3]
.Lstore:
strb r5, [r0]
adds r0, r0, #1
adds r4, r4, #1
blt .Lcopy_loop
b .Lloop
.Ldata:
movs r5, #1
.Ldata_loop:
bl upkr_decode_bit
adcs r5, r5, r5
lsrs r4, r5, #8
beq .Ldata_loop
b .Lstore
.Lend:
add sp, sp, #FRAME_SIZE
pop { r4, r5, r6, r7, pc }
.type upkr_decode_length, %function
// r0 .. -length tmp (saved)
// r1 ..
// r2 ..
// r3 ..
// r4 .. -length (returned)
// r5 .. context index (saved)
// r6 .. (saved)
// r7 ..
upkr_decode_length:
push { r0, r5, r6, lr }
movs r0, #0
subs r4, r0, #1
.Lbit_loop:
adds r5, r5, #1
bl upkr_decode_bit
beq 1f
adds r5, r5, #1
bl upkr_decode_bit
beq 2f
adds r0, r0, r4
2:
lsls r4, r4, #1
b .Lbit_loop
1:
adds r4, r4, r0
pop { r0, r5, r6, pc }
.type upkr_decode_bit, %function
// r0 .. tmp / prob (saved)
// r1 .. in_ptr (modified)
// r2 .. state (modified)
// r3 .. scratch (saved)
// r4 ..
// r5 .. context index (preserved)
// r6 .. bit (returned)
// r7 .. probs ptr (preserved)
upkr_fill_state:
lsls r2, r2, #8
ldrb r6, [r1]
adds r1, r1, #1
orrs r2, r2, r6
upkr_decode_bit:
lsrs r6, r2, #12
beq upkr_fill_state
push { r0, r1, r3, lr }
ldrb r0, [r7, r5]
lsrs r3, r2, #8
uxtb r1, r2
subs r6, r1, r0
blt 1f
subs r1, r2, r0
rsbs r0, r0, #0
1:
muls r3, r3, r0
adds r2, r1, r3
rsbs r3, r0, #0
uxtb r3, r3
lsrs r3, r3, #4
adcs r0, r0, r3
cmp r6, #0
blt 1f
rsbs r0, r0, #0
1:
strb r0, [r7, r5]
lsrs r6, r6, #31
pop { r0, r1, r3, pc }

View File

@@ -0,0 +1,142 @@
.section .text
#define FRAME_SIZE (256+32*4+4)
// x8 prob array ptr
// x9 prev was literal
// x10 out ptr
// x11 in ptr
// x12 offset
// x13 state
.global upkr_unpack
.type upkr_unpack, %function
upkr_unpack:
mv t4, ra
mv x17, x8
mv t6, x9
li x13, FRAME_SIZE
li x9, 128
1:
addi sp, sp, -1
sb x9, 0(sp)
addi x13, x13, -1
bnez x13, 1b
.Lmainloop:
li x14, 0
jal upkr_decode_bit
beqz x15, .Lliteral
slli x14, x14, 8
beqz x9, .Lread_offset_inc_x14
jal upkr_decode_bit
bnez x15, .Lread_offset
.Lfinished_offset:
addi x14, x14, 64
jal t3, upkr_decode_number
1:
add x14, x10, t0
lbu x14, (x14)
.Lstore_byte:
sb x14, (x10)
addi x10, x10, 1
addi x9, x9, 1
blt x9, x0, 1b
j .Lmainloop
.Lliteral:
jal upkr_decode_bit
addi x14, x14, -1
slli x14, x14, 1
add x14, x14, x15
srli x9, x14, 8
beqz x9, .Lliteral
j .Lstore_byte
.Lread_offset_inc_x14:
addi x14, x14, 1
.Lread_offset:
jal t3, upkr_decode_number
addi t0, x9, 1
bnez t0, .Lfinished_offset
.Ldone:
addi sp, sp, FRAME_SIZE
mv x8, x17
mv x9, t6
jr t4
// x14 context index
// return: x9 negtive decoded number
upkr_decode_number:
mv t5, x14
li x9, 0
li x8, -1
1:
jal upkr_decode_bit
beqz x15, 1f
jal upkr_decode_bit
beqz x15, 2f
add x9, x9, x8
2:
slli x8, x8, 1
j 1b
1:
add x9, x9, x8
mv x14, t5
jr t3
upkr_load_byte:
lbu x15, 0(x11)
addi x11, x11, 1
slli x13, x13, 8
add x13, x13, x15
// x8 prob array ptr
// x11 in ptr
// x13 state
// x14 context index
// return:
// x14 context index + 1
// x15 decoded bit
upkr_decode_bit:
srli x15, x13, 12
beqz x15, upkr_load_byte
mv t1, x14
mv t2, x10
add x14, x14, sp
lbu x12, 0(x14)
andi x10, x13, 255
sltu x15, x10, x12
srli x13, x13, 8
beqz x15, .Lelse
mul x13, x13, x12
add x13, x13, x10
li x10, 256 + 8
sub x10, x10, x12
srli x10, x10, 4
add x12, x12, x10
j .Lendif
.Lelse:
li x16, 256
sub x16, x16, x12
mul x13, x13, x16
add x13, x13, x10
sub x13, x13, x12
addi x10, x12, 8
srli x10, x10, 4
sub x12, x12, x10
.Lendif:
sb x12, 0(x14)
addi x14, t1, 1
mv x10, t2
ret

5
c_unpacker/.gitignore vendored Normal file
View File

@@ -0,0 +1,5 @@
unpack
unpack_bitstream
unpack_debug
*.upk

10
c_unpacker/Makefile Normal file
View File

@@ -0,0 +1,10 @@
all: unpack unpack_bitstream
unpack: main.c unpack.c
cc -O2 -o unpack main.c unpack.c
unpack_bitstream: main.c unpack.c
cc -O2 -D UPKR_BITSTREAM -o unpack_bitstream main.c unpack.c
unpack_debug: main.c unpack.c
cc -g -o unpack_debug main.c unpack.c

View File

@@ -0,0 +1,33 @@
int upkr_decode_bit(int context_index) {
#ifdef UPKR_BITSTREAM
while(upkr_state < 32768) {
if(upkr_bits_left == 0) {
upkr_current_byte = *upkr_data_ptr++;
upkr_bits_left = 8;
}
upkr_state = (upkr_state << 1) + (upkr_current_byte & 1);
upkr_current_byte >>= 1;
--upkr_bits_left;
}
#else
while(upkr_state < 4096) {
upkr_state = (upkr_state << 8) | *upkr_data_ptr++;
}
#endif
int prob = upkr_probs[context_index];
int bit = (upkr_state & 255) < prob ? 1 : 0;
if(bit) {
prob = 256 - prob;
}
upkr_state -= prob * ((upkr_state >> 8) + (bit ^ 1));
prob -= (prob + 8) >> 4;
if(bit) {
prob = -prob;
}
upkr_probs[context_index] = prob;
return bit;
}

26
c_unpacker/main.c Normal file
View File

@@ -0,0 +1,26 @@
#include <stdio.h>
#include <stdlib.h>
void* upkr_unpack(void* destination, void* compressed_data);
int main(int argn, char** argv) {
void* input_buffer = malloc(1024*1024);
void* output_buffer = malloc(1024*1024);
FILE* in_file = fopen(argv[1], "rb");
int in_size = fread(input_buffer, 1, 1024*1024, in_file);
fclose(in_file);
printf("Compressed size: %d\n", in_size);
void* end_ptr = upkr_unpack(output_buffer, input_buffer);
int out_size = (char*)end_ptr - (char*)output_buffer;
printf("Uncompressed size: %d\n", out_size);
FILE* out_file = fopen(argv[2], "wb");
fwrite(output_buffer, 1, out_size, out_file);
fclose(out_file);
return 0;
}

4
c_unpacker/readme.txt Normal file
View File

@@ -0,0 +1,4 @@
a very simple unpacker in c, as a reference for people wanting to implement their own unpacker.
absolutely not production ready, it makes no effort to ensure the output buffer can actually
hold the uncompressed data.
!!! Never run on untrusted input !!!

96
c_unpacker/unpack.c Normal file
View File

@@ -0,0 +1,96 @@
typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned long u32;
u8* upkr_data_ptr;
u8 upkr_probs[1 + 255 + 1 + 2*32 + 2*32];
#ifdef UPKR_BITSTREAM
u16 upkr_state;
u8 upkr_current_byte;
int upkr_bits_left;
#else
u32 upkr_state;
#endif
int upkr_decode_bit(int context_index) {
#ifdef UPKR_BITSTREAM
while(upkr_state < 32768) {
if(upkr_bits_left == 0) {
upkr_current_byte = *upkr_data_ptr++;
upkr_bits_left = 8;
}
upkr_state = (upkr_state << 1) + (upkr_current_byte & 1);
upkr_current_byte >>= 1;
--upkr_bits_left;
}
#else
while(upkr_state < 4096) {
upkr_state = (upkr_state << 8) | *upkr_data_ptr++;
}
#endif
int prob = upkr_probs[context_index];
int bit = (upkr_state & 255) < prob ? 1 : 0;
if(bit) {
upkr_state = prob * (upkr_state >> 8) + (upkr_state & 255);
prob += (256 - prob + 8) >> 4;
} else {
upkr_state = (256 - prob) * (upkr_state >> 8) + (upkr_state & 255) - prob;
prob -= (prob + 8) >> 4;
}
upkr_probs[context_index] = prob;
return bit;
}
int upkr_decode_length(int context_index) {
int length = 0;
int bit_pos = 0;
while(upkr_decode_bit(context_index)) {
length |= upkr_decode_bit(context_index + 1) << bit_pos++;
context_index += 2;
}
return length | (1 << bit_pos);
}
void* upkr_unpack(void* destination, void* compressed_data) {
upkr_data_ptr = (u8*)compressed_data;
upkr_state = 0;
#ifdef UPKR_BITSTREAM
upkr_bits_left = 0;
#endif
for(int i = 0; i < sizeof(upkr_probs); ++i)
upkr_probs[i] = 128;
u8* write_ptr = (u8*)destination;
int prev_was_match = 0;
int offset = 0;
for(;;) {
if(upkr_decode_bit(0)) {
if(prev_was_match || upkr_decode_bit(256)) {
offset = upkr_decode_length(257) - 1;
if(offset == 0) {
break;
}
}
int length = upkr_decode_length(257 + 64);
while(length--) {
*write_ptr = write_ptr[-offset];
++write_ptr;
}
prev_was_match = 1;
} else {
int byte = 1;
while(byte < 256) {
int bit = upkr_decode_bit(byte);
byte = (byte << 1) + bit;
}
*write_ptr++ = byte;
prev_was_match = 0;
}
}
return write_ptr;
}

View File

@@ -1,4 +1,7 @@
use crate::rans::{PROB_BITS, ONE_PROB}; use crate::{
rans::{ONE_PROB, PROB_BITS},
Config,
};
const INIT_PROB: u16 = 1 << (PROB_BITS - 1); const INIT_PROB: u16 = 1 << (PROB_BITS - 1);
const UPDATE_RATE: u32 = 4; const UPDATE_RATE: u32 = 4;
@@ -6,7 +9,9 @@ const UPDATE_ADD: u32 = 8;
#[derive(Clone)] #[derive(Clone)]
pub struct ContextState { pub struct ContextState {
contexts: Vec<u16>, contexts: Vec<u8>,
invert_bit_encoding: bool,
simplified_prob_update: bool,
} }
pub struct Context<'a> { pub struct Context<'a> {
@@ -15,9 +20,11 @@ pub struct Context<'a> {
} }
impl ContextState { impl ContextState {
pub fn new(size: usize) -> ContextState { pub fn new(size: usize, config: &Config) -> ContextState {
ContextState { ContextState {
contexts: vec![INIT_PROB; size], contexts: vec![INIT_PROB as u8; size],
invert_bit_encoding: config.invert_bit_encoding,
simplified_prob_update: config.simplified_prob_update,
} }
} }
@@ -28,15 +35,26 @@ impl ContextState {
impl<'a> Context<'a> { impl<'a> Context<'a> {
pub fn prob(&self) -> u16 { pub fn prob(&self) -> u16 {
self.state.contexts[self.index] self.state.contexts[self.index] as u16
} }
pub fn update(&mut self, bit: bool) { pub fn update(&mut self, bit: bool) {
let old = self.state.contexts[self.index]; let old = self.state.contexts[self.index];
self.state.contexts[self.index] = if bit {
old + ((ONE_PROB - old as u32 + UPDATE_ADD) >> UPDATE_RATE) as u16 self.state.contexts[self.index] = if self.state.simplified_prob_update {
let offset = if bit ^ self.state.invert_bit_encoding {
ONE_PROB as i32 >> UPDATE_RATE
} else {
0
};
(offset + old as i32 - ((old as i32 + UPDATE_ADD as i32) >> UPDATE_RATE)) as u8
} else { } else {
old - ((old + UPDATE_ADD as u16) >> UPDATE_RATE) if bit ^ self.state.invert_bit_encoding {
old + ((ONE_PROB - old as u32 + UPDATE_ADD) >> UPDATE_RATE) as u8
} else {
old - ((old as u32 + UPDATE_ADD) >> UPDATE_RATE) as u8
}
}; };
} }
} }

View File

@@ -1,12 +1,16 @@
use crate::lz;
use crate::match_finder::MatchFinder; use crate::match_finder::MatchFinder;
use crate::rans::RansCoder; use crate::rans::RansCoder;
use crate::ProgressCallback; use crate::ProgressCallback;
use crate::{lz, Config};
pub fn pack(data: &[u8], mut progress_callback: Option<ProgressCallback>) -> Vec<u8> { pub fn pack(
data: &[u8],
config: &Config,
mut progress_callback: Option<ProgressCallback>,
) -> Vec<u8> {
let mut match_finder = MatchFinder::new(data); let mut match_finder = MatchFinder::new(data);
let mut rans_coder = RansCoder::new(); let mut rans_coder = RansCoder::new(config);
let mut state = lz::CoderState::new(); let mut state = lz::CoderState::new(config);
let mut pos = 0; let mut pos = 0;
while pos < data.len() { while pos < data.len() {
@@ -22,7 +26,7 @@ pub fn pack(data: &[u8], mut progress_callback: Option<ProgressCallback>) -> Vec
offset: offset as u32, offset: offset as u32,
len: m.length as u32, len: m.length as u32,
} }
.encode(&mut rans_coder, &mut state); .encode(&mut rans_coder, &mut state, config);
pos += m.length; pos += m.length;
encoded_match = true; encoded_match = true;
} }
@@ -41,7 +45,7 @@ pub fn pack(data: &[u8], mut progress_callback: Option<ProgressCallback>) -> Vec
offset: offset as u32, offset: offset as u32,
len: length as u32, len: length as u32,
} }
.encode(&mut rans_coder, &mut state); .encode(&mut rans_coder, &mut state, config);
pos += length; pos += length;
encoded_match = true; encoded_match = true;
} }
@@ -49,11 +53,11 @@ pub fn pack(data: &[u8], mut progress_callback: Option<ProgressCallback>) -> Vec
} }
if !encoded_match { if !encoded_match {
lz::Op::Literal(data[pos]).encode(&mut rans_coder, &mut state); lz::Op::Literal(data[pos]).encode(&mut rans_coder, &mut state, config);
pos += 1; pos += 1;
} }
} }
lz::encode_eof(&mut rans_coder, &mut state); lz::encode_eof(&mut rans_coder, &mut state, config);
rans_coder.finish() rans_coder.finish()
} }

View File

@@ -2,17 +2,65 @@ mod context_state;
mod greedy_packer; mod greedy_packer;
mod lz; mod lz;
mod match_finder; mod match_finder;
mod rans;
mod parsing_packer; mod parsing_packer;
mod rans;
pub use lz::unpack; pub use lz::unpack;
pub type ProgressCallback<'a> = &'a mut dyn FnMut(usize); pub type ProgressCallback<'a> = &'a mut dyn FnMut(usize);
pub fn pack(data: &[u8], level: u8, progress_callback: Option<ProgressCallback>) -> Vec<u8> { pub struct Config {
if level == 0 { pub use_bitstream: bool,
greedy_packer::pack(data, progress_callback) pub parity_contexts: usize,
} else {
parsing_packer::pack(data, level, progress_callback) pub invert_bit_encoding: bool,
pub is_match_bit: bool,
pub new_offset_bit: bool,
pub continue_value_bit: bool,
pub bitstream_is_big_endian: bool,
pub simplified_prob_update: bool,
pub no_repeated_offsets: bool,
}
impl Default for Config {
fn default() -> Config {
Config {
use_bitstream: false,
parity_contexts: 1,
invert_bit_encoding: false,
is_match_bit: true,
new_offset_bit: true,
continue_value_bit: true,
bitstream_is_big_endian: false,
simplified_prob_update: false,
no_repeated_offsets: false,
}
} }
} }
pub fn pack(
data: &[u8],
level: u8,
config: Config,
progress_callback: Option<ProgressCallback>,
) -> Vec<u8> {
if level == 0 {
greedy_packer::pack(data, &config, progress_callback)
} else {
parsing_packer::pack(data, level, &config, progress_callback)
}
}
pub fn compressed_size(mut data: &[u8]) -> f32 {
let mut state = 0;
while state < 4096 {
state = (state << 8) | data[0] as u32;
data = &data[1..];
}
data.len() as f32 + (state as f32).log2() / 8.
}

112
src/lz.rs
View File

@@ -1,5 +1,6 @@
use crate::context_state::ContextState; use crate::context_state::ContextState;
use crate::rans::{EntropyCoder, RansDecoder}; use crate::rans::{EntropyCoder, RansDecoder};
use crate::Config;
#[derive(Copy, Clone, Debug)] #[derive(Copy, Clone, Debug)]
pub enum Op { pub enum Op {
@@ -8,40 +9,66 @@ pub enum Op {
} }
impl Op { impl Op {
pub fn encode(&self, coder: &mut dyn EntropyCoder, state: &mut CoderState) { pub fn encode(&self, coder: &mut dyn EntropyCoder, state: &mut CoderState, config: &Config) {
let literal_base = state.pos % state.parity_contexts * 256;
match self { match self {
&Op::Literal(lit) => { &Op::Literal(lit) => {
encode_bit(coder, state, 0, false); encode_bit(coder, state, literal_base, !config.is_match_bit);
let mut context_index = 1; let mut context_index = 1;
for i in (0..8).rev() { for i in (0..8).rev() {
let bit = (lit >> i) & 1 != 0; let bit = (lit >> i) & 1 != 0;
encode_bit(coder, state, context_index, bit); encode_bit(coder, state, literal_base + context_index, bit);
context_index = (context_index << 1) | bit as usize; context_index = (context_index << 1) | bit as usize;
} }
state.prev_was_match = false; state.prev_was_match = false;
state.pos += 1;
} }
&Op::Match { offset, len } => { &Op::Match { offset, len } => {
encode_bit(coder, state, 0, true); encode_bit(coder, state, literal_base, config.is_match_bit);
if !state.prev_was_match { if !state.prev_was_match && !config.no_repeated_offsets {
encode_bit(coder, state, 256, offset != state.last_offset); encode_bit(
coder,
state,
256 * state.parity_contexts,
(offset != state.last_offset) == config.new_offset_bit,
);
} else { } else {
assert!(offset != state.last_offset); assert!(offset != state.last_offset || config.no_repeated_offsets);
} }
if offset != state.last_offset { if offset != state.last_offset || config.no_repeated_offsets {
encode_length(coder, state, 257, offset + 1); encode_length(
coder,
state,
256 * state.parity_contexts + 1,
offset + 1,
config,
);
state.last_offset = offset; state.last_offset = offset;
} }
encode_length(coder, state, 257 + 64, len); encode_length(coder, state, 256 * state.parity_contexts + 65, len, config);
state.prev_was_match = true; state.prev_was_match = true;
state.pos += len as usize;
} }
} }
} }
} }
pub fn encode_eof(coder: &mut dyn EntropyCoder, state: &mut CoderState) { pub fn encode_eof(coder: &mut dyn EntropyCoder, state: &mut CoderState, config: &Config) {
encode_bit(coder, state, 0, true); encode_bit(
encode_bit(coder, state, 256, true); coder,
encode_length(coder, state, 257, 1); state,
state.pos % state.parity_contexts * 256,
config.is_match_bit,
);
if !state.prev_was_match {
encode_bit(
coder,
state,
256 * state.parity_contexts,
config.new_offset_bit,
);
}
encode_length(coder, state, 256 * state.parity_contexts + 1, 1, config);
} }
fn encode_bit( fn encode_bit(
@@ -58,32 +85,37 @@ fn encode_length(
state: &mut CoderState, state: &mut CoderState,
context_start: usize, context_start: usize,
mut value: u32, mut value: u32,
config: &Config,
) { ) {
assert!(value >= 1); assert!(value >= 1);
let mut context_index = context_start; let mut context_index = context_start;
while value >= 2 { while value >= 2 {
encode_bit(coder, state, context_index, true); encode_bit(coder, state, context_index, config.continue_value_bit);
encode_bit(coder, state, context_index + 1, value & 1 != 0); encode_bit(coder, state, context_index + 1, value & 1 != 0);
context_index += 2; context_index += 2;
value >>= 1; value >>= 1;
} }
encode_bit(coder, state, context_index, false); encode_bit(coder, state, context_index, !config.continue_value_bit);
} }
#[derive(Clone)] #[derive(Clone)]
pub struct CoderState { pub struct CoderState {
contexts: ContextState, contexts: ContextState,
last_offset: u32, last_offset: u32,
prev_was_match: bool prev_was_match: bool,
pos: usize,
parity_contexts: usize,
} }
impl CoderState { impl CoderState {
pub fn new() -> CoderState { pub fn new(config: &Config) -> CoderState {
CoderState { CoderState {
contexts: ContextState::new(1 + 255 + 1 + 64 + 64), contexts: ContextState::new((1 + 255) * config.parity_contexts + 1 + 64 + 64, config),
last_offset: 0, last_offset: 0,
prev_was_match: false prev_was_match: false,
pos: 0,
parity_contexts: config.parity_contexts,
} }
} }
@@ -92,9 +124,9 @@ impl CoderState {
} }
} }
pub fn unpack(packed_data: &[u8]) -> Vec<u8> { pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
let mut decoder = RansDecoder::new(packed_data); let mut decoder = RansDecoder::new(packed_data, &config);
let mut contexts = ContextState::new(1 + 255 + 1 + 64 + 64); let mut contexts = ContextState::new((1 + 255) * config.parity_contexts + 1 + 64 + 64, &config);
let mut result = vec![]; let mut result = vec![];
let mut offset = 0; let mut offset = 0;
let mut prev_was_match = false; let mut prev_was_match = false;
@@ -103,10 +135,13 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
decoder: &mut RansDecoder, decoder: &mut RansDecoder,
contexts: &mut ContextState, contexts: &mut ContextState,
mut context_index: usize, mut context_index: usize,
config: &Config,
) -> usize { ) -> usize {
let mut length = 0; let mut length = 0;
let mut bit_pos = 0; let mut bit_pos = 0;
while decoder.decode_with_context(&mut contexts.context_mut(context_index)) { while decoder.decode_with_context(&mut contexts.context_mut(context_index))
== config.continue_value_bit
{
length |= (decoder.decode_with_context(&mut contexts.context_mut(context_index + 1)) length |= (decoder.decode_with_context(&mut contexts.context_mut(context_index + 1))
as usize) as usize)
<< bit_pos; << bit_pos;
@@ -117,14 +152,32 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
} }
loop { loop {
if decoder.decode_with_context(&mut contexts.context_mut(0)) { let literal_base = result.len() % config.parity_contexts * 256;
if prev_was_match || decoder.decode_with_context(&mut contexts.context_mut(256)) { if decoder.decode_with_context(&mut contexts.context_mut(literal_base))
offset = decode_length(&mut decoder, &mut contexts, 257) - 1; == config.is_match_bit
{
if config.no_repeated_offsets
|| prev_was_match
|| decoder
.decode_with_context(&mut contexts.context_mut(256 * config.parity_contexts))
== config.new_offset_bit
{
offset = decode_length(
&mut decoder,
&mut contexts,
256 * config.parity_contexts + 1,
&config,
) - 1;
if offset == 0 { if offset == 0 {
break; break;
} }
} }
let length = decode_length(&mut decoder, &mut contexts, 257 + 64); let length = decode_length(
&mut decoder,
&mut contexts,
256 * config.parity_contexts + 65,
&config,
);
for _ in 0..length { for _ in 0..length {
result.push(result[result.len() - offset]); result.push(result[result.len() - offset]);
} }
@@ -133,7 +186,8 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
let mut context_index = 1; let mut context_index = 1;
let mut byte = 0; let mut byte = 0;
for i in (0..8).rev() { for i in (0..8).rev() {
let bit = decoder.decode_with_context(&mut contexts.context_mut(context_index)); let bit = decoder
.decode_with_context(&mut contexts.context_mut(literal_base + context_index));
context_index = (context_index << 1) | bit as usize; context_index = (context_index << 1) | bit as usize;
byte |= (bit as u8) << i; byte |= (bit as u8) << i;
} }

View File

@@ -1,60 +1,151 @@
use anyhow::{bail, Result}; use anyhow::Result;
use std::ffi::OsStr;
use std::io::prelude::*; use std::io::prelude::*;
use std::process;
use std::{fs::File, path::PathBuf}; use std::{fs::File, path::PathBuf};
fn main() -> Result<()> { fn main() -> Result<()> {
let mut args = pico_args::Arguments::from_env(); let mut config = upkr::Config::default();
let mut reverse = false;
let mut unpack = false;
let mut level = 2;
let mut infile: Option<PathBuf> = None;
let mut outfile: Option<PathBuf> = None;
match args.subcommand()?.as_ref().map(|s| s.as_str()) { let mut parser = lexopt::Parser::from_env();
None => print_help(), while let Some(arg) = parser.next()? {
Some("pack") => { use lexopt::prelude::*;
let level = args.opt_value_from_str(["-l", "--level"])?.unwrap_or(2u8); match arg {
Short('b') | Long("bitstream") => config.use_bitstream = true,
Short('p') | Long("parity") => config.parity_contexts = parser.value()?.parse()?,
Short('r') | Long("reverse") => reverse = true,
Long("invert-is-match-bit") => config.is_match_bit = false,
Long("invert-new-offset-bit") => config.new_offset_bit = false,
Long("invert-continue-value-bit") => config.continue_value_bit = false,
Long("invert-bit-encoding") => config.invert_bit_encoding = true,
Long("simplified-prob-update") => config.simplified_prob_update = true,
Long("big-endian-bitstream") => {
config.use_bitstream = true;
config.bitstream_is_big_endian = true;
}
Long("no-repeated-offsets") => config.no_repeated_offsets = true,
let infile = args.free_from_os_str::<PathBuf, bool>(|s| Ok(s.into()))?; Long("z80") => {
let outfile = args.free_from_os_str::<PathBuf, bool>(|s| Ok(s.into()))?; config.use_bitstream = true;
config.bitstream_is_big_endian = true;
config.invert_bit_encoding = true;
config.simplified_prob_update = true;
}
Long("x86") => {
config.use_bitstream = true;
config.continue_value_bit = false;
config.is_match_bit = false;
}
let mut data = vec![]; Short('u') | Long("unpack") => unpack = true,
File::open(infile)?.read_to_end(&mut data)?; Short('l') | Long("level") => level = parser.value()?.parse()?,
Short('h') | Long("help") => print_help(0),
let mut pb = pbr::ProgressBar::new(data.len() as u64); Value(val) if infile.is_none() => infile = Some(val.try_into()?),
pb.set_units(pbr::Units::Bytes); Value(val) if outfile.is_none() => outfile = Some(val.try_into()?),
let packed_data = upkr::pack( _ => return Err(arg.unexpected().into()),
&data,
level,
Some(&mut |pos| {
pb.set(pos as u64);
}),
);
pb.finish();
println!(
"Compressed {} bytes to {} bytes ({}%)",
data.len(),
packed_data.len(),
packed_data.len() as f32 * 100. / data.len() as f32
);
File::create(outfile)?.write_all(&packed_data)?;
} }
Some("unpack") => { }
let infile = args.free_from_os_str::<PathBuf, bool>(|s| Ok(s.into()))?;
let outfile = args.free_from_os_str::<PathBuf, bool>(|s| Ok(s.into()))?;
let mut data = vec![]; let infile = infile.unwrap_or_else(|| print_help(1));
File::open(infile)?.read_to_end(&mut data)?; let outfile = outfile.unwrap_or_else(|| {
let packed_data = upkr::unpack(&data); let mut name = infile.clone();
File::create(outfile)?.write_all(&packed_data)?; if unpack {
if name.extension().filter(|&e| e == "upk").is_some() {
name.set_extension("");
} else {
name.set_extension("bin");
}
} else {
let mut filename = name
.file_name()
.unwrap_or_else(|| OsStr::new(""))
.to_os_string();
filename.push(".upk");
name.set_file_name(filename);
} }
Some(other) => { name
bail!("Unknown subcommand '{}'", other); });
if config.parity_contexts != 1 && config.parity_contexts != 2 && config.parity_contexts != 4 {
eprintln!("--parity has to be 1, 2, or 4");
process::exit(1);
}
if !unpack {
let mut data = vec![];
File::open(infile)?.read_to_end(&mut data)?;
if reverse {
data.reverse();
} }
let mut pb = pbr::ProgressBar::new(data.len() as u64);
pb.set_units(pbr::Units::Bytes);
let mut packed_data = upkr::pack(
&data,
level,
config,
Some(&mut |pos| {
pb.set(pos as u64);
}),
);
pb.finish();
if reverse {
packed_data.reverse();
}
println!(
"Compressed {} bytes to {} bytes ({}%)",
data.len(),
packed_data.len(),
packed_data.len() as f32 * 100. / data.len() as f32
);
File::create(outfile)?.write_all(&packed_data)?;
} else {
let mut data = vec![];
File::open(infile)?.read_to_end(&mut data)?;
if reverse {
data.reverse();
}
let mut unpacked_data = upkr::unpack(&data, config);
if reverse {
unpacked_data.reverse();
}
File::create(outfile)?.write_all(&unpacked_data)?;
} }
Ok(()) Ok(())
} }
fn print_help() { fn print_help(exit_code: i32) -> ! {
eprintln!("Usage:"); eprintln!("Usage:");
eprintln!(" upkr pack [-l level(0-9)] <infile> <outfile>"); eprintln!(" upkr [-l level(0-9)] [config options] <infile> [<outfile>]");
eprintln!(" upkr unpack <infile> <outfile>"); eprintln!(" upkr -u [config options] <infile> [<outfile>]");
std::process::exit(1); eprintln!();
eprintln!(" -l, --level N compression level 0-9");
eprintln!(" -u, --unpack unpack infile");
eprintln!();
eprintln!("Config presets for specific unpackers:");
eprintln!(" --z80 --big-endian-bitstream --invert-bit-encoding --simplified-prob-update");
eprintln!(" --x86 --bitstream --invert-is-match-bit --invert-continue-value-bit");
eprintln!();
eprintln!("Config options (need to match when packing/unpacking):");
eprintln!(" -b, --bitstream bitstream mode");
eprintln!(" -p, --parity N use N (2/4) parity contexts");
eprintln!(" -r, --reverse reverse input & output");
eprintln!();
eprintln!("Config options to tailor output to specific optimized unpackers:");
eprintln!(" --invert-is-match-bit");
eprintln!(" --invert-new-offset-bit");
eprintln!(" --invert-continue-value-bit");
eprintln!(" --invert-bit-encoding");
eprintln!(" --simplified-prob-update");
eprintln!(" --big-endian-bitstream (implies --bitstream)");
eprintln!(" --no-repeated-offsets");
process::exit(exit_code);
} }

View File

@@ -6,19 +6,24 @@ use crate::match_finder::MatchFinder;
use crate::rans::{CostCounter, RansCoder}; use crate::rans::{CostCounter, RansCoder};
use crate::{lz, ProgressCallback}; use crate::{lz, ProgressCallback};
pub fn pack(data: &[u8], level: u8, progress_cb: Option<ProgressCallback>) -> Vec<u8> { pub fn pack(
let mut parse = parse(data, Config::from_level(level), progress_cb); data: &[u8],
level: u8,
config: &crate::Config,
progress_cb: Option<ProgressCallback>,
) -> Vec<u8> {
let mut parse = parse(data, Config::from_level(level), config, progress_cb);
let mut ops = vec![]; let mut ops = vec![];
while let Some(link) = parse { while let Some(link) = parse {
ops.push(link.op); ops.push(link.op);
parse = link.prev.clone(); parse = link.prev.clone();
} }
let mut state = lz::CoderState::new(); let mut state = lz::CoderState::new(config);
let mut coder = RansCoder::new(); let mut coder = RansCoder::new(config);
for op in ops.into_iter().rev() { for op in ops.into_iter().rev() {
op.encode(&mut coder, &mut state); op.encode(&mut coder, &mut state, config);
} }
lz::encode_eof(&mut coder, &mut state); lz::encode_eof(&mut coder, &mut state, config);
coder.finish() coder.finish()
} }
@@ -38,6 +43,7 @@ type Arrivals = HashMap<usize, Vec<Arrival>>;
fn parse( fn parse(
data: &[u8], data: &[u8],
config: Config, config: Config,
encoding_config: &crate::Config,
mut progress_cb: Option<ProgressCallback>, mut progress_cb: Option<ProgressCallback>,
) -> Option<Rc<Parse>> { ) -> Option<Rc<Parse>> {
let mut match_finder = MatchFinder::new(data) let mut match_finder = MatchFinder::new(data)
@@ -102,6 +108,7 @@ fn parse(
length: usize, length: usize,
arrival: &Arrival, arrival: &Arrival,
max_arrivals: usize, max_arrivals: usize,
config: &crate::Config,
) { ) {
cost_counter.reset(); cost_counter.reset();
let mut state = arrival.state.clone(); let mut state = arrival.state.clone();
@@ -109,7 +116,7 @@ fn parse(
offset: offset as u32, offset: offset as u32,
len: length as u32, len: length as u32,
}; };
op.encode(cost_counter, &mut state); op.encode(cost_counter, &mut state, config);
add_arrival( add_arrival(
arrivals, arrivals,
pos + length, pos + length,
@@ -129,13 +136,13 @@ fn parse(
0, 0,
Arrival { Arrival {
parse: None, parse: None,
state: lz::CoderState::new(), state: lz::CoderState::new(encoding_config),
cost: 0.0, cost: 0.0,
}, },
max_arrivals, max_arrivals,
); );
let cost_counter = &mut CostCounter::new(); let cost_counter = &mut CostCounter::new(encoding_config);
let mut best_per_offset = HashMap::new(); let mut best_per_offset = HashMap::new();
for pos in 0..data.len() { for pos in 0..data.len() {
let match_length = |offset: usize| { let match_length = |offset: usize| {
@@ -185,6 +192,7 @@ fn parse(
m.length, m.length,
&arrival, &arrival,
max_arrivals, max_arrivals,
encoding_config,
); );
if m.length >= config.greedy_size { if m.length >= config.greedy_size {
break 'arrival_loop; break 'arrival_loop;
@@ -208,6 +216,7 @@ fn parse(
length, length,
&arrival, &arrival,
max_arrivals, max_arrivals,
encoding_config,
); );
found_last_offset |= offset as u32 == arrival.state.last_offset(); found_last_offset |= offset as u32 == arrival.state.last_offset();
if offset < near_matches.len() { if offset < near_matches.len() {
@@ -228,6 +237,7 @@ fn parse(
length, length,
&arrival, &arrival,
max_arrivals, max_arrivals,
encoding_config,
); );
} }
} }
@@ -235,7 +245,7 @@ fn parse(
cost_counter.reset(); cost_counter.reset();
let mut state = arrival.state; let mut state = arrival.state;
let op = lz::Op::Literal(data[pos]); let op = lz::Op::Literal(data[pos]);
op.encode(cost_counter, &mut state); op.encode(cost_counter, &mut state, encoding_config);
add_arrival( add_arrival(
&mut arrivals, &mut arrivals,
pos + 1, pos + 1,

View File

@@ -1,6 +1,5 @@
use crate::context_state::Context; use crate::{context_state::Context, Config};
const L_BITS: u32 = 12;
pub const PROB_BITS: u32 = 8; pub const PROB_BITS: u32 = 8;
pub const ONE_PROB: u32 = 1 << PROB_BITS; pub const ONE_PROB: u32 = 1 << PROB_BITS;
@@ -13,43 +12,93 @@ pub trait EntropyCoder {
} }
} }
pub struct RansCoder(Vec<u16>); pub struct RansCoder {
bits: Vec<u16>,
use_bitstream: bool,
bitstream_is_big_endian: bool,
invert_bit_encoding: bool,
}
impl EntropyCoder for RansCoder { impl EntropyCoder for RansCoder {
fn encode_bit(&mut self, bit: bool, prob: u16) { fn encode_bit(&mut self, bit: bool, prob: u16) {
assert!(prob < 32768); assert!(prob < 32768);
self.0.push(prob | ((bit as u16) << 15)); self.bits
.push(prob | (((bit ^ self.invert_bit_encoding) as u16) << 15));
} }
} }
impl RansCoder { impl RansCoder {
pub fn new() -> RansCoder { pub fn new(config: &Config) -> RansCoder {
RansCoder(Vec::new()) RansCoder {
bits: Vec::new(),
use_bitstream: config.use_bitstream,
bitstream_is_big_endian: config.bitstream_is_big_endian,
invert_bit_encoding: config.invert_bit_encoding,
}
} }
pub fn finish(self) -> Vec<u8> { pub fn finish(self) -> Vec<u8> {
let mut buffer = vec![]; let mut buffer = vec![];
let mut state = 1 << L_BITS; let l_bits: u32 = if self.use_bitstream { 15 } else { 12 };
let mut state = 1 << l_bits;
const MAX_STATE_FACTOR: u32 = 1 << (L_BITS + 8 - PROB_BITS); let mut byte = 0u8;
for step in self.0.into_iter().rev() { let mut bit = if self.bitstream_is_big_endian { 0 } else { 8 };
let mut flush_state: Box<dyn FnMut(&mut u32)> = if self.use_bitstream {
if self.bitstream_is_big_endian {
Box::new(|state: &mut u32| {
byte |= ((*state & 1) as u8) << bit;
bit += 1;
if bit == 8 {
buffer.push(byte);
byte = 0;
bit = 0;
}
*state >>= 1;
})
} else {
Box::new(|state: &mut u32| {
bit -= 1;
byte |= ((*state & 1) as u8) << bit;
if bit == 0 {
buffer.push(byte);
byte = 0;
bit = 8;
}
*state >>= 1;
})
}
} else {
Box::new(|state: &mut u32| {
buffer.push(*state as u8);
*state >>= 8;
})
};
let num_flush_bits = if self.use_bitstream { 1 } else { 8 };
let max_state_factor: u32 = 1 << (l_bits + num_flush_bits - PROB_BITS);
for step in self.bits.into_iter().rev() {
let prob = step as u32 & 32767; let prob = step as u32 & 32767;
let (start, prob) = if step & 32768 != 0 { let (start, prob) = if step & 32768 != 0 {
(0, prob) (0, prob)
} else { } else {
(prob, ONE_PROB - prob) (prob, ONE_PROB - prob)
}; };
let max_state = MAX_STATE_FACTOR * prob; let max_state = max_state_factor * prob;
while state >= max_state { while state >= max_state {
buffer.push(state as u8); flush_state(&mut state);
state >>= 8;
} }
state = ((state / prob) << PROB_BITS) + (state % prob) + start; state = ((state / prob) << PROB_BITS) + (state % prob) + start;
} }
while state > 0 { while state > 0 {
buffer.push(state as u8); flush_state(&mut state);
state >>= 8; }
drop(flush_state);
if self.use_bitstream && byte != 0 {
buffer.push(byte);
} }
buffer.reverse(); buffer.reverse();
@@ -60,10 +109,11 @@ impl RansCoder {
pub struct CostCounter { pub struct CostCounter {
cost: f64, cost: f64,
log2_table: Vec<f64>, log2_table: Vec<f64>,
invert_bit_encoding: bool,
} }
impl CostCounter { impl CostCounter {
pub fn new() -> CostCounter { pub fn new(config: &Config) -> CostCounter {
let log2_table = (0..ONE_PROB) let log2_table = (0..ONE_PROB)
.map(|prob| { .map(|prob| {
let inv_prob = ONE_PROB as f64 / prob as f64; let inv_prob = ONE_PROB as f64 / prob as f64;
@@ -73,6 +123,7 @@ impl CostCounter {
CostCounter { CostCounter {
cost: 0.0, cost: 0.0,
log2_table, log2_table,
invert_bit_encoding: config.invert_bit_encoding,
} }
} }
@@ -87,7 +138,7 @@ impl CostCounter {
impl EntropyCoder for CostCounter { impl EntropyCoder for CostCounter {
fn encode_bit(&mut self, bit: bool, prob: u16) { fn encode_bit(&mut self, bit: bool, prob: u16) {
let prob = if bit { let prob = if bit ^ self.invert_bit_encoding {
prob as u32 prob as u32
} else { } else {
ONE_PROB - prob as u32 ONE_PROB - prob as u32
@@ -99,14 +150,26 @@ impl EntropyCoder for CostCounter {
pub struct RansDecoder<'a> { pub struct RansDecoder<'a> {
data: &'a [u8], data: &'a [u8],
state: u32, state: u32,
use_bitstream: bool,
byte: u8,
bits_left: u8,
invert_bit_encoding: bool,
bitstream_is_big_endian: bool,
} }
const PROB_MASK: u32 = ONE_PROB - 1; const PROB_MASK: u32 = ONE_PROB - 1;
const L: u32 = 1 << L_BITS;
impl<'a> RansDecoder<'a> { impl<'a> RansDecoder<'a> {
pub fn new(data: &'a [u8]) -> RansDecoder<'a> { pub fn new(data: &'a [u8], config: &Config) -> RansDecoder<'a> {
RansDecoder { data, state: 0 } RansDecoder {
data,
state: 0,
use_bitstream: config.use_bitstream,
byte: 0,
bits_left: 0,
invert_bit_encoding: config.invert_bit_encoding,
bitstream_is_big_endian: config.bitstream_is_big_endian,
}
} }
pub fn decode_with_context(&mut self, context: &mut Context) -> bool { pub fn decode_with_context(&mut self, context: &mut Context) -> bool {
@@ -117,9 +180,27 @@ impl<'a> RansDecoder<'a> {
pub fn decode_bit(&mut self, prob: u16) -> bool { pub fn decode_bit(&mut self, prob: u16) -> bool {
let prob = prob as u32; let prob = prob as u32;
while self.state < L { if self.use_bitstream {
self.state = (self.state << 8) | self.data[0] as u32; while self.state < 32768 {
self.data = &self.data[1..]; if self.bits_left == 0 {
self.byte = self.data[0];
self.data = &self.data[1..];
self.bits_left = 8;
}
if self.bitstream_is_big_endian {
self.state = (self.state << 1) | (self.byte >> 7) as u32;
self.byte <<= 1;
} else {
self.state = (self.state << 1) | (self.byte & 1) as u32;
self.byte >>= 1;
}
self.bits_left -= 1;
}
} else {
while self.state < 4096 {
self.state = (self.state << 8) | self.data[0] as u32;
self.data = &self.data[1..];
}
} }
let bit = (self.state & PROB_MASK) < prob; let bit = (self.state & PROB_MASK) < prob;
@@ -131,6 +212,6 @@ impl<'a> RansDecoder<'a> {
}; };
self.state = prob * (self.state >> PROB_BITS) + (self.state & PROB_MASK) - start; self.state = prob * (self.state >> PROB_BITS) + (self.state & PROB_MASK) - start;
bit bit ^ self.invert_bit_encoding
} }
} }