vox_codegen/targets/swift/

wire.rs

//! Swift wire type code generation.
//!
//! Generates a complete Swift source file containing all wire protocol types
//! (Message, MessagePayload, etc.) with encode/decode methods. The generated
//! code is driven by facet `Shape`s from `vox_types::message`.
//!
//! The only special-cased type is `Payload` (`ShapeKind::Opaque`), which maps to
//! `OpaquePayload` with both length-prefixed and trailing byte handling.
//! Everything else is normal struct/enum codegen.

use facet_core::{Field, ScalarType, Shape};
use heck::ToLowerCamelCase;
use vox_types::{
    EnumInfo, ShapeKind, StructInfo, VariantKind, classify_shape, classify_variant,
    extract_schemas, is_bytes,
};

/// A wire type to generate Swift code for.
pub struct WireType {
    /// The Swift name for this type (e.g. "Message", "HelloV7")
    pub swift_name: String,
    /// The facet Shape for this type
    pub shape: &'static Shape,
}

/// Generate a complete Swift wire types file.
pub fn generate_wire_types(types: &[WireType]) -> String {
    let mut out = String::new();
    out.push_str("// @generated by vox-codegen\n");
    out.push_str("// DO NOT EDIT — regenerate with `cargo xtask codegen --swift-wire`\n\n");
    out.push_str("import Foundation\n\n");

    // Preamble: error type + helpers + OpaquePayload
    out.push_str(&generate_preamble());

    // Metadata typealias (Metadata is Vec<MetadataEntry> in Rust, we alias for convenience)
    out.push_str("public typealias Metadata = [MetadataEntry]\n\n");

    // Generate each type
    for wt in types {
        out.push_str(&generate_one_type(&wt.swift_name, wt.shape, types));
        out.push('\n');
    }

    // Generate factory methods extension on Message
    out.push_str(&generate_factory_methods(types));

    if let Some(root) = types
        .iter()
        .find(|wt| wt.swift_name == "Message")
        .or_else(|| types.last())
    {
        let extracted = extract_schemas(root.shape).expect("wire schema extraction should succeed");
        let cbor_bytes = facet_cbor::to_vec(&extracted.schemas)
            .expect("wire schema CBOR serialization should succeed");
        let body = cbor_bytes
            .iter()
            .map(|b| b.to_string())
            .collect::<Vec<_>>()
            .join(", ");
        out.push('\n');
        out.push_str(&format!(
            "public let wireMessageSchemasCbor: [UInt8] = [{body}]\n"
        ));
    }

    out
}

/// Generate the static preamble: WireError, helpers, OpaquePayload.
fn generate_preamble() -> String {
    let mut out = String::new();

    // Error type
    out.push_str("public enum WireError: Error, Equatable {\n");
    out.push_str("    case truncated\n");
    out.push_str("    case unknownVariant(UInt64)\n");
    out.push_str("    case overflow\n");
    out.push_str("    case invalidUtf8\n");
    out.push_str("    case trailingBytes\n");
    out.push_str("}\n\n");

    // OpaquePayload
    out.push_str("public struct OpaquePayload: Sendable, Equatable {\n");
    out.push_str("    public var bytes: [UInt8]\n\n");
    out.push_str("    public init(_ bytes: [UInt8]) {\n");
    out.push_str("        self.bytes = bytes\n");
    out.push_str("    }\n\n");
    out.push_str("    func encode() -> [UInt8] {\n");
    out.push_str("        let len = UInt32(bytes.count)\n");
    out.push_str("        return [\n");
    out.push_str("            UInt8(truncatingIfNeeded: len),\n");
    out.push_str("            UInt8(truncatingIfNeeded: len >> 8),\n");
    out.push_str("            UInt8(truncatingIfNeeded: len >> 16),\n");
    out.push_str("            UInt8(truncatingIfNeeded: len >> 24),\n");
    out.push_str("        ] + bytes\n");
    out.push_str("    }\n\n");
    out.push_str("    static func decode(from data: Data, offset: inout Int) throws -> Self {\n");
    out.push_str("        guard offset + 4 <= data.count else { throw WireError.truncated }\n");
    out.push_str("        let start = data.startIndex + offset\n");
    out.push_str("        let len = Int(UInt32(data[start]) | (UInt32(data[start + 1]) << 8) | (UInt32(data[start + 2]) << 16) | (UInt32(data[start + 3]) << 24))\n");
    out.push_str("        offset += 4\n");
    out.push_str("        guard offset + len <= data.count else { throw WireError.truncated }\n");
    out.push_str("        let payloadStart = data.startIndex + offset\n");
    out.push_str("        let payloadEnd = payloadStart + len\n");
    out.push_str("        let payload = Array(data[payloadStart..<payloadEnd])\n");
    out.push_str("        offset += len\n");
    out.push_str("        return .init(payload)\n");
    out.push_str("    }\n\n");
    out.push_str("    /// Encode without a length prefix — for trailing fields only.\n");
    out.push_str("    func encodeTrailing() -> [UInt8] {\n");
    out.push_str("        bytes\n");
    out.push_str("    }\n\n");
    out.push_str("    /// Decode by consuming all remaining bytes — for trailing fields only.\n");
    out.push_str("    static func decodeTrailing(from data: Data, offset: inout Int) -> Self {\n");
    out.push_str("        let start = data.startIndex + offset\n");
    out.push_str("        let remaining = Array(data[start...])\n");
    out.push_str("        offset = data.count\n");
    out.push_str("        return .init(remaining)\n");
    out.push_str("    }\n");
    out.push_str("}\n\n");

    // Helper functions
    out.push_str("@inline(__always)\n");
    out.push_str(
        "private func decodeWireVarintU32(from data: Data, offset: inout Int) throws -> UInt32 {\n",
    );
    out.push_str("    let value = try decodeVarint(from: data, offset: &offset)\n");
    out.push_str("    guard value <= UInt64(UInt32.max) else {\n");
    out.push_str("        throw WireError.overflow\n");
    out.push_str("    }\n");
    out.push_str("    return UInt32(value)\n");
    out.push_str("}\n\n");

    out.push_str("@inline(__always)\n");
    out.push_str(
        "private func decodeWireString(from data: Data, offset: inout Int) throws -> String {\n",
    );
    out.push_str("    do {\n");
    out.push_str("        return try decodeString(from: data, offset: &offset)\n");
    out.push_str("    } catch PostcardError.invalidUtf8 {\n");
    out.push_str("        throw WireError.invalidUtf8\n");
    out.push_str("    } catch PostcardError.truncated {\n");
    out.push_str("        throw WireError.truncated\n");
    out.push_str("    } catch {\n");
    out.push_str("        throw error\n");
    out.push_str("    }\n");
    out.push_str("}\n\n");

    out.push_str("@inline(__always)\n");
    out.push_str(
        "private func decodeWireBytes(from data: Data, offset: inout Int) throws -> Data {\n",
    );
    out.push_str("    do {\n");
    out.push_str("        return try decodeBytes(from: data, offset: &offset)\n");
    out.push_str("    } catch PostcardError.truncated {\n");
    out.push_str("        throw WireError.truncated\n");
    out.push_str("    } catch {\n");
    out.push_str("        throw error\n");
    out.push_str("    }\n");
    out.push_str("}\n\n");

    out
}

/// Get the wire Swift type for a field's shape, taking the `WireType` list into account.
fn swift_wire_type(shape: &'static Shape, _field: Option<&Field>, types: &[WireType]) -> String {
    if is_bytes(shape) {
        return "[UInt8]".into();
    }

    match classify_shape(shape) {
        ShapeKind::Scalar(scalar) => swift_scalar_type(scalar),
        ShapeKind::List { element } | ShapeKind::Slice { element } => {
            format!("[{}]", swift_wire_type(element, None, types))
        }
        ShapeKind::Option { inner } => {
            format!("{}?", swift_wire_type(inner, None, types))
        }
        ShapeKind::Array { element, .. } => {
            format!("[{}]", swift_wire_type(element, None, types))
        }
        ShapeKind::Struct(StructInfo {
            name: Some(name), ..
        }) => lookup_wire_name(name, types),
        ShapeKind::Enum(EnumInfo {
            name: Some(name), ..
        }) => lookup_wire_name(name, types),
        ShapeKind::Pointer { pointee } => swift_wire_type(pointee, _field, types),
        ShapeKind::Opaque => "OpaquePayload".into(),
        ShapeKind::TupleStruct { fields } if fields.len() == 1 => {
            swift_wire_type(fields[0].shape(), None, types)
        }
        _ => "Any /* unsupported */".into(),
    }
}

fn swift_scalar_type(scalar: ScalarType) -> String {
    match scalar {
        ScalarType::Bool => "Bool".into(),
        ScalarType::U8 => "UInt8".into(),
        ScalarType::U16 => "UInt16".into(),
        ScalarType::U32 => "UInt32".into(),
        ScalarType::U64 | ScalarType::USize => "UInt64".into(),
        ScalarType::I8 => "Int8".into(),
        ScalarType::I16 => "Int16".into(),
        ScalarType::I32 => "Int32".into(),
        ScalarType::I64 | ScalarType::ISize => "Int64".into(),
        ScalarType::F32 => "Float".into(),
        ScalarType::F64 => "Double".into(),
        ScalarType::Char | ScalarType::Str | ScalarType::String | ScalarType::CowStr => {
            "String".into()
        }
        ScalarType::Unit => "Void".into(),
        _ => "Data".into(),
    }
}

/// Look up the wire name for a Rust type name. If a matching WireType exists, use its
/// swift_name; otherwise fall back to the Rust name.
fn lookup_wire_name(rust_name: &str, types: &[WireType]) -> String {
    for wt in types {
        // Match by checking the reflected type identifier for the underlying Rust shape.
        if wt.shape.type_identifier.ends_with(rust_name) {
            return wt.swift_name.clone();
        }
    }
    rust_name.to_string()
}

/// Generate a single wire type (struct or enum).
fn generate_one_type(swift_name: &str, shape: &'static Shape, types: &[WireType]) -> String {
    match classify_shape(shape) {
        ShapeKind::Struct(StructInfo { fields, .. }) => {
            generate_struct(swift_name, fields, types, swift_name == "Message")
        }
        ShapeKind::Enum(EnumInfo { variants, .. }) => generate_enum(swift_name, variants, types),
        _ => format!("// Unsupported shape for {swift_name}\n"),
    }
}

/// Generate a Swift struct with encode/decode methods.
fn generate_struct(name: &str, fields: &[Field], types: &[WireType], is_top_level: bool) -> String {
    let mut out = String::new();

    // Struct definition
    out.push_str(&format!("public struct {name}: Sendable, Equatable {{\n"));
    for f in fields {
        let field_name = f.name.to_lower_camel_case();
        let field_type = swift_wire_type(f.shape(), Some(f), types);
        out.push_str(&format!("    public var {field_name}: {field_type}\n"));
    }

    // Initializer
    if !fields.is_empty() {
        out.push_str("\n    public init(");
        for (i, f) in fields.iter().enumerate() {
            if i > 0 {
                out.push_str(", ");
            }
            let field_name = f.name.to_lower_camel_case();
            let field_type = swift_wire_type(f.shape(), Some(f), types);
            out.push_str(&format!("{field_name}: {field_type}"));
        }
        out.push_str(") {\n");
        for f in fields {
            let field_name = f.name.to_lower_camel_case();
            out.push_str(&format!("        self.{field_name} = {field_name}\n"));
        }
        out.push_str("    }\n");
    }

    // Encode method
    let vis = if is_top_level { "public " } else { "" };
    out.push_str(&format!("\n    {vis}func encode() -> [UInt8] {{\n"));
    if fields.is_empty() {
        out.push_str("        []\n");
    } else if fields.len() == 1 {
        let f = &fields[0];
        let expr = encode_field_expr(f, types);
        out.push_str(&format!("        {expr}\n"));
    } else {
        out.push_str("        var out: [UInt8] = []\n");
        for f in fields {
            let expr = encode_field_expr(f, types);
            out.push_str(&format!("        out += {expr}\n"));
        }
        out.push_str("        return out\n");
    }
    out.push_str("    }\n");

    // Decode method (with offset)
    out.push_str(&format!(
        "\n    {vis}static func decode(from data: Data, offset: inout Int) throws -> Self {{\n"
    ));
    for f in fields {
        let field_name = f.name.to_lower_camel_case();
        let decode_expr = decode_field_expr(f, types);
        out.push_str(&format!("        let {field_name} = {decode_expr}\n"));
    }
    let field_names: Vec<String> = fields
        .iter()
        .map(|f| {
            let n = f.name.to_lower_camel_case();
            format!("{n}: {n}")
        })
        .collect();
    out.push_str(&format!(
        "        return .init({})\n",
        field_names.join(", ")
    ));
    out.push_str("    }\n");

    // Top-level Message gets an extra decode(from:) without offset that checks trailing bytes
    if is_top_level {
        out.push_str("\n    public static func decode(from data: Data) throws -> Self {\n");
        out.push_str("        var offset = 0\n");
        out.push_str("        let result = try decode(from: data, offset: &offset)\n");
        out.push_str("        guard offset == data.count else {\n");
        out.push_str("            throw WireError.trailingBytes\n");
        out.push_str("        }\n");
        out.push_str("        return result\n");
        out.push_str("    }\n");
    }

    out.push_str("}\n");
    out
}

/// Generate a Swift enum with varint-discriminanted encode/decode methods.
fn generate_enum(name: &str, variants: &[facet_core::Variant], types: &[WireType]) -> String {
    let mut out = String::new();

    // Enum definition
    out.push_str(&format!("public enum {name}: Sendable, Equatable {{\n"));
    for v in variants {
        let variant_name = v.name.to_lower_camel_case();
        match classify_variant(v) {
            VariantKind::Unit => {
                out.push_str(&format!("    case {variant_name}\n"));
            }
            VariantKind::Newtype { inner } => {
                let inner_type = swift_wire_type(inner, v.data.fields.first(), types);
                out.push_str(&format!("    case {variant_name}({inner_type})\n"));
            }
            VariantKind::Tuple { fields } => {
                let field_types: Vec<String> = fields
                    .iter()
                    .map(|f| swift_wire_type(f.shape(), Some(f), types))
                    .collect();
                out.push_str(&format!(
                    "    case {variant_name}({})\n",
                    field_types.join(", ")
                ));
            }
            VariantKind::Struct { fields } => {
                let field_decls: Vec<String> = fields
                    .iter()
                    .map(|f| {
                        format!(
                            "{}: {}",
                            f.name.to_lower_camel_case(),
                            swift_wire_type(f.shape(), Some(f), types)
                        )
                    })
                    .collect();
                out.push_str(&format!(
                    "    case {variant_name}({})\n",
                    field_decls.join(", ")
                ));
            }
        }
    }

    // Encode
    out.push_str("\n    func encode() -> [UInt8] {\n");
    out.push_str("        switch self {\n");
    for (i, v) in variants.iter().enumerate() {
        let variant_name = v.name.to_lower_camel_case();
        match classify_variant(v) {
            VariantKind::Unit => {
                out.push_str(&format!(
                    "        case .{variant_name}:\n            return encodeVarint(UInt64({i}))\n"
                ));
            }
            VariantKind::Newtype { inner } => {
                let encode = encode_shape_expr(inner, "val", v.data.fields.first(), types);
                out.push_str(&format!(
                    "        case .{variant_name}(let val):\n            return encodeVarint(UInt64({i})) + {encode}\n"
                ));
            }
            VariantKind::Tuple { fields } => {
                let bindings: Vec<String> = (0..fields.len()).map(|j| format!("f{j}")).collect();
                let binding_str = bindings
                    .iter()
                    .map(|b| format!("let {b}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                let encodes: Vec<String> = fields
                    .iter()
                    .enumerate()
                    .map(|(j, f)| encode_shape_expr(f.shape(), &format!("f{j}"), Some(f), types))
                    .collect();
                out.push_str(&format!(
                    "        case .{variant_name}({binding_str}):\n            return encodeVarint(UInt64({i})) + {}\n",
                    encodes.join(" + ")
                ));
            }
            VariantKind::Struct { fields } => {
                let bindings: Vec<String> = fields
                    .iter()
                    .map(|f| f.name.to_lower_camel_case())
                    .collect();
                let binding_str = bindings
                    .iter()
                    .map(|b| format!("let {b}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                let encodes: Vec<String> = fields
                    .iter()
                    .map(|f| {
                        encode_shape_expr(f.shape(), &f.name.to_lower_camel_case(), Some(f), types)
                    })
                    .collect();
                out.push_str(&format!(
                    "        case .{variant_name}({binding_str}):\n            return encodeVarint(UInt64({i})) + {}\n",
                    encodes.join(" + ")
                ));
            }
        }
    }
    out.push_str("        }\n");
    out.push_str("    }\n");

    // Decode
    out.push_str("\n    static func decode(from data: Data, offset: inout Int) throws -> Self {\n");
    out.push_str("        let disc = try decodeVarint(from: data, offset: &offset)\n");
    out.push_str("        switch disc {\n");
    for (i, v) in variants.iter().enumerate() {
        let variant_name = v.name.to_lower_camel_case();
        out.push_str(&format!("        case {i}:\n"));
        match classify_variant(v) {
            VariantKind::Unit => {
                out.push_str(&format!("            return .{variant_name}\n"));
            }
            VariantKind::Newtype { inner } => {
                let decode = decode_shape_expr(inner, v.data.fields.first(), types);
                out.push_str(&format!("            return .{variant_name}({decode})\n"));
            }
            VariantKind::Tuple { fields } => {
                for (j, f) in fields.iter().enumerate() {
                    let decode = decode_shape_expr(f.shape(), Some(f), types);
                    out.push_str(&format!("            let f{j} = {decode}\n"));
                }
                let args: Vec<String> = (0..fields.len()).map(|j| format!("f{j}")).collect();
                out.push_str(&format!(
                    "            return .{variant_name}({})\n",
                    args.join(", ")
                ));
            }
            VariantKind::Struct { fields } => {
                for f in fields {
                    let field_name = f.name.to_lower_camel_case();
                    let decode = decode_shape_expr(f.shape(), Some(f), types);
                    out.push_str(&format!("            let {field_name} = {decode}\n"));
                }
                let args: Vec<String> = fields
                    .iter()
                    .map(|f| {
                        let n = f.name.to_lower_camel_case();
                        format!("{n}: {n}")
                    })
                    .collect();
                out.push_str(&format!(
                    "            return .{variant_name}({})\n",
                    args.join(", ")
                ));
            }
        }
    }
    out.push_str("        default:\n");
    out.push_str("            throw WireError.unknownVariant(disc)\n");
    out.push_str("        }\n");
    out.push_str("    }\n");

    out.push_str("}\n");
    out
}

/// Generate an encode expression for a struct field.
fn encode_field_expr(field: &Field, types: &[WireType]) -> String {
    let field_name = field.name.to_lower_camel_case();
    encode_shape_expr(field.shape(), &field_name, Some(field), types)
}

/// Generate an encode expression for a shape with a given value expression.
fn encode_shape_expr(
    shape: &'static Shape,
    value: &str,
    field: Option<&Field>,
    types: &[WireType],
) -> String {
    let is_trailing = field.is_some_and(|f| f.has_builtin_attr("trailing"));

    // Opaque type → OpaquePayload
    if matches!(classify_shape(shape), ShapeKind::Opaque) {
        return if is_trailing {
            format!("{value}.encodeTrailing()")
        } else {
            format!("{value}.encode()")
        };
    }

    if is_bytes(shape) {
        return format!("encodeBytes({value})");
    }

    match classify_shape(shape) {
        ShapeKind::Scalar(scalar) => encode_scalar(scalar, value),
        ShapeKind::List { element } | ShapeKind::Slice { element } => {
            let inner = encode_element_closure(element, types);
            format!("encodeVec({value}, encoder: {inner})")
        }
        ShapeKind::Option { inner } => {
            let inner_closure = encode_element_closure(inner, types);
            format!("encodeOption({value}, encoder: {inner_closure})")
        }
        ShapeKind::Struct(StructInfo { .. }) | ShapeKind::Enum(EnumInfo { .. }) => {
            format!("{value}.encode()")
        }
        ShapeKind::Pointer { pointee } => encode_shape_expr(pointee, value, field, types),
        ShapeKind::TupleStruct { fields } if fields.len() == 1 => {
            encode_shape_expr(fields[0].shape(), value, field, types)
        }
        _ => format!("[] /* unsupported encode for {value} */"),
    }
}

fn encode_scalar(scalar: ScalarType, value: &str) -> String {
    match scalar {
        ScalarType::Bool => format!("encodeBool({value})"),
        ScalarType::U8 => format!("encodeU8({value})"),
        ScalarType::I8 => format!("encodeI8({value})"),
        ScalarType::U16 => format!("encodeU16({value})"),
        ScalarType::I16 => format!("encodeI16({value})"),
        ScalarType::U32 => format!("encodeVarint(UInt64({value}))"),
        ScalarType::I32 => format!("encodeI32({value})"),
        ScalarType::U64 | ScalarType::USize => format!("encodeVarint({value})"),
        ScalarType::I64 | ScalarType::ISize => format!("encodeI64({value})"),
        ScalarType::F32 => format!("encodeF32({value})"),
        ScalarType::F64 => format!("encodeF64({value})"),
        ScalarType::Char | ScalarType::Str | ScalarType::String | ScalarType::CowStr => {
            format!("encodeString({value})")
        }
        _ => "[] /* unsupported scalar */".to_string(),
    }
}

/// Generate an encode closure for use with encodeVec.
fn encode_element_closure(shape: &'static Shape, _types: &[WireType]) -> String {
    if is_bytes(shape) {
        return "{ encodeBytes($0) }".into();
    }

    match classify_shape(shape) {
        ShapeKind::Scalar(scalar) => {
            let expr = encode_scalar(scalar, "$0");
            format!("{{ {expr} }}")
        }
        ShapeKind::Struct(StructInfo { .. }) | ShapeKind::Enum(EnumInfo { .. }) => {
            "{ $0.encode() }".into()
        }
        ShapeKind::List { element } | ShapeKind::Slice { element } => {
            let inner = encode_element_closure(element, _types);
            format!("{{ encodeVec($0, encoder: {inner}) }}")
        }
        ShapeKind::Opaque => "{ $0.encode() }".into(),
        ShapeKind::Pointer { pointee } => encode_element_closure(pointee, _types),
        _ => "{ _ in [] }".into(),
    }
}

/// Generate a decode expression for a struct field.
fn decode_field_expr(field: &Field, types: &[WireType]) -> String {
    decode_shape_expr(field.shape(), Some(field), types)
}

/// Generate a decode expression for a shape.
fn decode_shape_expr(shape: &'static Shape, field: Option<&Field>, types: &[WireType]) -> String {
    let is_trailing = field.is_some_and(|f| f.has_builtin_attr("trailing"));

    // Opaque type → OpaquePayload
    if matches!(classify_shape(shape), ShapeKind::Opaque) {
        return if is_trailing {
            "OpaquePayload.decodeTrailing(from: data, offset: &offset)".into()
        } else {
            "try OpaquePayload.decode(from: data, offset: &offset)".into()
        };
    }

    if is_bytes(shape) {
        return "Array(try decodeWireBytes(from: data, offset: &offset))".into();
    }

    match classify_shape(shape) {
        ShapeKind::Scalar(scalar) => decode_scalar(scalar),
        ShapeKind::List { element } | ShapeKind::Slice { element } => {
            let inner = decode_element_closure(element, types);
            format!("try decodeVec(from: data, offset: &offset, decoder: {inner})")
        }
        ShapeKind::Option { inner } => {
            let inner_closure = decode_element_closure(inner, types);
            format!("try decodeOption(from: data, offset: &offset, decoder: {inner_closure})")
        }
        ShapeKind::Struct(StructInfo {
            name: Some(name), ..
        }) => {
            let swift_name = lookup_wire_name(name, types);
            format!("try {swift_name}.decode(from: data, offset: &offset)")
        }
        ShapeKind::Enum(EnumInfo {
            name: Some(name), ..
        }) => {
            let swift_name = lookup_wire_name(name, types);
            format!("try {swift_name}.decode(from: data, offset: &offset)")
        }
        ShapeKind::Pointer { pointee } => decode_shape_expr(pointee, field, types),
        ShapeKind::TupleStruct { fields } if fields.len() == 1 => {
            decode_shape_expr(fields[0].shape(), field, types)
        }
        _ => "nil /* unsupported decode */".into(),
    }
}

fn decode_scalar(scalar: ScalarType) -> String {
    match scalar {
        ScalarType::Bool => "try decodeBool(from: data, offset: &offset)".into(),
        ScalarType::U8 => "try decodeU8(from: data, offset: &offset)".into(),
        ScalarType::I8 => "try decodeI8(from: data, offset: &offset)".into(),
        ScalarType::U16 => "try decodeU16(from: data, offset: &offset)".into(),
        ScalarType::I16 => "try decodeI16(from: data, offset: &offset)".into(),
        ScalarType::U32 => "try decodeWireVarintU32(from: data, offset: &offset)".into(),
        ScalarType::I32 => "try decodeI32(from: data, offset: &offset)".into(),
        ScalarType::U64 | ScalarType::USize => {
            "try decodeVarint(from: data, offset: &offset)".into()
        }
        ScalarType::I64 | ScalarType::ISize => "try decodeI64(from: data, offset: &offset)".into(),
        ScalarType::F32 => "try decodeF32(from: data, offset: &offset)".into(),
        ScalarType::F64 => "try decodeF64(from: data, offset: &offset)".into(),
        ScalarType::Char | ScalarType::Str | ScalarType::String | ScalarType::CowStr => {
            "try decodeWireString(from: data, offset: &offset)".into()
        }
        _ => "nil /* unsupported scalar decode */".into(),
    }
}

/// Generate a decode closure for use with decodeVec.
fn decode_element_closure(shape: &'static Shape, types: &[WireType]) -> String {
    if is_bytes(shape) {
        return "{ data, off in Array(try decodeWireBytes(from: data, offset: &off)) }".into();
    }

    match classify_shape(shape) {
        ShapeKind::Scalar(scalar) => {
            let expr = decode_scalar_with_params(scalar, "data", "off");
            format!("{{ data, off in {expr} }}")
        }
        ShapeKind::Struct(StructInfo {
            name: Some(name), ..
        }) => {
            let swift_name = lookup_wire_name(name, types);
            format!("{{ data, off in try {swift_name}.decode(from: data, offset: &off) }}")
        }
        ShapeKind::Enum(EnumInfo {
            name: Some(name), ..
        }) => {
            let swift_name = lookup_wire_name(name, types);
            format!("{{ data, off in try {swift_name}.decode(from: data, offset: &off) }}")
        }
        ShapeKind::List { element } | ShapeKind::Slice { element } => {
            let inner = decode_element_closure(element, types);
            format!("{{ data, off in try decodeVec(from: data, offset: &off, decoder: {inner}) }}")
        }
        ShapeKind::Opaque => {
            "{ data, off in try OpaquePayload.decode(from: data, offset: &off) }".into()
        }
        ShapeKind::Pointer { pointee } => decode_element_closure(pointee, types),
        _ => "{ _, _ in throw WireError.truncated }".into(),
    }
}

fn decode_scalar_with_params(scalar: ScalarType, data: &str, offset: &str) -> String {
    match scalar {
        ScalarType::Bool => format!("try decodeBool(from: {data}, offset: &{offset})"),
        ScalarType::U8 => format!("try decodeU8(from: {data}, offset: &{offset})"),
        ScalarType::I8 => format!("try decodeI8(from: {data}, offset: &{offset})"),
        ScalarType::U16 => format!("try decodeU16(from: {data}, offset: &{offset})"),
        ScalarType::I16 => format!("try decodeI16(from: {data}, offset: &{offset})"),
        ScalarType::U32 => format!("try decodeWireVarintU32(from: {data}, offset: &{offset})"),
        ScalarType::I32 => format!("try decodeI32(from: {data}, offset: &{offset})"),
        ScalarType::U64 | ScalarType::USize => {
            format!("try decodeVarint(from: {data}, offset: &{offset})")
        }
        ScalarType::I64 | ScalarType::ISize => {
            format!("try decodeI64(from: {data}, offset: &{offset})")
        }
        ScalarType::F32 => format!("try decodeF32(from: {data}, offset: &{offset})"),
        ScalarType::F64 => format!("try decodeF64(from: {data}, offset: &{offset})"),
        ScalarType::Char | ScalarType::Str | ScalarType::String | ScalarType::CowStr => {
            format!("try decodeWireString(from: {data}, offset: &{offset})")
        }
        _ => "nil /* unsupported scalar */".to_string(),
    }
}

/// Generate convenience factory methods on Message.
fn generate_factory_methods(types: &[WireType]) -> String {
    // Find the MessagePayload shape to inspect variants
    let payload_wt = types.iter().find(|wt| wt.swift_name == "MessagePayload");
    let payload_wt = match payload_wt {
        Some(wt) => wt,
        None => return String::new(),
    };

    let variants = match classify_shape(payload_wt.shape) {
        ShapeKind::Enum(EnumInfo { variants, .. }) => variants,
        _ => return String::new(),
    };

    let mut out = String::new();
    out.push_str("public extension Message {\n");

    for v in variants {
        let variant_name = v.name.to_lower_camel_case();
        if let VariantKind::Newtype { inner } = classify_variant(v) {
            let inner_swift = swift_wire_type(inner, None, types);
            // Control messages use connId=0.
            let is_control = matches!(
                v.name,
                "Hello" | "HelloYourself" | "ProtocolError" | "Ping" | "Pong"
            );

            if is_control {
                out.push_str(&format!(
                    "    static func {variant_name}(_ value: {inner_swift}) -> Message {{\n"
                ));
                out.push_str(&format!(
                    "        Message(connectionId: 0, payload: .{variant_name}(value))\n"
                ));
                out.push_str("    }\n\n");
            } else {
                out.push_str(&format!(
                    "    static func {variant_name}(connId: UInt64, _ value: {inner_swift}) -> Message {{\n"
                ));
                out.push_str(&format!(
                    "        Message(connectionId: connId, payload: .{variant_name}(value))\n"
                ));
                out.push_str("    }\n\n");
            }
        }
    }

    // Additional ergonomic factory methods that flatten nested structs.
    // These match the existing hand-coded API.
    out.push_str("    static func protocolError(description: String) -> Message {\n");
    out.push_str("        Message(connectionId: 0, payload: .protocolError(.init(description: description)))\n");
    out.push_str("    }\n\n");

    out.push_str("    static func connectionOpen(connId: UInt64, settings: ConnectionSettings, metadata: [MetadataEntry]) -> Message {\n");
    out.push_str("        Message(connectionId: connId, payload: .connectionOpen(.init(connectionSettings: settings, metadata: metadata)))\n");
    out.push_str("    }\n\n");

    out.push_str("    static func connectionAccept(connId: UInt64, settings: ConnectionSettings, metadata: [MetadataEntry]) -> Message {\n");
    out.push_str("        Message(connectionId: connId, payload: .connectionAccept(.init(connectionSettings: settings, metadata: metadata)))\n");
    out.push_str("    }\n\n");

    out.push_str("    static func connectionReject(connId: UInt64, metadata: [MetadataEntry]) -> Message {\n");
    out.push_str("        Message(connectionId: connId, payload: .connectionReject(.init(metadata: metadata)))\n");
    out.push_str("    }\n\n");

    out.push_str(
        "    static func connectionClose(connId: UInt64, metadata: [MetadataEntry]) -> Message {\n",
    );
    out.push_str("        Message(connectionId: connId, payload: .connectionClose(.init(metadata: metadata)))\n");
    out.push_str("    }\n\n");

    out.push_str("    static func request(\n");
    out.push_str("        connId: UInt64,\n");
    out.push_str("        requestId: UInt64,\n");
    out.push_str("        methodId: UInt64,\n");
    out.push_str("        metadata: [MetadataEntry],\n");
    out.push_str("        schemas: [UInt8] = [],\n");
    out.push_str("        payload: [UInt8]\n");
    out.push_str("    ) -> Message {\n");
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .requestMessage(\n");
    out.push_str("                .init(\n");
    out.push_str("                    id: requestId,\n");
    out.push_str("                    body: .call(.init(methodId: methodId, metadata: metadata, args: .init(payload), schemas: schemas))\n");
    out.push_str("                ))\n");
    out.push_str("        )\n");
    out.push_str("    }\n\n");

    out.push_str("    static func response(\n");
    out.push_str("        connId: UInt64,\n");
    out.push_str("        requestId: UInt64,\n");
    out.push_str("        metadata: [MetadataEntry],\n");
    out.push_str("        schemas: [UInt8] = [],\n");
    out.push_str("        payload: [UInt8]\n");
    out.push_str("    ) -> Message {\n");
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .requestMessage(\n");
    out.push_str("                .init(\n");
    out.push_str("                    id: requestId,\n");
    out.push_str("                    body: .response(.init(metadata: metadata, ret: .init(payload), schemas: schemas))\n");
    out.push_str("                ))\n");
    out.push_str("        )\n");
    out.push_str("    }\n\n");

    out.push_str("    static func cancel(connId: UInt64, requestId: UInt64, metadata: [MetadataEntry] = []) -> Message {\n");
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .requestMessage(\n");
    out.push_str("                .init(\n");
    out.push_str("                    id: requestId,\n");
    out.push_str("                    body: .cancel(.init(metadata: metadata))\n");
    out.push_str("                ))\n");
    out.push_str("        )\n");
    out.push_str("    }\n\n");

    out.push_str(
        "    static func data(connId: UInt64, channelId: UInt64, payload: [UInt8]) -> Message {\n",
    );
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .channelMessage(.init(id: channelId, body: .item(.init(item: .init(payload)))))\n");
    out.push_str("        )\n");
    out.push_str("    }\n\n");

    out.push_str("    static func close(connId: UInt64, channelId: UInt64, metadata: [MetadataEntry] = []) -> Message {\n");
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .channelMessage(.init(id: channelId, body: .close(.init(metadata: metadata))))\n");
    out.push_str("        )\n");
    out.push_str("    }\n\n");

    out.push_str("    static func reset(connId: UInt64, channelId: UInt64, metadata: [MetadataEntry] = []) -> Message {\n");
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .channelMessage(.init(id: channelId, body: .reset(.init(metadata: metadata))))\n");
    out.push_str("        )\n");
    out.push_str("    }\n\n");

    out.push_str(
        "    static func credit(connId: UInt64, channelId: UInt64, bytes: UInt32) -> Message {\n",
    );
    out.push_str("        Message(\n");
    out.push_str("            connectionId: connId,\n");
    out.push_str("            payload: .channelMessage(.init(id: channelId, body: .grantCredit(.init(additional: bytes))))\n");
    out.push_str("        )\n");
    out.push_str("    }\n");

    out.push_str("}\n");
    out
}