WAT: Add concrete type support

Author: kateinoigakukun

Sources/WAT/Encoder.swift

@@ -211,15 +211,16 @@ struct ElementExprCollector: AnyInstructionVisitor {
 
 extension WAT.WatParser.ElementDecl {
     func encode(to encoder: inout Encoder, wat: inout Wat) throws {
-        func isMemory64(tableIndex: Int) -> Bool {
+        func isMemory64(tableIndex: Int) throws -> Bool {
             guard tableIndex < wat.tablesMap.count else { return false }
-            return wat.tablesMap[tableIndex].type.limits.isMemory64
+            return try wat.tablesMap[tableIndex].type.resolve(wat.types).limits.isMemory64
         }
 
         var flags: UInt32 = 0
         var tableIndex: UInt32? = nil
         var isPassive = false
         var hasTableIndex = false
+        let type = try type.resolve(wat.types)
         switch self.mode {
         case .active(let table, _):
             let index: Int?
@@ -250,7 +251,7 @@ extension WAT.WatParser.ElementDecl {
         try collector.parse(indices: indices, wat: &wat)
         var useExpression: Bool {
             // if all instructions are ref.func, use function indices representation
-            return !collector.isAllRefFunc || self.type != .funcRef
+            return !collector.isAllRefFunc || type != .funcRef
         }
         if useExpression {
             // use expression
@@ -269,7 +270,7 @@ extension WAT.WatParser.ElementDecl {
                 try encoder.writeInstruction(lexer: &lexer, wat: &wat)
             case .synthesized(let offset):
                 var exprEncoder = ExpressionEncoder()
-                if isMemory64(tableIndex: Int(tableIndex ?? 0)) {
+                if try isMemory64(tableIndex: Int(tableIndex ?? 0)) {
                     try exprEncoder.visitI64Const(value: Int64(offset))
                 } else {
                     try exprEncoder.visitI32Const(value: Int32(offset))
@@ -341,6 +342,7 @@ extension Export: WasmEncodable {
 
 extension WatParser.GlobalDecl {
     func encode(to encoder: inout Encoder, wat: inout Wat) throws {
+        let type = try self.type.resolve(wat.types)
         encoder.encode(type)
         guard case var .definition(expr) = kind else {
             fatalError("imported global declaration should not be encoded here")
@@ -575,10 +577,11 @@ func encode(module: inout Wat, options: EncodeOptions) throws -> [UInt8] {
                     // Encode locals
                     var localsEntries: [(type: ValueType, count: UInt32)] = []
                     for local in locals {
-                        if localsEntries.last?.type == local.type {
+                        let type = try local.type.resolve(module.types)
+                        if localsEntries.last?.type == type {
                             localsEntries[localsEntries.count - 1].count += 1
                         } else {
-                            localsEntries.append((type: local.type, count: 1))
+                            localsEntries.append((type: type, count: 1))
                         }
                     }
                     exprEncoder.encoder.encodeVector(localsEntries) { local, encoder in
@@ -622,9 +625,9 @@ func encode(module: inout Wat, options: EncodeOptions) throws -> [UInt8] {
     // Section 4: Table section
     let tables = module.tablesMap.definitions()
     if !tables.isEmpty {
-        encoder.section(id: 0x04) { encoder in
-            encoder.encodeVector(tables) { table, encoder in
-                table.type.encode(to: &encoder)
+        try encoder.section(id: 0x04) { encoder in
+            try encoder.encodeVector(tables) { table, encoder in
+                try table.type.resolve(module.types).encode(to: &encoder)
             }
         }
     }

Sources/WAT/NameMapping.swift

@@ -21,8 +21,12 @@ protocol ImportableModuleFieldDecl {
     var importNames: WatParser.ImportNames? { get }
 }
 
+protocol NameToIndexResolver {
+    func resolveIndex(use: Parser.IndexOrId) throws -> Int
+}
+
 /// A map of module field declarations indexed by their name
-struct NameMapping<Decl: NamedModuleFieldDecl> {
+struct NameMapping<Decl: NamedModuleFieldDecl>: NameToIndexResolver {
     private var decls: [Decl] = []
     private var nameToIndex: [String: Int] = [:]
 
@@ -94,15 +98,21 @@ extension NameMapping where Decl: ImportableModuleFieldDecl {
     }
 }
 
+typealias TypesNameMapping = NameMapping<TypesMap.NamedResolvedType>
+
 /// A map of unique function types indexed by their name or type signature
 struct TypesMap {
-    private var nameMapping = NameMapping<WatParser.FunctionTypeDecl>()
+    struct NamedResolvedType: NamedModuleFieldDecl {
+        let id: Name?
+        let type: WatParser.FunctionType
+    }
+    private(set) var nameMapping = NameMapping<NamedResolvedType>()
     /// Tracks the earliest index for each function type
     private var indices: [FunctionType: Int] = [:]
 
     /// Adds a new function type to the mapping
     @discardableResult
-    mutating func add(_ decl: WatParser.FunctionTypeDecl) throws -> Int {
+    mutating func add(_ decl: NamedResolvedType) throws -> Int {
         try nameMapping.add(decl)
         // Normalize the function type signature without parameter names
         if let existing = indices[decl.type.signature] {
@@ -120,7 +130,7 @@ struct TypesMap {
             return existing
         }
         return try add(
-            WatParser.FunctionTypeDecl(
+            NamedResolvedType(
                 id: nil,
                 type: WatParser.FunctionType(signature: signature, parameterNames: [])
             )
@@ -170,10 +180,10 @@ struct TypesMap {
     mutating func resolveBlockType(use: WatParser.TypeUse) throws -> BlockType {
         switch (use.index, use.inline) {
         case let (indexOrId?, inline):
-            let (type, index) = try resolveAndCheck(use: indexOrId, inline: inline)
+            let (type, index) = try resolveAndCheck(use: indexOrId, inline: inline?.resolve(nameMapping))
             return try resolveBlockType(signature: type.signature, resolveSignatureIndex: { _ in index })
         case (nil, let inline?):
-            return try resolveBlockType(signature: inline.signature)
+            return try resolveBlockType(signature: inline.resolve(nameMapping).signature)
         case (nil, nil): return .empty
         }
     }
@@ -183,7 +193,7 @@ struct TypesMap {
         case let (indexOrId?, _):
             return try nameMapping.resolveIndex(use: indexOrId)
         case (nil, let inline):
-            let inline = inline?.signature ?? WasmTypes.FunctionType(parameters: [], results: [])
+            let inline = try inline?.resolve(nameMapping).signature ?? WasmTypes.FunctionType(parameters: [], results: [])
             return try addAnonymousSignature(inline)
         }
     }
@@ -209,16 +219,16 @@ struct TypesMap {
     mutating func resolve(use: WatParser.TypeUse) throws -> (type: WatParser.FunctionType, index: Int) {
         switch (use.index, use.inline) {
         case let (indexOrId?, inline):
-            return try resolveAndCheck(use: indexOrId, inline: inline)
+            return try resolveAndCheck(use: indexOrId, inline: inline?.resolve(nameMapping))
         case (nil, let inline):
             // If no index and no inline type, then it's a function type with no parameters or results
-            let inline = inline ?? WatParser.FunctionType(signature: WasmTypes.FunctionType(parameters: [], results: []), parameterNames: [])
+            let inline = try inline?.resolve(nameMapping) ?? WatParser.FunctionType(signature: WasmTypes.FunctionType(parameters: [], results: []), parameterNames: [])
             // Check if the inline type already exists
             if let index = indices[inline.signature] {
                 return (inline, index)
             }
             // Add inline type to the index space if it doesn't already exist
-            let index = try add(WatParser.FunctionTypeDecl(id: nil, type: inline))
+            let index = try add(NamedResolvedType(id: nil, type: inline))
             return (inline, index)
         }
     }
@@ -233,11 +243,11 @@ extension TypesMap: Collection {
         nameMapping.index(after: i)
     }
 
-    subscript(position: Int) -> WatParser.FunctionTypeDecl {
+    subscript(position: Int) -> NamedResolvedType {
         return nameMapping[position]
     }
 
-    func makeIterator() -> NameMapping<WatParser.FunctionTypeDecl>.Iterator {
+    func makeIterator() -> NameMapping<NamedResolvedType>.Iterator {
         return nameMapping.makeIterator()
     }
 }

Sources/WAT/Parser/ExpressionParser.swift

@@ -2,7 +2,7 @@ import WasmParser
 import WasmTypes
 
 struct ExpressionParser<Visitor: InstructionVisitor> {
-    typealias LocalsMap = NameMapping<WatParser.LocalDecl>
+    typealias LocalsMap = NameMapping<WatParser.ResolvedLocalDecl>
     private struct LabelStack {
         private var stack: [String?] = []
 
@@ -43,10 +43,11 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
         type: WatParser.FunctionType,
         locals: [WatParser.LocalDecl],
         lexer: Lexer,
-        features: WasmFeatureSet
+        features: WasmFeatureSet,
+        typeMap: TypesNameMapping
     ) throws {
         self.parser = Parser(lexer)
-        self.locals = try Self.computeLocals(type: type, locals: locals)
+        self.locals = try Self.computeLocals(type: type, locals: locals, typeMap: typeMap)
         self.features = features
     }
 
@@ -56,13 +57,17 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
         self.features = features
     }
 
-    static func computeLocals(type: WatParser.FunctionType, locals: [WatParser.LocalDecl]) throws -> LocalsMap {
+    static func computeLocals(
+        type: WatParser.FunctionType,
+        locals: [WatParser.LocalDecl],
+        typeMap: TypesNameMapping
+    ) throws -> LocalsMap {
         var localsMap = LocalsMap()
         for (name, type) in zip(type.parameterNames, type.signature.parameters) {
-            try localsMap.add(WatParser.LocalDecl(id: name, type: type))
+            try localsMap.add(.init(id: name, type: type))
         }
         for local in locals {
-            try localsMap.add(local)
+            try localsMap.add(local.resolve(typeMap))
         }
         return localsMap
     }
@@ -261,8 +266,10 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
         case "select":
             // Special handling for "select", which have two variants 1. with type, 2. without type
             let results = try withWatParser({ try $0.results() })
+            let types = wat.types
             return { visitor in
                 if let type = results.first {
+                    let type = try type.resolve(types)
                     return try visitor.visitTypedSelect(type: type)
                 } else {
                     return try visitor.visitSelect()
@@ -338,7 +345,9 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
     }
 
     private mutating func blockType(wat: inout Wat) throws -> BlockType {
-        let results = try withWatParser({ try $0.results() })
+        let results = try withWatParser {
+            try $0.results().map { try $0.resolve(wat.types) }
+        }
         if !results.isEmpty {
             return try wat.types.resolveBlockType(results: results)
         }