| 1 | use crate::io::bitreader::{BitReader, BitReaderError, BitReaderResult}; |
| 2 | |
| 3 | #[derive(Debug)] |
| 4 | pub enum UintCodeType { |
| 5 | UnaryOnes, |
| 6 | UnaryZeroes, |
| 7 | Unary012, |
| 8 | Unary210, |
| 9 | LimitedUnary(u32, u32), |
| 10 | LimitedZeroes(u32), |
| 11 | LimitedOnes(u32), |
| 12 | Golomb(u8), |
| 13 | Rice(u8), |
| 14 | Gamma, |
| 15 | GammaP, |
| 16 | } |
| 17 | |
| 18 | pub enum IntCodeType { |
| 19 | Golomb(u8), |
| 20 | Rice(u8), |
| 21 | Gamma, |
| 22 | GammaP, |
| 23 | } |
| 24 | |
| 25 | pub trait IntCodeReader { |
| 26 | fn read_code(&mut self, t: UintCodeType) -> BitReaderResult<u32>; |
| 27 | fn read_code_signed(&mut self, t: IntCodeType) -> BitReaderResult<i32>; |
| 28 | } |
| 29 | |
| 30 | fn read_unary(br: &mut BitReader, terminator: u32) -> BitReaderResult<u32> { |
| 31 | let mut res: u32 = 0; |
| 32 | loop { |
| 33 | if br.read(1)? == terminator { return Ok(res); } |
| 34 | res += 1; |
| 35 | } |
| 36 | } |
| 37 | |
| 38 | fn read_unary_lim(br: &mut BitReader, len: u32, terminator: u32) -> BitReaderResult<u32> { |
| 39 | let mut res: u32 = 0; |
| 40 | loop { |
| 41 | if br.read(1)? == terminator { return Ok(res); } |
| 42 | res += 1; |
| 43 | if res == len { return Ok(res); } |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | fn read_unary210(br: &mut BitReader) -> BitReaderResult<u32> { |
| 48 | let val = read_unary_lim(br, 2, 0)?; |
| 49 | Ok(2 - val) |
| 50 | } |
| 51 | |
| 52 | fn read_golomb(br: &mut BitReader, m: u8) -> BitReaderResult<u32> { |
| 53 | if m == 0 { return Err(BitReaderError::InvalidValue); } |
| 54 | let nbits = (8 - m.leading_zeros()) as u8; |
| 55 | if (m & (m - 1)) == 0 { return read_rice(br, nbits); } |
| 56 | let cutoff = u32::from((1 << nbits) - m); |
| 57 | let pfx = read_unary(br, 0)?; |
| 58 | let tail = br.read(nbits - 1)?; |
| 59 | if tail < cutoff { |
| 60 | let res = pfx * u32::from(m) + tail; |
| 61 | Ok (res) |
| 62 | } else { |
| 63 | let add = br.read(1)?; |
| 64 | let res = pfx * u32::from(m) + (tail - cutoff) * 2 + add + cutoff; |
| 65 | Ok (res) |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | fn read_rice(br: &mut BitReader, k: u8) -> BitReaderResult<u32> { |
| 70 | let pfx = read_unary(br, 1)?; |
| 71 | let ret = (pfx << k) + br.read(k)?; |
| 72 | Ok(ret) |
| 73 | } |
| 74 | |
| 75 | fn read_gamma(br: &mut BitReader) -> BitReaderResult<u32> { |
| 76 | let mut ret = 1; |
| 77 | while br.read(1)? != 1 { |
| 78 | ret = (ret << 1) | br.read(1)?; |
| 79 | } |
| 80 | Ok(ret - 1) |
| 81 | } |
| 82 | |
| 83 | fn read_gammap(br: &mut BitReader) -> BitReaderResult<u32> { |
| 84 | let pfx = read_unary(br, 1)?; |
| 85 | if pfx > 32 { return Err(BitReaderError::InvalidValue); } |
| 86 | let ret = (1 << pfx) + br.read(pfx as u8)?; |
| 87 | Ok(ret) |
| 88 | } |
| 89 | |
| 90 | fn uval_to_sval0mp(uval: u32) -> i32 { |
| 91 | if (uval & 1) != 0 { -((uval >> 1) as i32) } |
| 92 | else { (uval >> 1) as i32 } |
| 93 | } |
| 94 | |
| 95 | fn uval_to_sval0pm(uval: u32) -> i32 { |
| 96 | if (uval & 1) != 0 { ((uval + 1) >> 1) as i32 } |
| 97 | else { -((uval >> 1) as i32) } |
| 98 | } |
| 99 | |
| 100 | impl<'a> IntCodeReader for BitReader<'a> { |
| 101 | #[inline(always)] |
| 102 | fn read_code(&mut self, t: UintCodeType) -> BitReaderResult<u32> { |
| 103 | match t { |
| 104 | UintCodeType::UnaryOnes => read_unary(self, 0), |
| 105 | UintCodeType::UnaryZeroes => read_unary(self, 1), |
| 106 | UintCodeType::LimitedZeroes(len) => read_unary_lim(self, len, 1), |
| 107 | UintCodeType::LimitedOnes(len) => read_unary_lim(self, len, 0), |
| 108 | UintCodeType::LimitedUnary(len, term) => read_unary_lim(self, len, term), |
| 109 | UintCodeType::Unary012 => read_unary_lim(self, 2, 0), |
| 110 | UintCodeType::Unary210 => read_unary210(self), |
| 111 | UintCodeType::Golomb(m) => read_golomb(self, m), |
| 112 | UintCodeType::Rice(k) => read_rice(self, k), |
| 113 | UintCodeType::Gamma => read_gamma(self), |
| 114 | UintCodeType::GammaP => read_gammap(self), |
| 115 | } |
| 116 | } |
| 117 | #[allow(unused_variables)] |
| 118 | fn read_code_signed(&mut self, t: IntCodeType) -> BitReaderResult<i32> { |
| 119 | let uval = |
| 120 | match t { |
| 121 | IntCodeType::Golomb(m) => read_golomb(self, m)?, |
| 122 | IntCodeType::Rice(k) => read_rice(self, k)?, |
| 123 | IntCodeType::Gamma => read_gamma(self)?, |
| 124 | IntCodeType::GammaP => read_gammap(self)?, |
| 125 | }; |
| 126 | match t { |
| 127 | IntCodeType::Golomb(m) => Ok(uval_to_sval0mp(uval)), |
| 128 | IntCodeType::Rice(k) => Ok(uval_to_sval0mp(uval)), |
| 129 | IntCodeType::Gamma => Ok(uval_to_sval0pm(uval)), |
| 130 | IntCodeType::GammaP => Ok(uval_to_sval0pm(uval)), |
| 131 | } |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | #[cfg(test)] |
| 136 | mod test { |
| 137 | use super::*; |
| 138 | use crate::io::bitreader::*; |
| 139 | |
| 140 | #[test] |
| 141 | fn int_codes() { |
| 142 | const GDATA: [u8; 6] = [0b000_001_01, 0b0_0110_011, 0b1_1000_100, 0b1_1010_101, 0b10_10111_1, 0b1000_0000]; |
| 143 | let src = &GDATA; |
| 144 | let mut br = BitReader::new(src, BitReaderMode::BE); |
| 145 | for i in 0..11 { |
| 146 | assert_eq!(br.read_code(UintCodeType::Golomb(5)).unwrap(), i); |
| 147 | } |
| 148 | } |
| 149 | } |