Simplify nibble handling

A bit unexpectedly, nibble handling shows up in the profiler mainly
because the current impl is tuned towards slicing while the most common
operation is prefix comparison - since the code is simple, might has
well get rid of some of the excess fat by always aliging the nibbles to
the byte buffer.

nimbus/db/aristo/aristo_desc/desc_nibbles.nim

@@ -8,16 +8,23 @@
 # at your option. This file may not be copied, modified, or distributed
 # except according to those terms.
 
-import stew/[arraybuf, arrayops]
+{.push raises: [], gcsafe, inline.}
+
+import stew/[arraybuf, arrayops, bitops2, endians2, staticfor]
 
 export arraybuf
 
 type
   NibblesBuf* = object
     ## Allocation-free type for storing up to 64 4-bit nibbles, as seen in the
     ## Ethereum MPT
-    bytes: array[32, byte]
-    ibegin, iend: int8
+    limbs: array[4, uint64]
+      # Each limb holds 16 nibbles in big endian order - for buffers shorter
+      # 64 nibbles we make sure the last limb holding any data is zero-padded
+      # (so as to avoid UB on uninitialized reads) - for example a buffer
+      # holding one nibble will have one fully initialized limb and 3
+      # uninitialized limbs.
+    iend: uint8
       # Where valid nibbles can be found - we use indices here to avoid copies
       # wen slicing - iend not inclusive
 
@@ -26,136 +33,200 @@ type
 func high*(T: type NibblesBuf): int =
   63
 
-func fromBytes*(T: type NibblesBuf, bytes: openArray[byte]): T =
-  result.iend = 2 * (int8 result.bytes.copyFrom(bytes))
-
 func nibble*(T: type NibblesBuf, nibble: byte): T =
-  result.bytes[0] = nibble shl 4
+  result.limbs[0] = uint64(nibble) shl (64 - 4)
   result.iend = 1
 
-template `[]`*(r: NibblesBuf, i: int): byte =
-  let pos = r.ibegin + i
-  if (pos and 1) != 0:
-    (r.bytes[pos shr 1] and 0xf)
+template limb(i: int | uint8): uint8 =
+  # In which limb can nibble i be found?
+  uint8(i) shr 4 # shr 4 = div 16 = 16 nibbles per limb
+
+template shift(i: int | uint8): int =
+  # How many bits to shift to find nibble i within its limb?
+  60 - ((i mod 16) shl 2) # shl 2 = 4 bits per nibble
+
+func `[]`*(r: NibblesBuf, i: int): byte =
+  let
+    ilimb = i.limb
+    ishift = i.shift
+  byte((r.limbs[ilimb] shr ishift) and 0x0f)
+
+func `[]=`*(r: var NibblesBuf, i: int, v: byte) =
+  let
+    ilimb = i.limb
+    ishift = i.shift
+
+  r.limbs[ilimb] =
+    (uint64(v and 0x0f) shl ishift) or ((r.limbs[ilimb] and not (0x0f'u64 shl ishift)))
+
+func fromBytes*(T: type NibblesBuf, bytes: openArray[byte]): T {.noinit.} =
+  if bytes.len >= 32:
+    result.iend = 64
+    staticFor i, 0 ..< result.limbs.len:
+      const pos = i * 8 # 16 nibbles per limb, 2 nibbles per byte
+      result.limbs[i] = uint64.fromBytesBE(bytes.toOpenArray(pos, pos + 7))
   else:
-    (r.bytes[pos shr 1] shr 4)
-
-template `[]=`*(r: NibblesBuf, i: int, v: byte) =
-  let pos = r.ibegin + i
-  r.bytes[pos shr 1] =
-    if (pos and 1) != 0:
-      (v and 0x0f) or (r.bytes[pos shr 1] and 0xf0)
-    else:
-      (v shl 4) or (r.bytes[pos shr 1] and 0x0f)
+    let blen = uint8(bytes.len)
+    result.iend = blen * 2
+
+    block done:
+      staticFor i, 0 ..< result.limbs.len:
+        const pos = i * 8
+        if pos + 7 < blen:
+          result.limbs[i] = uint64.fromBytesBE(bytes.toOpenArray(pos, pos + 7))
+        else:
+          if pos < blen:
+            var tmp = 0'u64
+            var shift = 56'u8
+            for j in uint8(pos) ..< blen:
+              tmp = tmp or uint64(bytes[j]) shl shift
+              shift -= 8
+
+            result.limbs[i] = tmp
+          break done
 
 func len*(r: NibblesBuf): int =
-  r.iend - r.ibegin
-
-func `==`*(lhs, rhs: NibblesBuf): bool =
-  if lhs.len == rhs.len:
-    for i in 0 ..< lhs.len:
-      if lhs[i] != rhs[i]:
-        return false
-    return true
-  else:
-    return false
+  int(r.iend)
 
 func `$`*(r: NibblesBuf): string =
   result = newStringOfCap(64)
   for i in 0 ..< r.len:
     const chars = "0123456789abcdef"
     result.add chars[r[i]]
 
+func `==`*(lhs, rhs: NibblesBuf): bool =
+  if lhs.iend != rhs.iend:
+    return false
+
+  let last = (lhs.iend + 15).limb # Last limb containing any data
+  staticFor i, 0 ..< lhs.limbs.len:
+    if uint8(i) < last and lhs.limbs[i] != rhs.limbs[i]:
+      return false
+  true
+
+func sharedPrefixLen*(lhs, rhs: NibblesBuf): int =
+  let len = min(lhs.iend, rhs.iend)
+  staticFor i, 0 ..< lhs.limbs.len:
+    const pos = i * 16
+
+    if (pos + 16) >= len or lhs.limbs[i] != rhs.limbs[i]:
+      return
+        if pos < len:
+          let mask =
+            if len - pos >= 16:
+              0'u64
+            else:
+              (not 0'u64) shr ((len - pos) * 4)
+          pos + leadingZeros((lhs.limbs[i] xor rhs.limbs[i]) or mask) shr 2
+        else:
+          pos
+
+  64
+
+func startsWith*(lhs, rhs: NibblesBuf): bool =
+  sharedPrefixLen(lhs, rhs) == rhs.len
+
 func slice*(r: NibblesBuf, ibegin: int, iend = -1): NibblesBuf {.noinit.} =
-  result.bytes = r.bytes
-  result.ibegin = r.ibegin + ibegin.int8
   let e =
     if iend < 0:
-      min(64, r.iend + iend + 1)
+      min(64, r.len + iend + 1)
     else:
-      min(64, r.ibegin + iend)
-  doAssert ibegin >= 0 and e <= result.bytes.len * 2
-  result.iend = e.int8
+      min(64, iend)
+
+  # With noinit, we have to be careful not to read result.bytes
+  result.iend = uint8(e - ibegin)
+
+  var ilimb = ibegin.limb
+  let shift = (ibegin mod 16) shl 2
+  block done:
+    staticFor i, 0 ..< result.limbs.len:
+      if uint8(i) >= (result.iend + 15).limb:
+        break done
+
+      var cur = r.limbs[ilimb]
+      ilimb += 1
+      var next =
+        if shift != 0 and ilimb < uint8 r.limbs.len:
+          r.limbs[ilimb]
+        else:
+          0'u64
+
+      result.limbs[i] = (cur shl shift) or (next shr (64 - shift))
 
 func replaceSuffix*(r: NibblesBuf, suffix: NibblesBuf): NibblesBuf =
-  for i in 0 ..< r.len - suffix.len:
-    result[i] = r[i]
+  result.limbs = r.limbs
+  # TODO unroll this loop!
   for i in 0 ..< suffix.len:
     result[i + r.len - suffix.len] = suffix[i]
-  result.iend = min(64, r.len + suffix.len).int8
+  result.iend = r.iend
 
-template writeFirstByte(nibbleCountExpr) {.dirty.} =
-  let nibbleCount = nibbleCountExpr
-  var oddnessFlag = (nibbleCount and 1) != 0
-  result.setLen((nibbleCount div 2) + 1)
-  result[0] = byte((int(isLeaf) * 2 + int(oddnessFlag)) shl 4)
-  var writeHead = 0
+func toHexPrefix*(r: NibblesBuf, isLeaf = false): HexPrefixBuf {.noinit.} =
+  # We'll adjust to the actual length below, but this hack allows us to write
+  # full limbs
 
-template writeNibbles(r) {.dirty.} =
-  for i in 0 ..< r.len:
-    let nextNibble = r[i]
-    if oddnessFlag:
-      result[writeHead] = result[writeHead] or nextNibble
-    else:
-      inc writeHead
-      result[writeHead] = nextNibble shl 4
-    oddnessFlag = not oddnessFlag
+  result.setLen(33)
+  let
+    limbs = (r.iend + 15).limb
+    isOdd = (r.iend and 1) > 0
 
-func toHexPrefix*(r: NibblesBuf, isLeaf = false): HexPrefixBuf =
-  writeFirstByte(r.len)
-  writeNibbles(r)
+  result[0] = (byte(isLeaf) * 2 + byte(isOdd)) shl 4
 
-func toHexPrefix*(r1, r2: NibblesBuf, isLeaf = false): HexPrefixBuf =
-  writeFirstByte(r1.len + r2.len)
-  writeNibbles(r1)
-  writeNibbles(r2)
+  if isOdd:
+    result[0] = result[0] or byte(r.limbs[0] shr 60)
 
-func sharedPrefixLen*(lhs, rhs: NibblesBuf): int =
-  result = 0
-  while result < lhs.len and result < rhs.len:
-    if lhs[result] != rhs[result]:
-      break
-    inc result
+    staticFor i, 0 ..< r.limbs.len:
+      if i < limbs:
+        let next =
+          when i == r.limbs.high:
+            0'u64
+          else:
+            r.limbs[i + 1]
+        let limb = r.limbs[i] shl 4 or next shr 60
 
-func startsWith*(lhs, rhs: NibblesBuf): bool =
-  sharedPrefixLen(lhs, rhs) == rhs.len
+        const pos = i * 8 + 1
+        assign(result.data.toOpenArray(pos, pos + 7), limb.toBytesBE())
+  else:
+    staticFor i, 0 ..< r.limbs.len:
+      if i < limbs:
+        let limb = r.limbs[i]
+        const pos = i * 8 + 1
+        assign(result.data.toOpenArray(pos, pos + 7), limb.toBytesBE())
+
+  result.setLen(int((r.iend shr 1) + 1))
 
 func fromHexPrefix*(
     T: type NibblesBuf, r: openArray[byte]
 ): tuple[isLeaf: bool, nibbles: NibblesBuf] {.noinit.} =
-  result.nibbles.ibegin = 0
-
   if r.len > 0:
     result.isLeaf = (r[0] and 0x20) != 0
     let hasOddLen = (r[0] and 0x10) != 0
 
-    result.nibbles.iend =
-      if hasOddLen:
-        result.nibbles.bytes[0] = r[0] shl 4
-
-        let bytes = min(31, r.len - 1)
-        for j in 0 ..< bytes:
-          result.nibbles.bytes[j] = result.nibbles.bytes[j] or r[j + 1] shr 4
-          result.nibbles.bytes[j + 1] = r[j + 1] shl 4
-
-        int8(bytes) * 2 + 1
-      else:
-        let bytes = min(32, r.len - 1)
-        assign(result.nibbles.bytes.toOpenArray(0, bytes - 1), r.toOpenArray(1, bytes))
-        int8(bytes) * 2
+    if hasOddLen:
+      let bytes = uint8(min(31, r.len - 1))
+      result.nibbles =
+        NibblesBuf.nibble(r[0] and 0x0f) &
+        NibblesBuf.fromBytes(r.toOpenArray(1, int bytes))
+    else:
+      result.nibbles = NibblesBuf.fromBytes(r.toOpenArray(1, r.high()))
   else:
     result.isLeaf = false
     result.nibbles.iend = 0
 
-func `&`*(a, b: NibblesBuf): NibblesBuf {.noinit.} =
-  result.ibegin = 0
-  for i in 0 ..< a.len:
-    result[i] = a[i]
+func `&`*(a, b: NibblesBuf): NibblesBuf =
+  result.iend = min(64'u8, a.iend + b.iend)
+
+  block adone:
+    staticFor i, 0 ..< result.limbs.len:
+      if uint8(i) >= ((a.iend + 15).limb):
+        break adone
 
+      result.limbs[i] = a.limbs[i]
+
+  # TODO unroll this loop too!
   for i in 0 ..< b.len:
     result[i + a.len] = b[i]
 
-  result.iend = int8(min(64, a.len + b.len))
-
-template getBytes*(a: NibblesBuf): array[32, byte] =
-  a.bytes
+func getBytes*(a: NibblesBuf): array[32, byte] =
+  staticFor i, 0 ..< a.limbs.len:
+    const pos = i * 8
+    assign(result.toOpenArray(pos, pos + 7), a.limbs[i].toBytesBE)

nimbus/db/aristo/aristo_fetch.nim

@@ -26,12 +26,9 @@ import
 proc retrieveLeaf(
     db: AristoDbRef;
     root: VertexID;
-    path: openArray[byte];
+    path: Hash32;
       ): Result[VertexRef,AristoError] =
-  if path.len == 0:
-    return err(FetchPathInvalid)
-
-  for step in stepUp(NibblesBuf.fromBytes(path), root, db):
+  for step in stepUp(NibblesBuf.fromBytes(path.data), root, db):
     let vtx = step.valueOr:
       if error in HikeAcceptableStopsNotFound:
         return err(FetchPathNotFound)
@@ -68,7 +65,7 @@ proc retrieveAccountLeaf(
   # Updated payloads are stored in the layers so if we didn't find them there,
   # it must have been in the database
   let
-    leafVtx = db.retrieveLeaf(VertexID(1), accPath.data).valueOr:
+    leafVtx = db.retrieveLeaf(VertexID(1), accPath).valueOr:
       if error == FetchPathNotFound:
         db.accLeaves.put(accPath, nil)
       return err(error)
@@ -168,7 +165,7 @@ proc retrieveStoragePayload(
 
   # Updated payloads are stored in the layers so if we didn't find them there,
   # it must have been in the database
-  let leafVtx = db.retrieveLeaf(? db.fetchStorageIdImpl(accPath), stoPath.data).valueOr:
+  let leafVtx = db.retrieveLeaf(? db.fetchStorageIdImpl(accPath), stoPath).valueOr:
     if error == FetchPathNotFound:
       db.stoLeaves.put(mixPath, nil)
     return err(error)

nimbus/db/aristo/aristo_merge.nim

@@ -41,7 +41,7 @@ proc layersPutLeaf(
 proc mergePayloadImpl(
     db: AristoDbRef, # Database, top layer
     root: VertexID, # MPT state root
-    path: openArray[byte], # Leaf item to add to the database
+    path: Hash32, # Leaf item to add to the database
     leaf: Opt[VertexRef],
     payload: LeafPayload, # Payload value
 ): Result[(VertexRef, VertexRef, VertexRef), AristoError] =
@@ -51,7 +51,7 @@ proc mergePayloadImpl(
   ## accordingly.
   ##
   var
-    path = NibblesBuf.fromBytes(path)
+    path = NibblesBuf.fromBytes(path.data)
     cur = root
     (vtx, _) = db.getVtxRc((root, cur)).valueOr:
       if error != GetVtxNotFound:
@@ -189,7 +189,7 @@ proc mergeAccountRecord*(
   let
     pyl =  LeafPayload(pType: AccountData, account: accRec)
     updated = db.mergePayloadImpl(
-        VertexID(1), accPath.data, db.cachedAccLeaf(accPath), pyl).valueOr:
+        VertexID(1), accPath, db.cachedAccLeaf(accPath), pyl).valueOr:
       if error == MergeNoAction:
         return ok false
       return err(error)
@@ -230,7 +230,7 @@ proc mergeStorageData*(
     # Call merge
     pyl = LeafPayload(pType: StoData, stoData: stoData)
     updated = db.mergePayloadImpl(
-        useID.vid, stoPath.data, db.cachedStoLeaf(mixPath), pyl).valueOr:
+        useID.vid, stoPath, db.cachedStoLeaf(mixPath), pyl).valueOr:
       if error == MergeNoAction:
         assert stoID.isValid         # debugging only
         return ok()