leo_passes/common/
item_walkers.rs

1// Copyright (C) 2019-2026 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17//! Iterators over the functions and composites defined in a `Program`, `Library`, or `Stub`,
18//! plus `items_at_path` for filtering a `Location`-keyed map at a given module path.
19//!
20//! Used by passes like monomorphization, function inlining, and option lowering to seed their
21//! lookup maps with every Leo-AST-backed definition reachable from the top-level program —
22//! including those hidden behind `FromLeo` and `FromLibrary` stubs. `FromAleo` stubs are skipped
23//! because their bodies live in Aleo bytecode and are not walkable as AST.
24
25use indexmap::IndexMap;
26use leo_ast::{Composite, Function, Library, Location, Program, Stub};
27use leo_span::Symbol;
28
29/// Collects items from `map` whose `Location` lies at `path_prefix` within `program` — i.e. the
30/// item's path is exactly `path_prefix` followed by a single final segment. Yields
31/// `(final_segment, value)` pairs. Pass `&[]` for top-level items.
32///
33/// For example, given `program = foo` and entries keyed by
34/// `foo::bar`, `foo::baz::qux`, `foo::baz::quux`, `other::bar`:
35/// - `path_prefix = &[]` yields `(bar, ...)` — only the `foo::bar` top-level entry.
36/// - `path_prefix = &[baz]` yields `(qux, ...)` and `(quux, ...)` — items directly inside
37///   `foo::baz`, not `other::bar` (wrong program) and not `foo::bar` (path differs).
38pub fn items_at_path<'a, V: Clone + 'a>(
39    map: &'a IndexMap<Location, V>,
40    program: Symbol,
41    path_prefix: &'a [Symbol],
42) -> impl Iterator<Item = (Symbol, V)> + 'a {
43    map.iter().filter_map(move |(loc, v)| {
44        loc.path
45            .split_last()
46            .filter(|(_, rest)| *rest == path_prefix && loc.program == program)
47            .map(|(last, _)| (*last, v.clone()))
48    })
49}
50
51/// Yields `(Location, &Function)` for every function defined in `program`, including those in
52/// nested modules. `program.stubs` is not traversed.
53pub fn program_functions(program: &Program) -> impl Iterator<Item = (Location, &Function)> {
54    let scopes = program.program_scopes.iter().flat_map(|(_, scope)| {
55        let prog = scope.program_id.as_symbol();
56        scope.functions.iter().map(move |(name, f)| (Location::new(prog, vec![*name]), f))
57    });
58    let modules = program.modules.iter().flat_map(module_functions);
59    scopes.chain(modules)
60}
61
62/// Yields `(Location, &Composite)` for every composite defined in `program`, including those in
63/// nested modules.
64pub fn program_composites(program: &Program) -> impl Iterator<Item = (Location, &Composite)> {
65    let scopes = program.program_scopes.iter().flat_map(|(_, scope)| {
66        let prog = scope.program_id.as_symbol();
67        scope.composites.iter().map(move |(name, c)| (Location::new(prog, vec![*name]), c))
68    });
69    let modules = program.modules.iter().flat_map(module_composites);
70    scopes.chain(modules)
71}
72
73/// Yields `(Location, &Function)` for every function defined in `library`, including those in
74/// nested modules.
75pub fn library_functions(library: &Library) -> impl Iterator<Item = (Location, &Function)> {
76    let name = library.name;
77    let top = library.functions.iter().map(move |(sym, f)| (Location::new(name, vec![*sym]), f));
78    let modules = library.modules.iter().flat_map(module_functions);
79    top.chain(modules)
80}
81
82/// Yields `(Location, &Composite)` for every composite defined in `library`, including those in
83/// nested modules.
84pub fn library_composites(library: &Library) -> impl Iterator<Item = (Location, &Composite)> {
85    let name = library.name;
86    let top = library.structs.iter().map(move |(sym, c)| (Location::new(name, vec![*sym]), c));
87    let modules = library.modules.iter().flat_map(module_composites);
88    top.chain(modules)
89}
90
91/// Yields `(Location, &Function)` for every function defined in `stub` whose body is available as
92/// Leo AST. `FromAleo` stubs are skipped because their bodies live in Aleo bytecode.
93pub fn stub_functions(stub: &Stub) -> Box<dyn Iterator<Item = (Location, &Function)> + '_> {
94    match stub {
95        Stub::FromLeo { program, .. } => Box::new(program_functions(program)),
96        Stub::FromLibrary { library, .. } => Box::new(library_functions(library)),
97        Stub::FromAleo { .. } => Box::new(std::iter::empty()),
98    }
99}
100
101/// Yields `(Location, &Composite)` for every composite defined in `stub` whose body is available
102/// as Leo AST. `FromAleo` stubs are skipped.
103pub fn stub_composites(stub: &Stub) -> Box<dyn Iterator<Item = (Location, &Composite)> + '_> {
104    match stub {
105        Stub::FromLeo { program, .. } => Box::new(program_composites(program)),
106        Stub::FromLibrary { library, .. } => Box::new(library_composites(library)),
107        Stub::FromAleo { .. } => Box::new(std::iter::empty()),
108    }
109}
110
111fn module_functions<'a>(
112    (path, m): (&'a Vec<Symbol>, &'a leo_ast::Module),
113) -> impl Iterator<Item = (Location, &'a Function)> {
114    m.functions.iter().map(move |(name, f)| {
115        let full: Vec<Symbol> = path.iter().copied().chain(std::iter::once(*name)).collect();
116        (Location::new(m.unit_name, full), f)
117    })
118}
119
120fn module_composites<'a>(
121    (path, m): (&'a Vec<Symbol>, &'a leo_ast::Module),
122) -> impl Iterator<Item = (Location, &'a Composite)> {
123    m.composites.iter().map(move |(name, c)| {
124        let full: Vec<Symbol> = path.iter().copied().chain(std::iter::once(*name)).collect();
125        (Location::new(m.unit_name, full), c)
126    })
127}
leo_passes/common/item_walkers.rs

leo_passes/common/
item_walkers.rs
