From 7da36bf4598997516bcbffb37595b7bbd31df232 Mon Sep 17 00:00:00 2001
From: web3-developer <51288821+web3-developer@users.noreply.github.com>
Date: Fri, 3 May 2024 22:55:51 +0800
Subject: [PATCH] Minor refactor to packNibbles and unpackNibbles and add
 additional tests. (#2162)

---
 fluffy/network/state/state_content.nim        | 15 +++--
 .../test_state_content_keys.nim               | 58 +++++++++++++++++++
 2 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/fluffy/network/state/state_content.nim b/fluffy/network/state/state_content.nim
index 8a1e62e6d..7614be7cd 100644
--- a/fluffy/network/state/state_content.nim
+++ b/fluffy/network/state/state_content.nim
@@ -214,18 +214,17 @@ func packNibbles*(nibbles: seq[byte]): Nibbles =
   if nibbles.len() == 0:
     return Nibbles(@[byte(0x00)])
 
-  let isOddLength = (nibbles.len() %% 2 == 1)
+  let isEvenLength = nibbles.len() mod 2 == 0
 
   var
-    output = newSeq[byte]()
-    highNibble = not isOddLength
+    output = newSeqOfCap[byte](nibbles.len() div 2 + 1)
+    highNibble = isEvenLength
     currentByte: byte = 0
 
-  if isOddLength:
-    currentByte = 0x10
-
-  if not isOddLength:
+  if isEvenLength:
     output.add(0x00)
+  else:
+    currentByte = 0x10
 
   for i, nibble in nibbles:
     if highNibble:
@@ -240,7 +239,7 @@ func packNibbles*(nibbles: seq[byte]): Nibbles =
 func unpackNibbles*(nibbles: Nibbles): seq[byte] =
   doAssert(nibbles.len() <= MAX_PACKED_NIBBLES_LEN, "Packed nibbles length is too long")
 
-  var output = newSeq[byte]()
+  var output = newSeqOfCap[byte](nibbles.len * 2)
 
   for i, pair in nibbles:
     if i == 0 and pair == 0x00:
diff --git a/fluffy/tests/state_network_tests/test_state_content_keys.nim b/fluffy/tests/state_network_tests/test_state_content_keys.nim
index ab26dc9de..2164bb094 100644
--- a/fluffy/tests/state_network_tests/test_state_content_keys.nim
+++ b/fluffy/tests/state_network_tests/test_state_content_keys.nim
@@ -14,6 +14,45 @@ import
 const testVectorDir = "./vendor/portal-spec-tests/tests/mainnet/state/serialization/"
 
 suite "State Content Keys":
+  test "Encode/decode empty nibbles":
+    const
+      expected = "00"
+      nibbles: seq[byte] = @[]
+      packedNibbles = packNibbles(nibbles)
+      unpackedNibbles = unpackNibbles(packedNibbles)
+
+    let encoded = SSZ.encode(packedNibbles)
+
+    check:
+      encoded.toHex() == expected
+      unpackedNibbles == nibbles
+
+  test "Encode/decode zero nibble":
+    const
+      expected = "10"
+      nibbles: seq[byte] = @[0]
+      packedNibbles = packNibbles(nibbles)
+      unpackedNibbles = unpackNibbles(packedNibbles)
+
+    let encoded = SSZ.encode(packedNibbles)
+
+    check:
+      encoded.toHex() == expected
+      unpackedNibbles == nibbles
+
+  test "Encode/decode one nibble":
+    const
+      expected = "11"
+      nibbles: seq[byte] = @[1]
+      packedNibbles = packNibbles(nibbles)
+      unpackedNibbles = unpackNibbles(packedNibbles)
+
+    let encoded = SSZ.encode(packedNibbles)
+
+    check:
+      encoded.toHex() == expected
+      unpackedNibbles == nibbles
+
   test "Encode/decode even nibbles":
     const
       expected = "008679e8ed"
@@ -40,6 +79,25 @@ suite "State Content Keys":
       encoded.toHex() == expected
       unpackedNibbles == nibbles
 
+  test "Encode/decode max length nibbles":
+    const
+      expected = "008679e8eda65bd257638cf8cf09b8238888947cc3c0bea2aa2cc3f1c4ac7a3002"
+      nibbles: seq[byte] =
+        @[
+          8, 6, 7, 9, 0xe, 8, 0xe, 0xd, 0xa, 6, 5, 0xb, 0xd, 2, 5, 7, 6, 3, 8, 0xc, 0xf,
+          8, 0xc, 0xf, 0, 9, 0xb, 8, 2, 3, 8, 8, 8, 8, 9, 4, 7, 0xc, 0xc, 3, 0xc, 0,
+          0xb, 0xe, 0xa, 2, 0xa, 0xa, 2, 0xc, 0xc, 3, 0xf, 1, 0xc, 4, 0xa, 0xc, 7, 0xa,
+          3, 0, 0, 2,
+        ]
+      packedNibbles = packNibbles(nibbles)
+      unpackedNibbles = unpackNibbles(packedNibbles)
+
+    let encoded = SSZ.encode(packedNibbles)
+
+    check:
+      encoded.toHex() == expected
+      unpackedNibbles == nibbles
+
   test "Encode/decode AccountTrieNodeKey":
     const file = testVectorDir & "account_trie_node_key.yaml"