leo_passes/option_lowering/
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;
18
19use leo_ast::*;
20use leo_span::{Span, Symbol};
21
22use indexmap::IndexMap;
23
24pub struct OptionLoweringVisitor<'a> {
25    pub state: &'a mut CompilerState,
26    // The name of the current program scope
27    pub program: Symbol,
28    // The path to the current module. This should be an empty vector for the root.
29    pub module: Vec<Symbol>,
30    // The name of the current function, if we're inside one.
31    pub function: Option<Symbol>,
32    // The newly created structs. Each struct correspond to a converted optional type. All these
33    // structs are to be inserted in the program scope.
34    pub new_structs: IndexMap<Symbol, Composite>,
35    // The reconstructed structs. These are the new versions of the existing structs in the program.
36    pub reconstructed_structs: IndexMap<Vec<Symbol>, Composite>,
37}
38
39impl OptionLoweringVisitor<'_> {
40    /// Enter module scope with path `module`, execute `func`, and then return to the parent module.
41    pub fn in_module_scope<T>(&mut self, module: &[Symbol], func: impl FnOnce(&mut Self) -> T) -> T {
42        let parent_module = self.module.clone();
43        self.module = module.to_vec();
44        let result = func(self);
45        self.module = parent_module;
46        result
47    }
48
49    /// Wraps an expression of a given type in an `Optional<T>`-like struct representing `Some(value)`.
50    ///
51    /// This function creates a struct expression that encodes an optional value with `is_some = true`
52    /// and the provided expression as the `val` field. It also ensures that the type is fully
53    /// reconstructed, which guarantees that all necessary struct definitions are available and registered.
54    ///
55    /// # Parameters
56    /// - `expr`: The expression to wrap as the value of the optional.
57    /// - `ty`: The type of the expression.
58    ///
59    /// # Returns
60    /// - An `Expression` representing the constructed `Optional<T>` struct instance.
61    pub fn wrap_optional_value(&mut self, expr: Expression, ty: Type) -> Expression {
62        let is_some_expr = Expression::Literal(Literal {
63            span: Span::default(),
64            id: self.state.node_builder.next_id(),
65            variant: LiteralVariant::Boolean(true),
66        });
67
68        // Fully lower the type before proceeding. This also ensures that all required structs
69        // are actually registered in `self.new_structs`.
70        let lowered_inner_type = self.reconstruct_type(ty).0;
71
72        let struct_name = crate::make_optional_struct_symbol(&lowered_inner_type);
73        let struct_expr = StructExpression {
74            path: Path::from(Identifier::new(struct_name, self.state.node_builder.next_id())).into_absolute(),
75            const_arguments: vec![],
76            members: vec![
77                StructVariableInitializer {
78                    identifier: Identifier::new(Symbol::intern("is_some"), self.state.node_builder.next_id()),
79                    expression: Some(is_some_expr),
80                    span: Span::default(),
81                    id: self.state.node_builder.next_id(),
82                },
83                StructVariableInitializer {
84                    identifier: Identifier::new(Symbol::intern("val"), self.state.node_builder.next_id()),
85                    expression: Some(expr),
86                    span: Span::default(),
87                    id: self.state.node_builder.next_id(),
88                },
89            ],
90            span: Span::default(),
91            id: self.state.node_builder.next_id(),
92        };
93
94        struct_expr.into()
95    }
96
97    /// Constructs an `Optional<T>`-like struct representing `None` for a given inner type.
98    ///
99    /// The returned struct expression sets `is_some = false`, and provides a zero value for the `val`
100    /// field, where "zero" is defined according to the type:
101    /// numeric types use literal zero, structs are recursively zero-initialized, etc.
102    ///
103    /// This function assumes that all required struct types are already reconstructed and registered.
104    ///
105    /// # Parameters
106    /// - `inner_ty`: The type `T` inside the `Optional<T>`.
107    ///
108    /// # Returns
109    /// - An `Expression` representing the constructed `Optional<T>` struct instance with `None`.
110    pub fn wrap_none(&mut self, inner_ty: &Type) -> Expression {
111        let is_some_expr = Expression::Literal(Literal {
112            span: Span::default(),
113            id: self.state.node_builder.next_id(),
114            variant: LiteralVariant::Boolean(false),
115        });
116
117        // Fully lower the type before proceeding. This also ensures that all required structs
118        // are actually registered in `self.new_structs`.
119        let lowered_inner_type = self.reconstruct_type(inner_ty.clone()).0;
120
121        // Even though the `val` field of the struct will not be used as long as `is_some` is
122        // `false`, we still have to set it to something. We choose "zero", whatever "zero" means
123        // for each type.
124
125        // Instead of relying on the symbol table (which does not get updated in this pass), we rely on the set of
126        // reconstructed structs which is produced for all program scopes and all modules before doing anything else.
127        let reconstructed_structs = &self.reconstructed_structs;
128        let struct_lookup = |sym: &[Symbol]| {
129            reconstructed_structs
130                .get(sym) // check the new version of existing structs
131                .or_else(|| self.new_structs.get(sym.last().unwrap())) // check the newly produced structs
132                .expect("must exist by construction")
133                .members
134                .iter()
135                .map(|mem| (mem.identifier.name, mem.type_.clone()))
136                .collect()
137        };
138
139        let zero_val_expr =
140            zero_value_expression(&lowered_inner_type, Span::default(), &self.state.node_builder, &struct_lookup)
141                .expect("");
142
143        let struct_name = crate::make_optional_struct_symbol(&lowered_inner_type);
144
145        let struct_expr = StructExpression {
146            path: Path::from(Identifier::new(struct_name, self.state.node_builder.next_id())).into_absolute(),
147            const_arguments: vec![],
148            members: vec![
149                StructVariableInitializer {
150                    identifier: Identifier::new(Symbol::intern("is_some"), self.state.node_builder.next_id()),
151                    expression: Some(is_some_expr.clone()),
152                    span: Span::default(),
153                    id: self.state.node_builder.next_id(),
154                },
155                StructVariableInitializer {
156                    identifier: Identifier::new(Symbol::intern("val"), self.state.node_builder.next_id()),
157                    expression: Some(zero_val_expr.clone()),
158                    span: Span::default(),
159                    id: self.state.node_builder.next_id(),
160                },
161            ],
162            span: Span::default(),
163            id: self.state.node_builder.next_id(),
164        };
165
166        struct_expr.into()
167    }
168}
169
170#[allow(clippy::type_complexity)]
171fn zero_value_expression(
172    ty: &Type,
173    span: Span,
174    node_builder: &NodeBuilder,
175    struct_lookup: &dyn Fn(&[Symbol]) -> Vec<(Symbol, Type)>,
176) -> Option<Expression> {
177    let id = node_builder.next_id();
178
179    match ty {
180        // Numeric types
181        Type::Integer(IntegerType::I8) => Some(Literal::integer(IntegerType::I8, "0".to_string(), span, id).into()),
182        Type::Integer(IntegerType::I16) => Some(Literal::integer(IntegerType::I16, "0".to_string(), span, id).into()),
183        Type::Integer(IntegerType::I32) => Some(Literal::integer(IntegerType::I32, "0".to_string(), span, id).into()),
184        Type::Integer(IntegerType::I64) => Some(Literal::integer(IntegerType::I64, "0".to_string(), span, id).into()),
185        Type::Integer(IntegerType::I128) => Some(Literal::integer(IntegerType::I128, "0".to_string(), span, id).into()),
186        Type::Integer(IntegerType::U8) => Some(Literal::integer(IntegerType::U8, "0".to_string(), span, id).into()),
187        Type::Integer(IntegerType::U16) => Some(Literal::integer(IntegerType::U16, "0".to_string(), span, id).into()),
188        Type::Integer(IntegerType::U32) => Some(Literal::integer(IntegerType::U32, "0".to_string(), span, id).into()),
189        Type::Integer(IntegerType::U64) => Some(Literal::integer(IntegerType::U64, "0".to_string(), span, id).into()),
190        Type::Integer(IntegerType::U128) => Some(Literal::integer(IntegerType::U128, "0".to_string(), span, id).into()),
191
192        // Boolean
193        Type::Boolean => Some(Literal::boolean(false, span, id).into()),
194
195        // Field, Group, Scalar
196        Type::Field => Some(Literal::field("0".to_string(), span, id).into()),
197        Type::Group => Some(Literal::group("0".to_string(), span, id).into()),
198        Type::Scalar => Some(Literal::scalar("0".to_string(), span, id).into()),
199
200        // Structs (composite types)
201        Type::Composite(composite_type) => {
202            let path = &composite_type.path;
203            let members = struct_lookup(&path.absolute_path());
204
205            let struct_members = members
206                .into_iter()
207                .map(|(symbol, member_type)| {
208                    let member_id = node_builder.next_id();
209                    let zero_expr = zero_value_expression(&member_type, span, node_builder, struct_lookup)?;
210
211                    Some(StructVariableInitializer {
212                        span,
213                        id: member_id,
214                        identifier: Identifier::new(symbol, node_builder.next_id()),
215                        expression: Some(zero_expr),
216                    })
217                })
218                .collect::<Option<Vec<_>>>()?;
219
220            Some(Expression::Struct(StructExpression {
221                span,
222                id,
223                path: path.clone(),
224                const_arguments: composite_type.const_arguments.clone(),
225                members: struct_members,
226            }))
227        }
228
229        // Arrays
230        Type::Array(array_type) => {
231            let element_ty = &array_type.element_type;
232
233            let element_expr = zero_value_expression(element_ty, span, node_builder, struct_lookup)?;
234
235            Some(Expression::Repeat(
236                RepeatExpression { span, id, expr: element_expr, count: *array_type.length.clone() }.into(),
237            ))
238        }
239
240        // Other types are not expected or supported just yet
241        _ => None,
242    }
243}
leo_passes/option_lowering/visitor.rs

leo_passes/option_lowering/
visitor.rs
