From f477eb6877da4fb039e680a8fb16158ef34f97d9 Mon Sep 17 00:00:00 2001
From: Zahary Karadjov <zahary@gmail.com>
Date: Fri, 16 Aug 2019 18:50:04 +0200
Subject: [PATCH] Add a `sszList` helper for writing tests involving SSZ lists

---
 beacon_chain/ssz.nim |  7 +++++--
 tests/test_ssz.nim   | 16 ++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/beacon_chain/ssz.nim b/beacon_chain/ssz.nim
index 1c865e6e4..3337009d7 100644
--- a/beacon_chain/ssz.nim
+++ b/beacon_chain/ssz.nim
@@ -48,7 +48,7 @@ type
 
   Chunk = array[bytesPerChunk, byte]
 
-  TypeWithMaxLen[T; maxLen: static int64] = distinct T
+  TypeWithMaxLen*[T; maxLen: static int64] = distinct T
 
   SizePrefixed*[T] = distinct T
   SszMaxSizeExceeded* = object of SerializationError
@@ -84,13 +84,16 @@ when false:
   # TODO: Nim can't handle yet this simpler definition. File an issue.
   template valueOf[T; N](x: TypeWithMaxLen[T, N]): auto = T(x)
 else:
-  proc unwrapImpl[T; N](x: ptr TypeWithMaxLen[T, N]): ptr T =
+  proc unwrapImpl[T; N: static int64](x: ptr TypeWithMaxLen[T, N]): ptr T =
     cast[ptr T](x)
 
   template valueOf(x: TypeWithMaxLen): auto =
     let xaddr = unsafeAddr x
     unwrapImpl(xaddr)[]
 
+template sszList*(x: seq|array, maxLen: static int64): auto =
+  TypeWithMaxLen[type(x), maxLen](x)
+
 template toSszType*(x: auto): auto =
   mixin toSszType
 
diff --git a/tests/test_ssz.nim b/tests/test_ssz.nim
index 869a5a03d..e3108ab9f 100644
--- a/tests/test_ssz.nim
+++ b/tests/test_ssz.nim
@@ -68,6 +68,9 @@ type
   Baz = object
     i: uint64
 
+proc toDigest[N: static int](x: array[N, byte]): Eth2Digest =
+  result.data[0 .. N-1] = x
+
 suite "SSZ Navigation":
   test "simple object fields":
     var foo = Foo(bar: Bar(b: "bar", baz: Baz(i: 10'u64)))
@@ -81,3 +84,16 @@ suite "SSZ Navigation":
     let mountedBar = mountedFoo.bar
     check mountedBar.baz.i == 10'u64
 
+  test "lists with max size":
+    let a = [byte 0x01, 0x02, 0x03].toDigest
+    let b = [byte 0x04, 0x05, 0x06].toDigest
+    let c = [byte 0x07, 0x08, 0x09].toDigest
+
+    let leaves = sszList(@[a, b, c], int64(1 shl 3))
+    let root = hashTreeRoot(leaves)
+    check $root == "5248085B588FAB1DD1E03F3CD62201602B12E6560665935964F46E805977E8C5"
+
+    let leaves2 = sszList(@[a, b, c], int64(1 shl 10))
+    let root2 = hashTreeRoot(leaves2)
+    check $root2 == "9FB7D518368DC14E8CC588FB3FD2749BEEF9F493FEF70AE34AF5721543C67173"
+