Optimize BitStreamReader, speed up EV Whitelist loading

components/packed_ct_ev_whitelist/bit_stream_reader.cc

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "components/packed_ct_ev_whitelist/bit_stream_reader.h"
 
 #include "base/big_endian.h"
 #include "base/logging.h"
 #include "base/numerics/safe_conversions.h"
-

[Ryan Sleevi, 2015/04/09 01:33:57]

Don't delete this newline

[ian fette, 2015/04/09 16:22:13]

Done.

 namespace packed_ct_ev_whitelist {
 namespace internal {
 
 BitStreamReader::BitStreamReader(const base::StringPiece& source)
     : source_(source), current_byte_(0), current_bit_(7) {
   DCHECK_LT(source_.length(), UINT32_MAX);
 }
 
 bool BitStreamReader::ReadUnaryEncoding(uint64_t* out) {
   if (BitsLeft() == 0)
     return false;
 
   *out = 0;
   while ((BitsLeft() > 0) && ReadBit())
     ++(*out);
 
   return true;
 }
 
 bool BitStreamReader::ReadBits(uint8_t num_bits, uint64_t* out) {
   DCHECK_LE(num_bits, 64);
 
   if (BitsLeft() < num_bits)
     return false;
 
   *out = 0;
-  for (uint8_t i = 0; i < num_bits; ++i)
-    (*out) |= (static_cast<uint64_t>(ReadBit()) << (num_bits - (i + 1)));
+
+  for (; num_bits && (current_bit_ != 7); --num_bits)
+    (*out) |= (static_cast<uint64_t>(ReadBit() << (num_bits - 1)));
+  for (; num_bits / 8; num_bits -= 8) {
+    (*out) |= (static_cast<uint64_t>(source_.data()[current_byte_]) & 0xff)
+              << (num_bits - 8);

[Ryan Sleevi, 2015/04/09 01:33:57]

I'm not terribly keen on this, which duplicates the logic of ReadBit()

Seems like there should be an internal ReadByte helper that's responsible for
checking pre-conditions and updating current_byte_

[ian fette, 2015/04/09 02:43:06]

One important precondition is whether the stream is currently byte-aligned or
not. Hence the first for loop looking at current_bit_. I'm not sure if there's a
good way to avoid duplicating that logic, as even if you had a ReadByte() helper
you wouldn't want to call it unless you were byte-aligned, otherwise you're
going to just end up reading bit-by-bit. Could add a bool IsByteAligned() helper
if you want to avoid touching current_bit_ but it seems a bit overkill IMO.

The second for loop does indeed mirror some of the logic of ReadBit (namely,
adjusting current_byte as you mention, and touching source_.data() directly).
This could be moved off into a helper function, which would hopefully get
inlined, but then we're going to have to decide whether that helper function
needs to check that it's byte aligned first. I presume it would (unless it were
private?) at which point we're double checking preconditions. Not end of the
world but the whole point was to improve the performance here :) Those
preconditions would also get checked then 8 times for a 64-bit read (common). 

Personally, I'd prefer to leave it in-line there, but if you want it in a
separate function, I'd propose making it
a) private
b) Only DCHECK() the preconditions (that current_bit_ == 7 and BitsLeft() >= 8)
c) return a uint8_t and update current_byte_ 

The only downside would be if someone expects ReadByte() to work regardless of
whether the stream is currently byte-aligned (or at the very least to check
whether that condition is met). Could either handle both cases (in which case
you're having an extra switch statement run 8 times for a 64-bit read) or just
document that being byte-aligned is a precondition.

Thoughts?

[Ryan Sleevi, 2015/04/09 02:57:16]

Right, I was just talking in the context of the second loop (e.g. precondition
is that things are aligned). The first loop uses ReadBit(), such that
current_bit_/current_byte_ are always set.

I'm not worried about the precondition cost - DCHECK optimizes out in official
builds (and has for a long time) - it's just to guard against developer mistake,
which would happen if someone adjusted the first loop without adjusting the
second.
Yup, exactly what I was trying to suggest.
Exactly this.

From a design perspective, my grumpiness towards the current approach came from
the fact that ReadBits() didn't mutate internal state directly, but instead
delegated through ReadBit, so I wasn't keen to spread the mutation 'up'.

There's definitely room for 'cleverness', given the read patters (read 64 bits,
read 0-7 bits, read 47 bits, read 0-7 bits, read 47 bits, etc) where a
copy-mask-shift approach at the 128-bit level (so we can span the two 64-bits)
would totally work, and might yield better loop unrolling.

But that does feel like extreme microptimizations for the task, given the
discussion on the bug.

[ian fette, 2015/04/09 16:22:13]

Done.

+    current_byte_++;
+  }
+  for (; num_bits; --num_bits)
+    (*out) |= (static_cast<uint64_t>(ReadBit() << (num_bits - 1)));
 
   return true;
 }
 
 uint64_t BitStreamReader::BitsLeft() const {
   if (current_byte_ == source_.length())
     return 0;
   DCHECK_GT(source_.length(), current_byte_);
   return (source_.length() - (current_byte_ + 1)) * 8 + current_bit_ + 1;
 }
 
 uint8_t BitStreamReader::ReadBit() {
   DCHECK_GT(BitsLeft(), 0u);
   DCHECK(current_bit_ < 8 && current_bit_ >= 0);
   uint8_t res =
       (source_.data()[current_byte_] & (1 << current_bit_)) >> current_bit_;
   current_bit_--;
   if (current_bit_ < 0) {
     current_byte_++;
     current_bit_ = 7;
   }
 
   return res;
 }
 
 }  // namespace internal
 }  // namespace packed_ct_ev_whitelist