Convenience varints API going from int straight to openarray

Author: zah

stew/varints.nim

@@ -58,6 +58,17 @@ type
     when defined(debug):
       state: VarintState
 
+  VarintBuffer* = object
+    bytes*: array[10, byte]
+    totalBytesWritten*: int
+
+func append*(buf: var VarintBuffer, b: byte) =
+  buf.bytes[buf.totalBytesWritten] = b
+  inc buf.totalBytesWritten
+
+template writtenBytes*(buf: VarintBuffer): auto =
+  buf.bytes.toOpenArray(0, buf.totalBytesWritten - 1)
+
 func maxBits(T: type VarintParser): uint8 {.compileTime.} =
   when T.flavour == ProtoBuf:
     uint8(sizeof(T.IntType) * 8)
@@ -199,3 +210,12 @@ func vsizeof*(x: SomeInteger): int {.inline.} =
   if x == 0: 1
   else: (log2trunc(x) + 1 + 7 - 1) div 7
 
+template varintBytes*(x: SomeInteger,
+                      flavour: static VarintFlavour = ProtoBuf): untyped =
+  var buf: VarintBuffer
+  buf.appendVarint(x, flavour)
+  # TODO: toOpenArray doesn't work here for some reason, so we must
+  # use the less optimal approach of allocating a sequence copy.
+  # buf.bytes.toOpenArray(0, buf.totalBytesWritten - 1)
+  buf.bytes[0 .. buf.totalBytesWritten - 1]
+

tests/test_varints.nim

@@ -25,42 +25,31 @@ const edgeValues = {
   0xFFFF_FFFF_FFFF_FFFF'u64 : "ffffffffffffffffff01"
 }
 
-type
-  PseudoStream = object
-    bytes: array[12, byte]
-    bytesWritten: int
-
-func append(s: var PseudoStream, b: byte) =
-  s.bytes[s.bytesWritten] = b
-  inc s.bytesWritten
-
-template writtenData(s: PseudoStream): auto =
-  s.bytes.toOpenArray(0, s.bytesWritten - 1)
-
 suite "varints":
   template roundtipTest(val) =
-    var s {.inject.}: PseudoStream
+    var s {.inject.}: VarintBuffer
     s.appendVarint val
 
     var roundtripVal: uint64
     let bytesRead = readVarint(s.bytes, roundtripVal)
 
     check:
       val == roundtripVal
-      bytesRead == s.bytesWritten
+      bytesRead == s.totalBytesWritten
       bytesRead == vsizeof(val)
 
   test "[ProtoBuf] Success edge cases test":
     for pair in edgeValues:
       let (val, hex) = pair
       roundtipTest val
       check:
-        s.bytesWritten == hex.len div 2
-        toHex(s.writtenData) == hex
+        s.totalBytesWritten == hex.len div 2
+        toHex(s.writtenBytes) == hex
+        toHex(val.varintBytes) == hex
 
   test "[ProtoBuf] random values":
-      for i in 0..10000:
-        let val = rand(0'u64 .. 0xFFFF_FFFF_FFFF_FFFE'u64)
-        roundtipTest val
+    for i in 0..10000:
+      let val = rand(0'u64 .. 0xFFFF_FFFF_FFFF_FFFE'u64)
+      roundtipTest val
 
   # TODO Migrate the rest of the LibP2P test cases