From 810c8bf275069cbb21d8d3acbe2355a7c95815c9 Mon Sep 17 00:00:00 2001
From: Daira Hopwood <daira@jacaranda.org>
Date: Mon, 30 Aug 2021 19:42:51 +0100
Subject: [PATCH] Simplify parse_compact_size usage.

Signed-off-by: Daira Hopwood <daira@jacaranda.org>
---
 unified_address.py |  8 ++------
 zc_utils.py        | 12 ++++++------
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/unified_address.py b/unified_address.py
index 81165fd..a1dec6a 100644
--- a/unified_address.py
+++ b/unified_address.py
@@ -58,14 +58,10 @@ def decode_unified(addr_str):
     result = {}
     while len(rest) > 0:
         if s == 0:
-            (receiver_type, receiver_type_len) = parse_compact_size(rest)
-            rest = rest[receiver_type_len:]
-
+            (receiver_type, rest) = parse_compact_size(rest)
             s = 1
         elif s == 1:
-            (receiver_len, receiver_len_len) = parse_compact_size(rest)
-            rest = rest[receiver_len_len:]
-
+            (receiver_len, rest) = parse_compact_size(rest)
             expected_len = {0: 20, 1: 20, 2: 43, 3: 43}.get(receiver_type)
             if expected_len is not None:
                 assert receiver_len == expected_len, "incorrect receiver length"
diff --git a/zc_utils.py b/zc_utils.py
index b9ff34d..fcd1ca5 100644
--- a/zc_utils.py
+++ b/zc_utils.py
@@ -14,22 +14,22 @@ def parse_compact_size(rest):
     assert len(rest) >= 1
     b = rest[0]
     if b < 253:
-        return (b, 1)
+        return (b, rest[1:])
     elif b == 253:
         assert len(rest) >= 3
-        return (struct.unpack('<H', rest[1:3])[0], 3)
+        return (struct.unpack('<H', rest[1:3])[0], rest[3:])
     elif b == 254:
         assert len(rest) >= 5
-        return (struct.unpack('<I', rest[1:5])[0], 5)
+        return (struct.unpack('<I', rest[1:5])[0], rest[5:])
     else:
         assert len(rest) >= 9
-        return (struct.unpack('<Q', rest[1:9])[0], 9)
+        return (struct.unpack('<Q', rest[1:9])[0], rest[9:])
 
 
 def test_round_trip(n, encoding):
     assert write_compact_size(n) == encoding
-    assert parse_compact_size(encoding) == (n, len(encoding))
-    assert parse_compact_size(encoding + b'\x2a') == (n, len(encoding))
+    assert parse_compact_size(encoding) == (n, b'')
+    assert parse_compact_size(encoding + b'*') == (n, b'*')
     try:
         parse_compact_size(encoding[:-1])
     except AssertionError: