CM WAT normalization: remove ComponentEncoder.GroupedFields (#325)

MaxDesiatov · web-flow · commit e552c8c2bc19 · 2026-03-05T21:16:33.000Z
Following up on #323, `ComponentEncoder.GroupedFields` was responsible for normalization in the encoder. As we're shifting normalization to `ComponentWatParser`, which now accumulates fully populated `NormalizedDefinition`, `GroupedFields` is no longer needed for batching. As CM sections can be unordered, discontiguous, and interleaved, the encoder in the future will only have to iterate through `NormalizedDefinition` and flush a section when definition kind changes.
diff --git a/Sources/WAT/BinaryEncoding/ComponentEncoder.swift b/Sources/WAT/BinaryEncoding/ComponentEncoder.swift
@@ -22,24 +22,35 @@
         ) throws(WatParserError) -> [UInt8] {
             writeHeader()
 
-            // Collect metadata needed for encoding
-            let fields = try groupFields(component.fields)
+            // Compute metadata directly from component fields
+            let coreInstances: [(ComponentWatParser.CoreInstanceDef, Location)] = component.fields.compactMap { field in
+                if case .coreInstance(let def) = field.kind { return (def, field.location) }
+                return nil
+            }
+            let canons: [(ComponentWatParser.CanonDef, Location)] = component.fields.compactMap { field in
+                if case .canon(let def) = field.kind { return (def, field.location) }
+                return nil
+            }
+            let exports: [(ComponentWatParser.ExportDef, Location)] = component.fields.compactMap { field in
+                if case .componentExport(let def) = field.kind { return (def, field.location) }
+                return nil
+            }
 
             // Build core instance index mapping (parser index -> binary index)
             // This accounts for inline export instances that shift indices
-            let coreInstanceIndexMapping = buildCoreInstanceIndexMapping(component: component, fields: fields)
+            let coreInstanceIndexMapping = buildCoreInstanceIndexMapping(coreInstances: coreInstances)
 
             let coreFuncAliases = try collectCoreFuncAliases(
-                from: fields.canons,
+                from: canons,
                 component: component,
                 coreInstanceIndexMapping: coreInstanceIndexMapping
             )
-            let componentFuncAliases = try collectComponentFuncAliases(from: fields.canons, component: component)
-            let exportFuncAliases = try collectExportFuncAliases(from: fields.exports, component: component)
+            let componentFuncAliases = try collectComponentFuncAliases(from: canons, component: component)
+            let exportFuncAliases = try collectExportFuncAliases(from: exports, component: component)
 
             // Collect type indices that are exported
             var exportedTypeIndices = Set<Int>()
-            for (exportDef, _) in fields.exports {
+            for (exportDef, _) in exports {
                 if case .type(let indexOrId) = exportDef.descriptor {
                     if case .index(let idx, _) = indexOrId {
                         exportedTypeIndices.insert(Int(idx))
@@ -77,7 +88,7 @@
             }
 
             // Helper to flush accumulated core types as a batched section
-            var pendingCoreTypes: [UInt32] = []
+            var pendingCoreTypes: [ComponentWatParser.CoreTypeDef] = []
             func flushPendingCoreTypes() throws(WatParserError) {
                 if !pendingCoreTypes.isEmpty {
                     try encodeBatchedCoreTypes(
@@ -89,7 +100,7 @@
             }
 
             // Helper to flush accumulated core instances as a batched section
-            var pendingCoreInstances: [(CoreInstanceIndex, Location)] = []
+            var pendingCoreInstances: [(ComponentWatParser.CoreInstanceDef, Location)] = []
             func flushPendingCoreInstances() throws(WatParserError) {
                 if !pendingCoreInstances.isEmpty {
                     try encodeCoreInstances(
@@ -141,25 +152,25 @@
                 try flushPendingSectionsExcept(field.kind.sectionKind)
 
                 switch field.kind {
-                case .componentType(let typeIndex):
+                case .componentType(_, let typeIndex):
                     // Collect all unemitted types up to and including this type index
                     // This ensures anonymous dependency types (with lower indices) get batched
                     // with the types that reference them
-                    let maxIndex = Int(typeIndex.rawValue)
+                    let maxIndex = typeIndex
                     for (oldIndex, _) in typeIndexMapping.sorted(by: { $0.key < $1.key }) {
                         if oldIndex <= maxIndex && !emittedTypes.contains(oldIndex) && !pendingComponentTypes.contains(oldIndex) {
                             pendingComponentTypes.append(oldIndex)
                         }
                     }
 
-                case .coreModule(let index):
-                    try encodeSingleCoreModule(index, component: component, options: options)
+                case .coreModule(let moduleDef):
+                    try encodeSingleCoreModule(moduleDef, options: options)
 
-                case .coreInstance(let index):
-                    pendingCoreInstances.append((index, field.location))
+                case .coreInstance(let instanceDef):
+                    pendingCoreInstances.append((instanceDef, field.location))
 
-                case .coreType(let index):
-                    pendingCoreTypes.append(UInt32(index))
+                case .coreType(let coreTypeDef):
+                    pendingCoreTypes.append(coreTypeDef)
 
                 case .canon(let canonDef):
                     try emitRequiredTypesForCanon(
@@ -208,11 +219,11 @@
                         emittedExportFuncAliases: &emittedExportFuncAliases
                     )
 
-                case .component(let index):
-                    try encodeSingleComponent(index, component: component, options: options)
+                case .component(let nestedDef):
+                    try encodeSingleComponent(nestedDef, options: options)
 
-                case .instance(let index):
-                    try encodeSingleComponentInstance(index, component: component, location: field.location)
+                case .instance(let instanceDef):
+                    try encodeSingleComponentInstance(instanceDef, component: component, location: field.location)
 
                 case .alias(let alias):
                     pendingAliases.append(alias)
@@ -586,15 +597,10 @@
 
         // Encode a single core module as its own section
         private mutating func encodeSingleCoreModule(
-            _ moduleIndex: CoreModuleIndex,
-            component: ComponentWatParser.ComponentDef,
+            _ moduleDef: ComponentWatParser.ModuleDef,
             options: EncodeOptions
         ) throws(WatParserError) {
-            guard Int(moduleIndex.rawValue) < component.coreModulesMap.count else {
-                throw WatParserError("Invalid core module index \(moduleIndex.rawValue)", location: nil)
-            }
-
-            var moduleDef = component.coreModulesMap[Int(moduleIndex.rawValue)]
+            var moduleDef = moduleDef
             var moduleBytes = try WAT.encode(module: &moduleDef.wat, options: options)
 
             if options.nameSection, let moduleId = moduleDef.id {
@@ -740,17 +746,15 @@
 
         // Encode multiple core types in a single batched section
         private mutating func encodeBatchedCoreTypes(
-            _ typeIndices: [UInt32],
+            _ typeDefs: [ComponentWatParser.CoreTypeDef],
             component: ComponentWatParser.ComponentDef
         ) throws(WatParserError) {
-            guard !typeIndices.isEmpty else { return }
+            guard !typeDefs.isEmpty else { return }
 
             try underlying.section(id: 0x03) { encoder throws(WatParserError) in
-                encoder.writeUnsignedLEB128(UInt32(typeIndices.count))
+                encoder.writeUnsignedLEB128(UInt32(typeDefs.count))
 
-                for typeIndex in typeIndices {
-                    guard Int(typeIndex) < component.coreTypesMap.count else { continue }
-                    let typeDef = component.coreTypesMap[Int(typeIndex)]
+                for typeDef in typeDefs {
                     try Self.encodeCoreTypeContent(typeDef, component: component, encoder: &encoder)
                 }
             }
@@ -1201,15 +1205,9 @@
 
         // Encode a single nested component as its own section
         private mutating func encodeSingleComponent(
-            _ componentIndex: ComponentIndex,
-            component: ComponentWatParser.ComponentDef,
+            _ nestedComponent: ComponentWatParser.ComponentDef,
             options: EncodeOptions
         ) throws(WatParserError) {
-            guard Int(componentIndex.rawValue) < component.componentsMap.count else {
-                throw WatParserError("Invalid component index \(componentIndex.rawValue)", location: nil)
-            }
-
-            let nestedComponent = component.componentsMap[Int(componentIndex.rawValue)]
             var nestedEncoder = ComponentEncoder()
             _ = try nestedEncoder.encode(nestedComponent, options: options)
 
@@ -1220,16 +1218,10 @@
 
         // Encode a single component instance as its own section
         private mutating func encodeSingleComponentInstance(
-            _ instanceIndex: ComponentInstanceIndex,
+            _ instanceDef: ComponentWatParser.ComponentInstanceDef,
             component: ComponentWatParser.ComponentDef,
             location: Location
         ) throws(WatParserError) {
-            guard Int(instanceIndex.rawValue) < component.componentInstancesMap.count else {
-                throw WatParserError("Invalid component instance index \(instanceIndex.rawValue)", location: location)
-            }
-
-            let instanceDef = component.componentInstancesMap[Int(instanceIndex.rawValue)]
-
             guard let componentRef = instanceDef.componentRef else {
                 throw WatParserError("Component instance has no component reference", location: location)
             }
@@ -1255,57 +1247,16 @@
             }
         }
 
-        private func groupFields(_ fields: [ComponentWatParser.ComponentDefField]) throws(WatParserError) -> GroupedFields {
-            var result = GroupedFields()
-
-            for field in fields {
-                switch field.kind {
-                case .coreModule(let index):
-                    result.coreModules.append((index, field.location))
-                case .coreInstance(let index):
-                    result.coreInstances.append((index, field.location))
-                case .coreType(let index):
-                    result.coreTypes.append((index, field.location))
-                case .component(let index):
-                    result.components.append((index, field.location))
-                case .canon(let canonDef):
-                    result.canons.append((canonDef, field.location))
-                case .componentExport(let exportDef):
-                    result.exports.append((exportDef, field.location))
-                case .componentImport(let importDef):
-                    result.imports.append((importDef, field.location))
-                case .instance(let index):
-                    result.instances.append((index, field.location))
-                case .componentType:
-                    // Component types are handled separately in the main encoding loop
-                    // They don't need to be grouped for auxiliary calculations
-                    break
-                case .alias:
-                    // Aliases are handled in the main encoding loop with batching
-                    // They don't need to be grouped for auxiliary calculations
-                    break
-                }
-            }
-
-            return result
-        }
-
         /// Build a mapping from parser core instance index to binary index.
         /// This accounts for inline export instances that shift indices.
         private func buildCoreInstanceIndexMapping(
-            component: ComponentWatParser.ComponentDef,
-            fields: GroupedFields
+            coreInstances: [(ComponentWatParser.CoreInstanceDef, Location)]
         ) -> [Int: Int] {
             var mapping: [Int: Int] = [:]
             var binaryIndex = 0
 
             // Process core instances in order, counting inline export instances
-            for (instanceIndex, _) in fields.coreInstances {
-                let idx = Int(instanceIndex.rawValue)
-                guard idx < component.coreInstancesMap.count else { continue }
-
-                let instanceDef = component.coreInstancesMap[idx]
-
+            for (parserIndex, (instanceDef, _)) in coreInstances.enumerated() {
                 // Count inline export instances that will be emitted first
                 var inlineExportCount = 0
                 for arg in instanceDef.arguments {
@@ -1316,7 +1267,7 @@
 
                 // Inline exports come first, then the main instance
                 // So the main instance index is offset by inline export count
-                mapping[idx] = binaryIndex + inlineExportCount
+                mapping[parserIndex] = binaryIndex + inlineExportCount
                 binaryIndex += 1 + inlineExportCount  // 1 for main instance + inline exports
             }
 
@@ -1372,21 +1323,18 @@
         }
 
         private mutating func encodeCoreInstances(
-            _ instances: [(CoreInstanceIndex, Location)],
+            _ instances: [(ComponentWatParser.CoreInstanceDef, Location)],
             component: ComponentWatParser.ComponentDef
         ) throws(WatParserError) {
             // First pass: collect all inline export instances that need to be created
             // and compute the mapping from original to actual instance indices
             var inlineExportInstances: [(exports: [ComponentWatParser.CoreInstanceDef.Argument.Kind.Export], location: Location)] = []
-            var inlineExportInstanceMapping: [Int: Int] = [:]  // Map from (instance index, arg index) hash to inline instance index
-
-            for (instanceIndex, location) in instances {
-                guard Int(instanceIndex.rawValue) < component.coreInstancesMap.count else { continue }
-                let instanceDef = component.coreInstancesMap[Int(instanceIndex.rawValue)]
+            var inlineExportInstanceMapping: [Int: Int] = [:]  // Map from (instance position, arg index) hash to inline instance index
 
+            for (instancePosition, (instanceDef, location)) in instances.enumerated() {
                 for (argIndex, arg) in instanceDef.arguments.enumerated() {
                     if case .exports(let exports) = arg.kind {
-                        let key = Int(instanceIndex.rawValue) * 1000 + argIndex
+                        let key = instancePosition * 1000 + argIndex
                         inlineExportInstanceMapping[key] = inlineExportInstances.count
                         inlineExportInstances.append((exports: exports, location: location))
                     }
@@ -1414,13 +1362,7 @@
                 // Second: encode instantiate instances (form 0x00)
                 let inlineInstanceBaseIndex = inlineExportInstances.count
 
-                for (instanceIndex, location) in instances {
-                    guard Int(instanceIndex.rawValue) < component.coreInstancesMap.count else {
-                        throw WatParserError("Invalid core instance index \(instanceIndex.rawValue)", location: location)
-                    }
-
-                    let instanceDef = component.coreInstancesMap[Int(instanceIndex.rawValue)]
-
+                for (instancePosition, (instanceDef, location)) in instances.enumerated() {
                     encoder.output.append(0x00)  // instantiate form
 
                     let moduleIndex = try component.coreModulesMap.resolveIndex(use: instanceDef.moduleId)
@@ -1439,7 +1381,7 @@
                         case .exports:
                             // Reference the inline export instance we created earlier
                             encoder.output.append(0x12)
-                            let key = Int(instanceIndex.rawValue) * 1000 + argIndex
+                            let key = instancePosition * 1000 + argIndex
                             guard let inlineIndex = inlineExportInstanceMapping[key] else {
                                 throw WatParserError("Internal error: inline export instance not found", location: location)
                             }
@@ -1708,17 +1650,6 @@
         }
     }
 
-    private struct GroupedFields {
-        var coreModules: [(CoreModuleIndex, Location)] = []
-        var coreInstances: [(CoreInstanceIndex, Location)] = []
-        var coreTypes: [(UInt32, Location)] = []
-        var components: [(ComponentIndex, Location)] = []
-        var instances: [(ComponentInstanceIndex, Location)] = []
-        var canons: [(ComponentWatParser.CanonDef, Location)] = []
-        var exports: [(ComponentWatParser.ExportDef, Location)] = []
-        var imports: [(ComponentWatParser.ImportDef, Location)] = []
-    }
-
     private typealias CoreFuncAlias = (instanceIndex: Int, exportName: String)
     private typealias ComponentFuncAlias = (instanceIndex: Int, exportName: String)
 
diff --git a/Sources/WAT/Parser/ComponentWatParser.swift b/Sources/WAT/Parser/ComponentWatParser.swift
@@ -70,8 +70,8 @@
                     componentStack.append(ComponentDef(id: nestedId))
                     try parseComponentFields(&parser)  // Recursive call for nested components
                     let nested = componentStack.removeLast()
-                    let idx = try currentComponent.componentsMap.add(nested)
-                    currentComponent.fields.append(.init(location: location, kind: .component(.init(idx))))
+                    _ = try currentComponent.componentsMap.add(nested)
+                    currentComponent.fields.append(.init(location: location, kind: .component(nested)))
                 case "import":
                     try parseComponentImport(&parser, location: location)
                 case "export":
@@ -97,23 +97,23 @@
             switch coreKeyword {
             case "module":
                 let moduleDef = try self.parseModuleDef(&parser)
-                let index = try self.currentComponent.coreModulesMap.add(moduleDef)
+                _ = try self.currentComponent.coreModulesMap.add(moduleDef)
                 self.currentComponent.fields.append(
-                    .init(location: location, kind: .coreModule(.init(index)))
+                    .init(location: location, kind: .coreModule(moduleDef))
                 )
 
             case "instance":
                 let instanceDef = try self.parseCoreInstanceDef(&parser)
-                let index = try self.currentComponent.coreInstancesMap.add(instanceDef)
+                _ = try self.currentComponent.coreInstancesMap.add(instanceDef)
                 self.currentComponent.fields.append(
-                    .init(location: location, kind: .coreInstance(.init(index)))
+                    .init(location: location, kind: .coreInstance(instanceDef))
                 )
 
             case "type":
                 let coreTypeDef = try self.parseCoreTypeDef(&parser)
-                let resolvedId = try self.currentComponent.coreTypesMap.add(coreTypeDef)
+                _ = try self.currentComponent.coreTypesMap.add(coreTypeDef)
                 self.currentComponent.fields.append(
-                    .init(location: location, kind: .coreType(TypeIndex(resolvedId)))
+                    .init(location: location, kind: .coreType(coreTypeDef))
                 )
 
             case "func":
@@ -364,9 +364,9 @@
             try parser.expect(.rightParen)  // Close (instantiate ...)
 
             let instanceDef = ComponentInstanceDef(id: instanceId, componentRef: componentRef, arguments: arguments)
-            let index = try currentComponent.componentInstancesMap.add(instanceDef)
+            _ = try currentComponent.componentInstancesMap.add(instanceDef)
             currentComponent.fields.append(
-                .init(location: location, kind: .instance(.init(index)))
+                .init(location: location, kind: .instance(instanceDef))
             )
         }
 
@@ -589,7 +589,7 @@
             self.currentComponent.fields.append(
                 .init(
                     location: typeLocation,
-                    kind: .componentType(ComponentTypeIndex(rawValue: typeIndex))
+                    kind: .componentType(typeDef, typeIndex: typeIndex)
                 )
             )
 
diff --git a/Sources/WAT/Parser/NormalizedDefinition.swift b/Sources/WAT/Parser/NormalizedDefinition.swift
