From b52dab9fd769e996a9b734759d27c88985bc6e7b Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Tue, 15 Dec 2020 19:15:22 +0100
Subject: [PATCH] use stew/leb128 (#481)

* avoids multiple reallocations in readLp
* simplifies varint implementation
* remove vbuffer.length (unused)
---
 libp2p/multiaddress.nim    |   2 -
 libp2p/stream/lpstream.nim |  17 +--
 libp2p/varint.nim          | 246 ++++++++++++-------------------------
 libp2p/vbuffer.nim         |   1 -
 4 files changed, 89 insertions(+), 177 deletions(-)

diff --git a/libp2p/multiaddress.nim b/libp2p/multiaddress.nim
index f5a987371..456c86601 100644
--- a/libp2p/multiaddress.nim
+++ b/libp2p/multiaddress.nim
@@ -63,7 +63,6 @@ proc hash*(a: MultiAddress): Hash =
   var h: Hash = 0
   h = h !& hash(a.data.buffer)
   h = h !& hash(a.data.offset)
-  h = h !& hash(a.data.length)
   !$h
 
 proc ip4StB(s: string, vb: var VBuffer): bool =
@@ -458,7 +457,6 @@ proc trimRight(s: string, ch: char): string =
 proc shcopy*(m1: var MultiAddress, m2: MultiAddress) =
   shallowCopy(m1.data.buffer, m2.data.buffer)
   m1.data.offset = m2.data.offset
-  m1.data.length = m2.data.length
 
 proc protoCode*(ma: MultiAddress): MaResult[MultiCodec] =
   ## Returns MultiAddress ``ma`` protocol code.
diff --git a/libp2p/stream/lpstream.nim b/libp2p/stream/lpstream.nim
index 507efcaef..c433b3a38 100644
--- a/libp2p/stream/lpstream.nim
+++ b/libp2p/stream/lpstream.nim
@@ -11,7 +11,6 @@ import std/oids
 import stew/byteutils
 import chronicles, chronos, metrics
 import ../varint,
-       ../vbuffer,
        ../peerinfo,
        ../multiaddress
 
@@ -236,12 +235,16 @@ proc readLp*(s: LPStream, maxSize: int): Future[seq[byte]] {.async, gcsafe.} =
 method write*(s: LPStream, msg: seq[byte]): Future[void] {.base.} =
   doAssert(false, "not implemented!")
 
-proc writeLp*(s: LPStream, msg: string | seq[byte]): Future[void] {.gcsafe.} =
-  ## write length prefixed
-  var buf = initVBuffer()
-  buf.writeSeq(msg)
-  buf.finish()
-  s.write(buf.buffer)
+proc writeLp*(s: LPStream, msg: openArray[byte]): Future[void] =
+  ## Write `msg` with a varint-encoded length prefix
+  let vbytes = PB.toBytes(msg.len().uint64)
+  var buf = newSeqUninitialized[byte](msg.len() + vbytes.len)
+  buf[0..<vbytes.len] = vbytes.toOpenArray()
+  buf[vbytes.len..<buf.len] = msg
+  s.write(buf)
+
+proc writeLp*(s: LPStream, msg: string): Future[void] =
+  writeLp(s, msg.toOpenArrayByte(0, msg.high))
 
 proc write*(s: LPStream, pbytes: pointer, nbytes: int): Future[void] {.deprecated: "seq".} =
   s.write(@(toOpenArray(cast[ptr UncheckedArray[byte]](pbytes), 0, nbytes - 1)))
diff --git a/libp2p/varint.nim b/libp2p/varint.nim
index 5ace9fe30..7e10f3870 100644
--- a/libp2p/varint.nim
+++ b/libp2p/varint.nim
@@ -18,9 +18,8 @@
 
 {.push raises: [Defect].}
 
-import bitops, typetraits
-import stew/results
-export results
+import stew/[byteutils, leb128, results]
+export leb128, results
 
 type
   VarintError* {.pure.} = enum
@@ -56,50 +55,52 @@ type
   SomeVarint* = PBSomeVarint | LPSomeVarint
   SomeUVarint* = PBSomeUVarint | LPSomeUVarint
 
-proc vsizeof*(x: SomeUVarint): int {.inline.} =
+template toUleb(x: uint64): uint64 = x
+template toUleb(x: uint32): uint32 = x
+template toUleb(x: uint16): uint16 = x
+template toUleb(x: uint8): uint8 = x
+
+func toUleb(x: zint64): uint64 =
+  let v = cast[uint64](x)
+  (v shl 1) xor (0 - (v shr 63))
+
+func toUleb(x: zint32): uint32 =
+  let v = cast[uint32](x)
+  (v shl 1) xor (0 - (v shr 31))
+
+template toUleb(x: hint64): uint64 = cast[uint64](x)
+template toUleb(x: hint32): uint32 = cast[uint32](x)
+
+template toUleb(x: zint): uint64 =
+  when sizeof(x) == sizeof(zint64):
+    uint(toUleb(zint64(x)))
+  else:
+    uint(toUleb(zint32(x)))
+
+template toUleb(x: hint): uint =
+  when sizeof(x) == sizeof(hint64):
+    uint(toUleb(hint64(x)))
+  else:
+    uint(toUleb(hint32(x)))
+
+template fromUleb(x: uint64, T: type uint64): T = x
+template fromUleb(x: uint32, T: type uint32): T = x
+
+template fromUleb(x: uint64, T: type zint64): T =
+  cast[T]((x shr 1) xor (0 - (x and 1)))
+
+template fromUleb(x: uint32, T: type zint32): T =
+  cast[T]((x shr 1) xor (0 - (x and 1)))
+
+template fromUleb(x: uint64, T: type hint64): T =
+  cast[T](x)
+
+template fromUleb(x: uint32, T: type hint32): T =
+  cast[T](x)
+
+proc vsizeof*(x: SomeVarint): int {.inline.} =
   ## Returns number of bytes required to encode integer ``x`` as varint.
-  if x == type(x)(0):
-    1
-  else:
-    (fastLog2(x) + 1 + 7 - 1) div 7
-
-proc vsizeof*(x: PBSomeSVarint): int {.inline.} =
-  ## Returns number of bytes required to encode signed integer ``x``.
-  ##
-  ## Note: This procedure interprets signed integer as ProtoBuffer's
-  ## ``int32`` and ``int64`` integers.
-  when sizeof(x) == 8:
-    if int64(x) == 0'i64:
-      1
-    else:
-      (fastLog2(uint64(x)) + 1 + 7 - 1) div 7
-  else:
-    if int32(x) == 0'i32:
-      1
-    else:
-      (fastLog2(uint32(x)) + 1 + 7 - 1) div 7
-
-proc vsizeof*(x: PBZigVarint): int {.inline.} =
-  ## Returns number of bytes required to encode signed integer ``x``.
-  ##
-  ## Note: This procedure interprets signed integer as ProtoBuffer's
-  ## ``sint32`` and ``sint64`` integer.
-  when sizeof(x) == 8:
-    if int64(x) == 0'i64:
-      1
-    else:
-      if int64(x) < 0'i64:
-        vsizeof(not(uint64(x) shl 1))
-      else:
-        vsizeof(uint64(x) shl 1)
-  else:
-    if int32(x) == 0'i32:
-      1
-    else:
-      if int32(x) < 0'i32:
-        vsizeof(not(uint32(x) shl 1))
-      else:
-        vsizeof(uint32(x) shl 1)
+  Leb128.len(toUleb(x))
 
 proc getUVarint*[T: PB|LP](vtype: typedesc[T],
                            pbytes: openarray[byte],
@@ -126,39 +127,26 @@ proc getUVarint*[T: PB|LP](vtype: typedesc[T],
   ## LibP2P
   ## When decoding 5th byte of 32bit integer only 4 bits from byte will be
   ## decoded, all other bits will be ignored.
-  when vtype is PB:
-    const MaxBits = byte(sizeof(outval) * 8)
-  else:
-    when sizeof(outval) == 8:
-      const MaxBits = 63'u8
-    else:
-      const MaxBits = byte(sizeof(outval) * 8)
+  outlen = 0
+  outval = type(outval)(0)
 
-  var shift = 0'u8
-  outlen = 0
-  outval = type(outval)(0)
-  for i in 0..<len(pbytes):
-    let b = pbytes[i]
-    if shift >= MaxBits:
-      outlen = 0
-      outval = type(outval)(0)
+  let parsed = type(outval).fromBytes(pbytes, Leb128)
+
+  if parsed.len == 0:
+    return err(VarintError.Incomplete)
+  if parsed.len < 0:
+    return err(VarintError.Overflow)
+
+  when vtype is LP and sizeof(outval) == 8:
+    if parsed.val >= 0x8000_0000_0000_0000'u64:
       return err(VarintError.Overflow)
-    else:
-      outval = outval or (type(outval)(b and 0x7F'u8) shl shift)
-      shift += 7
-    inc(outlen)
-    if (b and 0x80'u8) == 0'u8:
-      # done, success
-      if outlen != vsizeof(outval):
-        outval = type(outval)(0)
-        outlen = 0
-        return err(VarintError.Overlong)
-      else:
-        return ok()
-  
-  outlen = 0
-  outval = type(outval)(0)
-  err(VarintError.Incomplete)
+
+  if vsizeof(parsed.val) != parsed.len:
+    return err(VarintError.Overlong)
+
+  (outval, outlen) = parsed
+
+  ok()
 
 proc putUVarint*[T: PB|LP](vtype: typedesc[T],
                            pbytes: var openarray[byte],
@@ -180,34 +168,20 @@ proc putUVarint*[T: PB|LP](vtype: typedesc[T],
   ## LibP2P
   ## Maximum encoded length of 63bit integer is 9 octets.
   ## Maximum encoded length of 32bit integer is 5 octets.
-  var buffer: array[10, byte]
-  var value = outval
-  var k = 0
+  when vtype is LP and sizeof(outval) == 8:
+    if (uint64(outval) and 0x8000_0000_0000_0000'u64) != 0'u64:
+      return err(VarintError.Overflow)
 
-  when vtype is LP:
-    if sizeof(outval) == 8:
-      if (uint64(outval) and 0x8000_0000_0000_0000'u64) != 0'u64:
-        return err(VarintError.Overflow)
-
-  if value <= type(outval)(0x7F):
-    buffer[0] = byte(outval and 0xFF)
-    inc(k)
-  else:
-    while value != type(outval)(0):
-      buffer[k] = byte((value and 0x7F) or 0x80)
-      value = value shr 7
-      inc(k)
-    buffer[k - 1] = buffer[k - 1] and 0x7F'u8
-
-  outlen = k
-  if len(pbytes) >= k:
-    copyMem(addr pbytes[0], addr buffer[0], k)
+  let bytes = toBytes(outval, Leb128)
+  outlen = len(bytes)
+  if len(pbytes) >= outlen:
+    pbytes[0..<outlen] = bytes.toOpenArray()
     ok()
   else:
     err(VarintError.Overrun)
 
 proc getSVarint*(pbytes: openarray[byte], outsize: var int,
-                 outval: var PBSomeSVarint): VarintResult[void] {.inline.} =
+                 outval: var (PBZigVarint | PBSomeSVarint)): VarintResult[void] {.inline.} =
   ## Decode signed integer (``int32`` or ``int64``) from buffer ``pbytes``
   ## and store it to ``outval``.
   ##
@@ -233,44 +207,11 @@ proc getSVarint*(pbytes: openarray[byte], outsize: var int,
 
   let res = PB.getUVarint(pbytes, outsize, value)
   if res.isOk():
-    outval = cast[type(outval)](value)
-  res
-
-proc getSVarint*(pbytes: openarray[byte], outsize: var int,
-                 outval: var PBZigVarint): VarintResult[void] {.inline.} =
-  ## Decode Google ProtoBuf's zigzag encoded signed integer (``sint32`` or
-  ## ``sint64`` ) from buffer ``pbytes`` and store it to ``outval``.
-  ##
-  ## On success ``outlen`` will be set to number of bytes processed while
-  ## decoding signed varint.
-  ##
-  ## If array ``pbytes`` is empty, ``Incomplete`` error will be returned.
-  ##
-  ## If there not enough bytes available in array ``pbytes`` to decode `signed
-  ## varint`, ``Incomplete`` error will be returned.
-  ##
-  ## If encoded value can produce integer overflow, ``Overflow`` error will be
-  ## returned.
-  ##
-  ## Note, when decoding 10th byte of 64bit integer only 1 bit from byte will be
-  ## decoded, all other bits will be ignored. When decoding 5th byte of 32bit
-  ## integer only 4 bits from byte will be decoded, all other bits will be
-  ## ignored.
-  when sizeof(outval) == 8:
-    var value: uint64
-  else:
-    var value: uint32
-
-  let res = PB.getUVarint(pbytes, outsize, value)
-  if res.isOk():
-    if (value and type(value)(1)) != type(value)(0):
-      outval = cast[type(outval)](not(value shr 1))
-    else:
-      outval = cast[type(outval)](value shr 1)
+    outval = fromUleb(value, type(outval))
   res
 
 proc putSVarint*(pbytes: var openarray[byte], outsize: var int,
-                 outval: PBZigVarint): VarintResult[void] {.inline.} =
+                 outval: (PBZigVarint | PBSomeSVarint)): VarintResult[void] {.inline.} =
   ## Encode signed integer ``outval`` using ProtoBuffer's zigzag encoding
   ## (``sint32`` or ``sint64``) and store it to array ``pbytes``.
   ##
@@ -283,38 +224,7 @@ proc putSVarint*(pbytes: var openarray[byte], outsize: var int,
   ##
   ## Maximum encoded length of 64bit integer is 10 octets.
   ## Maximum encoded length of 32bit integer is 5 octets.
-  when sizeof(outval) == 8:
-    var value: uint64 =
-      if int64(outval) < 0'i64:
-        not(uint64(outval) shl 1)
-      else:
-        uint64(outval) shl 1
-  else:
-    var value: uint32 =
-      if int32(outval) < 0'i32:
-        not(uint32(outval) shl 1)
-      else:
-        uint32(outval) shl 1
-  PB.putUVarint(pbytes, outsize, value)
-
-proc putSVarint*(pbytes: var openarray[byte], outsize: var int,
-                 outval: PBSomeSVarint): VarintResult[void] {.inline.} =
-  ## Encode signed integer ``outval`` (``int32`` or ``int64``) and store it to
-  ## array ``pbytes``.
-  ##
-  ## On success ``outlen`` will hold number of bytes (octets) used to encode
-  ## unsigned integer ``v``.
-  ##
-  ## If there not enough bytes available in buffer ``pbytes``, ``Incomplete``
-  ## error will be returned and ``outlen`` will be set to number of bytes
-  ## required.
-  ##
-  ## Maximum encoded length of 64bit integer is 10 octets.
-  ## Maximum encoded length of 32bit integer is 5 octets.
-  when sizeof(outval) == 8:
-    PB.putUVarint(pbytes, outsize, uint64(outval))
-  else:
-    PB.putUVarint(pbytes, outsize, uint32(outval))
+  PB.putUVarint(pbytes, outsize, toUleb(outval))
 
 template varintFatal(msg) =
   const m = msg
@@ -355,6 +265,9 @@ proc getVarint*[T: PB|LP](vtype: typedesc[T], pbytes: openarray[byte],
       varintFatal("LibP2P's varint do not support type [" &
                    typetraits.name(type(value)) & "]")
 
+template toBytes*(vtype: typedesc[PB], value: PBSomeVarint): auto =
+  toBytes(toUleb(value), Leb128)
+
 proc encodeVarint*(vtype: typedesc[PB],
                    value: PBSomeVarint): VarintResult[seq[byte]] {.inline.} =
   ## Encode integer to Google ProtoBuf's `signed/unsigned varint` and returns
@@ -374,7 +287,6 @@ proc encodeVarint*(vtype: typedesc[PB],
     ok(buffer)
   else:
     err(res.error())
-  
 
 proc encodeVarint*(vtype: typedesc[LP],
                    value: LPSomeVarint): VarintResult[seq[byte]] {.inline.} =
diff --git a/libp2p/vbuffer.nim b/libp2p/vbuffer.nim
index 6f315c7ba..4da8886ec 100644
--- a/libp2p/vbuffer.nim
+++ b/libp2p/vbuffer.nim
@@ -17,7 +17,6 @@ type
   VBuffer* = object
     buffer*: seq[byte]
     offset*: int
-    length*: int
 
 template isEmpty*(vb: VBuffer): bool =
   ## Returns ``true`` if buffer ``vb`` is empty.