From a1df4badabe7e80cb18917ce93a7514f09f99389 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Mon, 12 Sep 2022 14:30:31 +0200
Subject: [PATCH] Support distinct types

---
 contractabi/decoding.nim           | 4 ++++
 contractabi/encoding.nim           | 5 +++++
 contractabi/selector.nim           | 4 ++++
 tests/contractabi/testDecoding.nim | 6 ++++++
 tests/contractabi/testEncoding.nim | 5 +++++
 tests/contractabi/testSelector.nim | 4 ++++
 6 files changed, 28 insertions(+)

diff --git a/contractabi/decoding.nim b/contractabi/decoding.nim
index 7283d94..5931f25 100644
--- a/contractabi/decoding.nim
+++ b/contractabi/decoding.nim
@@ -1,3 +1,4 @@
+import std/typetraits
 import pkg/stint
 import pkg/stew/endians2
 import pkg/stew/byteutils
@@ -159,6 +160,9 @@ func decode[I,T](decoder: var AbiDecoder, _: type array[I,T]): ?!array[I,T] =
 func decode(decoder: var AbiDecoder, T: type string): ?!T =
   success string.fromBytes(?decoder.read(seq[byte]))
 
+func decode[T: distinct](decoder: var AbiDecoder, _: type T): ?!T =
+  success T(?decoder.read(distinctBase T))
+
 func readOffset(decoder: var AbiDecoder): ?!int =
   let offset = ?decoder.read(uint64)
   success decoder.currentTuple.start + offset.int
diff --git a/contractabi/encoding.nim b/contractabi/encoding.nim
index 88dd7c8..416a366 100644
--- a/contractabi/encoding.nim
+++ b/contractabi/encoding.nim
@@ -1,3 +1,4 @@
+import std/typetraits
 import pkg/stint
 import pkg/upraises
 import pkg/stew/byteutils
@@ -131,6 +132,10 @@ func encode(encoder: var AbiEncoder, tupl: tuple) =
     encoder.write(element)
   encoder.finishTuple()
 
+func encode(encoder: var AbiEncoder, value: distinct) =
+  type Base = distinctBase(typeof(value))
+  encoder.write(Base(value))
+
 func finish(encoder: var AbiEncoder): Tuple =
   doAssert encoder.stack.len == 1, "not all tuples were finished"
   doAssert encoder.stack[0].bytes.len mod 32 == 0, "encoding invariant broken"
diff --git a/contractabi/selector.nim b/contractabi/selector.nim
index f8fdbe8..f015640 100644
--- a/contractabi/selector.nim
+++ b/contractabi/selector.nim
@@ -1,4 +1,5 @@
 import std/strutils
+import std/typetraits
 import pkg/nimcrypto
 import pkg/stint
 import pkg/stew/byteutils
@@ -60,6 +61,9 @@ func solidityType*(Tuple: type tuple): string =
     names.add(solidityType(typeof parameter))
   "(" & names.join(",") & ")"
 
+func solidityType*(T: distinct type): string =
+  solidityType(distinctBase T)
+
 func signature*(function: string, Parameters: type tuple = ()): string =
   function & solidityType(Parameters)
 
diff --git a/tests/contractabi/testDecoding.nim b/tests/contractabi/testDecoding.nim
index 8db1023..92d4e1b 100644
--- a/tests/contractabi/testDecoding.nim
+++ b/tests/contractabi/testDecoding.nim
@@ -3,6 +3,9 @@ import pkg/questionable/results
 import contractabi
 import ./examples
 
+type SomeDistinctType = distinct uint16
+func `==`*(a, b: SomeDistinctType): bool {.borrow.}
+
 suite "ABI decoding":
 
   proc checkDecode[T](value: T) =
@@ -140,3 +143,6 @@ suite "ABI decoding":
 
   test "decodes strings":
     checkDecode("hello!☺")
+
+  test "decodes distinct types as their base type":
+    checkDecode(SomeDistinctType(0xAABB'u16))
diff --git a/tests/contractabi/testEncoding.nim b/tests/contractabi/testEncoding.nim
index 9f7d844..cd6f520 100644
--- a/tests/contractabi/testEncoding.nim
+++ b/tests/contractabi/testEncoding.nim
@@ -138,6 +138,11 @@ suite "ABI encoding":
   test "encodes strings as UTF-8 byte sequence":
     check AbiEncoder.encode("hello!☺") == AbiEncoder.encode("hello!☺".toBytes)
 
+  test "encodes distinct types as their base type":
+    type SomeDistinctType = distinct uint16
+    let value = 0xAABB'u16
+    check AbiEncoder.encode(SomeDistinctType(value)) == AbiEncoder.encode(value)
+
   test "can determine whether types are dynamic or static":
     check static AbiEncoder.isStatic(uint8)
     check static AbiEncoder.isDynamic(seq[byte])
diff --git a/tests/contractabi/testSelector.nim b/tests/contractabi/testSelector.nim
index a329daa..d84f499 100644
--- a/tests/contractabi/testSelector.nim
+++ b/tests/contractabi/testSelector.nim
@@ -35,6 +35,10 @@ suite "function selector":
     check solidityType((Address, string, bool)) == "(address,string,bool)"
     check solidityType(SomeEnum) == "uint8"
 
+  test "translates distinct type into base type":
+    type SomeDistinctType = distinct uint16
+    check solidityType(SomeDistinctType) == solidityType(uint16)
+
   test "calculates solidity function selector":
     check $selector("transfer", (Address, UInt256)) == "0xa9059cbb"
     check $selector("transferFrom", (Address, Address, UInt256)) == "0x23b872dd"