treesitter_types/codegen/

type_mapper.rs

use super::grammar_ir::{FieldInfo, NodeType, TypeRef};
use super::name_mangler;
use proc_macro2::Ident;

/// What Rust type to emit for a top-level node.
#[derive(Debug, Clone)]
pub enum TypeDecision {
    /// A named node with fields and/or children → emit a struct.
    Struct(StructDef),
    /// A named terminal node (no fields, no children, no subtypes) → emit a leaf struct.
    LeafStruct(LeafStructDef),
    /// A supertype node (has `subtypes`) → emit an enum.
    SupertypeEnum(SupertypeEnumDef),
}

#[derive(Debug, Clone)]
pub struct StructDef {
    pub type_name: Ident,
    pub kind: String,
    pub fields: Vec<FieldDef>,
    pub children: Option<ChildrenDef>,
}

#[derive(Debug, Clone)]
pub struct LeafStructDef {
    pub type_name: Ident,
    pub kind: String,
}

#[derive(Debug, Clone)]
pub struct SupertypeEnumDef {
    pub type_name: Ident,
    pub kind: String,
    pub variants: Vec<VariantDef>,
}

#[derive(Debug, Clone)]
pub struct FieldDef {
    pub field_name: Ident,
    /// The original field name string (for `child_by_field_name`).
    pub raw_field_name: String,
    pub field_type: FieldType,
}

#[derive(Debug, Clone)]
pub struct ChildrenDef {
    pub field_type: FieldType,
}

/// How to represent a field's type in Rust.
#[derive(Debug, Clone)]
pub enum FieldType {
    /// A single concrete type (e.g., `Identifier<'tree>`).
    Direct(TypeReference),
    /// Wrapped in `Option<T>`.
    Optional(TypeReference),
    /// Wrapped in `Vec<T>`.
    Repeated(TypeReference),
}

/// Reference to a Rust type, which may be a single named type or an alternation enum.
#[derive(Debug, Clone)]
pub enum TypeReference {
    /// A single named node type.
    Named(Ident),
    /// An alternation enum generated for this field.
    Alternation(AlternationEnumDef),
}

#[derive(Debug, Clone)]
pub struct AlternationEnumDef {
    pub type_name: Ident,
    pub variants: Vec<VariantDef>,
}

#[derive(Debug, Clone)]
pub struct VariantDef {
    pub variant_name: Ident,
    pub kind: String,
    pub named: bool,
    /// True if this variant references a supertype (has subtypes in the grammar).
    /// Supertypes are abstract — tree-sitter never emits them as node kinds at runtime.
    pub is_supertype: bool,
}

/// Maps all `NodeType` entries into `TypeDecision`s.
pub fn map_types(nodes: &[NodeType]) -> Vec<TypeDecision> {
    // Collect the set of type names that are supertypes
    let supertype_kinds: std::collections::HashSet<&str> = nodes
        .iter()
        .filter(|n| n.subtypes.is_some())
        .map(|n| n.type_name.as_str())
        .collect();

    nodes
        .iter()
        .filter(|n| n.named)
        .map(|n| map_node(n, &supertype_kinds))
        .collect()
}

fn map_node(node: &NodeType, supertype_kinds: &std::collections::HashSet<&str>) -> TypeDecision {
    let raw_kind = &node.type_name;

    // Supertype nodes (e.g., _expression, statement) → enum
    if let Some(subtypes) = &node.subtypes {
        return TypeDecision::SupertypeEnum(SupertypeEnumDef {
            type_name: supertype_ident(raw_kind),
            kind: raw_kind.clone(),
            variants: subtypes
                .iter()
                .map(|tr| make_variant_def(tr, supertype_kinds))
                .collect(),
        });
    }

    // Named node with fields or children → struct
    if !node.fields.is_empty() || node.children.is_some() {
        let type_name = name_mangler::type_ident(raw_kind);
        let fields = node
            .fields
            .iter()
            .map(|(field_name, field_info)| {
                let parent_name = type_name.to_string();
                map_field(field_name, field_info, &parent_name, supertype_kinds)
            })
            .collect();
        let children = node.children.as_ref().map(|c| {
            let parent_name = type_name.to_string();
            map_children(c, &parent_name, supertype_kinds)
        });
        return TypeDecision::Struct(StructDef {
            type_name,
            kind: raw_kind.clone(),
            fields,
            children,
        });
    }

    // Leaf node (named, no fields, no children, no subtypes)
    TypeDecision::LeafStruct(LeafStructDef {
        type_name: name_mangler::type_ident(raw_kind),
        kind: raw_kind.clone(),
    })
}

fn make_variant_def(tr: &TypeRef, supertype_kinds: &std::collections::HashSet<&str>) -> VariantDef {
    VariantDef {
        variant_name: name_mangler::variant_name(&tr.type_name, tr.named),
        kind: tr.type_name.clone(),
        named: tr.named,
        is_supertype: supertype_kinds.contains(tr.type_name.as_str()),
    }
}

fn map_field(
    field_name: &str,
    field_info: &FieldInfo,
    parent_name: &str,
    supertype_kinds: &std::collections::HashSet<&str>,
) -> FieldDef {
    let type_ref = map_type_reference(&field_info.types, parent_name, field_name, supertype_kinds);
    let field_type = match (field_info.required, field_info.multiple) {
        (_, true) => FieldType::Repeated(type_ref),
        (false, false) => FieldType::Optional(type_ref),
        (true, false) => FieldType::Direct(type_ref),
    };

    FieldDef {
        field_name: name_mangler::field_ident(field_name),
        raw_field_name: field_name.to_owned(),
        field_type,
    }
}

fn map_children(
    children: &FieldInfo,
    parent_name: &str,
    supertype_kinds: &std::collections::HashSet<&str>,
) -> ChildrenDef {
    let type_ref = map_type_reference(&children.types, parent_name, "children", supertype_kinds);
    let field_type = match (children.required, children.multiple) {
        (_, true) => FieldType::Repeated(type_ref),
        (false, false) => FieldType::Optional(type_ref),
        (true, false) => FieldType::Direct(type_ref),
    };
    ChildrenDef { field_type }
}

fn map_type_reference(
    types: &[TypeRef],
    parent_name: &str,
    field_name: &str,
    supertype_kinds: &std::collections::HashSet<&str>,
) -> TypeReference {
    // Filter to only named types for the type reference.
    // Anonymous nodes (punctuation) in field types are unusual but can appear in alternations.
    if types.len() == 1 && types[0].named {
        TypeReference::Named(name_mangler::type_ident(&types[0].type_name))
    } else {
        // Multiple types or contains anonymous → alternation enum
        let enum_name = format!(
            "{}{}",
            parent_name,
            name_mangler::to_pascal_case(field_name)
        );
        TypeReference::Alternation(AlternationEnumDef {
            type_name: quote::format_ident!("{}", enum_name),
            variants: types
                .iter()
                .map(|tr| make_variant_def(tr, supertype_kinds))
                .collect(),
        })
    }
}

/// Supertype nodes start with `_` (e.g., `_expression`). Strip the prefix for the type name.
fn supertype_ident(kind: &str) -> Ident {
    let stripped = kind.strip_prefix('_').unwrap_or(kind);
    name_mangler::type_ident(stripped)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codegen::grammar_ir::parse_node_types;

    #[test]
    fn test_leaf_node_maps_to_leaf_struct() {
        let nodes = parse_node_types(r#"[{"type": "identifier", "named": true}]"#).unwrap();
        let decisions = map_types(&nodes);
        assert_eq!(decisions.len(), 1);
        assert!(
            matches!(&decisions[0], TypeDecision::LeafStruct(def) if def.type_name == "Identifier")
        );
    }

    #[test]
    fn test_unnamed_nodes_are_skipped() {
        let nodes = parse_node_types(r#"[{"type": ".", "named": false}]"#).unwrap();
        let decisions = map_types(&nodes);
        assert!(decisions.is_empty());
    }

    #[test]
    fn test_node_with_fields_maps_to_struct() {
        let json = r#"[{
            "type": "import_spec",
            "named": true,
            "fields": {
                "path": {
                    "multiple": false,
                    "required": true,
                    "types": [{"type": "interpreted_string_literal", "named": true}]
                }
            }
        }]"#;
        let nodes = parse_node_types(json).unwrap();
        let decisions = map_types(&nodes);
        assert_eq!(decisions.len(), 1);
        let TypeDecision::Struct(def) = &decisions[0] else {
            panic!("expected Struct");
        };
        assert_eq!(def.type_name.to_string(), "ImportSpec");
        assert_eq!(def.fields.len(), 1);
        assert_eq!(def.fields[0].field_name.to_string(), "path");
        assert!(matches!(&def.fields[0].field_type, FieldType::Direct(_)));
    }

    #[test]
    fn test_optional_field() {
        let json = r#"[{
            "type": "import_spec",
            "named": true,
            "fields": {
                "name": {
                    "multiple": false,
                    "required": false,
                    "types": [{"type": "identifier", "named": true}]
                }
            }
        }]"#;
        let nodes = parse_node_types(json).unwrap();
        let decisions = map_types(&nodes);
        let TypeDecision::Struct(def) = &decisions[0] else {
            panic!("expected Struct");
        };
        assert!(matches!(&def.fields[0].field_type, FieldType::Optional(_)));
    }

    #[test]
    fn test_repeated_field() {
        let json = r#"[{
            "type": "block",
            "named": true,
            "fields": {
                "statements": {
                    "multiple": true,
                    "required": false,
                    "types": [{"type": "statement", "named": true}]
                }
            }
        }]"#;
        let nodes = parse_node_types(json).unwrap();
        let decisions = map_types(&nodes);
        let TypeDecision::Struct(def) = &decisions[0] else {
            panic!("expected Struct");
        };
        assert!(matches!(&def.fields[0].field_type, FieldType::Repeated(_)));
    }

    #[test]
    fn test_alternation_field() {
        let json = r#"[{
            "type": "import_spec",
            "named": true,
            "fields": {
                "name": {
                    "multiple": false,
                    "required": false,
                    "types": [
                        {"type": ".", "named": false},
                        {"type": "identifier", "named": true}
                    ]
                }
            }
        }]"#;
        let nodes = parse_node_types(json).unwrap();
        let decisions = map_types(&nodes);
        let TypeDecision::Struct(def) = &decisions[0] else {
            panic!("expected Struct");
        };
        let FieldType::Optional(TypeReference::Alternation(alt)) = &def.fields[0].field_type else {
            panic!("expected Optional(Alternation)");
        };
        assert_eq!(alt.type_name.to_string(), "ImportSpecName");
        assert_eq!(alt.variants.len(), 2);
        assert_eq!(alt.variants[0].variant_name.to_string(), "Dot");
        assert_eq!(alt.variants[1].variant_name.to_string(), "Identifier");
    }

    #[test]
    fn test_supertype_maps_to_enum() {
        let json = r#"[{
            "type": "_expression",
            "named": true,
            "subtypes": [
                {"type": "binary_expression", "named": true},
                {"type": "identifier", "named": true}
            ]
        }]"#;
        let nodes = parse_node_types(json).unwrap();
        let decisions = map_types(&nodes);
        let TypeDecision::SupertypeEnum(def) = &decisions[0] else {
            panic!("expected SupertypeEnum");
        };
        assert_eq!(def.type_name.to_string(), "Expression");
        assert_eq!(def.variants.len(), 2);
    }

    #[test]
    fn test_node_with_children() {
        let json = r#"[{
            "type": "import_spec_list",
            "named": true,
            "fields": {},
            "children": {
                "multiple": true,
                "required": false,
                "types": [{"type": "import_spec", "named": true}]
            }
        }]"#;
        let nodes = parse_node_types(json).unwrap();
        let decisions = map_types(&nodes);
        let TypeDecision::Struct(def) = &decisions[0] else {
            panic!("expected Struct");
        };
        assert!(def.children.is_some());
        assert!(matches!(
            &def.children.as_ref().unwrap().field_type,
            FieldType::Repeated(_)
        ));
    }
}