Skip to content

feat(passes): leverage snarkVM aggregate ternary opcodes #29349

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mohammadfawaz wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

feat(passes): leverage snarkVM aggregate ternary opcodes #29349

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
287 changes: 78 additions & 209 deletions Cargo.lock
View file
Open in desktop

Large diffs are not rendered by default.

11 changes: 6 additions & 5 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,14 +69,15 @@ serial_test = "3.1"
sha2 = "0.11"
similar = "3.1"
# We enable "test_consensus_heights" to give developers freedom to set custom consensus version upgrade heights and "test_targets" because the devnode's genesis block was created with the same feature flag.
# Uses the published `4.7.1` release, which ships V15 `view fn` support (snarkVM PRs #3238, #3257, #3253).
# Pinned to the `mohammadfawaz/complex_ternary` branch (snarkVM PR #3222), rebased on staging so it
# ships V15 `view fn` support (snarkVM PRs #3238, #3257, #3253) alongside the aggregate `ternary` opcodes.
# `history` is required to compile `Process::evaluate_view_with_history`, the public-facing view entry point.
snarkvm = { version = "4.7.1", features = [ "test_consensus_heights", "dev_skip_checks", "test_targets", "history" ] }
# Wasm-friendly snarkVM subsets pinned at the same version as the umbrella crate; used by
snarkvm = { git = "https://github.com/ProvableHQ/snarkVM", branch = "mohammadfawaz/complex_ternary", features = [ "test_consensus_heights", "dev_skip_checks", "test_targets", "history" ] }
# Wasm-friendly snarkVM subsets pinned at the same branch as the umbrella crate; used by
# `leo-ast` and `leo-disassembler` on `wasm32-unknown-unknown` where the full `snarkvm`
# crate pulls in native-only ledger/package code.
snarkvm-console = { version = "4.7.1", default-features = false, features = [ "account", "program", "types", "dev_skip_checks", "test_consensus_heights", "test_targets", "wasm" ] }
snarkvm-synthesizer-program = { version = "4.7.1", features = [ "wasm" ] }
snarkvm-console = { git = "https://github.com/ProvableHQ/snarkVM", branch = "mohammadfawaz/complex_ternary", default-features = false, features = [ "account", "program", "types", "dev_skip_checks", "test_consensus_heights", "test_targets", "wasm" ] }
snarkvm-synthesizer-program = { git = "https://github.com/ProvableHQ/snarkVM", branch = "mohammadfawaz/complex_ternary", features = [ "wasm" ] }
sys-info = "0.9"
tempfile = "3.20"
thiserror = "2.0"
Expand Down
67 changes: 28 additions & 39 deletions crates/passes/src/flattening/ast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,21 +52,17 @@ impl AstReconstructor for FlatteningVisitor<'_> {
(CompositeExpression { members, ..input }.into(), statements)
}

/// Reconstructs ternary expressions over arrays, composites, and tuples, accumulating any statements that are generated.
/// This is necessary because Aleo instructions does not support ternary expressions over composite data types.
/// For example, the ternary expression `cond ? (a, b) : (c, d)` is flattened into the following:
/// Reconstructs ternary expressions. Tuple operands are destructured element-wise and record
/// operands are destructured field-wise, since the Aleo `ternary` instruction only accepts
/// plaintext operands. Ternaries over literal, array, and plaintext struct types are left
/// intact; `ConsensusVersion::V15` accepts these as direct operands (snarkVM PR #3222).
///
/// Example: `cond ? (a, b) : (c, d)` flattens to:
/// ```leo
/// let var$0 = cond ? a : c;
/// let var$1 = cond ? b : d;
/// (var$0, var$1)
/// ```
/// For composites, the ternary expression `cond ? a : b`, where `a` and `b` are both composites `Foo { bar: u8, baz: u8 }`, is flattened into the following:
/// ```leo
/// let var$0 = cond ? a.bar : b.bar;
/// let var$1 = cond ? a.baz : b.baz;
/// let var$2 = Foo { bar: var$0, baz: var$1 };
/// var$2
/// ```
fn reconstruct_ternary(
&mut self,
input: TernaryExpression,
Expand Down Expand Up @@ -94,39 +90,32 @@ impl AstReconstructor for FlatteningVisitor<'_> {
}

match &if_true_type {
Type::Array(if_true_type) => self.ternary_array(
if_true_type,
&input.condition,
&as_identifier(input.if_true),
&as_identifier(input.if_false),
),
Type::Composite(if_true_type) => {
let composite_location = if_true_type.path.expect_global_location();
// Get the composite definitions.
let composite_path = if_true_type.path.clone();
let if_true_type = self
.state
.symbol_table
.lookup_struct(self.program, composite_location)
.or_else(|| self.state.symbol_table.lookup_record(self.program, composite_location))
.expect("This definition should exist")
.clone();

self.ternary_composite(
&composite_path,
&if_true_type,
&input.condition,
&as_identifier(input.if_true),
&as_identifier(input.if_false),
)
}
Type::Tuple(if_true_type) => {
self.ternary_tuple(if_true_type, &input.condition, &input.if_true, &input.if_false)
}
Type::Composite(composite_type) => {
// The Aleo `ternary` instruction accepts plaintext structs but not records, so we
// destructure record ternaries field-wise and leave struct ternaries intact.
let composite_location = composite_type.path.expect_global_location();
if let Some(record) = self.state.symbol_table.lookup_record(self.program, composite_location) {
let composite_path = composite_type.path.clone();
let record = record.clone();
self.ternary_composite(
&composite_path,
&record,
&input.condition,
&as_identifier(input.if_true),
&as_identifier(input.if_false),
)
} else {
assert!(matches!(&input.if_true, Expression::Path(..)));
assert!(matches!(&input.if_false, Expression::Path(..)));

(input.into(), Default::default())
}
}
_ => {
// There's nothing to be done - SSA has guaranteed that `if_true` and `if_false` are identifiers,
// so there's not even any point in reconstructing them.

// SSA guarantees both branches are identifiers, so nothing else needs rewriting.
assert!(matches!(&input.if_true, Expression::Path(..)));
assert!(matches!(&input.if_false, Expression::Path(..)));

Expand Down
20 changes: 16 additions & 4 deletions crates/passes/src/flattening/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,22 @@
// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.

//! The flattening pass traverses the AST after the SSA pass and converts into a sequential code.
//! The pass flattens `ConditionalStatement`s into a sequence of `AssignStatement`s.
//! The pass rewrites `ReturnStatement`s into `AssignStatement`s and consolidates the returned values as a single `ReturnStatement` at the end of the function.
//! 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.
//! Note that this transformation is not applied to async functions.
//!
//! In transition (offchain) functions the pass:
//! - Flattens `ConditionalStatement`s into a sequence of `AssignStatement`s guarded by ternaries.
//! - Rewrites `ReturnStatement`s into `AssignStatement`s and consolidates the returned values into
//! a single `ReturnStatement` at the end of the function.
//! - Guards asserts with the active path condition and any early-return guard.
//!
//! In finalize (onchain) functions the pass preserves `ConditionalStatement`s, `ReturnStatement`s,
//! and `AssertStatement`s verbatim, since finalize code can use Aleo branch instructions directly.
//!
//! In both contexts the pass rewrites ternary expressions whose branches are tuples into
//! per-element ternary expressions and ternaries over records into per-field ternaries followed
//! by a composite init, since the Aleo `ternary` instruction only accepts plaintext operands.
//! Ternaries over literal, array, and plaintext struct types are left intact;
//! `ConsensusVersion::V15` accepts these as direct operands of the `ternary` instruction
//! (snarkVM PR #3222).
//!
//! Consider the following Leo code, output by the SSA pass.
//! ```leo
Expand Down
114 changes: 6 additions & 108 deletions crates/passes/src/flattening/visitor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,6 @@
use crate::CompilerState;

use leo_ast::{
ArrayAccess,
ArrayExpression,
ArrayType,
AstReconstructor,
BinaryExpression,
BinaryOperation,
Expand All @@ -30,7 +27,6 @@ use leo_ast::{
CompositeType,
Expression,
Identifier,
IntegerType,
Literal,
Member,
MemberAccess,
Expand Down Expand Up @@ -364,95 +360,7 @@ impl FlatteningVisitor<'_> {
}
}

// For use in `ternary_array`.
fn make_array_access_definition(
&mut self,
i: usize,
identifier: Identifier,
array_type: &ArrayType,
) -> (Identifier, Statement) {
let index =
Literal::integer(IntegerType::U32, i.to_string(), Default::default(), self.state.node_builder.next_id());
self.state.type_table.insert(index.id(), Type::Integer(IntegerType::U32));
let access: Expression = ArrayAccess {
array: Path::from(identifier).to_local().into(),
index: index.into(),
span: Default::default(),
id: self.state.node_builder.next_id(),
}
.into();
self.state.type_table.insert(access.id(), array_type.element_type().clone());
self.unique_simple_definition(access)
}

pub fn ternary_array(
&mut self,
array: &ArrayType,
condition: &Expression,
first: &Identifier,
second: &Identifier,
) -> (Expression, Vec<Statement>) {
// Initialize a vector to accumulate any statements generated.
let mut statements = Vec::new();
// For each array element, construct a new ternary expression.
let elements = (0..array.length.as_u32().expect("length should be known at this point") as usize)
.map(|i| {
// Create an assignment statement for the first access expression.
let (first, stmt) = self.make_array_access_definition(i, *first, array);
statements.push(stmt);
// Create an assignment statement for the second access expression.
let (second, stmt) = self.make_array_access_definition(i, *second, array);
statements.push(stmt);

// Recursively reconstruct the ternary expression.
let ternary = TernaryExpression {
condition: condition.clone(),
// Access the member of the first expression.
if_true: Path::from(first).to_local().into(),
// Access the member of the second expression.
if_false: Path::from(second).to_local().into(),
span: Default::default(),
id: self.state.node_builder.next_id(),
};
self.state.type_table.insert(ternary.id(), array.element_type().clone());

let (expression, stmts) = self.reconstruct_ternary(ternary, &());

// Accumulate any statements generated.
statements.extend(stmts);

expression
})
.collect();

// Construct the array expression.
let (expr, stmts) = self.reconstruct_array(
ArrayExpression {
elements,
span: Default::default(),
id: {
// Create a node ID for the array expression.
let id = self.state.node_builder.next_id();
// Set the type of the node ID.
self.state.type_table.insert(id, Type::Array(array.clone()));
id
},
},
&(),
);

// Accumulate any statements generated.
statements.extend(stmts);

// Create a new assignment statement for the array expression.
let (identifier, statement) = self.unique_simple_definition(expr);

statements.push(statement);

(Path::from(identifier).to_local().into(), statements)
}

// For use in `ternary_composite`.
/// Construct an access expression `inner.name` bound to a fresh identifier.
fn make_composite_access_definition(
&mut self,
inner: Identifier,
Expand All @@ -470,6 +378,9 @@ impl FlatteningVisitor<'_> {
self.unique_simple_definition(expr)
}

/// Destructure a ternary over a composite (record) type into per-field ternaries followed by a
/// composite init. This is only used for records — plaintext structs are passed to the Aleo
/// `ternary` instruction directly.
pub fn ternary_composite(
&mut self,
composite_path: &Path,
Expand All @@ -478,9 +389,7 @@ impl FlatteningVisitor<'_> {
first: &Identifier,
second: &Identifier,
) -> (Expression, Vec<Statement>) {
// Initialize a vector to accumulate any statements generated.
let mut statements = Vec::new();
// For each composite member, construct a new ternary expression.
let members = composite
.members
.iter()
Expand All @@ -489,7 +398,6 @@ impl FlatteningVisitor<'_> {
statements.push(stmt);
let (second, stmt) = self.make_composite_access_definition(*second, *identifier, type_.clone());
statements.push(stmt);
// Recursively reconstruct the ternary expression.
let ternary = TernaryExpression {
condition: condition.clone(),
if_true: Path::from(first).to_local().into(),
Expand All @@ -499,8 +407,6 @@ impl FlatteningVisitor<'_> {
};
self.state.type_table.insert(ternary.id(), type_.clone());
let (expression, stmts) = self.reconstruct_ternary(ternary, &());

// Accumulate any statements generated.
statements.extend(stmts);

CompositeFieldInitializer {
Expand All @@ -515,32 +421,24 @@ impl FlatteningVisitor<'_> {
let (expr, stmts) = self.reconstruct_composite_init(
CompositeExpression {
path: composite_path.clone(),
const_arguments: Vec::new(), // All const arguments should have been resolved by now
const_arguments: Vec::new(),
members,
span: Default::default(),
id: {
// Create a new node ID for the comopsite expression.
let id = self.state.node_builder.next_id();
// Set the type of the node ID.
self.state.type_table.insert(
id,
Type::Composite(CompositeType {
path: composite_path.clone(),
const_arguments: Vec::new(), // all const generics should have been resolved by now
}),
Type::Composite(CompositeType { path: composite_path.clone(), const_arguments: Vec::new() }),
);
id
},
},
&(),
);

// Accumulate any statements generated.
statements.extend(stmts);

// Create a new assignment statement for the composite expression.
let (identifier, statement) = self.unique_simple_definition(expr);

statements.push(statement);

(Path::from(identifier).to_local().into(), statements)
Expand Down
11 changes: 5 additions & 6 deletions tests/expectations/compiler/array/array_write.out
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,11 @@ function with_tuples:
function in_conditional:
input r0 as T.private;
input r1 as boolean.private;
ternary r1 1u8 r0.x into r2;
ternary r1 r0.z.x 2u8 into r3;
ternary r1 r0.z.y 3u32 into r4;
cast r3 r4 into r5 as S;
cast r2 r0.y r5 into r6 as T;
output r6 as T.private;
cast 2u8 3u32 into r2 as S;
ternary r1 1u8 r0.x into r3;
ternary r1 r0.z r2 into r4;
cast r3 r0.y r4 into r5 as T;
output r5 as T.private;

function f_record:
async f_record into r0;
Expand Down
9 changes: 5 additions & 4 deletions tests/expectations/compiler/cei/blind_storage_read_in_helper.out
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,10 +51,11 @@ finalize do_work:
await r0;
contains counter__[false] into r1;
get.or_use counter__[false] 0u32 into r2;
ternary r1 r2 0u32 into r3;
cast r1 r3 into r4 as Optional__JzunLORyB8U;
ternary r4.is_some r4.val 0u32 into r5;
set r5 into balances[0u32];
cast true r2 into r3 as Optional__JzunLORyB8U;
cast false 0u32 into r4 as Optional__JzunLORyB8U;
ternary r1 r3 r4 into r5;
ternary r5.is_some r5.val 0u32 into r6;
set r6 into balances[0u32];

constructor:
assert.eq edition 0u16;
11 changes: 6 additions & 5 deletions tests/expectations/compiler/cei/class2_storage_stale_read_warn.out
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,13 @@ finalize do_work:
input r1 as u32.public;
contains counter__[false] into r2;
get.or_use counter__[false] 0u32 into r3;
ternary r2 r3 0u32 into r4;
cast r2 r4 into r5 as Optional__JzunLORyB8U;
ternary r5.is_some r5.val 0u32 into r6;
cast true r3 into r4 as Optional__JzunLORyB8U;
cast false 0u32 into r5 as Optional__JzunLORyB8U;
ternary r2 r4 r5 into r6;
ternary r6.is_some r6.val 0u32 into r7;
await r0;
add r6 r1 into r7;
set r7 into counter__[false];
add r7 r1 into r8;
set r8 into counter__[false];

constructor:
assert.eq edition 0u16;
Loading
Loading