Fix encoding of heap types and leave concrete heap types unimplemented

Author: kateinoigakukun

Sources/WAT/BinaryInstructionEncoder.swift

@@ -23,7 +23,7 @@ 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
@@ -172,7 +172,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)
     }

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 self.heapType {
-        case .funcRef: encoder.output.append(isNullable ? 0x70 : 0x71)
-        case .externRef: encoder.output.append(0x6F)
+        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 .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))
         }
     }
 }
@@ -511,7 +528,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) }

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() {

Sources/WAT/Parser/ExpressionParser.swift

@@ -361,15 +361,17 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
         return UInt32(index)
     }
 
-    private mutating func refKind() throws -> ReferenceType {
-        if try parser.take(.id) {
-            return .funcRef  // not sure about this.
-        } else if try parser.takeKeyword("func") {
+    /// https://webassembly.github.io/function-references/core/text/types.html#text-heaptype
+    private mutating func heapType(wat: inout Wat) throws -> HeapType {
+        if try parser.takeKeyword("func") {
             return .funcRef
         } else if try parser.takeKeyword("extern") {
             return .externRef
+        } else if let id = try parser.takeIndexOrId() {
+            let (decl, index) = try wat.types.resolve(use: id)
+            return .concrete(typeIndex: UInt32(index))
         }
-        throw WatParserError("expected \"func\" or \"extern\"", location: parser.lexer.location())
+        throw WatParserError("expected \"func\", \"extern\" or type index", location: parser.lexer.location())
     }
 
     private mutating func memArg(defaultAlign: UInt32) throws -> MemArg {
@@ -513,8 +515,8 @@ extension ExpressionParser {
     mutating func visitF64Const(wat: inout Wat) throws -> IEEE754.Float64 {
         return try parser.expectFloat64()
     }
-    mutating func visitRefNull(wat: inout Wat) throws -> ReferenceType {
-        return try refKind()
+    mutating func visitRefNull(wat: inout Wat) throws -> HeapType {
+        return try heapType(wat: &wat)
     }
     mutating func visitRefFunc(wat: inout Wat) throws -> UInt32 {
         return try functionIndex(wat: &wat)

Sources/WAT/Parser/WastParser.swift

@@ -54,31 +54,25 @@ struct WastParser {
     }
 
     struct ConstExpressionCollector: WastConstInstructionVisitor {
-        let addValue: (Value) -> Void
+        let addValue: (WastConstValue) -> Void
 
         mutating func visitI32Const(value: Int32) throws { addValue(.i32(UInt32(bitPattern: value))) }
         mutating func visitI64Const(value: Int64) throws { addValue(.i64(UInt64(bitPattern: value))) }
         mutating func visitF32Const(value: IEEE754.Float32) throws { addValue(.f32(value.bitPattern)) }
         mutating func visitF64Const(value: IEEE754.Float64) throws { addValue(.f64(value.bitPattern)) }
         mutating func visitRefFunc(functionIndex: UInt32) throws {
-            addValue(.ref(.function(FunctionAddress(functionIndex))))
+            addValue(.refFunc(functionIndex: functionIndex))
         }
-        mutating func visitRefNull(type: ReferenceType) throws {
-            let value: Reference
-            switch type.heapType {
-            case .externRef: value = .extern(nil)
-            case .funcRef: value = .function(nil)
-            }
-            addValue(.ref(value))
+        mutating func visitRefNull(type: HeapType) throws {
+            addValue(.refNull(type))
         }
-
-        mutating func visitRefExtern(value: UInt32) throws {
-            addValue(.ref(.extern(ExternAddress(value))))
+        func visitRefExtern(value: UInt32) throws {
+            addValue(.refExtern(value: value))
         }
     }
 
-    mutating func constExpression() throws -> [Value] {
-        var values: [Value] = []
+    mutating func argumentValues() throws -> [WastConstValue] {
+        var values: [WastConstValue] = []
         var collector = ConstExpressionCollector(addValue: { values.append($0) })
         var exprParser = ExpressionParser<ConstExpressionCollector>(lexer: parser.lexer, features: features)
         while try exprParser.parseWastConstInstruction(visitor: &collector) {}
@@ -137,16 +131,26 @@ public enum WastExecute {
     }
 }
 
+public enum WastConstValue {
+    case i32(UInt32)
+    case i64(UInt64)
+    case f32(UInt32)
+    case f64(UInt64)
+    case refNull(HeapType)
+    case refFunc(functionIndex: UInt32)
+    case refExtern(value: UInt32)
+}
+
 public struct WastInvoke {
     public let module: String?
     public let name: String
-    public let args: [Value]
+    public let args: [WastConstValue]
 
     static func parse(wastParser: inout WastParser) throws -> WastInvoke {
         try wastParser.parser.expectKeyword("invoke")
         let module = try wastParser.parser.takeId()
         let name = try wastParser.parser.expectString()
-        let args = try wastParser.constExpression()
+        let args = try wastParser.argumentValues()
         try wastParser.parser.expect(.rightParen)
         let invoke = WastInvoke(module: module?.value, name: name, args: args)
         return invoke
@@ -155,7 +159,7 @@ public struct WastInvoke {
 
 public enum WastExpectValue {
     /// A concrete value that is expected to be returned.
-    case value(Value)
+    case value(WastConstValue)
     /// A value that is expected to be a canonical NaN.
     /// Corresponds to `f32.const nan:canonical` in WAST.
     case f32CanonicalNaN

Sources/WasmKit/Execution/Errors.swift

@@ -139,6 +139,9 @@ extension TrapReason.Message {
     static func exportedFunctionNotFound(name: String, instance: Instance) -> Self {
         Self("exported function \(name) not found in instance \(instance)")
     }
+    static func unimplemented(feature: String) -> Self {
+        Self("\(feature) is not implemented yet")
+    }
 }
 
 struct ImportError: Error {

Sources/WasmKit/Execution/Instances.swift

@@ -275,10 +275,12 @@ struct TableEntity /* : ~Copyable */ {
     init(_ tableType: TableType, resourceLimiter: any ResourceLimiter) throws {
         let emptyElement: Reference
         switch tableType.elementType.heapType {
-        case .funcRef:
+        case .abstract(.funcRef):
             emptyElement = .function(nil)
-        case .externRef:
+        case .abstract(.externRef):
             emptyElement = .extern(nil)
+        case .concrete:
+            throw Trap(.unimplemented(feature: "heap type other than `func` and `extern`"))
         }
 
         let numberOfElements = Int(tableType.limits.min)

Sources/WasmKit/Execution/Instructions/InstructionSupport.swift

@@ -104,12 +104,12 @@ extension Int32: InstructionImmediate {
 // MARK: - Immediate type extensions
 
 extension Instruction.RefNullOperand {
-    init(result: VReg, type: ReferenceType) {
-        self.init(result: result, rawType: type.heapType.rawValue)  // need to figure out rawType here
+    init(result: VReg, type: AbstractHeapType) {
+        self.init(result: result, rawType: type.rawValue)  // need to figure out rawType here
     }
 
-    var type: ReferenceType {
-        ReferenceType(isNullable: true, heapType: HeapType(rawValue: rawType)!)  // is this still a valid conversion?
+    var type: AbstractHeapType {
+        AbstractHeapType(rawValue: rawType).unsafelyUnwrapped
     }
 }
 

Sources/WasmKit/Execution/Instructions/Misc.swift

@@ -21,7 +21,7 @@ extension Execution {
 extension Execution {
     mutating func refNull(sp: Sp, immediate: Instruction.RefNullOperand) {
         let value: Value
-        switch immediate.type.heapType {
+        switch immediate.type {
         case .externRef:
             value = .ref(.extern(nil))
         case .funcRef:

Sources/WasmKit/Execution/UntypedValue.swift

@@ -115,10 +115,12 @@ struct UntypedValue: Equatable, Hashable {
             return Int(storage)
         }
         switch type.heapType {
-        case .funcRef:
+        case .abstract(.funcRef):
             return .function(decodeOptionalInt())
-        case .externRef:
+        case .abstract(.externRef):
             return .extern(decodeOptionalInt())
+        case .concrete:
+            fatalError("heap type other than `func` and `extern` is not implemented yet")
         }
     }