Skip to main content

brk_bindgen/generate/
tree.rs

1//! Shared tree generation helpers.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use brk_types::TreeNode;
6
7use crate::{
8    ClientMetadata, PatternBaseResult, PatternField, child_type_name, get_fields_with_child_info,
9};
10
11/// Build a child path by appending a child name to a parent path.
12/// Uses "/" as separator. If parent is empty, returns just the child name.
13#[inline]
14pub fn build_child_path(parent: &str, child: &str) -> String {
15    if parent.is_empty() {
16        child.to_string()
17    } else {
18        format!("{}/{}", parent, child)
19    }
20}
21
22/// Pre-computed context for a single child node.
23pub struct ChildContext<'a> {
24    /// The child's field name in the tree.
25    pub name: &'a str,
26    /// The child node.
27    pub node: &'a TreeNode,
28    /// The field info for this child (with type_param set for generic patterns).
29    pub field: PatternField,
30    /// Pattern analysis result.
31    pub base_result: PatternBaseResult,
32    /// Whether this is a leaf node.
33    pub is_leaf: bool,
34    /// Whether to use an inline type instead of a pattern type (only meaningful for branches).
35    pub should_inline: bool,
36    /// The type name to use for inline branches.
37    pub inline_type_name: String,
38}
39
40/// Context for generating a tree node, returned by `prepare_tree_node`.
41pub struct TreeNodeContext<'a> {
42    /// Pre-computed context for each child.
43    pub children: Vec<ChildContext<'a>>,
44}
45
46/// Prepare a tree node for generation.
47/// Returns None if the node should be skipped (not a branch, already generated,
48/// or matches a parameterizable pattern).
49///
50/// The `path` parameter is the tree path to this node (e.g., "distribution/utxoCohorts").
51/// It's used to look up pre-computed PatternBaseResult from the analysis phase.
52pub fn prepare_tree_node<'a>(
53    node: &'a TreeNode,
54    name: &str,
55    path: &str,
56    pattern_lookup: &BTreeMap<Vec<PatternField>, String>,
57    metadata: &ClientMetadata,
58    generated: &mut BTreeSet<String>,
59) -> Option<TreeNodeContext<'a>> {
60    let TreeNode::Branch(branch_children) = node else {
61        return None;
62    };
63
64    let fields_with_child_info = get_fields_with_child_info(branch_children, name, pattern_lookup);
65    let fields: Vec<PatternField> = fields_with_child_info
66        .iter()
67        .map(|(f, _)| f.clone())
68        .collect();
69
70    // Look up the pre-computed base result, or use a default that forces inlining
71    let base_result = metadata
72        .get_node_base(path)
73        .cloned()
74        .unwrap_or_else(PatternBaseResult::force_inline);
75
76    // Skip if this matches a parameterizable pattern AND has no outlier AND field parts match
77    let pattern_compatible = pattern_lookup
78        .get(&fields)
79        .and_then(|name| metadata.find_pattern(name))
80        .is_none_or(|p| {
81            p.is_suffix_mode() == base_result.is_suffix_mode
82                && p.field_parts_match(&base_result.field_parts)
83        });
84    if let Some(pattern_name) = pattern_lookup.get(&fields)
85        && pattern_name != name
86        && metadata.is_parameterizable(pattern_name)
87        && !base_result.has_outlier
88        && pattern_compatible
89    {
90        return None;
91    }
92
93    // Skip if already generated
94    if generated.contains(name) {
95        return None;
96    }
97    generated.insert(name.to_string());
98
99    // Build child contexts with pre-computed decisions
100    let children: Vec<ChildContext<'a>> = branch_children
101        .iter()
102        .zip(fields_with_child_info)
103        .map(|((child_name, child_node), (mut field, child_fields))| {
104            let is_leaf = matches!(child_node, TreeNode::Leaf(_));
105
106            // Set type_param for generic patterns so field_type_annotation works directly
107            if let Some(cf) = &child_fields {
108                field.type_param = metadata.get_type_param(cf).cloned();
109            }
110
111            // Build child path and look up its pre-computed base result
112            let child_path = build_child_path(path, child_name);
113            let base_result = metadata
114                .get_node_base(&child_path)
115                .cloned()
116                .unwrap_or_else(PatternBaseResult::force_inline);
117
118            // Single lookup for the child's matching pattern (avoids repeated scans)
119            let matching_pattern = child_fields
120                .as_ref()
121                .and_then(|cf| metadata.find_pattern_by_fields(cf));
122
123            let matches_any_pattern = matching_pattern.is_some();
124            let pattern_compatible = matching_pattern.is_none_or(|p| {
125                p.is_suffix_mode() == base_result.is_suffix_mode
126                    && p.field_parts_match(&base_result.field_parts)
127            });
128            let is_parameterizable =
129                matching_pattern.is_none_or(|p| metadata.is_parameterizable(&p.name));
130
131            // should_inline determines if we generate an inline struct type
132            let should_inline = !is_leaf
133                && (!matches_any_pattern
134                    || !pattern_compatible
135                    || !is_parameterizable
136                    || base_result.has_outlier);
137
138            let inline_type_name = if should_inline {
139                child_type_name(name, child_name)
140            } else {
141                String::new()
142            };
143
144            ChildContext {
145                name: child_name,
146                node: child_node,
147                field,
148                base_result,
149                is_leaf,
150                should_inline,
151                inline_type_name,
152            }
153        })
154        .collect();
155
156    Some(TreeNodeContext { children })
157}