mirror of
https://github.com/ProvableHQ/leo.git
synced 2024-09-11 05:45:30 +03:00
Update flattener
This commit is contained in:
Pranav Gaddamadugu
Pranav Gaddamadugu
parent
f494a891c4
commit
cacc2154ea
@ -25,11 +25,12 @@ use leo_ast::{
|
||||
ExpressionReconstructor,
|
||||
Member,
|
||||
MemberAccess,
|
||||
Node,
|
||||
Statement,
|
||||
StructExpression,
|
||||
StructVariableInitializer,
|
||||
TernaryExpression,
|
||||
TupleExpression,
|
||||
Type,
|
||||
};
|
||||
|
||||
// TODO: Clean up logic. To be done in a follow-up PR (feat/tuples)
|
||||
@ -133,68 +134,55 @@ impl ExpressionReconstructor for Flattener<'_> {
|
||||
let mut statements = Vec::new();
|
||||
match (*input.if_true, *input.if_false) {
|
||||
// If both expressions are identifiers which are arrays, construct ternary expressions for each of the members and an array expression for the result.
|
||||
(Expression::Identifier(first), Expression::Identifier(second))
|
||||
if self.arrays.contains_key(&first.name) && self.arrays.contains_key(&second.name) =>
|
||||
{
|
||||
let first_array = self.arrays.get(&first.name).unwrap().clone();
|
||||
let second_array = self.arrays.get(&second.name).unwrap();
|
||||
// Note that type checking guarantees that both expressions have the same same type. This is a sanity check.
|
||||
assert_eq!(&first_array, second_array);
|
||||
(Expression::Identifier(first), Expression::Identifier(second)) => {
|
||||
match (self.type_table.get(&first.id()), self.type_table.get(&second.id())) {
|
||||
(Some(Type::Array(first_type)), Some(Type::Array(second_type))) => {
|
||||
// Note that type checking guarantees that both expressions have the same same type. This is a sanity check.
|
||||
assert_eq!(first_type, second_type);
|
||||
self.ternary_array(&first_type, &input.condition, &first, &second)
|
||||
}
|
||||
(Some(Type::Identifier(first_type)), Some(Type::Identifier(second_type))) => {
|
||||
// Get the struct definitions.
|
||||
let first_type = self.symbol_table.lookup_struct(first_type.name).unwrap();
|
||||
let second_type = self.symbol_table.lookup_struct(second_type.name).unwrap();
|
||||
// Note that type checking guarantees that both expressions have the same same type. This is a sanity check.
|
||||
assert_eq!(first_type, second_type);
|
||||
self.ternary_struct(first_type, &input.condition, &first, &second)
|
||||
}
|
||||
(Some(Type::Tuple(first_type)), Some(Type::Tuple(second_type))) => {
|
||||
// Note that type checking guarantees that both expressions have the same same type. This is a sanity check.
|
||||
assert_eq!(first_type, second_type);
|
||||
self.ternary_tuple(&first_type, &input.condition, &first, &second)
|
||||
}
|
||||
_ => {
|
||||
// Reconstruct the true case.
|
||||
let (if_true, stmts) = self.reconstruct_expression(Expression::Identifier(first));
|
||||
statements.extend(stmts);
|
||||
|
||||
self.ternary_array(first_array, &input.condition, &first, &second)
|
||||
// Reconstruct the false case.
|
||||
let (if_false, stmts) = self.reconstruct_expression(Expression::Identifier(second));
|
||||
statements.extend(stmts);
|
||||
|
||||
let (identifier, statement) =
|
||||
self.unique_simple_assign_statement(Expression::Ternary(TernaryExpression {
|
||||
condition: input.condition,
|
||||
if_true: Box::new(if_true),
|
||||
if_false: Box::new(if_false),
|
||||
span: input.span,
|
||||
id: input.id,
|
||||
}));
|
||||
|
||||
// Accumulate the new assignment statement.
|
||||
statements.push(statement);
|
||||
|
||||
(Expression::Identifier(identifier), statements)
|
||||
}
|
||||
}
|
||||
}
|
||||
// If both expressions are identifiers which are structs, construct ternary expression for each of the members and a struct expression for the result.
|
||||
(Expression::Identifier(first), Expression::Identifier(second))
|
||||
if self.structs.contains_key(&first.name) && self.structs.contains_key(&second.name) =>
|
||||
{
|
||||
let first_struct = self.structs.get(&first.name).unwrap().clone();
|
||||
let second_struct = self.structs.get(&second.name).unwrap();
|
||||
// Note that type checking guarantees that both expressions have the same same type. This is a sanity check.
|
||||
assert_eq!(&first_struct, second_struct);
|
||||
|
||||
self.ternary_struct(first_struct, &input.condition, &first, &second)
|
||||
}
|
||||
// If both expressions are identifiers which map to tuples, construct ternary expression over the tuples.
|
||||
(Expression::Identifier(first), Expression::Identifier(second))
|
||||
if self.tuples.contains_key(&first.name) && self.tuples.contains_key(&second.name) =>
|
||||
{
|
||||
// Note that this unwrap is safe since we check that `self.tuples` contains the key.
|
||||
let first_tuple = self.tuples.get(&first.name).unwrap();
|
||||
// Note that this unwrap is safe since we check that `self.tuples` contains the key.
|
||||
let second_tuple = self.tuples.get(&second.name).unwrap();
|
||||
// Note that type checking guarantees that both expressions have the same same type.
|
||||
self.reconstruct_ternary(TernaryExpression {
|
||||
condition: input.condition,
|
||||
if_true: Box::new(Expression::Tuple(first_tuple.clone())),
|
||||
if_false: Box::new(Expression::Tuple(second_tuple.clone())),
|
||||
span: input.span,
|
||||
id: input.id,
|
||||
})
|
||||
}
|
||||
// Otherwise, create a new intermediate assignment for the ternary expression are return the assigned variable.
|
||||
// Note that a new assignment must be created to flattened nested ternary expressions.
|
||||
(if_true, if_false) => {
|
||||
// Reconstruct the true case.
|
||||
let (if_true, stmts) = self.reconstruct_expression(if_true);
|
||||
statements.extend(stmts);
|
||||
|
||||
// Reconstruct the false case.
|
||||
let (if_false, stmts) = self.reconstruct_expression(if_false);
|
||||
statements.extend(stmts);
|
||||
|
||||
let (identifier, statement) =
|
||||
self.unique_simple_assign_statement(Expression::Ternary(TernaryExpression {
|
||||
condition: input.condition,
|
||||
if_true: Box::new(if_true),
|
||||
if_false: Box::new(if_false),
|
||||
span: input.span,
|
||||
id: input.id,
|
||||
}));
|
||||
|
||||
// Accumulate the new assignment statement.
|
||||
statements.push(statement);
|
||||
|
||||
(Expression::Identifier(identifier), statements)
|
||||
(expr1, expr2) => {
|
||||
println!("expr1: {:?}", expr1);
|
||||
println!("expr2: {:?}", expr2);
|
||||
unreachable!("SSA guarantees that the subexpressions of a ternary expression are identifiers.")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -24,15 +24,6 @@ impl ProgramReconstructor for Flattener<'_> {
|
||||
// First, flatten the finalize block. This allows us to initialize self.finalizes correctly.
|
||||
// Note that this is safe since the finalize block is independent of the function body.
|
||||
let finalize = function.finalize.map(|finalize| {
|
||||
// Initialize `self.structs` with the finalize's input as necessary.
|
||||
self.structs = Default::default();
|
||||
for input in &finalize.input {
|
||||
if let Type::Identifier(struct_name) = input.type_() {
|
||||
// Note that this unwrap is safe since type checking guarantees that the struct exists.
|
||||
let struct_ = self.symbol_table.lookup_struct(struct_name.name).unwrap().clone();
|
||||
self.structs.insert(input.identifier().name, struct_);
|
||||
}
|
||||
}
|
||||
// Flatten the finalize block.
|
||||
let mut block = self.reconstruct_block(finalize.block).0;
|
||||
|
||||
@ -53,16 +44,6 @@ impl ProgramReconstructor for Flattener<'_> {
|
||||
}
|
||||
});
|
||||
|
||||
// Initialize `self.structs` with the function's input as necessary.
|
||||
self.structs = Default::default();
|
||||
for input in &function.input {
|
||||
if let Type::Identifier(struct_name) = input.type_() {
|
||||
// Note that this unwrap is safe since type checking guarantees that the struct exists.
|
||||
let struct_ = self.symbol_table.lookup_struct(struct_name.name).unwrap().clone();
|
||||
self.structs.insert(input.identifier().name, struct_);
|
||||
}
|
||||
}
|
||||
|
||||
// Flatten the function body.
|
||||
let mut block = self.reconstruct_block(function.block).0;
|
||||
|
||||
|
||||
@ -107,13 +107,25 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
variant: AssertVariant::Assert(Expression::Binary(BinaryExpression {
|
||||
op: BinaryOperation::Or,
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the binary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Update the type table with the type of the binary expression.
|
||||
self.type_table.insert(id, Type::Boolean);
|
||||
id
|
||||
},
|
||||
// Take the logical negation of the guard.
|
||||
left: Box::new(Expression::Unary(UnaryExpression {
|
||||
op: UnaryOperation::Not,
|
||||
receiver: Box::new(guard),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the unary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Update the type table with the type of the unary expression.
|
||||
self.type_table.insert(id, Type::Boolean);
|
||||
id
|
||||
},
|
||||
})),
|
||||
right: Box::new(match assert.variant {
|
||||
// If the assert statement is an `assert`, use the expression as is.
|
||||
@ -124,7 +136,13 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
op: BinaryOperation::Eq,
|
||||
right: Box::new(right),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the unary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Update the type table with the type of the unary expression.
|
||||
self.type_table.insert(id, Type::Boolean);
|
||||
id
|
||||
},
|
||||
}),
|
||||
// If the assert statement is an `assert_ne`, construct a new inequality expression.
|
||||
AssertVariant::AssertNeq(left, right) => Expression::Binary(BinaryExpression {
|
||||
@ -132,7 +150,13 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
op: BinaryOperation::Neq,
|
||||
right: Box::new(right),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the unary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Update the type table with the type of the unary expression.
|
||||
self.type_table.insert(id, Type::Boolean);
|
||||
id
|
||||
},
|
||||
}),
|
||||
}),
|
||||
})),
|
||||
@ -184,26 +208,34 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
elements: (0..tuple.length())
|
||||
.zip_eq(tuple.elements().iter())
|
||||
.map(|(i, type_)| {
|
||||
let identifier = Identifier::new(
|
||||
// Return the identifier as an expression.
|
||||
Expression::Identifier(Identifier::new(
|
||||
self.assigner.unique_symbol(lhs_identifier.name, format!("$index${i}$")),
|
||||
self.node_builder.next_id(),
|
||||
);
|
||||
|
||||
// If the output type is a struct, add it to `self.structs`.
|
||||
if let Type::Identifier(struct_name) = type_ {
|
||||
// Note that this unwrap is safe since type checking guarantees that the struct exists.
|
||||
let struct_ = self.symbol_table.lookup_struct(struct_name.name).unwrap().clone();
|
||||
self.structs.insert(identifier.name, struct_);
|
||||
}
|
||||
|
||||
Expression::Identifier(identifier)
|
||||
{
|
||||
// Construct a node ID for the identifier.
|
||||
let id = self.node_builder.next_id();
|
||||
// Update the type table with the type.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
))
|
||||
})
|
||||
.collect(),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Construct a node ID for the tuple expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Update the type table with the type.
|
||||
self.type_table.insert(id, Type::Tuple(tuple.clone()));
|
||||
id
|
||||
},
|
||||
};
|
||||
// Add the `tuple_expression` to `self.tuples`.
|
||||
self.tuples.insert(lhs_identifier.name, tuple_expression.clone());
|
||||
|
||||
// Update the type table with the type of the tuple expression.
|
||||
self.type_table.insert(tuple_expression.id, Type::Tuple(tuple.clone()));
|
||||
|
||||
// Construct a new assignment statement with a tuple expression on the lhs.
|
||||
(
|
||||
Statement::Assign(Box::new(AssignStatement {
|
||||
@ -216,23 +248,7 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
)
|
||||
}
|
||||
// Otherwise, reconstruct the assignment as is.
|
||||
type_ => {
|
||||
// If the function returns a struct, add it to `self.structs`.
|
||||
if let Type::Identifier(struct_name) = type_ {
|
||||
// Note that this unwrap is safe since type checking guarantees that the struct exists.
|
||||
let struct_ = self.symbol_table.lookup_struct(struct_name.name).unwrap().clone();
|
||||
self.structs.insert(lhs_identifier.name, struct_);
|
||||
};
|
||||
(
|
||||
Statement::Assign(Box::new(AssignStatement {
|
||||
place: Expression::Identifier(lhs_identifier),
|
||||
value: Expression::Call(call),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
})),
|
||||
statements,
|
||||
)
|
||||
}
|
||||
_ => (self.simple_assign_statement(lhs_identifier, Expression::Call(call)), statements),
|
||||
}
|
||||
}
|
||||
// If the `rhs` is an invocation of `get` or `get_or_use` on a mapping, then check if the value type is a struct.
|
||||
@ -268,22 +284,8 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
}
|
||||
_ => unreachable!("Type checking guarantee that `arguments[0]` is the name of the mapping."),
|
||||
};
|
||||
// If the value type is a struct, add it to `self.structs`.
|
||||
if let Type::Identifier(struct_name) = value_type {
|
||||
// Note that this unwrap is safe since type checking guarantees that the struct exists.
|
||||
let struct_ = self.symbol_table.lookup_struct(struct_name.name).unwrap().clone();
|
||||
self.structs.insert(lhs_identifier.name, struct_);
|
||||
}
|
||||
// Reconstruct the assignment.
|
||||
(
|
||||
Statement::Assign(Box::new(AssignStatement {
|
||||
place: Expression::Identifier(lhs_identifier),
|
||||
value,
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
})),
|
||||
statements,
|
||||
)
|
||||
(self.simple_assign_statement(lhs_identifier, value), statements)
|
||||
}
|
||||
(Expression::Identifier(identifier), expression) => {
|
||||
(self.simple_assign_statement(identifier, expression), statements)
|
||||
@ -300,7 +302,7 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
let function = self.symbol_table.lookup_fn_symbol(function_name).unwrap();
|
||||
|
||||
let output_type = match &function.output_type {
|
||||
Type::Tuple(tuple) => tuple.clone(),
|
||||
Type::Tuple(tuple_type) => tuple_type.clone(),
|
||||
_ => unreachable!("Type checking guarantees that the output type is a tuple."),
|
||||
};
|
||||
|
||||
@ -309,14 +311,13 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
Expression::Identifier(identifier) => identifier,
|
||||
_ => unreachable!("Type checking guarantees that a tuple element on the lhs is an identifier."),
|
||||
};
|
||||
// If the output type is a struct, add it to `self.structs`.
|
||||
if let Type::Identifier(struct_name) = type_ {
|
||||
// Note that this unwrap is safe since type checking guarantees that the struct exists.
|
||||
let struct_ = self.symbol_table.lookup_struct(struct_name.name).unwrap().clone();
|
||||
self.structs.insert(identifier.name, struct_);
|
||||
}
|
||||
// Add the type of each identifier to the type table.
|
||||
self.type_table.insert(identifier.id, type_.clone());
|
||||
});
|
||||
|
||||
// Set the type of the tuple expression.
|
||||
self.type_table.insert(tuple.id, Type::Tuple(output_type.clone()));
|
||||
|
||||
(
|
||||
Statement::Assign(Box::new(AssignStatement {
|
||||
place: Expression::Tuple(tuple),
|
||||
@ -330,11 +331,14 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
// If the lhs is a tuple and the rhs is a tuple, create a new assign statement for each tuple element.
|
||||
(Expression::Tuple(lhs_tuple), Expression::Tuple(rhs_tuple)) => {
|
||||
statements.extend(lhs_tuple.elements.into_iter().zip(rhs_tuple.elements).map(|(lhs, rhs)| {
|
||||
let identifier = match &lhs {
|
||||
Expression::Identifier(identifier) => identifier,
|
||||
_ => unreachable!("Type checking guarantees that `lhs` is an identifier."),
|
||||
// Get the type of the rhs.
|
||||
let type_ = match self.type_table.get(&lhs.id()) {
|
||||
Some(type_) => type_.clone(),
|
||||
None => unreachable!("Type checking guarantees that the type of the lhs is in the type table."),
|
||||
};
|
||||
self.update_structs(identifier, &rhs);
|
||||
// Set the type of the lhs.
|
||||
self.type_table.insert(rhs.id(), type_);
|
||||
// Return the assign statement.
|
||||
Statement::Assign(Box::new(AssignStatement {
|
||||
place: lhs,
|
||||
value: rhs,
|
||||
@ -353,12 +357,14 @@ impl StatementReconstructor for Flattener<'_> {
|
||||
let rhs_tuple = self.tuples.get(&identifier.name).unwrap().clone();
|
||||
// Create a new assign statement for each tuple element.
|
||||
for (lhs, rhs) in lhs_tuple.elements.into_iter().zip(rhs_tuple.elements) {
|
||||
let identifier = match &lhs {
|
||||
Expression::Identifier(identifier) => identifier,
|
||||
_ => unreachable!("Type checking guarantees that `lhs` is an identifier."),
|
||||
// Get the type of the rhs.
|
||||
let type_ = match self.type_table.get(&lhs.id()) {
|
||||
Some(type_) => type_.clone(),
|
||||
None => unreachable!("Type checking guarantees that the type of the lhs is in the type table."),
|
||||
};
|
||||
self.update_structs(identifier, &rhs);
|
||||
|
||||
// Set the type of the lhs.
|
||||
self.type_table.insert(rhs.id(), type_);
|
||||
// Return the assign statement.
|
||||
statements.push(Statement::Assign(Box::new(AssignStatement {
|
||||
place: lhs,
|
||||
value: rhs,
|
||||
|
||||
@ -14,7 +14,7 @@
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
use crate::{Assigner, SymbolTable};
|
||||
use crate::{Assigner, SymbolTable, TypeTable};
|
||||
|
||||
use leo_ast::{
|
||||
AccessExpression,
|
||||
@ -31,6 +31,7 @@ use leo_ast::{
|
||||
Literal,
|
||||
Member,
|
||||
MemberAccess,
|
||||
Node,
|
||||
NodeBuilder,
|
||||
NonzeroNumber,
|
||||
ReturnStatement,
|
||||
@ -49,15 +50,17 @@ use leo_span::Symbol;
|
||||
|
||||
use indexmap::IndexMap;
|
||||
|
||||
// TODO: TypeTable should be placed behind a RefCell.
|
||||
|
||||
pub struct Flattener<'a> {
|
||||
/// The symbol table associated with the program.
|
||||
pub(crate) symbol_table: &'a SymbolTable,
|
||||
/// A mapping between node IDs and their types.
|
||||
pub(crate) type_table: &'a TypeTable,
|
||||
/// A counter used to generate unique node IDs.
|
||||
pub(crate) node_builder: &'a NodeBuilder,
|
||||
/// A struct used to construct (unique) assignment statements.
|
||||
pub(crate) assigner: &'a Assigner,
|
||||
/// The set of variables that are structs.
|
||||
pub(crate) structs: IndexMap<Symbol, Struct>,
|
||||
/// A stack of condition `Expression`s visited up to the current point in the AST.
|
||||
pub(crate) condition_stack: Vec<Expression>,
|
||||
/// A list containing tuples of guards and expressions associated `ReturnStatement`s.
|
||||
@ -67,21 +70,23 @@ pub struct Flattener<'a> {
|
||||
pub(crate) returns: Vec<(Option<Expression>, ReturnStatement)>,
|
||||
/// A mapping between variables and flattened tuple expressions.
|
||||
pub(crate) tuples: IndexMap<Symbol, TupleExpression>,
|
||||
/// A mapping from variables to array types.
|
||||
pub(crate) arrays: IndexMap<Symbol, ArrayType>,
|
||||
}
|
||||
|
||||
impl<'a> Flattener<'a> {
|
||||
pub(crate) fn new(symbol_table: &'a SymbolTable, node_builder: &'a NodeBuilder, assigner: &'a Assigner) -> Self {
|
||||
pub(crate) fn new(
|
||||
symbol_table: &'a SymbolTable,
|
||||
type_table: &'a TypeTable,
|
||||
node_builder: &'a NodeBuilder,
|
||||
assigner: &'a Assigner,
|
||||
) -> Self {
|
||||
Self {
|
||||
symbol_table,
|
||||
type_table,
|
||||
node_builder,
|
||||
assigner,
|
||||
structs: IndexMap::new(),
|
||||
condition_stack: Vec::new(),
|
||||
returns: Vec::new(),
|
||||
tuples: IndexMap::new(),
|
||||
arrays: IndexMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@ -97,12 +102,19 @@ impl<'a> Flattener<'a> {
|
||||
false => {
|
||||
let (first, rest) = self.condition_stack.split_first().unwrap();
|
||||
Some(rest.iter().cloned().fold(first.clone(), |acc, condition| {
|
||||
// Construct the binary expression.
|
||||
Expression::Binary(BinaryExpression {
|
||||
op: BinaryOperation::And,
|
||||
left: Box::new(acc),
|
||||
right: Box::new(condition),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the binary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, Type::Boolean);
|
||||
id
|
||||
},
|
||||
})
|
||||
}))
|
||||
}
|
||||
@ -136,11 +148,22 @@ impl<'a> Flattener<'a> {
|
||||
id: self.node_builder.next_id(),
|
||||
};
|
||||
let (value, stmts) = self.reconstruct_ternary(TernaryExpression {
|
||||
id: {
|
||||
// Create a new node ID for the ternary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Get the type of the node ID.
|
||||
let type_ = match self.type_table.get(&if_true.id()) {
|
||||
Some(type_) => type_,
|
||||
None => unreachable!("Type checking guarantees that all expressions have a type."),
|
||||
};
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_);
|
||||
id
|
||||
},
|
||||
condition: Box::new(guard),
|
||||
if_true: Box::new(if_true),
|
||||
if_false: Box::new(if_false),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
});
|
||||
statements.extend(stmts);
|
||||
|
||||
@ -167,11 +190,7 @@ impl<'a> Flattener<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Gets the type of the expression, if it's being tracked.
|
||||
pub(crate) fn get_type(&self, expression: &Expression) -> Option<Type> {}
|
||||
|
||||
/// A wrapper around `assigner.unique_simple_assign_statement` that updates `self.structs`.
|
||||
// TODO (@d0cd) Update to check for tuples and arrays
|
||||
pub(crate) fn unique_simple_assign_statement(&mut self, expr: Expression) -> (Identifier, Statement) {
|
||||
// Create a new variable for the expression.
|
||||
let name = self.assigner.unique_symbol("$var", "$");
|
||||
@ -180,20 +199,19 @@ impl<'a> Flattener<'a> {
|
||||
// Construct the assignment statement.
|
||||
let statement = self.simple_assign_statement(place, expr);
|
||||
|
||||
match &statement {
|
||||
Statement::Assign(assign) => {
|
||||
self.update_structs(&place, &assign.value);
|
||||
}
|
||||
_ => unreachable!("`assigner.unique_simple_assign_statement` always returns an assignment statement."),
|
||||
}
|
||||
(place, statement)
|
||||
}
|
||||
|
||||
/// A wrapper around `assigner.simple_assign_statement` that tracks the type of the lhs.
|
||||
pub(crate) fn simple_assign_statement(&mut self, lhs: Identifier, rhs: Expression) -> Statement {
|
||||
let statement = self.assigner.simple_assign_statement(lhs, rhs, self.node_builder.next_id());
|
||||
self.update_types(&statement);
|
||||
statement
|
||||
// Update the type table.
|
||||
let type_ = match self.type_table.get(&rhs.id()) {
|
||||
Some(type_) => type_,
|
||||
None => unreachable!("Type checking guarantees that all expressions have a type."),
|
||||
};
|
||||
self.type_table.insert(lhs.id(), type_);
|
||||
// Construct the statement.
|
||||
self.assigner.simple_assign_statement(lhs, rhs, self.node_builder.next_id())
|
||||
}
|
||||
|
||||
/// Folds a list of return statements into a single return statement and adds the produced statements to the block.
|
||||
@ -268,7 +286,7 @@ impl<'a> Flattener<'a> {
|
||||
|
||||
pub(crate) fn ternary_array(
|
||||
&mut self,
|
||||
array: ArrayType,
|
||||
array: &ArrayType,
|
||||
condition: &Expression,
|
||||
first: &Identifier,
|
||||
second: &Identifier,
|
||||
@ -286,10 +304,22 @@ impl<'a> Flattener<'a> {
|
||||
IntegerType::U32,
|
||||
i.to_string(),
|
||||
Default::default(),
|
||||
self.node_builder.next_id(),
|
||||
{
|
||||
// Create a new node ID for the literal.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, Type::Integer(IntegerType::U32));
|
||||
id
|
||||
},
|
||||
))),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the access expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, array.element_type().clone());
|
||||
id
|
||||
},
|
||||
})));
|
||||
statements.push(stmt);
|
||||
// Create an assignment statement for the second access expression.
|
||||
@ -300,10 +330,22 @@ impl<'a> Flattener<'a> {
|
||||
IntegerType::U32,
|
||||
i.to_string(),
|
||||
Default::default(),
|
||||
self.node_builder.next_id(),
|
||||
{
|
||||
// Create a new node ID for the literal.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, Type::Integer(IntegerType::U32));
|
||||
id
|
||||
},
|
||||
))),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the access expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, array.element_type().clone());
|
||||
id
|
||||
},
|
||||
})));
|
||||
statements.push(stmt);
|
||||
|
||||
@ -315,7 +357,13 @@ impl<'a> Flattener<'a> {
|
||||
// Access the member of the second expression.
|
||||
if_false: Box::new(Expression::Identifier(second)),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the ternary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, array.element_type().clone());
|
||||
id
|
||||
},
|
||||
});
|
||||
|
||||
// Accumulate any statements generated.
|
||||
@ -329,7 +377,13 @@ impl<'a> Flattener<'a> {
|
||||
let (expr, stmts) = self.reconstruct_array(ArrayExpression {
|
||||
elements,
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a node ID for the array expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, Type::Array(array.clone()));
|
||||
id
|
||||
},
|
||||
});
|
||||
|
||||
// Accumulate any statements generated.
|
||||
@ -338,9 +392,6 @@ impl<'a> Flattener<'a> {
|
||||
// Create a new assignment statement for the array expression.
|
||||
let (identifier, statement) = self.unique_simple_assign_statement(expr);
|
||||
|
||||
// Mark the lhs of the assignment as an array.
|
||||
self.arrays.insert(identifier.name, array);
|
||||
|
||||
statements.push(statement);
|
||||
|
||||
(Expression::Identifier(identifier), statements)
|
||||
@ -348,7 +399,7 @@ impl<'a> Flattener<'a> {
|
||||
|
||||
pub(crate) fn ternary_struct(
|
||||
&mut self,
|
||||
struct_: Struct,
|
||||
struct_: &Struct,
|
||||
condition: &Expression,
|
||||
first: &Identifier,
|
||||
second: &Identifier,
|
||||
@ -359,14 +410,20 @@ impl<'a> Flattener<'a> {
|
||||
let members = struct_
|
||||
.members
|
||||
.iter()
|
||||
.map(|Member { identifier, .. }| {
|
||||
.map(|Member { identifier, type_, .. }| {
|
||||
// Create an assignment statement for the first access expression.
|
||||
let (first, stmt) =
|
||||
self.unique_simple_assign_statement(Expression::Access(AccessExpression::Member(MemberAccess {
|
||||
inner: Box::new(Expression::Identifier(*first)),
|
||||
name: *identifier,
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the access expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
})));
|
||||
statements.push(stmt);
|
||||
// Create an assignment statement for the second access expression.
|
||||
@ -375,7 +432,13 @@ impl<'a> Flattener<'a> {
|
||||
inner: Box::new(Expression::Identifier(*second)),
|
||||
name: *identifier,
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the access expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
})));
|
||||
statements.push(stmt);
|
||||
// Recursively reconstruct the ternary expression.
|
||||
@ -386,7 +449,13 @@ impl<'a> Flattener<'a> {
|
||||
// Access the member of the second expression.
|
||||
if_false: Box::new(Expression::Identifier(second)),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the ternary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
});
|
||||
|
||||
// Accumulate any statements generated.
|
||||
@ -405,7 +474,13 @@ impl<'a> Flattener<'a> {
|
||||
name: struct_.identifier,
|
||||
members,
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the struct expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, Type::Identifier(struct_.identifier));
|
||||
id
|
||||
},
|
||||
});
|
||||
|
||||
// Accumulate any statements generated.
|
||||
@ -414,9 +489,6 @@ impl<'a> Flattener<'a> {
|
||||
// Create a new assignment statement for the struct expression.
|
||||
let (identifier, statement) = self.unique_simple_assign_statement(expr);
|
||||
|
||||
// Mark the lhs of the assignment as a struct.
|
||||
self.structs.insert(identifier.name, struct_);
|
||||
|
||||
statements.push(statement);
|
||||
|
||||
(Expression::Identifier(identifier), statements)
|
||||
@ -424,7 +496,7 @@ impl<'a> Flattener<'a> {
|
||||
|
||||
pub(crate) fn ternary_tuple(
|
||||
&mut self,
|
||||
tuple: TupleType,
|
||||
tuple_type: &TupleType,
|
||||
condition: &Expression,
|
||||
first: &Identifier,
|
||||
second: &Identifier,
|
||||
@ -432,15 +504,24 @@ impl<'a> Flattener<'a> {
|
||||
// Initialize a vector to accumulate any statements generated.
|
||||
let mut statements = Vec::new();
|
||||
// For each tuple element, construct a new ternary expression.
|
||||
let elements = (0..tuple.length())
|
||||
.map(|i| {
|
||||
let elements = tuple_type
|
||||
.elements()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, type_)| {
|
||||
// Create an assignment statement for the first access expression.
|
||||
let (first, stmt) =
|
||||
self.unique_simple_assign_statement(Expression::Access(AccessExpression::Tuple(TupleAccess {
|
||||
tuple: Box::new(Expression::Identifier(*first)),
|
||||
index: NonzeroNumber::from(i),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the access expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
})));
|
||||
statements.push(stmt);
|
||||
// Create an assignment statement for the second access expression.
|
||||
@ -449,7 +530,13 @@ impl<'a> Flattener<'a> {
|
||||
tuple: Box::new(Expression::Identifier(*second)),
|
||||
index: NonzeroNumber::from(i),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the access expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
})));
|
||||
statements.push(stmt);
|
||||
|
||||
@ -461,7 +548,13 @@ impl<'a> Flattener<'a> {
|
||||
// Access the member of the second expression.
|
||||
if_false: Box::new(Expression::Identifier(second)),
|
||||
span: Default::default(),
|
||||
id: self.node_builder.next_id(),
|
||||
id: {
|
||||
// Create a new node ID for the ternary expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, type_.clone());
|
||||
id
|
||||
},
|
||||
});
|
||||
|
||||
// Accumulate any statements generated.
|
||||
@ -472,7 +565,17 @@ impl<'a> Flattener<'a> {
|
||||
.collect();
|
||||
|
||||
// Construct the tuple expression.
|
||||
let tuple = TupleExpression { elements, span: Default::default(), id: self.node_builder.next_id() };
|
||||
let tuple = TupleExpression {
|
||||
elements,
|
||||
span: Default::default(),
|
||||
id: {
|
||||
// Create a new node ID for the tuple expression.
|
||||
let id = self.node_builder.next_id();
|
||||
// Set the type of the node ID.
|
||||
self.type_table.insert(id, Type::Tuple(tuple_type.clone()));
|
||||
id
|
||||
},
|
||||
};
|
||||
let (expr, stmts) = self.reconstruct_tuple(tuple.clone());
|
||||
|
||||
// Accumulate any statements generated.
|
||||
|
||||
@ -59,17 +59,17 @@ mod flatten_statement;
|
||||
pub mod flattener;
|
||||
pub use flattener::*;
|
||||
|
||||
use crate::{Assigner, Pass, SymbolTable};
|
||||
use crate::{Assigner, Pass, SymbolTable, TypeTable};
|
||||
|
||||
use leo_ast::{Ast, NodeBuilder, ProgramReconstructor};
|
||||
use leo_errors::Result;
|
||||
|
||||
impl<'a> Pass for Flattener<'a> {
|
||||
type Input = (Ast, &'a SymbolTable, &'a NodeBuilder, &'a Assigner);
|
||||
type Input = (Ast, &'a SymbolTable, &'a TypeTable, &'a NodeBuilder, &'a Assigner);
|
||||
type Output = Result<Ast>;
|
||||
|
||||
fn do_pass((ast, st, node_builder, assigner): Self::Input) -> Self::Output {
|
||||
let mut reconstructor = Flattener::new(st, node_builder, assigner);
|
||||
fn do_pass((ast, st, tt, node_builder, assigner): Self::Input) -> Self::Output {
|
||||
let mut reconstructor = Flattener::new(st, tt, node_builder, assigner);
|
||||
let program = reconstructor.reconstruct_program(ast.into_repr());
|
||||
|
||||
Ok(Ast::new(program))
|
||||
|
||||
Loading…
Reference in New Issue
Block a user