leo_passes/flattening/
mod.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
17//! The flattening pass traverses the AST after the SSA pass and converts into a sequential code.
18//! The pass flattens `ConditionalStatement`s into a sequence of `AssignStatement`s.
19//! The pass rewrites `ReturnStatement`s into `AssignStatement`s and consolidates the returned values as a single `ReturnStatement` at the end of the function.
20//! The pass rewrites ternary expressions over composite data types, into ternary expressions over the individual fields of the composite data type, followed by an expression constructing the composite data type.
21//! Note that this transformation is not applied to async functions.
22//!
23//! Consider the following Leo code, output by the SSA pass.
24//! ```leo
25//! function main(flag: u8, value: u8) -> u8 {
26//!     $var$0 = flag == 0u8;
27//!     if ($var$0) {
28//!         $var$1 = value + 1u8;
29//!         value$1 = $var$1;
30//!         return value$1;
31//!     } else {
32//!         $var$2 = value + 2u8;
33//!         value$2 = $var$2;
34//!     }
35//!     value$3 = $var$0 ? value$1 : value$2;
36//!     return value$3;
37//! }
38//! ```
39//!
40//! The flattening pass produces the following code.
41//! ```leo
42//! function main(flag: u8, value: u8) -> u8 {
43//!     $var$0 = flag == 0u8;
44//!     $var$1 = value + 1u8;
45//!     value$1 = $var$1;
46//!     $var$2 = value + 2u8;
47//!     value$2 = $var$2;
48//!     value$3 = $var$0 ? value$1 : value$2;
49//!     ret$4 = $var$0 ? value$1 : value$3;
50//!     return ret$4;
51//! }
52//! ```
53
54use crate::Pass;
55
56use leo_ast::ProgramReconstructor as _;
57use leo_errors::Result;
58use leo_span::Symbol;
59
60mod ast;
61
62mod program;
63
64mod visitor;
65use visitor::*;
66
67pub struct Flattening;
68
69impl Pass for Flattening {
70    type Input = ();
71    type Output = ();
72
73    const NAME: &str = "Flattening";
74
75    fn do_pass(_input: Self::Input, state: &mut crate::CompilerState) -> Result<Self::Output> {
76        let mut ast = std::mem::take(&mut state.ast);
77        let mut visitor = FlatteningVisitor {
78            state,
79            condition_stack: Vec::new(),
80            returns: Vec::new(),
81            program: Symbol::intern(""),
82            is_async: false,
83        };
84        ast.ast = visitor.reconstruct_program(ast.ast);
85        visitor.state.handler.last_err()?;
86        visitor.state.ast = ast;
87        Ok(())
88    }
89}