Function References Draft

Package.swift

@@ -70,7 +70,7 @@ let package = Package(
             ],
             exclude: ["CMakeLists.txt"]
         ),
-        .testTarget(name: "WasmParserTests", dependencies: ["WasmParser"]),
+        .testTarget(name: "WasmParserTests", dependencies: ["WasmParser", "WAT"]),
 
         .target(name: "WasmTypes", exclude: ["CMakeLists.txt"]),
 

README.md

@@ -70,6 +70,8 @@ and should work on the following platforms:
 |                      | [Memory64](https://github.com/WebAssembly/memory64/blob/main/proposals/memory64/Overview.md) | ✅ Implemented |
 |                      | [Tail call](https://github.com/WebAssembly/tail-call/blob/master/proposals/tail-call/Overview.md) | ✅ Implemented |
 |                      | [Threads and atomics](https://github.com/WebAssembly/threads/blob/master/proposals/threads/Overview.md) | 🚧 Parser implemented |
+|                      | [Typed Function References](https://github.com/WebAssembly/function-references/blob/main/proposals/function-references/Overview.md) | 📋 Todo  |
+|                      | [Garbage Collection](https://github.com/WebAssembly/gc/blob/main/proposals/gc/Overview.md) |  📋 Todo |
 | WASI | WASI Preview 1 | ✅ Implemented |
 
 

Sources/WAT/BinaryInstructionEncoder.swift

@@ -23,8 +23,9 @@ protocol BinaryInstructionEncoder: InstructionVisitor {
     mutating func encodeImmediates(relativeDepth: UInt32) throws
     mutating func encodeImmediates(table: UInt32) throws
     mutating func encodeImmediates(targets: BrTable) throws
-    mutating func encodeImmediates(type: ReferenceType) throws
+    mutating func encodeImmediates(type: HeapType) throws
     mutating func encodeImmediates(type: ValueType) throws
+    mutating func encodeImmediates(typeIndex: UInt32) throws
     mutating func encodeImmediates(value: IEEE754.Float32) throws
     mutating func encodeImmediates(value: IEEE754.Float64) throws
     mutating func encodeImmediates(value: Int32) throws
@@ -82,6 +83,14 @@ extension BinaryInstructionEncoder {
         try encodeInstruction([0x13])
         try encodeImmediates(typeIndex: typeIndex, tableIndex: tableIndex)
     }
+    mutating func visitCallRef(typeIndex: UInt32) throws {
+        try encodeInstruction([0x14])
+        try encodeImmediates(typeIndex: typeIndex)
+    }
+    mutating func visitReturnCallRef(typeIndex: UInt32) throws {
+        try encodeInstruction([0x15])
+        try encodeImmediates(typeIndex: typeIndex)
+    }
     mutating func visitDrop() throws { try encodeInstruction([0x1A]) }
     mutating func visitSelect() throws { try encodeInstruction([0x1B]) }
     mutating func visitTypedSelect(type: ValueType) throws {
@@ -171,7 +180,7 @@ extension BinaryInstructionEncoder {
         try encodeInstruction([0x44])
         try encodeImmediates(value: value)
     }
-    mutating func visitRefNull(type: ReferenceType) throws {
+    mutating func visitRefNull(type: HeapType) throws {
         try encodeInstruction([0xD0])
         try encodeImmediates(type: type)
     }
@@ -180,6 +189,15 @@ extension BinaryInstructionEncoder {
         try encodeInstruction([0xD2])
         try encodeImmediates(functionIndex: functionIndex)
     }
+    mutating func visitRefAsNonNull() throws { try encodeInstruction([0xD4]) }
+    mutating func visitBrOnNull(relativeDepth: UInt32) throws {
+        try encodeInstruction([0xD5])
+        try encodeImmediates(relativeDepth: relativeDepth)
+    }
+    mutating func visitBrOnNonNull(relativeDepth: UInt32) throws {
+        try encodeInstruction([0xD6])
+        try encodeImmediates(relativeDepth: relativeDepth)
+    }
     mutating func visitI32Eqz() throws { try encodeInstruction([0x45]) }
     mutating func visitCmp(_ cmp: Instruction.Cmp) throws {
         let opcode: [UInt8]

Sources/WAT/Encoder.swift

@@ -120,16 +120,33 @@ extension ValueType: WasmEncodable {
         case .f32: encoder.output.append(0x7D)
         case .f64: encoder.output.append(0x7C)
         case .v128: encoder.output.append(0x7B)
-        case .ref(let refType): refType.encode(to: &encoder)
+        case .ref(let refType): encoder.encode(refType)
         }
     }
 }
 
 extension ReferenceType: WasmEncodable {
+    func encode(to encoder: inout Encoder) {
+        switch (isNullable, heapType) {
+        // Use short form when available
+        case (true, .externRef): encoder.output.append(0x6F)
+        case (true, .funcRef): encoder.output.append(0x70)
+        default:
+            encoder.output.append(isNullable ? 0x63 : 0x64)
+            encoder.encode(heapType)
+        }
+    }
+}
+
+extension HeapType: WasmEncodable {
     func encode(to encoder: inout Encoder) {
         switch self {
-        case .funcRef: encoder.output.append(0x70)
-        case .externRef: encoder.output.append(0x6F)
+        case .abstract(.externRef): encoder.output.append(0x6F)
+        case .abstract(.funcRef): encoder.output.append(0x70)
+        case .concrete(let typeIndex):
+            // Note that the typeIndex is decoded as s33,
+            // so we need to encode it as signed.
+            encoder.writeSignedLEB128(Int64(typeIndex))
         }
     }
 }
@@ -194,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?
@@ -233,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
@@ -252,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))
@@ -324,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")
@@ -496,6 +515,7 @@ struct ExpressionEncoder: BinaryInstructionEncoder {
     mutating func encodeImmediates(functionIndex: UInt32) throws { encodeUnsigned(functionIndex) }
     mutating func encodeImmediates(globalIndex: UInt32) throws { encodeUnsigned(globalIndex) }
     mutating func encodeImmediates(localIndex: UInt32) throws { encodeUnsigned(localIndex) }
+    mutating func encodeImmediates(typeIndex: UInt32) throws { encodeUnsigned(typeIndex) }
     mutating func encodeImmediates(memarg: WasmParser.MemArg) throws {
         encodeUnsigned(UInt(memarg.align))
         encodeUnsigned(memarg.offset)
@@ -510,7 +530,7 @@ struct ExpressionEncoder: BinaryInstructionEncoder {
         encodeUnsigned(targets.defaultIndex)
     }
     mutating func encodeImmediates(type: WasmTypes.ValueType) throws { encoder.encode(type) }
-    mutating func encodeImmediates(type: WasmTypes.ReferenceType) throws { encoder.encode(type) }
+    mutating func encodeImmediates(type: WasmTypes.HeapType) throws { encoder.encode(type) }
     mutating func encodeImmediates(value: Int32) throws { encodeSigned(value) }
     mutating func encodeImmediates(value: Int64) throws { encodeSigned(value) }
     mutating func encodeImmediates(value: WasmParser.IEEE754.Float32) throws { encodeFixedWidth(value.bitPattern) }
@@ -557,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
@@ -604,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/ParseTextInstruction.swift

@@ -49,6 +49,12 @@ func parseTextInstruction<V: InstructionVisitor>(keyword: String, expressionPars
     case "return_call_indirect":
         let (typeIndex, tableIndex) = try expressionParser.visitReturnCallIndirect(wat: &wat)
         return { return try $0.visitReturnCallIndirect(typeIndex: typeIndex, tableIndex: tableIndex) }
+    case "call_ref":
+        let (typeIndex) = try expressionParser.visitCallRef(wat: &wat)
+        return { return try $0.visitCallRef(typeIndex: typeIndex) }
+    case "return_call_ref":
+        let (typeIndex) = try expressionParser.visitReturnCallRef(wat: &wat)
+        return { return try $0.visitReturnCallRef(typeIndex: typeIndex) }
     case "drop": return { return try $0.visitDrop() }
     case "select": return { return try $0.visitSelect() }
     case "local.get":
@@ -160,6 +166,13 @@ func parseTextInstruction<V: InstructionVisitor>(keyword: String, expressionPars
     case "ref.func":
         let (functionIndex) = try expressionParser.visitRefFunc(wat: &wat)
         return { return try $0.visitRefFunc(functionIndex: functionIndex) }
+    case "ref.as_non_null": return { return try $0.visitRefAsNonNull() }
+    case "br_on_null":
+        let (relativeDepth) = try expressionParser.visitBrOnNull(wat: &wat)
+        return { return try $0.visitBrOnNull(relativeDepth: relativeDepth) }
+    case "br_on_non_null":
+        let (relativeDepth) = try expressionParser.visitBrOnNonNull(wat: &wat)
+        return { return try $0.visitBrOnNonNull(relativeDepth: relativeDepth) }
     case "i32.eqz": return { return try $0.visitI32Eqz() }
     case "i32.eq": return { return try $0.visitCmp(.i32Eq) }
     case "i32.ne": return { return try $0.visitCmp(.i32Ne) }

Sources/WAT/Parser.swift

@@ -111,6 +111,8 @@ internal struct Parser {
         return String(decoding: bytes, as: UTF8.self)
     }
 
+    /// Parse a `u32` raw index or a symbolic `id` identifier.
+    /// https://webassembly.github.io/function-references/core/text/modules.html#indices
     mutating func takeIndexOrId() throws -> IndexOrId? {
         let location = lexer.location()
         if let index: UInt32 = try takeUnsignedInt() {