ProvableHQ/leo
1
1
Fork 0
mirror of https://github.com/ProvableHQ/leo.git synced 2024-11-23 15:15:47 +03:00
Code Issues Packages Projects Releases Wiki Activity

Redesign flattening

Browse Source
This commit is contained in:
Pranav Gaddamadugu 2023-10-13 11:11:53 -04:00 committed by Pranav Gaddamadugu
parent cacc2154ea
commit 992f0b83de
3 changed files with 10 additions and 275 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
compiler/passes/src/flattening/flatten_expression.rs
View File

@ -30,6 +30,7 @@ use leo_ast::{
StructExpression,
StructVariableInitializer,
TernaryExpression,
TupleAccess,
Type,
};
@ -38,60 +39,6 @@ use leo_ast::{
impl ExpressionReconstructor for Flattener<'_> {
type AdditionalOutput = Vec<Statement>;
/// Replaces a tuple access expression with the appropriate expression.
fn reconstruct_access(&mut self, input: AccessExpression) -> (Expression, Self::AdditionalOutput) {
let mut statements = Vec::new();
(
match input {
AccessExpression::Array(array) => Expression::Access(AccessExpression::Array(ArrayAccess {
array: Box::new(self.reconstruct_expression(*array.array).0),
index: Box::new(self.reconstruct_expression(*array.index).0),
span: array.span,
id: array.id,
})),
AccessExpression::AssociatedFunction(function) => {
Expression::Access(AccessExpression::AssociatedFunction(AssociatedFunction {
ty: function.ty,
name: function.name,
arguments: function
.arguments
.into_iter()
.map(|arg| self.reconstruct_expression(arg).0)
.collect(),
span: function.span,
id: function.id,
}))
}
AccessExpression::Member(member) => Expression::Access(AccessExpression::Member(MemberAccess {
inner: Box::new(self.reconstruct_expression(*member.inner).0),
name: member.name,
span: member.span,
id: member.id,
})),
AccessExpression::Tuple(tuple) => {
// Reconstruct the tuple expression.
let (expr, stmts) = self.reconstruct_expression(*tuple.tuple);
// Accumulate any statements produced.
statements.extend(stmts);
// Lookup the expression in the tuple map.
match expr {
Expression::Identifier(identifier) => {
// Note that this unwrap is safe since TYC guarantees that all tuples are declared and indices are valid.
self.tuples.get(&identifier.name).unwrap().elements[tuple.index.value()].clone()
}
_ => unreachable!("SSA guarantees that subexpressions are identifiers or literals."),
}
}
AccessExpression::AssociatedConstant(access) => {
Expression::Access(AccessExpression::AssociatedConstant(access))
}
},
statements,
)
}
/// Reconstructs a struct init expression, flattening any tuples in the expression.
fn reconstruct_struct_init(&mut self, input: StructExpression) -> (Expression, Self::AdditionalOutput) {
let mut statements = Vec::new();

210
compiler/passes/src/flattening/flatten_statement.rs
View File

@ -173,136 +173,11 @@ impl StatementReconstructor for Flattener<'_> {
fn reconstruct_assign(&mut self, assign: AssignStatement) -> (Statement, Self::AdditionalOutput) {
// Flatten the rhs of the assignment.
let (value, mut statements) = self.reconstruct_expression(assign.value);
match (assign.place, value.clone()) {
// If the lhs is an identifier and the rhs is a tuple, then add the tuple to `self.tuples`.
(Expression::Identifier(identifier), Expression::Tuple(tuple)) => {
self.tuples.insert(identifier.name, tuple);
// Note that tuple assignments are removed from the AST.
(Statement::dummy(Default::default(), self.node_builder.next_id()), statements)
}
// If the lhs is an identifier and the rhs is an identifier that is a tuple, then add it to `self.tuples`.
(Expression::Identifier(lhs_identifier), Expression::Identifier(rhs_identifier))
if self.tuples.contains_key(&rhs_identifier.name) =>
{
// Lookup the entry in `self.tuples` and add it for the lhs of the assignment.
// Note that the `unwrap` is safe since the match arm checks that the entry exists.
self.tuples.insert(lhs_identifier.name, self.tuples.get(&rhs_identifier.name).unwrap().clone());
// Note that tuple assignments are removed from the AST.
(Statement::dummy(Default::default(), self.node_builder.next_id()), statements)
}
// If the lhs is an identifier and the rhs is a function call that produces a tuple, then add it to `self.tuples`.
(Expression::Identifier(lhs_identifier), Expression::Call(call)) => {
// Retrieve the entry in the symbol table for the function call.
// Note that this unwrap is safe since type checking ensures that the function exists.
let function_name = match call.function.borrow() {
Expression::Identifier(rhs_identifier) => rhs_identifier.name,
_ => unreachable!("Parsing guarantees that `function` is an identifier."),
};
let function = self.symbol_table.lookup_fn_symbol(function_name).unwrap();
match &function.output_type {
// If the function returns a tuple, reconstruct the assignment and add an entry to `self.tuples`.
Type::Tuple(tuple) => {
// Create a new tuple expression with unique identifiers for each index of the lhs.
let tuple_expression = TupleExpression {
elements: (0..tuple.length())
.zip_eq(tuple.elements().iter())
.map(|(i, type_)| {
// Return the identifier as an expression.
Expression::Identifier(Identifier::new(
self.assigner.unique_symbol(lhs_identifier.name, format!("$index${i}$")),
{
// 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: {
// 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 {
place: Expression::Tuple(tuple_expression),
value: Expression::Call(call),
span: Default::default(),
id: self.node_builder.next_id(),
})),
statements,
)
}
// Otherwise, reconstruct the assignment as is.
_ => (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.
// Note that the parser rewrites `.get` and `.get_or_use` to `Mapping::get` and `Mapping::get_or_use` respectively.
(
Expression::Identifier(lhs_identifier),
Expression::Access(AccessExpression::AssociatedFunction(AssociatedFunction {
ty: Type::Identifier(Identifier { name: sym::Mapping, .. }),
name: Identifier { name: sym::get, .. },
arguments,
..
})),
)
| (
Expression::Identifier(lhs_identifier),
Expression::Access(AccessExpression::AssociatedFunction(AssociatedFunction {
ty: Type::Identifier(Identifier { name: sym::Mapping, .. }),
name: Identifier { name: sym::get_or_use, .. },
arguments,
..
})),
) => {
// Get the value type of the mapping.
let value_type = match arguments[0] {
Expression::Identifier(identifier) => {
// Retrieve the entry in the symbol table for the mapping.
// Note that this unwrap is safe since type checking ensures that the mapping exists.
let variable = self.symbol_table.lookup_variable(identifier.name).unwrap();
match &variable.type_ {
Type::Mapping(mapping_type) => &*mapping_type.value,
_ => unreachable!("Type checking guarantee that `arguments[0]` is a mapping."),
}
}
_ => unreachable!("Type checking guarantee that `arguments[0]` is the name of the mapping."),
};
// Reconstruct the assignment.
(self.simple_assign_statement(lhs_identifier, value), statements)
}
(Expression::Identifier(identifier), expression) => {
(self.simple_assign_statement(identifier, expression), statements)
}
// If the lhs is a tuple and the rhs is a function call, then return the reconstructed statement.
(Expression::Tuple(tuple), Expression::Call(call)) => {
// Retrieve the entry in the symbol table for the function call.
// Note that this unwrap is safe since type checking ensures that the function exists.
let function_name = match call.function.borrow() {
Expression::Identifier(rhs_identifier) => rhs_identifier.name,
_ => unreachable!("Parsing guarantees that `function` is an identifier."),
};
let function = self.symbol_table.lookup_fn_symbol(function_name).unwrap();
let output_type = match &function.output_type {
Type::Tuple(tuple_type) => tuple_type.clone(),
match (assign.place, &value) {
(Expression::Identifier(identifier), _) => (self.simple_assign_statement(identifier, value), statements),
(Expression::Tuple(tuple), expression) => {
let output_type = match &self.type_table.get(&expression.id()) {
Some(Type::Tuple(tuple_type)) => tuple_type.clone(),
_ => unreachable!("Type checking guarantees that the output type is a tuple."),
};
@ -321,67 +196,13 @@ impl StatementReconstructor for Flattener<'_> {
(
Statement::Assign(Box::new(AssignStatement {
place: Expression::Tuple(tuple),
value: Expression::Call(call),
value,
span: Default::default(),
id: self.node_builder.next_id(),
})),
statements,
)
}
// 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)| {
// 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."),
};
// 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,
span: Default::default(),
id: self.node_builder.next_id(),
}))
}));
(Statement::dummy(Default::default(), self.node_builder.next_id()), statements)
}
// If the lhs is a tuple and the rhs is an identifier that is a tuple, create a new assign statement for each tuple element.
(Expression::Tuple(lhs_tuple), Expression::Identifier(identifier))
if self.tuples.contains_key(&identifier.name) =>
{
// Lookup the entry in `self.tuples`.
// Note that the `unwrap` is safe since the match arm checks that the entry exists.
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) {
// 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."),
};
// 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,
span: Default::default(),
id: self.node_builder.next_id(),
})));
}
(Statement::dummy(Default::default(), self.node_builder.next_id()), statements)
}
// If the lhs of an assignment is a tuple, then the rhs can be one of the following:
// - A function call that produces a tuple. (handled above)
// - A tuple. (handled above)
// - An identifier that is a tuple. (handled above)
// - A ternary expression that produces a tuple. (handled when the rhs is flattened above)
(Expression::Tuple(_), _) => {
unreachable!("`Type checking guarantees that the rhs of an assignment to a tuple is a tuple.`")
}
_ => unreachable!("`AssignStatement`s can only have `Identifier`s or `Tuple`s on the left hand side."),
}
}
@ -442,12 +263,10 @@ impl StatementReconstructor for Flattener<'_> {
unreachable!("`ConsoleStatement`s should not be in the AST at this phase of compilation.")
}
/// Static single assignment converts definition statements into assignment statements.
fn reconstruct_definition(&mut self, _definition: DefinitionStatement) -> (Statement, Self::AdditionalOutput) {
unreachable!("`DefinitionStatement`s should not exist in the AST at this phase of compilation.")
}
// TODO: Error message requesting the user to enable loop-unrolling.
fn reconstruct_iteration(&mut self, _input: IterationStatement) -> (Statement, Self::AdditionalOutput) {
unreachable!("`IterationStatement`s should not be in the AST at this phase of compilation.");
}
@ -458,22 +277,9 @@ impl StatementReconstructor for Flattener<'_> {
// Construct the associated guard.
let guard = self.construct_guard();
// Note that SSA guarantees that `input.expression` is either a literal, identifier, or unit expression.
match input.expression {
// If the input is an identifier that maps to a tuple,
// construct a `ReturnStatement` with the tuple and add it to `self.returns`
Expression::Identifier(identifier) if self.tuples.contains_key(&identifier.name) => {
// Note that the `unwrap` is safe since the match arm checks that the entry exists in `self.tuples`.
let tuple = self.tuples.get(&identifier.name).unwrap().clone();
self.returns.push((guard, ReturnStatement {
span: input.span,
expression: Expression::Tuple(tuple),
finalize_arguments: input.finalize_arguments,
id: input.id,
}));
}
// Otherwise, add the expression directly.
_ => self.returns.push((guard, input)),
Expression::Unit(_) | Expression::Identifier(_) => self.returns.push((guard, input)),
_ => unreachable!("SSA guarantees that the expression is always an identifier or unit expression."),
};
(Statement::dummy(Default::default(), self.node_builder.next_id()), Default::default())

20
compiler/passes/src/flattening/flattener.rs
View File

@ -46,11 +46,6 @@ use leo_ast::{
Type,
UnitExpression,
};
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.
@ -68,8 +63,6 @@ pub struct Flattener<'a> {
/// Note that returns are inserted in the order they are encountered during a pre-order traversal of the AST.
/// Note that type checking guarantees that there is at most one return in a basic block.
pub(crate) returns: Vec<(Option<Expression>, ReturnStatement)>,
/// A mapping between variables and flattened tuple expressions.
pub(crate) tuples: IndexMap<Symbol, TupleExpression>,
}
impl<'a> Flattener<'a> {
@ -79,15 +72,7 @@ impl<'a> Flattener<'a> {
node_builder: &'a NodeBuilder,
assigner: &'a Assigner,
) -> Self {
Self {
symbol_table,
type_table,
node_builder,
assigner,
condition_stack: Vec::new(),
returns: Vec::new(),
tuples: IndexMap::new(),
}
Self { symbol_table, type_table, node_builder, assigner, condition_stack: Vec::new(), returns: Vec::new() }
}
/// Clears the state associated with `ReturnStatements`, returning the ones that were previously stored.
@ -584,9 +569,6 @@ impl<'a> Flattener<'a> {
// Create a new assignment statement for the tuple expression.
let (identifier, statement) = self.unique_simple_assign_statement(expr);
// Mark the lhs of the assignment as a tuple.
self.tuples.insert(identifier.name, tuple);
statements.push(statement);
(Expression::Identifier(identifier), statements)