ec.rs 25.5 KB
Newer Older
Guillaume Martres's avatar
Guillaume Martres committed
1
// Copyright (c) 2001-2016, Alliance for Open Media. All rights reserved
2
// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
Guillaume Martres's avatar
Guillaume Martres committed
3 4 5 6 7 8 9
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
Thomas Daede's avatar
Thomas Daede committed
10

Guillaume Martres's avatar
Guillaume Martres committed
11
#![allow(non_camel_case_types)]
Michael Bebenita's avatar
Michael Bebenita committed
12 13 14
#![cfg_attr(feature = "cargo-clippy", allow(cast_lossless))]
#![cfg_attr(feature = "cargo-clippy", allow(identity_op))]
#![cfg_attr(feature = "cargo-clippy", allow(needless_range_loop))]
Thomas Daede's avatar
Thomas Daede committed
15

16
use bitstream_io::{BitWriter, BigEndian};
17
use std;
18
use std::io;
19
use util::ILog;
20

21
pub const OD_BITRES: u8 = 3;
22 23
const EC_PROB_SHIFT: u32 = 6;
const EC_MIN_PROB: u32 = 4;
Monty Montgomery's avatar
Monty Montgomery committed
24 25
type ec_window = u32;

Monty Montgomery's avatar
Monty Montgomery committed
26 27 28 29 30 31 32
/// Public trait interface to a bitstream Writer: a Counter can be
/// used to count bits for cost analysis without actually storing
/// anything (using a new::WriterCounter() as a Writer), to record
/// tokens for later writing (using a new::WriterRecorder() as a
/// Writer) to write actual final bits out using a range encoder
/// (using a new::WriterEncoder() as a Writer).  A WriterRecorder's
/// contents can be replayed into a WriterEncoder.
Monty Montgomery's avatar
Monty Montgomery committed
33 34 35 36 37 38 39 40 41 42 43 44 45
pub trait Writer {
  /// Write a symbol s, using the passed in cdf reference; leaves cdf unchanged
  fn symbol(&mut self, s: u32, cdf: &[u16]);
  /// Write a symbol s, using the passed in cdf reference; updates the referenced cdf.
  fn symbol_with_update(&mut self, s: u32, cdf: &mut [u16]);
  /// Write a bool using passed in probability
  fn bool(&mut self, val: bool, f: u16);
  /// Write a single bit with flat proability
  fn bit(&mut self, bit: u16);
  /// Write literal bits with flat probability
  fn literal(&mut self, bits: u8, s: u32);
  /// Write passed level as a golomb code
  fn write_golomb(&mut self, level: u16);
46
  /// Return current length of range-coded bitstream in integer bits
Monty Montgomery's avatar
Monty Montgomery committed
47
  fn tell(&mut self) -> u32;
48
  /// Return currrent length of range-coded bitstream in fractional
Monty Montgomery's avatar
Monty Montgomery committed
49 50
  /// bits with OD_BITRES decimal precision
  fn tell_frac(&mut self) -> u32;
Luca Barbato's avatar
Luca Barbato committed
51
  /// Save current point in coding/recording to a checkpoint
Monty Montgomery's avatar
Monty Montgomery committed
52
  fn checkpoint(&mut self) -> WriterCheckpoint;
Luca Barbato's avatar
Luca Barbato committed
53
  /// Restore saved position in coding/recording from a checkpoint
Monty Montgomery's avatar
Monty Montgomery committed
54
  fn rollback(&mut self, &WriterCheckpoint);
Thomas Daede's avatar
Thomas Daede committed
55 56
}

Monty Montgomery's avatar
Monty Montgomery committed
57 58 59 60 61 62
/// StorageBackend is an internal trait used to tie a specific Writer
/// implementation's storage to the generic Writer.  It would be
/// private, but Rust is deprecating 'private trait in a public
/// interface' support.
pub trait StorageBackend {
  /// Store partially-computed range code into given storage backend
63
  fn store(&mut self, fl: u16, fh: u16, nms: u16);
Luca Barbato's avatar
Luca Barbato committed
64
  /// Return byte-length of encoded stream to date
Monty Montgomery's avatar
Monty Montgomery committed
65
  fn stream_bytes(&mut self) -> usize;
Luca Barbato's avatar
Luca Barbato committed
66
  /// Backend implementation of checkpoint to pass through Writer interface
Monty Montgomery's avatar
Monty Montgomery committed
67
  fn checkpoint(&mut self) -> WriterCheckpoint;
Luca Barbato's avatar
Luca Barbato committed
68
  /// Backend implementation of rollback to pass through Writer interface
Monty Montgomery's avatar
Monty Montgomery committed
69 70
  fn rollback(&mut self, &WriterCheckpoint);
}
Guillaume Martres's avatar
Guillaume Martres committed
71

Monty Montgomery's avatar
Monty Montgomery committed
72 73
#[derive(Debug, Clone)]
pub struct WriterBase<S> {
Michael Bebenita's avatar
Michael Bebenita committed
74
  /// The number of values in the current range.
Monty Montgomery's avatar
Monty Montgomery committed
75
  rng: u16,
Michael Bebenita's avatar
Michael Bebenita committed
76
  /// The number of bits of data in the current value.
Monty Montgomery's avatar
Monty Montgomery committed
77 78 79 80 81
  cnt: i16,
  /// Debug enable flag
  debug: bool,
  /// Use-specific storage
  s: S
Thomas Daede's avatar
Thomas Daede committed
82 83
}

Monty Montgomery's avatar
Monty Montgomery committed
84 85 86 87 88 89
#[derive(Debug, Clone)]
pub struct WriterCounter {
  /// Bytes that would be shifted out to date
  bytes: usize
}

Monty Montgomery's avatar
Monty Montgomery committed
90 91 92
#[derive(Debug, Clone)]
pub struct WriterRecorder {
  /// Storage for tokens
93
  storage: Vec<(u16, u16, u16)>,
Monty Montgomery's avatar
Monty Montgomery committed
94
  /// Bytes that would be shifted out to date
Monty Montgomery's avatar
Monty Montgomery committed
95
  bytes: usize
Monty Montgomery's avatar
Monty Montgomery committed
96
}
Michael Bebenita's avatar
Michael Bebenita committed
97

Monty Montgomery's avatar
Monty Montgomery committed
98 99 100 101 102
#[derive(Debug, Clone)]
pub struct WriterEncoder {
  /// A buffer for output bytes with their associated carry flags.
  precarry: Vec<u16>,
  /// The low end of the current range.
103
  low: ec_window
Monty Montgomery's avatar
Monty Montgomery committed
104
}
Michael Bebenita's avatar
Michael Bebenita committed
105

Monty Montgomery's avatar
Monty Montgomery committed
106 107
#[derive(Clone)]
pub struct WriterCheckpoint {
Luca Barbato's avatar
Luca Barbato committed
108
  /// Byte length coded/recorded to date
Monty Montgomery's avatar
Monty Montgomery committed
109
  stream_bytes: usize,
Luca Barbato's avatar
Luca Barbato committed
110
  /// To be defined by backend
Monty Montgomery's avatar
Monty Montgomery committed
111 112 113 114
  backend_var: usize,
  /// Saved number of values in the current range.
  rng: u16,
  /// Saved number of bits of data in the current value.
115
  cnt: i16
Monty Montgomery's avatar
Monty Montgomery committed
116
}
Michael Bebenita's avatar
Michael Bebenita committed
117

Monty Montgomery's avatar
Monty Montgomery committed
118 119
/// Constructor for a counting Writer
impl WriterCounter {
120 121
  pub fn new() -> WriterBase<WriterCounter> {
    WriterBase::new(WriterCounter { bytes: 0 })
Monty Montgomery's avatar
Monty Montgomery committed
122 123 124
  }
}

Monty Montgomery's avatar
Monty Montgomery committed
125 126
/// Constructor for a recording Writer
impl WriterRecorder {
127 128
  pub fn new() -> WriterBase<WriterRecorder> {
    WriterBase::new(WriterRecorder { storage: Vec::new(), bytes: 0 })
Michael Bebenita's avatar
Michael Bebenita committed
129
  }
Monty Montgomery's avatar
Monty Montgomery committed
130
}
Michael Bebenita's avatar
Michael Bebenita committed
131

Monty Montgomery's avatar
Monty Montgomery committed
132 133
/// Constructor for a encoding Writer
impl WriterEncoder {
134 135
  pub fn new() -> WriterBase<WriterEncoder> {
    WriterBase::new(WriterEncoder { precarry: Vec::new(), low: 0 })
Michael Bebenita's avatar
Michael Bebenita committed
136
  }
Monty Montgomery's avatar
Monty Montgomery committed
137
}
Michael Bebenita's avatar
Michael Bebenita committed
138

Monty Montgomery's avatar
Monty Montgomery committed
139 140 141
/// The Counter stores nothing we write to it, it merely counts the
/// bit usage like in an Encoder for cost analysis.
impl StorageBackend for WriterBase<WriterCounter> {
142 143
  fn store(&mut self, fl: u16, fh: u16, nms: u16) {
    let (_l, r) = self.lr_compute(fl, fh, nms);
Luca Barbato's avatar
Luca Barbato committed
144
    let d = 16 - r.ilog();
Monty Montgomery's avatar
Monty Montgomery committed
145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160
    let mut c = self.cnt;
    let mut s = c + (d as i16);

    if s >= 0 {
      c += 16;
      if s >= 8 {
        self.s.bytes += 1;
        c -= 8;
      }
      self.s.bytes += 1;
      s = c + (d as i16) - 24;
    }
    self.rng = r << d;
    self.cnt = s;
  }
  fn stream_bytes(&mut self) -> usize {
161
    self.s.bytes
Monty Montgomery's avatar
Monty Montgomery committed
162 163 164 165 166 167
  }
  fn checkpoint(&mut self) -> WriterCheckpoint {
    WriterCheckpoint {
      stream_bytes: self.s.bytes,
      backend_var: 0,
      rng: self.rng,
168
      cnt: self.cnt
Monty Montgomery's avatar
Monty Montgomery committed
169 170 171 172 173 174 175 176 177
    }
  }
  fn rollback(&mut self, checkpoint: &WriterCheckpoint) {
    self.rng = checkpoint.rng;
    self.cnt = checkpoint.cnt;
    self.s.bytes = checkpoint.stream_bytes;
  }
}

Monty Montgomery's avatar
Monty Montgomery committed
178 179
/// The Recorder does not produce a range-coded bitstream, but it
/// still tracks the range coding progress like in an Encoder, as it
180
/// neds to be able to report bit costs for RDO decisions.  It stores a
Monty Montgomery's avatar
Monty Montgomery committed
181 182
/// pair of mostly-computed range coding values per token recorded.
impl StorageBackend for WriterBase<WriterRecorder> {
183 184
  fn store(&mut self, fl: u16, fh: u16, nms: u16) {
    let (_l, r) = self.lr_compute(fl, fh, nms);
Luca Barbato's avatar
Luca Barbato committed
185
    let d = 16 - r.ilog();
Monty Montgomery's avatar
Monty Montgomery committed
186 187
    let mut c = self.cnt;
    let mut s = c + (d as i16);
Michael Bebenita's avatar
Michael Bebenita committed
188

Monty Montgomery's avatar
Monty Montgomery committed
189 190 191 192 193 194 195 196
    if s >= 0 {
      c += 16;
      if s >= 8 {
        self.s.bytes += 1;
        c -= 8;
      }
      self.s.bytes += 1;
      s = c + (d as i16) - 24;
Michael Bebenita's avatar
Michael Bebenita committed
197
    }
Monty Montgomery's avatar
Monty Montgomery committed
198 199
    self.rng = r << d;
    self.cnt = s;
200
    self.s.storage.push((fl, fh, nms));
Michael Bebenita's avatar
Michael Bebenita committed
201
  }
Monty Montgomery's avatar
Monty Montgomery committed
202
  fn stream_bytes(&mut self) -> usize {
203
    self.s.bytes
Monty Montgomery's avatar
Monty Montgomery committed
204 205 206 207 208 209
  }
  fn checkpoint(&mut self) -> WriterCheckpoint {
    WriterCheckpoint {
      stream_bytes: self.s.bytes,
      backend_var: self.s.storage.len(),
      rng: self.rng,
210
      cnt: self.cnt
Monty Montgomery's avatar
Monty Montgomery committed
211
    }
Luca Barbato's avatar
Luca Barbato committed
212
  }
Monty Montgomery's avatar
Monty Montgomery committed
213 214 215 216 217
  fn rollback(&mut self, checkpoint: &WriterCheckpoint) {
    self.rng = checkpoint.rng;
    self.cnt = checkpoint.cnt;
    self.s.bytes = checkpoint.stream_bytes;
    self.s.storage.truncate(checkpoint.backend_var);
Michael Bebenita's avatar
Michael Bebenita committed
218
  }
Monty Montgomery's avatar
Monty Montgomery committed
219
}
Michael Bebenita's avatar
Michael Bebenita committed
220

Monty Montgomery's avatar
Monty Montgomery committed
221 222 223 224 225
/// An Encoder produces an actual range-coded bitstream from passed in
/// tokens.  It does not retain any information about the coded
/// tokens, only the resulting bitstream, and so it cannot be replayed
/// (only checkpointed and rolled back).
impl StorageBackend for WriterBase<WriterEncoder> {
226 227
  fn store(&mut self, fl: u16, fh: u16, nms: u16) {
    let (l, r) = self.lr_compute(fl, fh, nms);
Monty Montgomery's avatar
Monty Montgomery committed
228
    let mut low = l + self.s.low;
Michael Bebenita's avatar
Michael Bebenita committed
229
    let mut c = self.cnt;
Luca Barbato's avatar
Luca Barbato committed
230
    let d = 16 - r.ilog();
Michael Bebenita's avatar
Michael Bebenita committed
231 232 233 234 235 236
    let mut s = c + (d as i16);

    if s >= 0 {
      c += 16;
      let mut m = (1 << c) - 1;
      if s >= 8 {
Monty Montgomery's avatar
Monty Montgomery committed
237
        self.s.precarry.push((low >> c) as u16);
Michael Bebenita's avatar
Michael Bebenita committed
238 239 240 241
        low &= m;
        c -= 8;
        m >>= 8;
      }
Monty Montgomery's avatar
Monty Montgomery committed
242
      self.s.precarry.push((low >> c) as u16);
Michael Bebenita's avatar
Michael Bebenita committed
243 244 245
      s = c + (d as i16) - 24;
      low &= m;
    }
Monty Montgomery's avatar
Monty Montgomery committed
246 247
    self.s.low = low << d;
    self.rng = r << d;
Michael Bebenita's avatar
Michael Bebenita committed
248 249
    self.cnt = s;
  }
Monty Montgomery's avatar
Monty Montgomery committed
250
  fn stream_bytes(&mut self) -> usize {
251
    self.s.precarry.len()
Monty Montgomery's avatar
Monty Montgomery committed
252 253 254
  }
  fn checkpoint(&mut self) -> WriterCheckpoint {
    WriterCheckpoint {
255 256 257 258
      stream_bytes: self.s.precarry.len(),
      backend_var: self.s.low as usize,
      rng: self.rng,
      cnt: self.cnt
Michael Bebenita's avatar
Michael Bebenita committed
259 260
    }
  }
Monty Montgomery's avatar
Monty Montgomery committed
261 262 263 264 265
  fn rollback(&mut self, checkpoint: &WriterCheckpoint) {
    self.rng = checkpoint.rng;
    self.cnt = checkpoint.cnt;
    self.s.low = checkpoint.backend_var as ec_window;
    self.s.precarry.truncate(checkpoint.stream_bytes);
Michael Bebenita's avatar
Michael Bebenita committed
266
  }
Monty Montgomery's avatar
Monty Montgomery committed
267
}
Michael Bebenita's avatar
Michael Bebenita committed
268

Monty Montgomery's avatar
Monty Montgomery committed
269 270
/// A few local helper functions needed by the Writer that are not
/// part of the public interface.
271
impl<S> WriterBase<S> {
Monty Montgomery's avatar
Monty Montgomery committed
272 273 274
  /// Internal constructor called by the subtypes that implement the
  /// actual encoder and Recorder.
  fn new(storage: S) -> Self {
275
    WriterBase { rng: 0x8000, cnt: -9, debug: false, s: storage }
Michael Bebenita's avatar
Michael Bebenita committed
276
  }
277
  /// Compute low and range values from token cdf values and local state
278
  fn lr_compute(&mut self, fl: u16, fh: u16, nms: u16) -> (ec_window, u16) {
279 280 281
    let u: u32;
    let v: u32;
    let mut r = self.rng as u32;
Luca Barbato's avatar
Luca Barbato committed
282
    debug_assert!(32768 <= r);
283 284
    if fl < 32768 {
      u = (((r >> 8) * (fl as u32 >> EC_PROB_SHIFT)) >> (7 - EC_PROB_SHIFT))
285
        + EC_MIN_PROB * nms as u32;
286
      v = (((r >> 8) * (fh as u32 >> EC_PROB_SHIFT)) >> (7 - EC_PROB_SHIFT))
287
        + EC_MIN_PROB * (nms - 1) as u32;
288 289
      (r - u, (u - v) as u16)
    } else {
290 291 292
      r -= (((r >> 8) * (fh as u32 >> EC_PROB_SHIFT)) >> (7 - EC_PROB_SHIFT))
        + EC_MIN_PROB * (nms - 1) as u32;
      (0, r as u16)
293 294
    }
  }
Michael Bebenita's avatar
Michael Bebenita committed
295 296 297 298 299 300 301 302 303
  /// Given the current total integer number of bits used and the current value of
  /// rng, computes the fraction number of bits used to `OD_BITRES` precision.
  /// This is used by `od_ec_enc_tell_frac()` and `od_ec_dec_tell_frac()`.
  /// `nbits_total`: The number of whole bits currently used, i.e., the value
  ///                returned by `od_ec_enc_tell()` or `od_ec_dec_tell()`.
  /// `rng`: The current value of rng from either the encoder or decoder state.
  /// Return: The number of bits scaled by `2**OD_BITRES`.
  ///         This will always be slightly larger than the exact value (e.g., all
  ///         rounding error is in the positive direction).
Monty Montgomery's avatar
Monty Montgomery committed
304
  fn frac_compute(nbits_total: u32, mut rng: u32) -> u32 {
Michael Bebenita's avatar
Michael Bebenita committed
305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328
    // To handle the non-integral number of bits still left in the encoder/decoder
    //  state, we compute the worst-case number of bits of val that must be
    //  encoded to ensure that the value is inside the range for any possible
    //  subsequent bits.
    // The computation here is independent of val itself (the decoder does not
    //  even track that value), even though the real number of bits used after
    //  od_ec_enc_done() may be 1 smaller if rng is a power of two and the
    //  corresponding trailing bits of val are all zeros.
    // If we did try to track that special case, then coding a value with a
    //  probability of 1/(1 << n) might sometimes appear to use more than n bits.
    // This may help explain the surprising result that a newly initialized
    //  encoder or decoder claims to have used 1 bit.
    let nbits = nbits_total << OD_BITRES;
    let mut l = 0;
    let mut i = OD_BITRES;
    while i > 0 {
      i -= 1;
      rng = (rng * rng) >> 15;
      let b = rng >> 16;
      l = (l << 1) | b;
      rng >>= b;
    }
    nbits - l
  }
Monty Montgomery's avatar
Monty Montgomery committed
329
  /// Function to update the CDF for Writer calls that do so.
Michael Bebenita's avatar
Michael Bebenita committed
330 331
  fn update_cdf(cdf: &mut [u16], val: u32) {
    let nsymbs = cdf.len() - 1;
fbossen's avatar
fbossen committed
332 333
    let rate = 3 + (nsymbs >> 1).min(2) + (cdf[nsymbs] >> 4) as usize;
    cdf[nsymbs] += 1 - (cdf[nsymbs] >> 5);
Michael Bebenita's avatar
Michael Bebenita committed
334 335

    // Single loop (faster)
336
    for (i, v) in cdf[..nsymbs - 1].iter_mut().enumerate() {
fbossen's avatar
fbossen committed
337 338
      if i as u32 >= val {
        *v -= *v >> rate;
Michael Bebenita's avatar
Michael Bebenita committed
339
      } else {
fbossen's avatar
fbossen committed
340
        *v += (32768 - *v) >> rate;
Michael Bebenita's avatar
Michael Bebenita committed
341 342 343
      }
    }
  }
344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364

  #[cfg(debug)]
  fn print_backtrace(&self, s: u32) {
    let mut depth = 3;
    backtrace::trace(|frame| {
      let ip = frame.ip();

      depth -= 1;

      if depth == 0 {
        backtrace::resolve(ip, |symbol| {
          if let Some(name) = symbol.name() {
            eprintln!("Writing symbol {} from {}", s, name);
          }
        });
        false
      } else {
        true
      }
    });
  }
Monty Montgomery's avatar
Monty Montgomery committed
365
}
366

367
/// Replay implementation specific to the Recorder
Monty Montgomery's avatar
Monty Montgomery committed
368
impl WriterBase<WriterRecorder> {
369
  /// Replays the partially-computed range tokens out of the Recorder's
Monty Montgomery's avatar
Monty Montgomery committed
370 371
  /// storage and into the passed in Writer, which may be an Encoder
  /// or another Recorder.  Clears the Recorder after replay.
gibix's avatar
gibix committed
372
  pub fn replay(&mut self, dest: &mut dyn StorageBackend) {
Monty Montgomery's avatar
Monty Montgomery committed
373
    for i in 0..self.s.storage.len() {
374 375
      let (fl, fh, nms) = self.s.storage[i];
      dest.store(fl, fh, nms);
Guillaume Martres's avatar
Guillaume Martres committed
376
    }
377 378
    self.rng = 0x8000;
    self.cnt = -9;
Monty Montgomery's avatar
Monty Montgomery committed
379 380
    self.s.storage.truncate(0);
    self.s.bytes = 0;
Michael Bebenita's avatar
Michael Bebenita committed
381
  }
Monty Montgomery's avatar
Monty Montgomery committed
382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414
}

/// Done implementation specific to the Encoder
impl WriterBase<WriterEncoder> {
  /// Indicates that there are no more symbols to encode.  Flushes
  /// remaining state into coding and returns a vector containing the
  /// final bitstream.
  pub fn done(&mut self) -> Vec<u8> {
    // We output the minimum number of bits that ensures that the symbols encoded
    // thus far will be decoded correctly regardless of the bits that follow.
    let l = self.s.low;
    let mut c = self.cnt;
    let mut s = 10;
    let m = 0x3FFF;
    let mut e = ((l + m) & !m) | (m + 1);

    s += c;

    if s > 0 {
      let mut n = (1 << (c + 16)) - 1;

      loop {
        self.s.precarry.push((e >> (c + 16)) as u16);
        e &= n;
        s -= 8;
        c -= 8;
        n >>= 8;

        if s <= 0 {
          break;
        }
      }
    }
Monty Montgomery's avatar
Monty Montgomery committed
415

Monty Montgomery's avatar
Monty Montgomery committed
416 417 418 419 420 421 422 423 424 425 426
    let mut c = 0;
    let mut offs = self.s.precarry.len();
    let mut out = vec![0 as u8; offs];
    while offs > 0 {
      offs -= 1;
      c += self.s.precarry[offs];
      out[offs] = c as u8;
      c >>= 8;
    }

    out
Michael Bebenita's avatar
Michael Bebenita committed
427
  }
Monty Montgomery's avatar
Monty Montgomery committed
428
}
429

Monty Montgomery's avatar
Monty Montgomery committed
430
/// Generic/shared implementation for Writers with StorageBackends (ie, Encoders and Recorders)
431 432 433 434
impl<S> Writer for WriterBase<S>
where
  WriterBase<S>: StorageBackend
{
Monty Montgomery's avatar
Monty Montgomery committed
435 436 437 438 439 440
  /// Encode a single binary value.
  /// `val`: The value to encode (0 or 1).
  /// `f`: The probability that the val is one, scaled by 32768.
  fn bool(&mut self, val: bool, f: u16) {
    debug_assert!(0 < f);
    debug_assert!(f < 32768);
441
    self.symbol(if val { 1 } else { 0 }, &[f, 0]);
Monty Montgomery's avatar
Monty Montgomery committed
442 443 444 445 446 447 448 449 450 451 452 453
  }
  /// Encode a single boolean value.
  /// `val`: The value to encode (false or true).
  /// `f`: The probability that the val is true, scaled by 32768.
  fn bit(&mut self, bit: u16) {
    self.bool(bit == 1, 16384);
  }
  /// Encode a literal bitstring, bit by bit in MSB order, with flat
  /// probability.
  /// 'bits': Length of bitstring
  /// 's': Bit string to encode
  fn literal(&mut self, bits: u8, s: u32) {
Monty Montgomery's avatar
Monty Montgomery committed
454
    for bit in (0..bits).rev() {
Monty Montgomery's avatar
Monty Montgomery committed
455
      self.bit((1 & (s >> bit)) as u16);
Monty Montgomery's avatar
Monty Montgomery committed
456 457
    }
  }
Monty Montgomery's avatar
Monty Montgomery committed
458 459 460 461 462 463 464 465
  /// Encodes a symbol given a cumulative distribution function (CDF) table in Q15.
  /// `s`: The index of the symbol to encode.
  /// `cdf`: The CDF, such that symbol s falls in the range
  ///        `[s > 0 ? cdf[s - 1] : 0, cdf[s])`.
  ///       The values must be monotonically non-decreasing, and the last value
  ///       must be exactly 32768. There should be at most 16 values.
  fn symbol(&mut self, s: u32, cdf: &[u16]) {
    debug_assert!(cdf[cdf.len() - 1] == 0);
466
    let nms = cdf.len() - s as usize;
Monty Montgomery's avatar
Monty Montgomery committed
467 468 469 470
    let fl = if s > 0 { cdf[s as usize - 1] } else { 32768 };
    let fh = cdf[s as usize];
    debug_assert!(fh <= fl);
    debug_assert!(fl <= 32768);
471
    self.store(fl, fh, nms as u16);
Monty Montgomery's avatar
Monty Montgomery committed
472 473 474
  }
  /// Encodes a symbol given a cumulative distribution function (CDF)
  /// table in Q15, then updates the CDF probabilities to relect we've
Luca Barbato's avatar
Luca Barbato committed
475
  /// written one more symbol 's'.
Monty Montgomery's avatar
Monty Montgomery committed
476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494
  /// `s`: The index of the symbol to encode.
  /// `cdf`: The CDF, such that symbol s falls in the range
  ///        `[s > 0 ? cdf[s - 1] : 0, cdf[s])`.
  ///       The values must be monotonically non-decreasing, and the last value
  ///       must be exactly 32768. There should be at most 16 values.
  fn symbol_with_update(&mut self, s: u32, cdf: &mut [u16]) {
    let nsymbs = cdf.len() - 1;
    #[cfg(debug)]
    {
      if self.debug {
        self.print_backtrace(s);
      }
    }
    self.symbol(s, &cdf[..nsymbs]);
    Self::update_cdf(cdf, s);
  }
  /// Encode a golomb to the bitstream.
  /// 'level': passed in value to encode
  fn write_golomb(&mut self, level: u16) {
Michael Bebenita's avatar
Michael Bebenita committed
495 496 497
    let x = level + 1;
    let mut i = x;
    let mut length = 0;
498

Michael Bebenita's avatar
Michael Bebenita committed
499 500 501
    while i != 0 {
      i >>= 1;
      length += 1;
502
    }
503
    debug_assert!(length > 0);
504

Michael Bebenita's avatar
Michael Bebenita committed
505 506
    for _ in 0..length - 1 {
      self.bit(0);
507
    }
Guillaume Martres's avatar
Guillaume Martres committed
508

Michael Bebenita's avatar
Michael Bebenita committed
509 510
    for i in (0..length).rev() {
      self.bit((x >> i) & 0x01);
511
    }
Michael Bebenita's avatar
Michael Bebenita committed
512
  }
Monty Montgomery's avatar
Monty Montgomery committed
513 514 515 516 517 518 519 520 521 522 523
  /// Returns the number of bits "used" by the encoded symbols so far.
  /// This same number can be computed in either the encoder or the
  /// decoder, and is suitable for making coding decisions.  The value
  /// will be the same whether using an Encoder or Recorder.
  /// Return: The integer number of bits.
  ///         This will always be slightly larger than the exact value (e.g., all
  ///          rounding error is in the positive direction).
  fn tell(&mut self) -> u32 {
    // The 10 here counteracts the offset of -9 baked into cnt, and adds 1 extra
    // bit, which we reserve for terminating the stream.
    (((self.stream_bytes() * 8) as i32) + (self.cnt as i32) + 10) as u32
Michael Bebenita's avatar
Michael Bebenita committed
524
  }
Monty Montgomery's avatar
Monty Montgomery committed
525 526 527 528 529 530 531 532 533
  /// Returns the number of bits "used" by the encoded symbols so far.
  /// This same number can be computed in either the encoder or the
  /// decoder, and is suitable for making coding decisions. The value
  /// will be the same whether using an Encoder or Recorder.
  /// Return: The number of bits scaled by `2**OD_BITRES`.
  ///         This will always be slightly larger than the exact value (e.g., all
  ///          rounding error is in the positive direction).
  fn tell_frac(&mut self) -> u32 {
    Self::frac_compute(self.tell(), self.rng as u32)
Michael Bebenita's avatar
Michael Bebenita committed
534
  }
Monty Montgomery's avatar
Monty Montgomery committed
535 536 537 538
  /// Save current point in coding/recording to a checkpoint that can
  /// be restored later.  A WriterCheckpoint can be generated for an
  /// Encoder or Recorder, but can only be used to rollback the Writer
  /// instance from which it was generated.
539
  fn checkpoint(&mut self) -> WriterCheckpoint {
Monty Montgomery's avatar
Monty Montgomery committed
540
    StorageBackend::checkpoint(self)
Michael Bebenita's avatar
Michael Bebenita committed
541
  }
Monty Montgomery's avatar
Monty Montgomery committed
542
  /// Roll back a given Writer to the state saved in the WriterCheckpoint
Luca Barbato's avatar
Luca Barbato committed
543
  /// 'wc': Saved Writer state/posiiton to restore
544 545
  fn rollback(&mut self, wc: &WriterCheckpoint) {
    StorageBackend::rollback(self, wc)
Michael Bebenita's avatar
Michael Bebenita committed
546
  }
Guillaume Martres's avatar
Guillaume Martres committed
547 548
}

549
pub trait BCodeWriter {
Michael Bebenita's avatar
Michael Bebenita committed
550 551 552 553 554 555 556 557 558 559 560 561
  fn recenter_nonneg(&mut self, r: u16, v: u16) -> u16;
  fn recenter_finite_nonneg(&mut self, n: u16, r: u16, v: u16) -> u16;
  fn write_quniform(&mut self, n: u16, v: u16) -> Result<(), std::io::Error>;
  fn write_subexpfin(
    &mut self, n: u16, k: u16, v: u16
  ) -> Result<(), std::io::Error>;
  fn write_refsubexpfin(
    &mut self, n: u16, k: u16, r: i16, v: i16
  ) -> Result<(), std::io::Error>;
  fn write_s_refsubexpfin(
    &mut self, n: u16, k: u16, r: i16, v: i16
  ) -> Result<(), std::io::Error>;
562 563
}

564
impl<W: io::Write> BCodeWriter for BitWriter<W, BigEndian> {
Michael Bebenita's avatar
Michael Bebenita committed
565 566 567
  fn recenter_nonneg(&mut self, r: u16, v: u16) -> u16 {
    /* Recenters a non-negative literal v around a reference r */
    if v > (r << 1) {
Michael Bebenita's avatar
Michael Bebenita committed
568
      v
Michael Bebenita's avatar
Michael Bebenita committed
569
    } else if v >= r {
Michael Bebenita's avatar
Michael Bebenita committed
570
      (v - r) << 1
Michael Bebenita's avatar
Michael Bebenita committed
571
    } else {
Michael Bebenita's avatar
Michael Bebenita committed
572
      ((r - v) << 1) - 1
Michael Bebenita's avatar
Michael Bebenita committed
573 574 575 576
    }
  }
  fn recenter_finite_nonneg(&mut self, n: u16, r: u16, v: u16) -> u16 {
    /* Recenters a non-negative literal v in [0, n-1] around a
577
           reference r also in [0, n-1] */
Michael Bebenita's avatar
Michael Bebenita committed
578
    if (r << 1) <= n {
Michael Bebenita's avatar
Michael Bebenita committed
579
      self.recenter_nonneg(r, v)
Michael Bebenita's avatar
Michael Bebenita committed
580
    } else {
Michael Bebenita's avatar
Michael Bebenita committed
581
      self.recenter_nonneg(n - 1 - r, n - 1 - v)
Michael Bebenita's avatar
Michael Bebenita committed
582 583 584 585 586 587 588 589 590 591
    }
  }
  fn write_quniform(&mut self, n: u16, v: u16) -> Result<(), std::io::Error> {
    /* Encodes a value v in [0, n-1] quasi-uniformly */
    if n <= 1 {
      return Ok(());
    };
    let l = 31 ^ ((n - 1) + 1).leading_zeros();
    let m = (1 << l) - n;
    if v < m {
Michael Bebenita's avatar
Michael Bebenita committed
592
      self.write(l - 1, v)
Michael Bebenita's avatar
Michael Bebenita committed
593 594
    } else {
      self.write(l - 1, m + ((v - m) >> 1))?;
Michael Bebenita's avatar
Michael Bebenita committed
595
      self.write_bit(((v - m) & 1) != 0)
Michael Bebenita's avatar
Michael Bebenita committed
596 597 598 599 600 601 602 603 604
    }
  }
  fn write_subexpfin(
    &mut self, n: u16, k: u16, v: u16
  ) -> Result<(), std::io::Error> {
    /* Finite subexponential code that codes a symbol v in [0, n-1] with parameter k */
    let mut i = 0;
    let mut mk = 0;
    loop {
Michael Bebenita's avatar
Michael Bebenita committed
605
      let b = if i > 0 { k + i - 1 } else { k };
Michael Bebenita's avatar
Michael Bebenita committed
606 607 608 609 610 611 612
      let a = 1 << b;
      if n <= mk + 3 * a {
        return self.write_quniform(n - mk, v - mk);
      } else {
        let t = v >= mk + a;
        self.write_bit(t)?;
        if t {
Michael Bebenita's avatar
Michael Bebenita committed
613
          i += 1;
Michael Bebenita's avatar
Michael Bebenita committed
614
          mk += a;
615
        } else {
Michael Bebenita's avatar
Michael Bebenita committed
616
          return self.write(b as u32, v - mk);
617
        }
Michael Bebenita's avatar
Michael Bebenita committed
618
      }
619
    }
Michael Bebenita's avatar
Michael Bebenita committed
620 621 622 623 624
  }
  fn write_refsubexpfin(
    &mut self, n: u16, k: u16, r: i16, v: i16
  ) -> Result<(), std::io::Error> {
    /* Finite subexponential code that codes a symbol v in [0, n-1] with
625 626
           parameter k based on a reference ref also in [0, n-1].
           Recenters symbol around r first and then uses a finite subexponential code. */
Michael Bebenita's avatar
Michael Bebenita committed
627
    let recentered_v = self.recenter_finite_nonneg(n, r as u16, v as u16);
Michael Bebenita's avatar
Michael Bebenita committed
628
    self.write_subexpfin(n, k, recentered_v)
Michael Bebenita's avatar
Michael Bebenita committed
629 630 631 632 633
  }
  fn write_s_refsubexpfin(
    &mut self, n: u16, k: u16, r: i16, v: i16
  ) -> Result<(), std::io::Error> {
    /* Signed version of the above function */
Michael Bebenita's avatar
Michael Bebenita committed
634
    self.write_refsubexpfin(
Michael Bebenita's avatar
Michael Bebenita committed
635 636 637 638
      (n << 1) - 1,
      k,
      r + (n - 1) as i16,
      v + (n - 1) as i16
Michael Bebenita's avatar
Michael Bebenita committed
639
    )
Michael Bebenita's avatar
Michael Bebenita committed
640
  }
641 642
}

Luca Barbato's avatar
Luca Barbato committed
643 644
#[cfg(test)]
mod test {
Michael Bebenita's avatar
Michael Bebenita committed
645
  use super::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
646

647 648
  const WINDOW_SIZE: i16 = 32;
  const LOTS_OF_BITS: i16 = 0x4000;
Michael Bebenita's avatar
Michael Bebenita committed
649

650
  #[derive(Debug)]
Michael Bebenita's avatar
Michael Bebenita committed
651
  struct Reader<'a> {
652 653 654 655 656
    buf: &'a [u8],
    bptr: usize,
    dif: ec_window,
    rng: u16,
    cnt: i16,
Michael Bebenita's avatar
Michael Bebenita committed
657 658 659 660
  }

  impl<'a> Reader<'a> {
    fn new(buf: &'a [u8]) -> Self {
661 662 663 664 665 666 667 668 669 670
      let mut r = Reader {
        buf,
        bptr: 0,
        dif: (1 << (WINDOW_SIZE - 1)) - 1,
        rng: 0x8000,
        cnt: -15,
      };
      r.refill();
      r
    }
Michael Bebenita's avatar
Michael Bebenita committed
671

672 673 674 675 676 677 678 679 680 681 682 683 684
    fn refill(&mut self) {
      let mut s = WINDOW_SIZE - 9 - (self.cnt + 15);
      while s >= 0 && self.bptr < self.buf.len() {
        assert!(s <= WINDOW_SIZE - 8);
        self.dif ^= (self.buf[self.bptr] as ec_window) << s;
        self.cnt += 8;
        s -= 8;
        self.bptr += 1;
      }
      if self.bptr >= self.buf.len() {
        self.cnt = LOTS_OF_BITS;
      }
    }
Michael Bebenita's avatar
Michael Bebenita committed
685

686 687 688 689 690 691 692 693 694 695 696 697
    fn normalize(&mut self, dif: ec_window, rng: u32) {
      assert!(rng <= 65536);
      let d = rng.leading_zeros() - 16;
      //let d = 16 - (32-rng.leading_zeros());
      //msb(rng) = 31-rng.leading_zeros();
      self.cnt -= d as i16;
      /*This is equivalent to shifting in 1's instead of 0's.*/
      self.dif = ((dif + 1) << d) - 1;
      self.rng = (rng << d) as u16;
      if self.cnt < 0 {
        self.refill()
      }
Michael Bebenita's avatar
Michael Bebenita committed
698 699 700
    }

    fn bool(&mut self, f: u32) -> bool {
701 702 703 704 705 706 707 708 709 710 711 712 713
      assert!(f < 32768);
      let r = self.rng as u32;
      assert!(self.dif >> (WINDOW_SIZE - 16) < r);
      assert!(32768 <= r);
      let v = ((r >> 8) * (f >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT)) + EC_MIN_PROB;
      let vw = v << (WINDOW_SIZE - 16);
      let (dif, rng, ret) = if self.dif >= vw {
        (self.dif - vw, r - v, false)
      } else {
        (self.dif, v, true)
      };
      self.normalize(dif, rng);
      ret
Michael Bebenita's avatar
Michael Bebenita committed
714 715
    }

Monty Montgomery's avatar
Monty Montgomery committed
716
    fn symbol(&mut self, icdf: &[u16]) -> i32 {
717 718 719 720 721 722 723 724 725 726 727 728 729 730 731
      let r = self.rng as u32;
      assert!(self.dif >> (WINDOW_SIZE - 16) < r);
      assert!(32768 <= r);
      let n = icdf.len() as u32 - 1;
      let c = self.dif >> (WINDOW_SIZE - 16);
      let mut v = self.rng as u32;
      let mut ret = 0i32;
      let mut u = v;
      v = (r >> 8) * (icdf[ret as usize] as u32 >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT);
      v += EC_MIN_PROB * (n - ret as u32);
      while c < v {
        u = v;
        ret += 1;
        v = (r >> 8) * (icdf[ret as usize] as u32 >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT);
        v += EC_MIN_PROB * (n - ret as u32);
Michael Bebenita's avatar
Michael Bebenita committed
732
      }
733 734 735 736 737
      assert!(v < u);
      assert!(u <= r);
      let new_dif = self.dif - (v << (WINDOW_SIZE - 16));
      self.normalize(new_dif, u - v);
      ret
Michael Bebenita's avatar
Michael Bebenita committed
738 739 740 741 742
    }
  }

  #[test]
  fn booleans() {
Monty Montgomery's avatar
Monty Montgomery committed
743
    let mut w = WriterEncoder::new();
Michael Bebenita's avatar
Michael Bebenita committed
744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767

    w.bool(false, 1);
    w.bool(true, 2);
    w.bool(false, 3);
    w.bool(true, 1);
    w.bool(true, 2);
    w.bool(false, 3);

    let b = w.done();

    let mut r = Reader::new(&b);

    assert_eq!(r.bool(1), false);
    assert_eq!(r.bool(2), true);
    assert_eq!(r.bool(3), false);
    assert_eq!(r.bool(1), true);
    assert_eq!(r.bool(2), true);
    assert_eq!(r.bool(3), false);
  }

  #[test]
  fn cdf() {
    let cdf = [7296, 3819, 1716, 0];

Monty Montgomery's avatar
Monty Montgomery committed
768
    let mut w = WriterEncoder::new();
Michael Bebenita's avatar
Michael Bebenita committed
769

Monty Montgomery's avatar
Monty Montgomery committed
770 771 772 773 774 775 776 777 778
    w.symbol(0, &cdf);
    w.symbol(0, &cdf);
    w.symbol(0, &cdf);
    w.symbol(1, &cdf);
    w.symbol(1, &cdf);
    w.symbol(1, &cdf);
    w.symbol(2, &cdf);
    w.symbol(2, &cdf);
    w.symbol(2, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
779 780 781 782 783

    let b = w.done();

    let mut r = Reader::new(&b);

Monty Montgomery's avatar
Monty Montgomery committed
784 785 786 787 788 789 790 791 792
    assert_eq!(r.symbol(&cdf), 0);
    assert_eq!(r.symbol(&cdf), 0);
    assert_eq!(r.symbol(&cdf), 0);
    assert_eq!(r.symbol(&cdf), 1);
    assert_eq!(r.symbol(&cdf), 1);
    assert_eq!(r.symbol(&cdf), 1);
    assert_eq!(r.symbol(&cdf), 2);
    assert_eq!(r.symbol(&cdf), 2);
    assert_eq!(r.symbol(&cdf), 2);
Michael Bebenita's avatar
Michael Bebenita committed
793 794 795 796 797 798
  }

  #[test]
  fn mixed() {
    let cdf = [7296, 3819, 1716, 0];

Monty Montgomery's avatar
Monty Montgomery committed
799
    let mut w = WriterEncoder::new();
Michael Bebenita's avatar
Michael Bebenita committed
800

Monty Montgomery's avatar
Monty Montgomery committed
801
    w.symbol(0, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
802
    w.bool(true, 2);
Monty Montgomery's avatar
Monty Montgomery committed
803
    w.symbol(0, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
804
    w.bool(true, 2);
Monty Montgomery's avatar
Monty Montgomery committed
805
    w.symbol(0, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
806
    w.bool(true, 2);
Monty Montgomery's avatar
Monty Montgomery committed
807
    w.symbol(1, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
808
    w.bool(true, 1);
Monty Montgomery's avatar
Monty Montgomery committed
809
    w.symbol(1, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
810
    w.bool(false, 2);
Monty Montgomery's avatar
Monty Montgomery committed
811 812 813 814
    w.symbol(1, &cdf);
    w.symbol(2, &cdf);
    w.symbol(2, &cdf);
    w.symbol(2, &cdf);
Michael Bebenita's avatar
Michael Bebenita committed
815 816 817 818 819

    let b = w.done();

    let mut r = Reader::new(&b);

Monty Montgomery's avatar
Monty Montgomery committed
820
    assert_eq!(r.symbol(&cdf), 0);
Michael Bebenita's avatar
Michael Bebenita committed
821
    assert_eq!(r.bool(2), true);
Monty Montgomery's avatar
Monty Montgomery committed
822
    assert_eq!(r.symbol(&cdf), 0);
Michael Bebenita's avatar
Michael Bebenita committed
823
    assert_eq!(r.bool(2), true);
Monty Montgomery's avatar
Monty Montgomery committed
824
    assert_eq!(r.symbol(&cdf), 0);
Michael Bebenita's avatar
Michael Bebenita committed
825
    assert_eq!(r.bool(2), true);
Monty Montgomery's avatar
Monty Montgomery committed
826
    assert_eq!(r.symbol(&cdf), 1);
Michael Bebenita's avatar
Michael Bebenita committed
827
    assert_eq!(r.bool(1), true);
Monty Montgomery's avatar
Monty Montgomery committed
828
    assert_eq!(r.symbol(&cdf), 1);
Michael Bebenita's avatar
Michael Bebenita committed
829
    assert_eq!(r.bool(2), false);
Monty Montgomery's avatar
Monty Montgomery committed
830 831 832 833
    assert_eq!(r.symbol(&cdf), 1);
    assert_eq!(r.symbol(&cdf), 2);
    assert_eq!(r.symbol(&cdf), 2);
    assert_eq!(r.symbol(&cdf), 2);
Michael Bebenita's avatar
Michael Bebenita committed
834
  }
Luca Barbato's avatar
Luca Barbato committed
835
}