leo_passes/monomorphization/
visitor.rs

1// Copyright (C) 2019-2025 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
17use crate::{CompilerState, Replacer};
18use indexmap::{IndexMap, IndexSet};
19use itertools::Itertools;
20use leo_ast::{
21    CallExpression,
22    Composite,
23    CompositeType,
24    Expression,
25    Function,
26    Identifier,
27    Path,
28    ProgramReconstructor,
29    StructExpression,
30};
31use leo_span::Symbol;
32
33pub struct MonomorphizationVisitor<'a> {
34    pub state: &'a mut CompilerState,
35    /// The main program.
36    pub program: Symbol,
37    /// A map to provide faster lookup of functions.
38    pub function_map: IndexMap<Vec<Symbol>, Function>,
39    /// A map to provide faster lookup of structs.
40    pub struct_map: IndexMap<Vec<Symbol>, Composite>,
41    /// A map of reconstructed functions in the current program scope.
42    pub reconstructed_functions: IndexMap<Vec<Symbol>, Function>,
43    /// A set of all functions that have been monomorphized at least once. This keeps track of the _original_ names of
44    /// the functions not the names of the monomorphized versions.
45    pub monomorphized_functions: IndexSet<Vec<Symbol>>,
46    /// A map of reconstructed functions in the current program scope.
47    pub reconstructed_structs: IndexMap<Vec<Symbol>, Composite>,
48    /// A set of all functions that have been monomorphized at least once. This keeps track of the _original_ names of
49    /// the functions not the names of the monomorphized versions.
50    pub monomorphized_structs: IndexSet<Vec<Symbol>>,
51    /// A vector of all the calls to const generic functions that have not been resolved.
52    pub unresolved_calls: Vec<CallExpression>,
53    /// A vector of all the struct expressions of const generic structs that have not been resolved.
54    pub unresolved_struct_exprs: Vec<StructExpression>,
55    /// A vector of all the struct type instantiations of const generic structs that have not been resolved.
56    pub unresolved_struct_types: Vec<CompositeType>,
57    /// Have we actually modified the program at all?
58    pub changed: bool,
59}
60
61impl MonomorphizationVisitor<'_> {
62    /// Monomorphizes a generic struct by substituting const parameters with concrete arguments and caching result.
63    /// Generates a unique name like `Foo::[1u32, 2u32]` based on the original name and the provided const arguments.
64    /// Replaces all const parameter references in the struct body with values, then resolves nested generics.
65    /// Assigns a new name and struct ID, clears const params, and stores the result to avoid reprocessing.
66    /// Panics if the original struct is not found in `reconstructed_structs` (should already be reconstructed).
67    ///
68    /// # Arguments
69    /// * `name` - Symbol of the original generic struct.
70    /// * `const_arguments` - Const values to substitute into the struct.
71    /// * Returns a `Symbol` for the newly monomorphized struct.
72    ///
73    /// Note: this functions already assumes that all provided const arguments are literals.
74    pub(crate) fn monomorphize_struct(&mut self, path: &Path, const_arguments: &Vec<Expression>) -> Path {
75        // Generate a unique name for the monomorphized struct based on const arguments.
76        //
77        // For `struct Foo::[x: u32, y: u32](..)`, the generated name would be `Foo::[1u32, 2u32]` for a struct
78        // expression that sets `x` to `1u32` and `y` to `2u32`. We know this name is safe to use because it's not a
79        // valid identifier in the user code.
80        //
81        // Later, we have to legalize these names because they are not valid Aleo identifiers. We do this in codegen.
82        let new_struct_path = path.clone().with_updated_last_symbol(leo_span::Symbol::intern(&format!(
83            "{}::[{}]",
84            path.identifier().name,
85            const_arguments.iter().format(", ")
86        )));
87
88        // Check if the new struct name is not already present in `reconstructed_structs`. This ensures that we do not
89        // add a duplicate definition for the same struct.
90        if self.reconstructed_structs.get(&new_struct_path.absolute_path()).is_none() {
91            let full_name = path.absolute_path();
92            // Look up the already reconstructed struct by name.
93            let struct_ = self
94                .reconstructed_structs
95                .get(&full_name)
96                .expect("Struct should already be reconstructed (post-order traversal).");
97
98            // Build mapping from const parameters to const argument values.
99            let const_param_map: IndexMap<_, _> =
100                struct_.const_parameters.iter().map(|param| param.identifier().name).zip_eq(const_arguments).collect();
101
102            // Function to replace path expressions with their corresponding const argument or keep them unchanged.
103            let replace_path = |expr: &Expression| match expr {
104                Expression::Path(path) => const_param_map
105                    .get(&path.identifier().name)
106                    .map_or(Expression::Path(path.clone()), |&expr| expr.clone()),
107                _ => expr.clone(),
108            };
109
110            let mut replacer = Replacer::new(replace_path, true /* refresh IDs */, self.state);
111
112            // Create a new version of `struct_` that has a new name, no const parameters, and a new struct ID.
113            //
114            // First, reconstruct the struct by changing all instances of const generic parameters to literals
115            // according to `const_param_map`.
116            let mut struct_ = replacer.reconstruct_struct(struct_.clone());
117
118            // Now, reconstruct the struct to actually monomorphize its content such as generic struct type
119            // instantiations.
120            struct_ = self.reconstruct_struct(struct_);
121            struct_.identifier = Identifier::new(new_struct_path.identifier().name, self.state.node_builder.next_id());
122            struct_.const_parameters = vec![];
123            struct_.id = self.state.node_builder.next_id();
124
125            // Keep track of the new struct in case other structs need it.
126            self.reconstructed_structs.insert(new_struct_path.absolute_path(), struct_);
127
128            // Now keep track of the struct we just monomorphized
129            self.monomorphized_structs.insert(full_name);
130        }
131
132        new_struct_path
133    }
134}
leo_passes/monomorphization/visitor.rs

leo_passes/monomorphization/
visitor.rs
