onnx-runtime-loader 0.1.0-dev.2

ONNX protobuf loader for the ORT 2.0 runtime: model parsing, weight mmap, and shape inference into onnx-runtime-ir
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
//! Model-local function inlining (ONNX function expansion) at load time.
//!
//! An ONNX `ModelProto` may declare reusable subgraphs as `FunctionProto`s in
//! `ModelProto.functions`. A node whose `(domain, op_type, overload)` matches a
//! declared function's `(domain, name, overload)` is a *function call*: it is
//! semantically equivalent to the function body with the call's actual inputs,
//! outputs, and attributes substituted in.
//!
//! Our executor only has kernels for primitive ops, so this module rewrites the
//! `ModelProto` at the proto level — **before** [`crate::graph_builder`] runs —
//! so the rest of the pipeline never sees a function call. Because the rewrite
//! is proto-level, the existing `NodeProto → IR` conversion (attributes,
//! control-flow subgraphs) is reused unchanged.
//!
//! ## Algorithm (standard ONNX function expansion)
//!
//! For each function-call node, we splice in a fresh copy of the matched
//! function body:
//!
//! 1. **Value remapping.** Formal `input[i]`/`output[j]` names are mapped to the
//!    call's actual argument names (positionally). Every *other* value name in
//!    the body (an intermediate result) is renamed to a globally-fresh unique
//!    name (`__fn{K}_{orig}`, bumped until unused) so instantiations never
//!    collide with each other or with pre-existing model names. The empty name
//!    `""` (ONNX "absent optional") is never remapped. A pass-through output
//!    whose formal name aliases an input is wired via a boundary `Identity`.
//!
//! 2. **Attribute binding.** A body-node attribute with a non-empty
//!    `ref_attr_name = A` is a reference to the function's formal attribute `A`.
//!    It is resolved from the call site (the call node's attribute `A`), else the
//!    function's declared default (`attribute_proto` entry named `A`), else — if
//!    `A` is a required attribute (`FunctionProto.attribute`) — an error; else
//!    the attribute is dropped. Literal (non-`ref`) attributes are kept as-is.
//!
//! 3. **Recursion + fixpoint.** A function body may call other functions; those
//!    calls are expanded recursively to a fixpoint. True recursion (a function
//!    that transitively calls itself) is rejected rather than looped forever.
//!
//! 4. **Control-flow subgraphs.** Function calls may appear inside If/Loop/Scan
//!    subgraph bodies, and function bodies may themselves contain control flow;
//!    both are handled by recursing into every node's `Graph`/`Graphs`
//!    attributes. Attribute binding and value remapping are scope-aware: nested
//!    `ref_attr_name` references are bound at every depth, and a subgraph's own
//!    locals (inputs, initializers, node outputs) shadow outer captures.
//!
//! ## Opset policy
//!
//! `FunctionProto.opset_import` domains/versions are merged into the model's
//! `opset_import`, taking the highest version per domain. Per the ONNX spec the
//! operator schemas for a shared domain must be compatible across the two opset
//! lists, so a version difference is not treated as a conflict; a domain the
//! model does not yet declare is added.
//!
//! ## Overload policy
//!
//! Matching is exact on the full `(domain, name, overload)` triple, so an
//! overload set is disambiguated by the node's `overload` field.

use std::borrow::Cow;
use std::collections::{HashMap, HashSet};

use crate::proto::onnx::{
    AttributeProto, FunctionProto, GraphProto, ModelProto, NodeProto, OperatorSetIdProto,
};
use crate::LoaderError;

/// Unique identity of a model-local function: `(domain, name, overload)`.
type FnKey = (String, String, String);

fn fn_key_of_function(f: &FunctionProto) -> FnKey {
    (f.domain.clone(), f.name.clone(), f.overload.clone())
}

fn fn_key_of_call(n: &NodeProto) -> FnKey {
    (n.domain.clone(), n.op_type.clone(), n.overload.clone())
}

/// Expand every call to a model-local function in `model` into the function's
/// body, so the returned `ModelProto`'s graph (and all nested subgraphs) contain
/// only calls to ops the runtime has kernels for.
///
/// When `model.functions` is empty this is a no-op and the input is borrowed
/// back unchanged (`Cow::Borrowed`). Otherwise a rewritten owned `ModelProto`
/// is returned with `functions` cleared and function opset imports merged in.
pub fn inline_functions(model: &ModelProto) -> Result<Cow<'_, ModelProto>, LoaderError> {
    if model.functions.is_empty() {
        return Ok(Cow::Borrowed(model));
    }

    let mut funcs: HashMap<FnKey, &FunctionProto> = HashMap::new();
    for f in &model.functions {
        funcs.insert(fn_key_of_function(f), f);
    }

    let graph = model
        .graph
        .as_ref()
        .ok_or_else(|| LoaderError::GraphBuild("ModelProto has no graph".into()))?;

    let mut counter: usize = 0;
    let mut stack: Vec<FnKey> = Vec::new();
    // Every value name already in use model-wide, so generated internal names
    // can be allocated to be globally fresh (BUG 4). Updated as inlining adds
    // new node outputs.
    let mut used: HashSet<String> = HashSet::new();
    collect_used_names(graph, &mut used);
    // Set when inlining synthesizes any default-domain (`""`/`ai.onnx`) node,
    // e.g. a boundary `Identity` alias, so we can guarantee the model declares a
    // default-domain opset import for it (BUG 3 regression).
    let mut synthesized_default = false;
    let new_graph = inline_graph(
        graph,
        &funcs,
        &mut counter,
        &mut stack,
        &mut used,
        &mut synthesized_default,
    )?;

    let mut out = model.clone();
    out.graph = Some(new_graph);
    out.opset_import = merged_opset_imports(model);
    if synthesized_default {
        ensure_default_opset_import(&mut out.opset_import);
    }
    out.functions.clear();
    Ok(Cow::Owned(out))
}

/// Conservative default `ai.onnx` opset version used only when inlining
/// synthesizes a default-domain node but the model (and its functions) declared
/// no default-domain opset import at all — a valid ONNX model that, e.g., only
/// called custom-domain functions. Any version ≥ 1 satisfies loader validation.
const DEFAULT_ONNX_OPSET_VERSION: i64 = 17;

/// The ONNX default operator-set domain has two equivalent spellings: the empty
/// string `""` and the explicit `"ai.onnx"`. They denote the SAME domain, so any
/// opset-import comparison/dedup must treat them as one.
fn is_default_domain(domain: &str) -> bool {
    domain.is_empty() || domain == "ai.onnx"
}

/// Canonical map key for an opset-import domain: every spelling of the default
/// domain collapses to a single key so duplicates cannot survive merging.
const DEFAULT_DOMAIN_KEY: &str = "";

fn domain_key(domain: &str) -> String {
    if is_default_domain(domain) {
        DEFAULT_DOMAIN_KEY.to_string()
    } else {
        domain.to_string()
    }
}

/// Ensure `imports` contains a default-domain (`""`/`ai.onnx`) opset entry so a
/// synthesized default-domain node (e.g. a boundary `Identity`) passes loader
/// validation. An existing default-domain import (under either spelling) is left
/// untouched — never downgraded, never duplicated.
fn ensure_default_opset_import(imports: &mut Vec<OperatorSetIdProto>) {
    let has_default = imports.iter().any(|o| is_default_domain(&o.domain));
    if !has_default {
        imports.push(OperatorSetIdProto {
            domain: String::new(),
            version: DEFAULT_ONNX_OPSET_VERSION,
        });
    }
}

/// Merge every function's `opset_import` into the model's, taking the highest
/// version per domain. Preserves the model's original import ordering, then
/// appends any domains introduced solely by functions (in first-seen order).
///
/// The default domain is canonicalized: `""` and `"ai.onnx"` collapse to a
/// SINGLE entry at the highest contributed version, so a model importing
/// `"ai.onnx"` plus a function (or synthesized Identity path) contributing `""`
/// never yields two logically-duplicate default-domain imports. The emitted
/// default entry keeps the model's original spelling if it declared one (so we
/// don't gratuitously rewrite `"ai.onnx"`→`""`); otherwise it is spelled `""`.
fn merged_opset_imports(model: &ModelProto) -> Vec<OperatorSetIdProto> {
    let mut order: Vec<String> = Vec::new();
    let mut best: HashMap<String, i64> = HashMap::new();
    // Preferred spelling for the emitted default-domain entry: the model's
    // original spelling if it imported the default domain, else `""`.
    let mut default_spelling: Option<String> = None;
    let mut note = |domain: &str, version: i64, from_model: bool| {
        if from_model && is_default_domain(domain) && default_spelling.is_none() {
            default_spelling = Some(domain.to_string());
        }
        let key = domain_key(domain);
        match best.entry(key.clone()) {
            std::collections::hash_map::Entry::Occupied(mut e) => {
                if version > *e.get() {
                    *e.get_mut() = version;
                }
            }
            std::collections::hash_map::Entry::Vacant(e) => {
                order.push(key);
                e.insert(version);
            }
        }
    };
    for o in &model.opset_import {
        note(&o.domain, o.version, true);
    }
    for f in &model.functions {
        for o in &f.opset_import {
            note(&o.domain, o.version, false);
        }
    }
    order
        .into_iter()
        .map(|key| {
            let version = best[&key];
            let domain = if key == DEFAULT_DOMAIN_KEY {
                default_spelling.clone().unwrap_or_default()
            } else {
                key
            };
            OperatorSetIdProto { domain, version }
        })
        .collect()
}

/// Rewrite `gp` so its node list contains no calls to any declared function.
/// Regular nodes are kept (with their control-flow subgraphs recursively
/// inlined); function-call nodes are replaced by their expanded bodies.
fn inline_graph(
    gp: &GraphProto,
    funcs: &HashMap<FnKey, &FunctionProto>,
    counter: &mut usize,
    stack: &mut Vec<FnKey>,
    used: &mut HashSet<String>,
    synthesized_default: &mut bool,
) -> Result<GraphProto, LoaderError> {
    let mut out = gp.clone();
    out.node = Vec::with_capacity(gp.node.len());
    for node in &gp.node {
        expand_node(node, funcs, counter, stack, used, synthesized_default, &mut out.node)?;
    }
    Ok(out)
}

/// Append the fully-inlined form of `node` to `sink`. If `node` calls a
/// function, its body (recursively inlined) is appended; otherwise the node is
/// appended with its subgraph attributes recursively inlined.
fn expand_node(
    node: &NodeProto,
    funcs: &HashMap<FnKey, &FunctionProto>,
    counter: &mut usize,
    stack: &mut Vec<FnKey>,
    used: &mut HashSet<String>,
    synthesized_default: &mut bool,
    sink: &mut Vec<NodeProto>,
) -> Result<(), LoaderError> {
    if let Some(func) = funcs.get(&fn_key_of_call(node)) {
        instantiate(node, func, funcs, counter, stack, used, synthesized_default, sink)?;
    } else {
        sink.push(inline_subgraph_attrs(
            node,
            funcs,
            counter,
            stack,
            used,
            synthesized_default,
        )?);
    }
    Ok(())
}

/// Return a copy of `node` whose `Graph`/`Graphs` attribute bodies have had any
/// function calls inside them inlined.
fn inline_subgraph_attrs(
    node: &NodeProto,
    funcs: &HashMap<FnKey, &FunctionProto>,
    counter: &mut usize,
    stack: &mut Vec<FnKey>,
    used: &mut HashSet<String>,
    synthesized_default: &mut bool,
) -> Result<NodeProto, LoaderError> {
    let mut out = node.clone();
    for attr in &mut out.attribute {
        if let Some(g) = attr.g.as_mut() {
            *g = inline_graph(g, funcs, counter, stack, used, synthesized_default)?;
        }
        for g in &mut attr.graphs {
            *g = inline_graph(g, funcs, counter, stack, used, synthesized_default)?;
        }
    }
    Ok(out)
}

/// Expand a single function call: substitute actual arguments and attributes
/// into a fresh copy of the function body, then recursively inline any calls the
/// body itself makes. Appends the resulting primitive nodes to `sink`.
#[allow(clippy::too_many_arguments)]
fn instantiate(
    call: &NodeProto,
    func: &FunctionProto,
    funcs: &HashMap<FnKey, &FunctionProto>,
    counter: &mut usize,
    stack: &mut Vec<FnKey>,
    used: &mut HashSet<String>,
    synthesized_default: &mut bool,
    sink: &mut Vec<NodeProto>,
) -> Result<(), LoaderError> {
    let key = fn_key_of_function(func);

    if stack.contains(&key) {
        let mut chain: Vec<String> = stack.iter().map(fmt_key).collect();
        chain.push(fmt_key(&key));
        return Err(LoaderError::RecursiveFunction {
            function: fmt_key(&key),
            chain: chain.join(" -> "),
        });
    }

    // Arity: passing *more* actuals than the function declares is illegal;
    // passing fewer is allowed (trailing optionals omitted, mapped to absent).
    if call.input.len() > func.input.len() {
        return Err(LoaderError::FunctionArityMismatch {
            function: fmt_key(&key),
            node: node_label(call),
            kind: "input",
            formal: func.input.len(),
            actual: call.input.len(),
        });
    }
    if call.output.len() > func.output.len() {
        return Err(LoaderError::FunctionArityMismatch {
            function: fmt_key(&key),
            node: node_label(call),
            kind: "output",
            formal: func.output.len(),
            actual: call.output.len(),
        });
    }

    let inst_id = *counter;
    *counter += 1;

    // The set of formal names actually produced by a body node. A formal output
    // that is *not* produced is a pass-through of an input (or otherwise-defined
    // value) and needs a boundary alias rather than a rename (BUG 3).
    let produced: HashSet<&str> = func
        .node
        .iter()
        .flat_map(|n| n.output.iter())
        .filter(|o| !o.is_empty())
        .map(String::as_str)
        .collect();

    // 1. Value remapping: formals -> actuals, everything else -> globally fresh.
    let mut rename: HashMap<String, String> = HashMap::new();
    // Boundary `Identity` aliases (src_actual -> dst_actual) for pass-through
    // outputs whose name aliases an input/other output (BUG 3).
    let mut aliases: Vec<(String, String)> = Vec::new();

    for (i, formal) in func.input.iter().enumerate() {
        if formal.is_empty() {
            continue;
        }
        let actual = call.input.get(i).cloned().unwrap_or_default();
        rename.insert(formal.clone(), actual);
    }
    for (j, formal) in func.output.iter().enumerate() {
        if formal.is_empty() {
            continue;
        }
        let actual = call.output.get(j).cloned().unwrap_or_default();
        if produced.contains(formal.as_str()) {
            // Genuinely produced by the body: consumers read the output actual.
            rename.insert(formal.clone(), actual);
        } else if let Some(src) = rename.get(formal) {
            // Pass-through: the formal is already bound (e.g. it is also an
            // input, or an earlier output). Keep body references reading the
            // source, and emit a boundary alias to the output actual.
            if !actual.is_empty() && src != &actual {
                aliases.push((src.clone(), actual));
            }
        } else {
            // Output not produced and not otherwise bound: map it directly.
            rename.insert(formal.clone(), actual);
        }
    }
    // Fresh, globally-unique names for internal (non-formal) body value names.
    for bn in &func.node {
        for name in bn.input.iter().chain(bn.output.iter()) {
            if name.is_empty() || rename.contains_key(name) {
                continue;
            }
            let fresh = fresh_name(name, inst_id, used);
            rename.insert(name.clone(), fresh);
        }
    }

    // 2. Attribute binding + value renaming for each body node.
    stack.push(key.clone());
    let result = (|| {
        let mut instantiated: Vec<NodeProto> = Vec::with_capacity(func.node.len());
        for (idx, bn) in func.node.iter().enumerate() {
            let mut nn = bn.clone();

            // Rename node name to a fresh unique one to avoid duplicate-name
            // collisions between instantiations.
            nn.name = if bn.name.is_empty() {
                format!("__fn{inst_id}_n{idx}")
            } else {
                format!("__fn{inst_id}_{}", bn.name)
            };

            // Bind attributes (resolve ref_attr_name against the call site) at
            // every depth, including nodes inside control-flow subgraphs (BUG 1).
            bind_node_attributes(&mut nn, call, func, &key)?;

            // Rename value references (inputs/outputs + captured names inside
            // any control-flow subgraph attributes, scope-aware).
            rename_value_refs(&mut nn, &rename);

            instantiated.push(nn);
        }

        // Boundary `Identity` aliases for pass-through outputs (BUG 3). Appended
        // last so their source values are already produced. `Identity` is a
        // default-domain op, so record that we synthesized one to guarantee the
        // model declares a default-domain opset import (BUG 3 regression).
        for (k, (src, dst)) in aliases.iter().enumerate() {
            *synthesized_default = true;
            instantiated.push(NodeProto {
                op_type: "Identity".to_string(),
                input: vec![src.clone()],
                output: vec![dst.clone()],
                name: format!("__fn{inst_id}_alias{k}"),
                ..Default::default()
            });
        }

        // 3. Recursively inline any function calls the body itself makes.
        let mut expanded: Vec<NodeProto> = Vec::new();
        for n in &instantiated {
            expand_node(n, funcs, counter, stack, used, synthesized_default, &mut expanded)?;
        }
        Ok::<Vec<NodeProto>, LoaderError>(expanded)
    })();
    stack.pop();

    sink.extend(result?);
    Ok(())
}

/// Bind a body node's attributes for a specific instantiation, recursing into
/// any control-flow subgraph so that `ref_attr_name` references carried by
/// nested nodes are resolved against the same call site (BUG 1).
fn bind_node_attributes(
    node: &mut NodeProto,
    call: &NodeProto,
    func: &FunctionProto,
    key: &FnKey,
) -> Result<(), LoaderError> {
    let mut bound: Vec<AttributeProto> = Vec::with_capacity(node.attribute.len());
    for attr in &node.attribute {
        if let Some(mut resolved) = bind_attribute(attr, call, func, key)? {
            if let Some(g) = resolved.g.as_mut() {
                for sub in &mut g.node {
                    bind_node_attributes(sub, call, func, key)?;
                }
            }
            for g in &mut resolved.graphs {
                for sub in &mut g.node {
                    bind_node_attributes(sub, call, func, key)?;
                }
            }
            bound.push(resolved);
        }
    }
    node.attribute = bound;
    Ok(())
}

/// Resolve a body-node attribute for a specific instantiation.
///
/// * Literal attribute (`ref_attr_name` empty): kept unchanged.
/// * Reference attribute (`ref_attr_name = A`): replaced by the call-site
///   attribute `A`, else the function's default for `A`, else dropped (if `A` is
///   optional) or an error (if `A` is required). The emitted attribute keeps the
///   body attribute's `name` and has `ref_attr_name` cleared.
///
/// Returns `Ok(None)` when the attribute should be omitted from the node.
fn bind_attribute(
    attr: &AttributeProto,
    call: &NodeProto,
    func: &FunctionProto,
    key: &FnKey,
) -> Result<Option<AttributeProto>, LoaderError> {
    if attr.ref_attr_name.is_empty() {
        return Ok(Some(attr.clone()));
    }
    let a = &attr.ref_attr_name;

    // Call-site value wins.
    if let Some(supplied) = call.attribute.iter().find(|ca| &ca.name == a) {
        let mut bound = supplied.clone();
        bound.name = attr.name.clone();
        bound.ref_attr_name.clear();
        return Ok(Some(bound));
    }
    // Otherwise the function's declared default, if any.
    if let Some(default) = func.attribute_proto.iter().find(|d| &d.name == a) {
        let mut bound = default.clone();
        bound.name = attr.name.clone();
        bound.ref_attr_name.clear();
        return Ok(Some(bound));
    }
    // No value and no default: an error if the attribute is required, else drop.
    if func.attribute.iter().any(|req| req == a) {
        return Err(LoaderError::MissingRequiredFunctionAttribute {
            function: fmt_key(key),
            node: node_label(call),
            attribute: a.clone(),
        });
    }
    Ok(None)
}

/// Apply `rename` to a node's value references: its inputs, its outputs, and any
/// value names captured inside its control-flow subgraph attributes. A name of
/// `""` (absent optional) is left untouched; a name absent from `rename` is left
/// as-is (subgraph-local names live in their own scope).
///
/// The node's own inputs/outputs live in the function-body scope, so they are
/// remapped directly. Subgraph attributes are remapped scope-aware
/// ([`rename_subgraph_refs`]).
fn rename_value_refs(node: &mut NodeProto, rename: &HashMap<String, String>) {
    for name in node.input.iter_mut().chain(node.output.iter_mut()) {
        if let Some(new) = rename.get(name.as_str()) {
            *name = new.clone();
        }
    }
    for attr in &mut node.attribute {
        if let Some(g) = attr.g.as_mut() {
            rename_subgraph_refs(g, rename);
        }
        for g in &mut attr.graphs {
            rename_subgraph_refs(g, rename);
        }
    }
}

/// Scope-aware renaming of outer-scope value captures inside a subgraph (BUG 2).
///
/// ONNX subgraphs have their own lexical scope. A subgraph's graph inputs,
/// initializers, and node outputs are *locals* that shadow any outer name, so
/// they must not be remapped. Only genuine captures of the enclosing scope —
/// node inputs, and `GraphProto.output` entries that directly name a captured
/// value — are rewritten to the outer actual. Shadowing is restored on descent
/// into deeper subgraphs by recomputing the local set at each level.
fn rename_subgraph_refs(gp: &mut GraphProto, rename: &HashMap<String, String>) {
    // Names locally bound in this subgraph shadow the outer scope.
    let mut locals: HashSet<&str> = HashSet::new();
    for i in &gp.input {
        if !i.name.is_empty() {
            locals.insert(i.name.as_str());
        }
    }
    for init in &gp.initializer {
        if !init.name.is_empty() {
            locals.insert(init.name.as_str());
        }
    }
    // Sparse initializers are also initializers (schema: GraphProto.
    // sparse_initializer), hence local bindings that shadow outer names.
    for sparse in &gp.sparse_initializer {
        if let Some(values) = &sparse.values
            && !values.name.is_empty()
        {
            locals.insert(values.name.as_str());
        }
    }    for n in &gp.node {
        for o in &n.output {
            if !o.is_empty() {
                locals.insert(o.as_str());
            }
        }
    }

    // Effective remap for this scope: outer captures minus anything shadowed.
    let effective: HashMap<String, String> = rename
        .iter()
        .filter(|(k, _)| !locals.contains(k.as_str()))
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();

    for n in &mut gp.node {
        for name in n.input.iter_mut().chain(n.output.iter_mut()) {
            if let Some(new) = effective.get(name.as_str()) {
                *name = new.clone();
            }
        }
        // Recurse into deeper subgraphs with this scope's effective map so a
        // name shadowed here stays shadowed, and is restored on the way out.
        for attr in &mut n.attribute {
            if let Some(g) = attr.g.as_mut() {
                rename_subgraph_refs(g, &effective);
            }
            for g in &mut attr.graphs {
                rename_subgraph_refs(g, &effective);
            }
        }
    }

    // A subgraph output that directly names a captured value must follow it.
    for o in &mut gp.output {
        if let Some(new) = effective.get(o.name.as_str()) {
            o.name = new.clone();
        }
    }
}

/// Collect every value name in use within `gp` (and its nested subgraphs):
/// graph inputs/outputs, initializers, value_info, and all node inputs/outputs.
/// Used to allocate globally-fresh generated names (BUG 4).
fn collect_used_names(gp: &GraphProto, used: &mut HashSet<String>) {
    for i in &gp.input {
        if !i.name.is_empty() {
            used.insert(i.name.clone());
        }
    }
    for o in &gp.output {
        if !o.name.is_empty() {
            used.insert(o.name.clone());
        }
    }
    for init in &gp.initializer {
        if !init.name.is_empty() {
            used.insert(init.name.clone());
        }
    }
    for sparse in &gp.sparse_initializer {
        if let Some(values) = &sparse.values
            && !values.name.is_empty()
        {
            used.insert(values.name.clone());
        }
    }
    for vi in &gp.value_info {
        if !vi.name.is_empty() {
            used.insert(vi.name.clone());
        }
    }
    for n in &gp.node {
        for name in n.input.iter().chain(n.output.iter()) {
            if !name.is_empty() {
                used.insert(name.clone());
            }
        }
        for attr in &n.attribute {
            if let Some(g) = &attr.g {
                collect_used_names(g, used);
            }
            for g in &attr.graphs {
                collect_used_names(g, used);
            }
        }
    }
}

/// Allocate a generated name for internal body value `base`, guaranteed unique
/// against every name already in use `used` (BUG 4). The chosen name is added to
/// `used` so subsequent allocations remain distinct.
fn fresh_name(base: &str, inst_id: usize, used: &mut HashSet<String>) -> String {
    let mut candidate = format!("__fn{inst_id}_{base}");
    let mut suffix = 0usize;
    while used.contains(&candidate) {
        suffix += 1;
        candidate = format!("__fn{inst_id}_{base}__{suffix}");
    }
    used.insert(candidate.clone());
    candidate
}

fn fmt_key(key: &FnKey) -> String {
    let (domain, name, overload) = key;
    let d = if domain.is_empty() { "ai.onnx" } else { domain };
    if overload.is_empty() {
        format!("{d}::{name}")
    } else {
        format!("{d}::{name}:{overload}")
    }
}

fn node_label(n: &NodeProto) -> String {
    if n.name.is_empty() {
        format!("<{}::{} (unnamed)>", n.domain, n.op_type)
    } else {
        n.name.clone()
    }
}