Refactor abi encoding and decoding

Author: mohanson

example/storage.py

@@ -28,9 +28,7 @@
 
 if args.action == 'set':
     user = pleth.wallet.Wallet(int(args.prikey, 0))
-    data = pleth.abi.function_selector('set', ['uint256']) + pleth.abi.argument_encoding([
-        pleth.abi.Uint256.encode(42),
-    ])
+    data = pleth.abi.function_selector('set', ['uint256']) + pleth.abi.argument_encoding([pleth.abi.Uint256], [42])
     hash = user.contract_exec(bytearray.fromhex(args.addr[2:]), 0, data)
     print(f'hash = 0x{hash.hex()}')
 

pleth/abi.py

@@ -15,11 +15,12 @@ def function_selector(name: str, args_type: list[str]) -> bytearray:
     return pleth.core.hash(bytearray(s.encode()))[:4]
 
 
-def argument_encoding(data: list[bytearray]) -> bytearray:
-    s = bytearray()
-    for e in data:
-        s.extend(e)
-    return s
+def argument_decoding(args_type: list, data: bytearray) -> list:
+    return Tuple(args_type).decode(io.BytesIO(data))
+
+
+def argument_encoding(args_type: list, args: list) -> bytearray:
+    return Tuple(args_type).encode(args)
 
 
 class Address:
@@ -32,6 +33,10 @@ def encode(cls, origin: bytearray) -> bytearray:
         assert len(origin) == 20
         return bytearray(12) + origin
 
+    @classmethod
+    def size(cls) -> int:
+        return 32
+
 
 class Bytes:
     @classmethod
@@ -46,6 +51,10 @@ def encode(cls, origin: bytearray) -> bytearray:
         padded = (length + 31) & -32
         return Uint256.encode(length) + origin + bytearray(padded - length)
 
+    @classmethod
+    def size(cls) -> int:
+        return 0
+
 
 class String:
     @classmethod
@@ -56,6 +65,53 @@ def decode(cls, reader: io.IOBase) -> str:
     def encode(cls, origin: str) -> bytearray:
         return Bytes.encode(bytearray(origin.encode()))
 
+    @classmethod
+    def size(cls) -> int:
+        return 0
+
+
+class Tuple:
+    def __init__(self, kype: list) -> None:
+        self.kype = kype
+        self.slen = 0
+        size_list = [k.size() for k in kype]
+        if all([s != 0 for s in size_list]):
+            self.slen = sum(size_list)
+
+    def decode(self, reader: io.IOBase) -> list:
+        muts = []
+        vals = []
+        for k in self.kype:
+            match k.size():
+                case 0:
+                    muts.append(len(vals))
+                    vals.append(Uint256.decode(reader))
+                case _:
+                    vals.append(k.decode(reader))
+        for i in muts:
+            vals[i] = self.kype[i].decode(reader)
+        return vals
+
+    def encode(self, origin: list) -> bytearray:
+        assert len(origin) == len(self.kype)
+        offs = 0
+        for k in self.kype:
+            offs += 32 if k.size() == 0 else k.size()
+        head = bytearray()
+        tail = bytearray()
+        for k, v in zip(self.kype, origin):
+            e = k.encode(v)
+            match k.size():
+                case 0:
+                    head.extend(Uint256.encode(offs + len(tail)))
+                    tail.extend(e)
+                case _:
+                    head.extend(e)
+        return head + tail
+
+    def size(self) -> int:
+        return self.slen
+
 
 class Uint256:
     @classmethod
@@ -67,3 +123,7 @@ def encode(cls, origin: int) -> bytearray:
         assert origin >= 0
         assert origin <= 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
         return bytearray(origin.to_bytes(32, 'big'))
+
+    @classmethod
+    def size(cls) -> int:
+        return 32

test/test_abi.py

@@ -0,0 +1,26 @@
+import io
+import pleth
+
+
+def test_tuple():
+    k = pleth.abi.Tuple([
+        pleth.abi.Uint256,
+        pleth.abi.Tuple([pleth.abi.Uint256, pleth.abi.Uint256]),
+        pleth.abi.String,
+        pleth.abi.String,
+    ])
+    v = [0x123, [0x456, 0x789], '1234567890', 'Hello, world!']
+    assert k.decode(io.BytesIO(bytes(k.encode(v)))) == v
+
+
+def test_tuple_erc20_name():
+    k = pleth.abi.Tuple([pleth.abi.String])
+    v = bytearray([
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x08,
+        0x55, 0x53, 0x44, 0x20, 0x43, 0x6f, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    ])
+    assert k.decode(io.BytesIO(v))[0] == 'USD Coin'