Support preset types in StarknetTypedData

Sources/Starknet/Data/TypedData/StarknetTypedData.swift

@@ -11,6 +11,7 @@ import Foundation
 public enum StarknetTypedDataError: Error, Equatable {
     case decodingError
     case invalidRevision(Felt)
+    case presetTypeRedefinition(String)
     case basicTypeRedefinition(String)
     case invalidTypeName(String)
     case danglingType(String)
@@ -73,15 +74,15 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
     public let domain: Domain
     public let message: [String: Element]
 
-    public var revision: Revision {
-        try! domain.resolveRevision()
-    }
+    public let revision: Revision
+    private let allTypes: [String: [TypeDeclarationWrapper]]
+    private let hashMethod: StarknetHashMethod
 
-    private var hashMethod: StarknetHashMethod {
-        switch revision {
-        case .v0: .pedersen
-        case .v1: .poseidon
-        }
+    fileprivate enum CodingKeys: CodingKey {
+        case types
+        case primaryType
+        case domain
+        case message
     }
 
     private func hashArray(_ values: [Felt]) -> Felt {
@@ -96,6 +97,18 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
         self.domain = domain
         self.message = message
 
+        self.revision = try domain.resolveRevision()
+
+        self.allTypes = self.types.merging(
+            Self.getPresetTypes(revision: self.revision),
+            uniquingKeysWith: { current, _ in current }
+        )
+
+        self.hashMethod = switch revision {
+        case .v0: .pedersen
+        case .v1: .poseidon
+        }
+
         try self.verifyTypes()
     }
 
@@ -128,7 +141,8 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
             throw StarknetTypedDataError.dependencyNotDefined(domain.separatorName)
         }
 
-        let basicTypes = getBasicTypes()
+        let basicTypes = Self.getBasicTypes(revision: revision)
+        let presetTypes = Self.getPresetTypes(revision: revision)
 
         let referencedTypes = try Set(types.values.flatMap { type in
             try type.flatMap { param in
@@ -154,6 +168,9 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
             guard !basicTypes.contains(typeName) else {
                 throw StarknetTypedDataError.basicTypeRedefinition(typeName)
             }
+            guard presetTypes[typeName] == nil else {
+                throw StarknetTypedDataError.presetTypeRedefinition(typeName)
+            }
             guard !typeName.isEmpty,
                   !typeName.isArray(),
                   !typeName.isEnum(),
@@ -193,13 +210,13 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
 
         while !toVisit.isEmpty {
             let currentType = toVisit.removeFirst()
-            let params = types[currentType] ?? []
+            let params = allTypes[currentType] ?? []
 
             try params.forEach { param in
                 let extractedTypes = try extractTypes(from: param).map { $0.strippingPointer() }
 
                 extractedTypes.forEach { extractedType in
-                    if types.keys.contains(extractedType), !dependencies.contains(extractedType) {
+                    if allTypes.keys.contains(extractedType), !dependencies.contains(extractedType) {
                         dependencies.append(extractedType)
                         toVisit.append(extractedType)
                     }
@@ -239,7 +256,7 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
             return "(\(enumTypes))"
         }
 
-        guard let params = types[dependency] else {
+        guard let params = allTypes[dependency] else {
             throw StarknetTypedDataError.dependencyNotDefined(dependency)
         }
 
@@ -265,7 +282,7 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
     }
 
     func encode(element: Element, forType typeName: String, context: Context? = nil) throws -> Felt {
-        if types.keys.contains(typeName) {
+        if allTypes.keys.contains(typeName) {
             let object = try unwrapObject(from: element)
 
             return try getStructHash(typeName: typeName, data: object)
@@ -312,7 +329,7 @@ public struct StarknetTypedData: Codable, Equatable, Hashable {
     private func encode(data: [String: Element], forType typeName: String) throws -> [Felt] {
         var values: [Felt] = []
 
-        guard let types = types[typeName] else {
+        guard let types = allTypes[typeName] else {
             throw StarknetTypedDataError.encodingError
         }
 
@@ -473,13 +490,36 @@ private extension StarknetTypedData {
     static let basicTypesV0: Set = ["felt", "bool", "string", "selector", "merkletree"]
     static let basicTypesV1: Set = basicTypesV0.union(["enum", "u128", "i128", "ContractAddress", "ClassHash", "timestamp", "shortstring"])
 
-    func getBasicTypes() -> Set<String> {
+    static let presetTypesV1 = [
+        "u256": [
+            StandardType(name: "low", type: "u128"),
+            StandardType(name: "high", type: "u128"),
+        ],
+        "TokenAmount": [
+            StandardType(name: "token_address", type: "ContractAddress"),
+            StandardType(name: "amount", type: "u256"),
+        ],
+        "NftId": [
+            StandardType(name: "collection_address", type: "ContractAddress"),
+            StandardType(name: "token_id", type: "u256"),
+        ],
+    ]
+
+    static func getBasicTypes(revision: Revision) -> Set<String> {
         switch revision {
         case .v0:
-            Self.basicTypesV0
+            basicTypesV0
         case .v1:
-            Self.basicTypesV1
+            basicTypesV1
+        }
+    }
+
+    static func getPresetTypes(revision: Revision) -> [String: [TypeDeclarationWrapper]] {
+        let types: [String: [any TypeDeclaration]] = switch revision {
+        case .v0: [:]
+        case .v1: Self.presetTypesV1
         }
+        return types.mapValues { $0.map { TypeDeclarationWrapper($0) } }
     }
 }
 
@@ -621,7 +661,7 @@ extension StarknetTypedData {
     private func getEnumVariants(context: Context) throws -> [TypeDeclarationWrapper] {
         let enumType: EnumType = try resolveType(context)
 
-        guard let variants = types[enumType.contains] else {
+        guard let variants = allTypes[enumType.contains] else {
             throw StarknetTypedDataError.dependencyNotDefined(enumType.contains)
         }
 
@@ -652,7 +692,7 @@ extension StarknetTypedData {
     private func resolveType<T: TypeDeclaration>(_ context: Context) throws -> T {
         let (parent, key) = (context.parent, context.key)
 
-        guard let parentType = types[parent] else {
+        guard let parentType = allTypes[parent] else {
             throw StarknetTypedDataError.parentNotDefined
         }
         guard let targetType = parentType.first(where: { $0.type.name == key }) else {

Tests/StarknetTests/Data/TypedDataTests.swift

@@ -22,6 +22,7 @@ final class TypedDataTests: XCTestCase {
         static let td = try! loadTypedDataFromFile(name: "typed_data_rev_1_example")
         static let tdFeltMerkleTree = try! loadTypedDataFromFile(name: "typed_data_rev_1_felt_merkletree_example")
         static let tdBasicTypes = try! loadTypedDataFromFile(name: "typed_data_rev_1_basic_types_example")
+        static let tdPresetTypes = try! loadTypedDataFromFile(name: "typed_data_rev_1_preset_types_example")
         static let tdEnum = try! loadTypedDataFromFile(name: "typed_data_rev_1_enum_example")
     }
 
@@ -95,23 +96,32 @@ final class TypedDataTests: XCTestCase {
     }
 
     func testTypesRedifintion() throws {
-        func testTypeRedifintion(_ type: String, _ revision: StarknetTypedData.Revision) throws {
+        func testBasicTypeRedifintion(_ type: String, _ revision: StarknetTypedData.Revision) throws {
             try XCTAssertThrowsError(makeTypedData(type, revision)) { error in
                 XCTAssertEqual(error as? StarknetTypedDataError, .basicTypeRedefinition(type))
             }
         }
+        func testPresetTypeRedifintion(_ type: String, _ revision: StarknetTypedData.Revision) throws {
+            try XCTAssertThrowsError(makeTypedData(type, revision)) { error in
+                XCTAssertEqual(error as? StarknetTypedDataError, .presetTypeRedefinition(type))
+            }
+        }
 
         let basicTypesV0 = [
             "felt", "bool", "string", "selector", "merkletree",
         ]
         let basicTypesV1 = basicTypesV0 + ["enum", "u128", "i128", "ContractAddress", "ClassHash", "timestamp", "shortstring"]
+        let presetTypesV1 = ["u256", "TokenAmount", "NftId"]
 
         try XCTAssertNoThrow(makeTypedData("myType", .v0))
         try basicTypesV0.forEach { type in
-            try testTypeRedifintion(type, .v0)
+            try testBasicTypeRedifintion(type, .v0)
         }
         try basicTypesV1.forEach { type in
-            try testTypeRedifintion(type, .v1)
+            try testBasicTypeRedifintion(type, .v1)
+        }
+        try presetTypesV1.forEach { type in
+            try testPresetTypeRedifintion(type, .v1)
         }
     }
 
@@ -287,6 +297,7 @@ final class TypedDataTests: XCTestCase {
             (Self.CasesRev1.td, "Person", "0x30f7aa21b8d67cb04c30f962dd29b95ab320cb929c07d1605f5ace304dadf34"),
             (Self.CasesRev1.td, "Mail", "0x560430bf7a02939edd1a5c104e7b7a55bbab9f35928b1cf5c7c97de3a907bd"),
             (Self.CasesRev1.tdBasicTypes, "Example", "0x1f94cd0be8b4097a41486170fdf09a4cd23aefbc74bb2344718562994c2c111"),
+            (Self.CasesRev1.tdPresetTypes, "Example", "0x1a25a8bb84b761090b1fadaebe762c4b679b0d8883d2bedda695ea340839a55"),
             (Self.CasesRev1.tdEnum, "Example", "0x380a54d417fb58913b904675d94a8a62e2abc3467f4b5439de0fd65fafdd1a8"),
             (Self.CasesRev1.tdFeltMerkleTree, "Example", "0x160b9c0e8a7c561f9c5d9e3cc2990a1b4d26e94aa319e9eb53e163cd06c71be"),
         ]
@@ -349,6 +360,12 @@ final class TypedDataTests: XCTestCase {
                 "message",
                 "0x391d09a51a31dd17f7270aaa9904688fbeeb9c56a7e2d15c5a6af32e981c730"
             ),
+            (
+                Self.CasesRev1.tdPresetTypes,
+                "Example",
+                "message",
+                "0x74fba3f77f8a6111a9315bac313bf75ecfa46d1234e0fda60312fb6a6517667"
+            ),
             (
                 Self.CasesRev1.tdEnum,
                 "Example",
@@ -416,6 +433,11 @@ final class TypedDataTests: XCTestCase {
                 "0xcd2a3d9f938e13cd947ec05abc7fe734df8dd826",
                 "0x2d80b87b8bc32068247c779b2ef0f15f65c9c449325e44a9df480fb01eb43ec"
             ),
+            (
+                Self.CasesRev1.tdPresetTypes,
+                "0xcd2a3d9f938e13cd947ec05abc7fe734df8dd826",
+                "0x185b339d5c566a883561a88fb36da301051e2c0225deb325c91bb7aa2f3473a"
+            ),
             (
                 Self.CasesRev1.tdEnum,
                 "0xcd2a3d9f938e13cd947ec05abc7fe734df8dd826",

Tests/StarknetTests/Resources/TypedData/typed_data_rev_1_preset_types_example.json

@@ -0,0 +1,37 @@
+{
+    "types": {
+        "StarknetDomain": [
+            { "name": "name", "type": "shortstring" },
+            { "name": "version", "type": "shortstring" },
+            { "name": "chainId", "type": "shortstring" },
+            { "name": "revision", "type": "shortstring" }
+        ],
+        "Example": [
+            { "name": "n0", "type": "TokenAmount" },
+            { "name": "n1", "type": "NftId" }
+        ]
+    },
+    "primaryType": "Example",
+    "domain": {
+        "name": "StarkNet Mail",
+        "version": "1",
+        "chainId": "1",
+        "revision": "1"
+    },
+    "message": {
+        "n0": {
+            "token_address": "0x049d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7",
+            "amount": {
+                "low": "0x3e8",
+                "high": "0x0"
+            }
+        },
+        "n1": {
+            "collection_address": "0x049d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7",
+            "token_id": {
+                "low": "0x3e8",
+                "high": "0x0"
+            }
+        }
+    }
+}