Bit Wrangling in Scala part deux (with Macros)

Hi all,

In my previous post, I talked about a simple library for extracting bit-level information from a sequence of bytes stored in ByteString.

Although the library works as expected, it is a little bit slower than bit shifting code that you would write by hand, even though the code is a bit difficult to read at 1st glance.

For example, let's say you want to read 7 bits at offset 9 within a byte string, then the following would be equivalent:

• int(bytes,9,7)
• (bytes(1) & 0x7f)

However, the 1st approach is a bit slower, because the code has got branches and loops, etc.. See the code below:

  /** Read an integer value at a given bit offset with a given length within a ByteString */
  def int(bytes: ByteString, offset: Int, length: Int): Int = {
    val startByte = offset / 8    val endByte   = (offset + length - 1) / 8    var i = startByte
    var value : Int = 0    var toRead : Int = 0    var readSoFar : Int = 0    while (i <= endByte) {
      if (i == startByte){
    val firstOffset = offset % 8    toRead = length.min(8 - firstOffset)
        value = int( bytes(i), firstOffset, toRead)
        readSoFar += toRead
      } else{
        value = (value << toRead)
        toRead = 8.min(length - readSoFar)
        val lastValue = int( bytes(i), 0, toRead)
        value = value | lastValue        readSoFar += toRead
      }
      i += 1    }
    value
  }

In order to speed up code, I decided to look at Scala Macros. With Macros, I was essentially able to shift most of the computation to compile time, and all you end up with an abstract syntax tree which captures the essence of the bit shifting computation. The key thing was to understand how to create and combine ASTs, but once I got that, the rest of the work was not too complicated. 

Here is the macro to extract boolean value at a given offset:

 /** Implementation of the boolean macro*/
  def booleanImpl(c: Context)(bytes: c.Expr[ByteString], offset: c.Expr[Int]): c.Expr[Boolean] = {
    import c.universe._
    implicit val cc: c.type = c    val Literal(Constant(o: Int)) = offset.tree
    val what = nthElement(c)(bytes.tree, o / 8)
    val masked = Apply(Select(what, op("&")), const(bitMasks(o % 8)))
    c.Expr(Apply(Select(masked, op("!=")), const(0)))
  }

Basically, the macro expects to find the offset value at compile time, so it has to be a constant literal. The macro constructs the AST by building the following sub-trees:

• A tree to select the nth element within a byte string
• A tree to apply the bit AND operator with a given mask to the tree above.
• And finally a tree to apply a boolean unequal to the tree create above. 

The key thing is that instead of manipulating an expression directly, you manipulate ASTs, but otherwise, the logic is the same. 

Speed-wise the macro version is faster that its non-macro equivalent, and it is as fast as a hand written version, but the benefit (IMHO)being that the code is easy to read and write. 

I've posted the macro code on my github repository. Have a look at: https://github.com/kafecho/BitWrangler