WAT: Add support for encoding the name section

Author: kateinoigakukun

Sources/WAT/Encoder.swift

@@ -535,7 +535,7 @@ struct ExpressionEncoder: InstructionEncoder {
     }
 }
 
-func encode(module: inout Wat) throws -> [UInt8] {
+func encode(module: inout Wat, options: EncodeOptions) throws -> [UInt8] {
     var encoder = Encoder()
     encoder.writeHeader()
 
@@ -675,5 +675,25 @@ func encode(module: inout Wat) throws -> [UInt8] {
         }
     }
 
+    // (Optional) Name Section
+    if !module.functionsMap.isEmpty, options.nameSection {
+        encoder.section(id: 0) { encoder in
+            encoder.encode("name")
+            // Subsection 1: Function names
+            encoder.section(id: 1) { encoder in
+                let functionNames = module.functionsMap.enumerated().compactMap { i, decl -> (Int, String)? in
+                    guard let name = decl.id else { return nil }
+                    return (i, name.value)
+                }
+                encoder.encodeVector(functionNames) { entry, encoder in
+                    let (index, name) = entry
+                    encoder.writeUnsignedLEB128(UInt(index))
+                    // Drop initial "$"
+                    encoder.encode(String(name.dropFirst()))
+                }
+            }
+        }
+    }
+
     return encoder.output
 }

Sources/WAT/WAT.swift

@@ -1,5 +1,19 @@
 import WasmParser
 
+/// Options for encoding a WebAssembly module into a binary format.
+public struct EncodeOptions {
+    /// Whether to include the name section.
+    public var nameSection: Bool
+
+    /// The default encoding options.
+    public static let `default` = EncodeOptions()
+
+    /// Creates a new encoding options instance.
+    public init(nameSection: Bool = false) {
+        self.nameSection = nameSection
+    }
+}
+
 /// Transforms a WebAssembly text format (WAT) string into a WebAssembly binary format byte array.
 /// - Parameter input: The WAT string to transform
 /// - Returns: The WebAssembly binary format byte array
@@ -17,9 +31,13 @@ import WasmParser
 /// )
 /// """)
 /// ```
-public func wat2wasm(_ input: String, features: WasmFeatureSet = .default) throws -> [UInt8] {
+public func wat2wasm(
+    _ input: String,
+    features: WasmFeatureSet = .default,
+    options: EncodeOptions = .default
+) throws -> [UInt8] {
     var wat = try parseWAT(input, features: features)
-    return try encode(module: &wat)
+    return try encode(module: &wat, options: options)
 }
 
 /// A WAT module representation.
@@ -65,8 +83,8 @@ public struct Wat {
     /// This method effectively consumes the module value, encoding it into a
     /// binary format byte array. If you need to encode the module multiple times,
     /// you should create a copy of the module value before encoding it.
-    public mutating func encode() throws -> [UInt8] {
-        try WAT.encode(module: &self)
+    public mutating func encode(options: EncodeOptions = .default) throws -> [UInt8] {
+        try WAT.encode(module: &self, options: options)
     }
 }
 

Tests/WATTests/EncoderTests.swift

@@ -78,7 +78,7 @@ class EncoderTests: XCTestCase {
                 assertEqual(watModule.id, expectedName)
                 switch watModule.source {
                 case .text(var watModule):
-                    moduleBytes = try encode(module: &watModule)
+                    moduleBytes = try encode(module: &watModule, options: .default)
                 case .binary(let bytes):
                     moduleBytes = bytes
                 case .quote(let watText):
@@ -135,4 +135,37 @@ class EncoderTests: XCTestCase {
             }
         #endif
     }
+
+    func testEncodeNameSection() throws {
+        let bytes = try wat2wasm(
+            """
+            (module
+                (func $foo)
+                (func)
+                (func $bar)
+            )
+            """,
+            options: EncodeOptions(nameSection: true)
+        )
+        
+        var parser = WasmParser.Parser(bytes: bytes)
+        var customSections: [CustomSection] = []
+        while let payload = try parser.parseNext() {
+            guard case .customSection(let section) = payload else {
+                continue
+            }
+            customSections.append(section)
+        }
+        let nameSection = customSections.first(where: { $0.name == "name" })
+        let nameParser = NameSectionParser(
+            stream: StaticByteStream(bytes: nameSection?.bytes ?? [])
+        )
+        let names = try nameParser.parseAll()
+        XCTAssertEqual(names.count, 1)
+        guard case .functions(let functionNames) = try XCTUnwrap(names.first) else {
+            XCTFail()
+            return
+        }
+        XCTAssertEqual(functionNames, [0: "foo", 2: "bar"])
+    }
 }