From 5f1dc751ca436e599c59df939344a641f935dd4d Mon Sep 17 00:00:00 2001
From: Zahary Karadjov <zahary@gmail.com>
Date: Sun, 8 Sep 2019 22:33:03 -0400
Subject: [PATCH] Convenience varints API going from int straight to openarray

---
 stew/varints.nim       | 20 ++++++++++++++++++++
 tests/test_varints.nim | 27 ++++++++-------------------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/stew/varints.nim b/stew/varints.nim
index d1f7d4d..f3c7385 100644
--- a/stew/varints.nim
+++ b/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]
+
diff --git a/tests/test_varints.nim b/tests/test_varints.nim
index bfa983a..303fd42 100644
--- a/tests/test_varints.nim
+++ b/tests/test_varints.nim
@@ -25,21 +25,9 @@ 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
@@ -47,7 +35,7 @@ suite "varints":
 
     check:
       val == roundtripVal
-      bytesRead == s.bytesWritten
+      bytesRead == s.totalBytesWritten
       bytesRead == vsizeof(val)
 
   test "[ProtoBuf] Success edge cases test":
@@ -55,12 +43,13 @@ suite "varints":
       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