From 9be5f770cbc60901e663fbb29ea512f841e33e6e Mon Sep 17 00:00:00 2001 From: ControlCplusControlV <44706811+ControlCplusControlV@users.noreply.github.com> Date: Mon, 16 Jan 2023 18:17:38 -0700 Subject: [PATCH 1/6] more progress on opts --- Cargo.toml | 1 + crates/miden-integration-tests/Cargo.toml | 2 +- crates/miden-integration-tests/src/lib.rs | 4 +- crates/miden-integration-tests/tests/lib.rs | 3 +- .../tests/optimization.rs | 89 +++++++++ .../tests/quickcheck_tests.rs | 2 +- crates/miden-integration-tests/tests/utils.rs | 40 ++-- crates/papyrus/src/ast_optimization.rs | 181 +++++++++++------- 8 files changed, 225 insertions(+), 97 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b32b031..59002e7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "crates/papyrus", + "crates/miden-integration-tests", "bin/scribe", "bin/repl", ] \ No newline at end of file diff --git a/crates/miden-integration-tests/Cargo.toml b/crates/miden-integration-tests/Cargo.toml index 6ec3cab..adc1d71 100644 --- a/crates/miden-integration-tests/Cargo.toml +++ b/crates/miden-integration-tests/Cargo.toml @@ -17,7 +17,7 @@ harness = true [dependencies] -scribe = {path = "../transpiler/"} +papyrus = { path = "../papyrus" } miden-assembly = { git = "http://github.com/maticnetwork/miden", branch = "next" } miden-processor = { git = "http://github.com/maticnetwork/miden", branch = "next" } miden-core = { git = "http://github.com/maticnetwork/miden", branch = "next" } diff --git a/crates/miden-integration-tests/src/lib.rs b/crates/miden-integration-tests/src/lib.rs index 8b13789..7a4432c 100644 --- a/crates/miden-integration-tests/src/lib.rs +++ b/crates/miden-integration-tests/src/lib.rs @@ -1 +1,3 @@ - +pub fn test() { + println!("Hello, world!"); +} \ No newline at end of file diff --git a/crates/miden-integration-tests/tests/lib.rs b/crates/miden-integration-tests/tests/lib.rs index 77e14a4..6179a5a 100644 --- a/crates/miden-integration-tests/tests/lib.rs +++ b/crates/miden-integration-tests/tests/lib.rs @@ -1,8 +1,7 @@ #[macro_use] extern crate quickcheck; mod future; -mod milestone_1; -mod milestone_2; mod quickcheck_tests; mod test; mod utils; +mod optimization; \ No newline at end of file diff --git a/crates/miden-integration-tests/tests/optimization.rs b/crates/miden-integration-tests/tests/optimization.rs index 9f2830b..247af08 100644 --- a/crates/miden-integration-tests/tests/optimization.rs +++ b/crates/miden-integration-tests/tests/optimization.rs @@ -1,5 +1,6 @@ use crate::utils::compile_example; use indoc::indoc; +use papyrus::ast_optimization::optimize_ast; #[test] fn optimization_basic_constant_replacement() { @@ -143,3 +144,91 @@ fn optimization_let_old_vars_die_v2() { "}, ); } + + +#[test] +fn test_for_loop_to_repeat_statement_optimization() { + //Conditional is lt + let source_code = " + let result := 0 + + for { let i := 0 } lt(i, 100) { i := add(i, 1) } + { + result := mul(result, 2) + } + "; + + let parsed = papyrus::parser::parse_yul_syntax(source_code); + let ast = optimize_ast(parsed); + + println!("{:#?}", ast); + + + + // insta::assert_debug_snapshot!(ast); + + + // //Conditional is lt, mul + // let source_code = " + // let result := 0 + + // for { let i := 0 } lt(i, 100) { i := mul(i, 2) } + // { + // result := mul(result, 2) + // } + // " + + // let parsed = parser::parse_yul_syntax(source_code); + // let ast = optimize_ast(parsed); + // insta::assert_debug_snapshot!(ast); + + + // //Conditional is gt, div + // let source_code = " + // let result := 0 + + // for { let i := 100 } gt(i, 0) { i := div(i, 2) } + // { + // result := mul(result, 2) + // } + // " + + // let parsed = parser::parse_yul_syntax(source_code); + // let ast = optimize_ast(parsed); + // insta::assert_debug_snapshot!(ast); + + + + // //Conditional is slt + // let source_code = " + // let result := 0 + + // for { let i := 100 } gt(i, 0) { i := sub(i, 1) } + // { + // result := mul(result, 2) + // } + // " + + // let parsed = parser::parse_yul_syntax(source_code); + // let ast = optimize_ast(parsed); + // insta::assert_debug_snapshot!(ast); + + + + // //Conditional is sgt + // let source_code = " + // let result := 0 + + // for { let i := 100 } gt(i, 0) { i := sub(i, 1) } + // { + // result := mul(result, 2) + // } + // " + + // let parsed = parser::parse_yul_syntax(source_code); + // let ast = optimize_ast(parsed); + // insta::assert_debug_snapshot!(ast); + + + +} \ No newline at end of file diff --git a/crates/miden-integration-tests/tests/quickcheck_tests.rs b/crates/miden-integration-tests/tests/quickcheck_tests.rs index 1f83502..7537a92 100644 --- a/crates/miden-integration-tests/tests/quickcheck_tests.rs +++ b/crates/miden-integration-tests/tests/quickcheck_tests.rs @@ -2,7 +2,7 @@ use crate::utils::{miden_to_u256, MidenResult}; use miden_core::StarkField; use quickcheck::{Arbitrary, Gen, TestResult}; use quickcheck_macros::quickcheck; -use scribe::{ +use papyrus::{ executor::execute, utils::{convert_u256_to_pushes, join_u32s_to_u256, load_all_procs, split_u256_to_u32s}, }; diff --git a/crates/miden-integration-tests/tests/utils.rs b/crates/miden-integration-tests/tests/utils.rs index cfbe178..905adc9 100644 --- a/crates/miden-integration-tests/tests/utils.rs +++ b/crates/miden-integration-tests/tests/utils.rs @@ -1,28 +1,31 @@ use colored::*; use miden_core::StarkField; use primitive_types::U256; -use scribe::ast_optimization::optimize_ast; -use scribe::executor; -use scribe::miden_generator; -use scribe::miden_generator::CompileOptions; -use scribe::parser; -use scribe::type_inference::infer_types; -use scribe::types::expressions_to_tree; -use scribe::types::YulFile; +use papyrus::ast_optimization::optimize_ast; +use papyrus::executor; +use papyrus::miden_generator; +use papyrus::miden_generator::CompileOptions; +use papyrus::parser; +use papyrus::type_inference::infer_types; +use papyrus::types::expressions_to_tree; +use papyrus::types::YulFile; use std::fs; pub enum MidenResult { U256(primitive_types::U256), U32(u32), } + +fn print_title(s: &str) { + let s1 = format!("=== {} ===", s).blue().bold(); + println!("{}", s1); + println!(" "); +} + //Function to display transpile Yul code and display each step of the transpilation process in the terminal. //This function is only used to demonstrate what Scribe does in a easy to read format. pub fn run_example(yul_code: &str, expected_output: MidenResult) { - fn print_title(s: &str) { - let s1 = format!("=== {} ===", s).blue().bold(); - println!("{}", s1); - println!(" "); - } + println!(); println!(); print_title("Input Yul"); @@ -38,7 +41,7 @@ pub fn run_example(yul_code: &str, expected_output: MidenResult) { println!("{}", expressions_to_tree(&ast)); println!(); - let (transpiler, miden_code) = miden_generator::transpile_program(ast, Default::default()); + let miden_code = miden_generator::transpile_program(ast, Default::default()); let mut trimmed_miden_code = miden_code .split('\n') // .skip_while(|line| *line != "# end std lib #") @@ -49,7 +52,7 @@ pub fn run_example(yul_code: &str, expected_output: MidenResult) { print_title("Generated Miden Assembly"); println!("{}", trimmed_miden_code); println!(); - println!("Estimated cost: {}", transpiler.cost); + // println!("Estimated cost: {}", transpiler.cost); println!(); fs::write(format!("./test_output.masm",), trimmed_miden_code) .expect("Unable to write Miden to file."); @@ -81,11 +84,6 @@ pub fn run_example(yul_code: &str, expected_output: MidenResult) { } pub fn compile_example(yul_code: &str, expected_output: &str) { - fn print_title(s: &str) { - let s1 = format!("=== {} ===", s).blue().bold(); - println!("{}", s1); - println!(" "); - } let parsed = parser::parse_yul_syntax(yul_code); @@ -93,7 +91,7 @@ pub fn compile_example(yul_code: &str, expected_output: &str) { let ast = infer_types(&ast); - let (_, miden_code) = miden_generator::transpile_program( + let miden_code = miden_generator::transpile_program( ast, CompileOptions { comments: false, diff --git a/crates/papyrus/src/ast_optimization.rs b/crates/papyrus/src/ast_optimization.rs index c94d7ba..5cab95e 100644 --- a/crates/papyrus/src/ast_optimization.rs +++ b/crates/papyrus/src/ast_optimization.rs @@ -8,9 +8,14 @@ use crate::types::*; pub fn optimize_ast(ast: Vec) -> Vec { // let mut assignment_visitor = VariableAssignmentVisitor::default(); // let ast = walk_ast(ast, &mut assignment_visitor); + + // let const_variables = assignment_visitor.get_const_variables(); // let ast = walk_ast(ast, &mut ConstVariableVisitor { const_variables }); + + + // walk_ast(ast, &mut ForLoopToRepeatVisitor {}) // TODO: fix optimizations ast @@ -73,77 +78,111 @@ impl VariableAssignmentVisitor { // add(i, 1) to i := add(1, i) will break this optimization. In the future we should support gt, // subtracting, etc. impl ExpressionVisitor for ForLoopToRepeatVisitor { - fn visit_expr(&mut self, _expr: Expr) -> Option { - todo!(); - // match &expr { - // Expr::ForLoop(ExprForLoop { - // init_block, - // conditional, - // after_block, - // interior_block, - // }) => { - // let start: Option; - // let iterator_identifier: Option; - // if let Some(first_expr) = (*init_block.exprs).first() { - // if let Expr::DeclareVariable(ExprDeclareVariable { identifier, rhs }) = - // first_expr - // { - // if let Some(Expr::Literal(value)) = rhs.clone().map(|e| *e) { - // start = Some(todo!("Need to get literal value here")); - // iterator_identifier = Some(identifier.to_string()); - // } else { - // return Some(expr); - // } - // } else { - // return Some(expr); - // } - // } else { - // return Some(expr); - // } - // - // if let Some(Expr::Assignment(assignment)) = (*after_block.exprs).first() { - // if *assignment - // == (ExprAssignment { - // typed_identifier: iterator_identifier.clone().unwrap(), - // rhs: Box::new(Expr::FunctionCall(ExprFunctionCall { - // function_name: "add".to_string(), - // exprs: Box::new(vec![ - // Expr::Variable(ExprVariableReference { - // identifier: iterator_identifier.clone().unwrap(), - // }), - // Expr::Literal(todo!("Need to get literal value here")), - // ]), - // })), - // }) - // {} - // } else { - // return Some(expr); - // } - // if let Expr::FunctionCall(ExprFunctionCall { - // function_name, - // exprs, - // }) = &**conditional - // { - // if function_name == "lt" - // && exprs[0] - // == Expr::Variable(ExprVariableReference { - // identifier: iterator_identifier.unwrap(), - // }) - // { - // if let Expr::Literal(value) = exprs[1] { - // return Some(Expr::Repeat(ExprRepeat { - // interior_block: interior_block.clone(), - // iterations: todo!("Get end value from literal"), - // })); - // } - // } - // } else { - // return Some(expr); - // } - // } - // _ => {} - // } - // Some(expr) + fn visit_expr(&mut self, expr: Expr) -> Option { + let preserve_value = Some(expr.clone()); + + if let Expr::ForLoop(ExprForLoop { + init_block, + conditional, + after_block, + interior_block, + }) = expr.clone() { + let mut i = 0; // Variable to keep track of where repeat begins + let init_value_name: String; + + + // Start out with the init block, and determine where iteration will start + /// This optimization will only be applied when the init block uses a number literal for the rhs + if init_block.exprs.len() != 1 { + return preserve_value; + } + + + + for expr in init_block.exprs { + if let Expr::Assignment(value) = expr.clone() { + if let Expr::Literal(value) = *value.rhs { + match value { + ExprLiteral::Number(n)=> { + i = n.value.0[3]; // TODO: make sure this is within u32 constraints + } + _ => { + return preserve_value; + } + } + } + } + } + + + let mut end_val = primitive_types::U256::zero(); + // With `i` now indicating where our iteration will start, we can check the conditional + if let Expr::FunctionCall(ExprFunctionCall{function_name, exprs, inferred_return_types, inferred_param_types }) = *conditional { + + if function_name == "lt"{ + if let Expr::Literal(ExprLiteral::Number(expr_literal_num)) = exprs.last().unwrap() { + end_val = expr_literal_num.value; + }; + + } else if function_name == "gt"{ + if let Expr::Literal(ExprLiteral::Number(expr_literal_num)) = exprs.last().unwrap() { + end_val = expr_literal_num.value; + }; + }else{ + return preserve_value; + } + + } + + // Determine iterator size + if after_block.exprs.len() != 1 { + return preserve_value; + } + let mut step = 0; + + for expr in after_block.exprs { + match expr { + Expr::Assignment(value) => { + if let Expr::FunctionCall(ExprFunctionCall { function_name, exprs, .. }) = *value.rhs { + if function_name == "add" { + + + + + for args in *exprs { + match args { + Expr::Literal(ExprLiteral::Number(expr_literal_num)) => { + step = expr_literal_num.value.0[3]; // TODO: make sure this is within u32 constraints + } + Expr::Variable(ExprVariableReference { identifier , ..}) => { + if identifier != "i" { // Should check identifier is the same as the one which appaered earlier + return preserve_value; + } + } + _ => return preserve_value + } + } + } + } + } + _ => return Some(expr) + } + } + + //TODO: + + // for { let i := 29 } lt(i, exponent) { i := add(i, 1) } + // { + // result := mul(result, base) + // } + + return Some(Expr::Repeat(ExprRepeat { + iterations: 53, + interior_block, + })); + }; + + Some(expr) } } From a5364d89b377460aa095ec168d7701d99daa88dc Mon Sep 17 00:00:00 2001 From: ControlCplusControlV <44706811+ControlCplusControlV@users.noreply.github.com> Date: Mon, 16 Jan 2023 19:04:56 -0700 Subject: [PATCH 2/6] changes --- crates/papyrus/src/ast_optimization.rs | 64 +++++++++++++++++++++++--- 1 file changed, 58 insertions(+), 6 deletions(-) diff --git a/crates/papyrus/src/ast_optimization.rs b/crates/papyrus/src/ast_optimization.rs index 5cab95e..f7a6c90 100644 --- a/crates/papyrus/src/ast_optimization.rs +++ b/crates/papyrus/src/ast_optimization.rs @@ -88,8 +88,13 @@ impl ExpressionVisitor for ForLoopToRepeatVisitor { interior_block, }) = expr.clone() { let mut i = 0; // Variable to keep track of where repeat begins - let init_value_name: String; + let mut init_value_variable_name: String = "".to_string(); + //get the init value name, + + //Get the end value + + // Get the step and see if it is mul, add etc, and get the repeat amount // Start out with the init block, and determine where iteration will start /// This optimization will only be applied when the init block uses a number literal for the rhs @@ -97,7 +102,13 @@ impl ExpressionVisitor for ForLoopToRepeatVisitor { return preserve_value; } - + if let Some(Expr::DeclareVariable(ExprDeclareVariable{typed_identifiers, rhs: _}))= init_block.exprs.first(){ + if let Some(TypedIdentifier{identifier, yul_type}) = typed_identifiers.first(){ + init_value_variable_name = identifier.to_owned(); + } + }else{ + return preserve_value; + }; for expr in init_block.exprs { if let Expr::Assignment(value) = expr.clone() { @@ -138,30 +149,71 @@ impl ExpressionVisitor for ForLoopToRepeatVisitor { if after_block.exprs.len() != 1 { return preserve_value; } - let mut step = 0; + let mut iterations = 0; // Final value for repeat statement for expr in after_block.exprs { match expr { Expr::Assignment(value) => { if let Expr::FunctionCall(ExprFunctionCall { function_name, exprs, .. }) = *value.rhs { - if function_name == "add" { - + match function_name.to_string().as_str() { + "add" => { + let mut step = 1; + for args in *exprs { + match args { + Expr::Literal(ExprLiteral::Number(expr_literal_num)) => { + step = expr_literal_num.value.0[3]; // TODO: make sure this is within u32 constraints + } + Expr::Variable(ExprVariableReference { identifier , ..}) => { + if identifier != init_value_variable_name { // Should check identifier is the same as the one which appaered earlier + return preserve_value; + } + } + _ => return preserve_value + } + iterations = ((end_val - i) / step).0[3]; // TODO: make sure this is within u32 constraints + } + } + "sub" => { + let mut step = 1; + let args = *exprs; + if let Some(Expr::Variable( ExprVariableReference { identifier , ..})) = args.first() { + if *identifier != init_value_variable_name { + return preserve_value; + } + } + if let Some(Expr::Literal( ExprLiteral::Number())) = args.first() { + if *identifier != init_value_variable_name { + return preserve_value; + } + } + + + } + "mul" => { + let step = 1; + let args = *exprs; for args in *exprs { match args { Expr::Literal(ExprLiteral::Number(expr_literal_num)) => { step = expr_literal_num.value.0[3]; // TODO: make sure this is within u32 constraints } Expr::Variable(ExprVariableReference { identifier , ..}) => { - if identifier != "i" { // Should check identifier is the same as the one which appaered earlier + if identifier != init_value_variable_name { // Should check identifier is the same as the one which appaered earlier return preserve_value; } } _ => return preserve_value } + } + } + "div" => { + let step = 1; + let args = *exprs; } + _ => return preserve_value, } } } From 4982282eb19ad731754cb117ba8a0c821b0fde25 Mon Sep 17 00:00:00 2001 From: ControlCplusControlV <44706811+ControlCplusControlV@users.noreply.github.com> Date: Fri, 20 Jan 2023 21:14:52 -0700 Subject: [PATCH 3/6] finalized everything --- crates/miden-integration-tests/src/lib.rs | 2 +- .../miden-integration-tests/test_output.masm | 12 +- .../miden-integration-tests/tests/future.rs | 2 +- crates/miden-integration-tests/tests/lib.rs | 2 +- .../tests/optimization.rs | 34 +- .../tests/quickcheck_tests.rs | 4 +- crates/miden-integration-tests/tests/test.rs | 2 +- crates/miden-integration-tests/tests/utils.rs | 7 +- crates/papyrus/src/ast_optimization.rs | 310 +++++++++++------- 9 files changed, 209 insertions(+), 166 deletions(-) diff --git a/crates/miden-integration-tests/src/lib.rs b/crates/miden-integration-tests/src/lib.rs index 7a4432c..9815d27 100644 --- a/crates/miden-integration-tests/src/lib.rs +++ b/crates/miden-integration-tests/src/lib.rs @@ -1,3 +1,3 @@ pub fn test() { println!("Hello, world!"); -} \ No newline at end of file +} diff --git a/crates/miden-integration-tests/test_output.masm b/crates/miden-integration-tests/test_output.masm index 7b1384b..38cd87d 100644 --- a/crates/miden-integration-tests/test_output.masm +++ b/crates/miden-integration-tests/test_output.masm @@ -2,12 +2,12 @@ use.std::math::u256 begin - # and() # - # u256 literal: 0 # - push.0 push.0 push.0 push.0 push.0 push.0 push.0 push.0 + # eq() # + # u256 literal: 2 # + push.2 push.0 push.0 push.0 push.0 push.0 push.0 push.0 - # u256 literal: 0 # - push.0 push.0 push.0 push.0 push.0 push.0 push.0 push.0 + # u256 literal: 2 # + push.2 push.0 push.0 push.0 push.0 push.0 push.0 push.0 - exec.u256::and + exec.u256::eq_unsafe end \ No newline at end of file diff --git a/crates/miden-integration-tests/tests/future.rs b/crates/miden-integration-tests/tests/future.rs index 20f861f..c834ee7 100644 --- a/crates/miden-integration-tests/tests/future.rs +++ b/crates/miden-integration-tests/tests/future.rs @@ -1,5 +1,5 @@ -use primitive_types::U256; use crate::utils::{run_example, MidenResult}; +use primitive_types::U256; #[ignore] #[test] diff --git a/crates/miden-integration-tests/tests/lib.rs b/crates/miden-integration-tests/tests/lib.rs index 6179a5a..9406445 100644 --- a/crates/miden-integration-tests/tests/lib.rs +++ b/crates/miden-integration-tests/tests/lib.rs @@ -1,7 +1,7 @@ #[macro_use] extern crate quickcheck; mod future; +mod optimization; mod quickcheck_tests; mod test; mod utils; -mod optimization; \ No newline at end of file diff --git a/crates/miden-integration-tests/tests/optimization.rs b/crates/miden-integration-tests/tests/optimization.rs index 247af08..ed4c6f2 100644 --- a/crates/miden-integration-tests/tests/optimization.rs +++ b/crates/miden-integration-tests/tests/optimization.rs @@ -145,7 +145,6 @@ fn optimization_let_old_vars_die_v2() { ); } - #[test] fn test_for_loop_to_repeat_statement_optimization() { //Conditional is lt @@ -158,18 +157,15 @@ fn test_for_loop_to_repeat_statement_optimization() { } "; - let parsed = papyrus::parser::parse_yul_syntax(source_code); + let parsed = papyrus::parser::parse_yul_syntax(source_code); let ast = optimize_ast(parsed); println!("{:#?}", ast); - - // insta::assert_debug_snapshot!(ast); - // //Conditional is lt, mul - // let source_code = " + // let source_code = " // let result := 0 // for { let i := 0 } lt(i, 100) { i := mul(i, 2) } @@ -178,13 +174,12 @@ fn test_for_loop_to_repeat_statement_optimization() { // } // " - // let parsed = parser::parse_yul_syntax(source_code); + // let parsed = parser::parse_yul_syntax(source_code); // let ast = optimize_ast(parsed); // insta::assert_debug_snapshot!(ast); - // //Conditional is gt, div - // let source_code = " + // let source_code = " // let result := 0 // for { let i := 100 } gt(i, 0) { i := div(i, 2) } @@ -193,14 +188,12 @@ fn test_for_loop_to_repeat_statement_optimization() { // } // " - // let parsed = parser::parse_yul_syntax(source_code); + // let parsed = parser::parse_yul_syntax(source_code); // let ast = optimize_ast(parsed); // insta::assert_debug_snapshot!(ast); - - // //Conditional is slt - // let source_code = " + // let source_code = " // let result := 0 // for { let i := 100 } gt(i, 0) { i := sub(i, 1) } @@ -208,15 +201,13 @@ fn test_for_loop_to_repeat_statement_optimization() { // result := mul(result, 2) // } // " - - // let parsed = parser::parse_yul_syntax(source_code); + + // let parsed = parser::parse_yul_syntax(source_code); // let ast = optimize_ast(parsed); // insta::assert_debug_snapshot!(ast); - - // //Conditional is sgt - // let source_code = " + // let source_code = " // let result := 0 // for { let i := 100 } gt(i, 0) { i := sub(i, 1) } @@ -225,10 +216,7 @@ fn test_for_loop_to_repeat_statement_optimization() { // } // " - // let parsed = parser::parse_yul_syntax(source_code); + // let parsed = parser::parse_yul_syntax(source_code); // let ast = optimize_ast(parsed); // insta::assert_debug_snapshot!(ast); - - - -} \ No newline at end of file +} diff --git a/crates/miden-integration-tests/tests/quickcheck_tests.rs b/crates/miden-integration-tests/tests/quickcheck_tests.rs index 7537a92..899ee12 100644 --- a/crates/miden-integration-tests/tests/quickcheck_tests.rs +++ b/crates/miden-integration-tests/tests/quickcheck_tests.rs @@ -1,11 +1,11 @@ use crate::utils::{miden_to_u256, MidenResult}; use miden_core::StarkField; -use quickcheck::{Arbitrary, Gen, TestResult}; -use quickcheck_macros::quickcheck; use papyrus::{ executor::execute, utils::{convert_u256_to_pushes, join_u32s_to_u256, load_all_procs, split_u256_to_u32s}, }; +use quickcheck::{Arbitrary, Gen, TestResult}; +use quickcheck_macros::quickcheck; #[derive(Clone, Debug)] struct U256(primitive_types::U256); diff --git a/crates/miden-integration-tests/tests/test.rs b/crates/miden-integration-tests/tests/test.rs index 815e64d..38033c2 100644 --- a/crates/miden-integration-tests/tests/test.rs +++ b/crates/miden-integration-tests/tests/test.rs @@ -1,5 +1,5 @@ -use primitive_types::U256; use crate::utils::{run_example, MidenResult}; +use primitive_types::U256; #[test] fn integration_math() { diff --git a/crates/miden-integration-tests/tests/utils.rs b/crates/miden-integration-tests/tests/utils.rs index 905adc9..6e17acd 100644 --- a/crates/miden-integration-tests/tests/utils.rs +++ b/crates/miden-integration-tests/tests/utils.rs @@ -1,6 +1,5 @@ use colored::*; use miden_core::StarkField; -use primitive_types::U256; use papyrus::ast_optimization::optimize_ast; use papyrus::executor; use papyrus::miden_generator; @@ -9,13 +8,13 @@ use papyrus::parser; use papyrus::type_inference::infer_types; use papyrus::types::expressions_to_tree; use papyrus::types::YulFile; +use primitive_types::U256; use std::fs; pub enum MidenResult { U256(primitive_types::U256), U32(u32), } - fn print_title(s: &str) { let s1 = format!("=== {} ===", s).blue().bold(); println!("{}", s1); @@ -25,7 +24,6 @@ fn print_title(s: &str) { //Function to display transpile Yul code and display each step of the transpilation process in the terminal. //This function is only used to demonstrate what Scribe does in a easy to read format. pub fn run_example(yul_code: &str, expected_output: MidenResult) { - println!(); println!(); print_title("Input Yul"); @@ -41,7 +39,7 @@ pub fn run_example(yul_code: &str, expected_output: MidenResult) { println!("{}", expressions_to_tree(&ast)); println!(); - let miden_code = miden_generator::transpile_program(ast, Default::default()); + let miden_code = miden_generator::transpile_program(ast, Default::default()); let mut trimmed_miden_code = miden_code .split('\n') // .skip_while(|line| *line != "# end std lib #") @@ -84,7 +82,6 @@ pub fn run_example(yul_code: &str, expected_output: MidenResult) { } pub fn compile_example(yul_code: &str, expected_output: &str) { - let parsed = parser::parse_yul_syntax(yul_code); let ast = optimize_ast(parsed); diff --git a/crates/papyrus/src/ast_optimization.rs b/crates/papyrus/src/ast_optimization.rs index f7a6c90..cdaf8e0 100644 --- a/crates/papyrus/src/ast_optimization.rs +++ b/crates/papyrus/src/ast_optimization.rs @@ -9,13 +9,9 @@ pub fn optimize_ast(ast: Vec) -> Vec { // let mut assignment_visitor = VariableAssignmentVisitor::default(); // let ast = walk_ast(ast, &mut assignment_visitor); - // let const_variables = assignment_visitor.get_const_variables(); // let ast = walk_ast(ast, &mut ConstVariableVisitor { const_variables }); - - - // walk_ast(ast, &mut ForLoopToRepeatVisitor {}) // TODO: fix optimizations ast @@ -81,158 +77,220 @@ impl ExpressionVisitor for ForLoopToRepeatVisitor { fn visit_expr(&mut self, expr: Expr) -> Option { let preserve_value = Some(expr.clone()); - if let Expr::ForLoop(ExprForLoop { - init_block, - conditional, - after_block, - interior_block, - }) = expr.clone() { - let mut i = 0; // Variable to keep track of where repeat begins - let mut init_value_variable_name: String = "".to_string(); - - //get the init value name, - - //Get the end value - - // Get the step and see if it is mul, add etc, and get the repeat amount + if let Expr::ForLoop(ExprForLoop { + init_block, + conditional, + after_block, + interior_block, + }) = expr.clone() + { + let mut i = 0; // Variable to keep track of where repeat begins + let mut init_value_variable_name: String = "".to_string(); + + // Get the step and see if it is mul, add etc, and get the repeat amount + + // Start out with the init block, and determine where iteration will start + /// This optimization will only be applied when the init block uses a number literal for the rhs + if init_block.exprs.len() != 1 { + return preserve_value; + } - // Start out with the init block, and determine where iteration will start - /// This optimization will only be applied when the init block uses a number literal for the rhs - if init_block.exprs.len() != 1 { - return preserve_value; + //get the init value name, + if let Some(Expr::DeclareVariable(ExprDeclareVariable { + typed_identifiers, + rhs: _, + })) = init_block.exprs.first() + { + if let Some(TypedIdentifier { + identifier, + yul_type, + }) = typed_identifiers.first() + { + init_value_variable_name = identifier.to_owned(); } - - if let Some(Expr::DeclareVariable(ExprDeclareVariable{typed_identifiers, rhs: _}))= init_block.exprs.first(){ - if let Some(TypedIdentifier{identifier, yul_type}) = typed_identifiers.first(){ - init_value_variable_name = identifier.to_owned(); - } - }else{ + } else { return preserve_value; - }; - - for expr in init_block.exprs { - if let Expr::Assignment(value) = expr.clone() { - if let Expr::Literal(value) = *value.rhs { - match value { - ExprLiteral::Number(n)=> { - i = n.value.0[3]; // TODO: make sure this is within u32 constraints - } - _ => { - return preserve_value; - } + }; + + for expr in init_block.exprs { + if let Expr::Assignment(value) = expr.clone() { + if let Expr::Literal(value) = *value.rhs { + match value { + ExprLiteral::Number(n) => { + i = n.value.0[3]; // TODO: make sure this is within u32 constraints + } + _ => { + return preserve_value; } } } } + } + //Get the end value - let mut end_val = primitive_types::U256::zero(); - // With `i` now indicating where our iteration will start, we can check the conditional - if let Expr::FunctionCall(ExprFunctionCall{function_name, exprs, inferred_return_types, inferred_param_types }) = *conditional { - - if function_name == "lt"{ - if let Expr::Literal(ExprLiteral::Number(expr_literal_num)) = exprs.last().unwrap() { - end_val = expr_literal_num.value; - }; - - } else if function_name == "gt"{ - if let Expr::Literal(ExprLiteral::Number(expr_literal_num)) = exprs.last().unwrap() { - end_val = expr_literal_num.value; - }; - }else{ - return preserve_value; - } - - } - - // Determine iterator size - if after_block.exprs.len() != 1 { + let mut end_val = 0; + // With `i` now indicating where our iteration will start, we can check the conditional + if let Expr::FunctionCall(ExprFunctionCall { + function_name, + exprs, + inferred_return_types, + inferred_param_types, + }) = *conditional + { + if function_name == "lt" { + if let Expr::Literal(ExprLiteral::Number(expr_literal_num)) = + exprs.last().unwrap() + { + end_val = expr_literal_num.value.0[3]; + }; + } else if function_name == "gt" { + if let Expr::Literal(ExprLiteral::Number(expr_literal_num)) = + exprs.last().unwrap() + { + end_val = expr_literal_num.value.0[3]; + }; + } else { return preserve_value; } - let mut iterations = 0; // Final value for repeat statement - - for expr in after_block.exprs { - match expr { - Expr::Assignment(value) => { - if let Expr::FunctionCall(ExprFunctionCall { function_name, exprs, .. }) = *value.rhs { - match function_name.to_string().as_str() { - "add" => { - let mut step = 1; - for args in *exprs { - match args { - Expr::Literal(ExprLiteral::Number(expr_literal_num)) => { - step = expr_literal_num.value.0[3]; // TODO: make sure this is within u32 constraints - } - Expr::Variable(ExprVariableReference { identifier , ..}) => { - if identifier != init_value_variable_name { // Should check identifier is the same as the one which appaered earlier - return preserve_value; - } + } + + // Determine iterator size + if after_block.exprs.len() != 1 { + return preserve_value; + } + let mut iterations: u32 = 0; // Final value for repeat statement + + for expr in after_block.exprs { + match expr { + Expr::Assignment(value) => { + if let Expr::FunctionCall(ExprFunctionCall { + function_name, + exprs, + .. + }) = *value.rhs + { + match function_name.to_string().as_str() { + "add" => { + let mut step = 1; + for args in *exprs { + match args { + Expr::Literal(ExprLiteral::Number( + expr_literal_num, + )) => { + step = expr_literal_num.value.0[3]; + // TODO: make sure this is within u32 constraints + } + Expr::Variable(ExprVariableReference { + identifier, + .. + }) => { + if identifier != init_value_variable_name { + // Should check identifier is the same as the one which appaered earlier + return preserve_value; } - _ => return preserve_value } - iterations = ((end_val - i) / step).0[3]; // TODO: make sure this is within u32 constraints - } + _ => return preserve_value, + } + iterations = ((end_val - i) / step).try_into().unwrap(); + // TODO: make sure this is within u32 constraints } - "sub" => { - let mut step = 1; - let args = *exprs; - if let Some(Expr::Variable( ExprVariableReference { identifier , ..})) = args.first() { - if *identifier != init_value_variable_name { - return preserve_value; - } + } + "sub" => { + let mut step = 1; + let args = *exprs; + if let Some(Expr::Variable(ExprVariableReference { + identifier, + .. + })) = args.first() + { + if *identifier != init_value_variable_name { + return preserve_value; } + } + if let Some(Expr::Literal(ExprLiteral::Number(num))) = + args.last() + { + step = num.value.0[3]; + } else { + return preserve_value; + } - if let Some(Expr::Literal( ExprLiteral::Number())) = args.first() { - if *identifier != init_value_variable_name { - return preserve_value; + iterations = ((i - end_val) / step).try_into().unwrap(); + // Assumes no overflow + } + "mul" => { + let mut step = 1; + for args in *exprs { + match args { + Expr::Literal(ExprLiteral::Number( + expr_literal_num, + )) => { + step = expr_literal_num.value.0[3]; + // TODO: make sure this is within u32 constraints + } + Expr::Variable(ExprVariableReference { + identifier, + .. + }) => { + if identifier != init_value_variable_name { + // Should check identifier is the same as the one which appaered earlier + return preserve_value; + } } + _ => return preserve_value, } - - - + let mut j = i.clone(); + while j < end_val { + j *= step; + iterations += 1; + } } - "mul" => { - let step = 1; - let args = *exprs; - for args in *exprs { - match args { - Expr::Literal(ExprLiteral::Number(expr_literal_num)) => { - step = expr_literal_num.value.0[3]; // TODO: make sure this is within u32 constraints - } - Expr::Variable(ExprVariableReference { identifier , ..}) => { - if identifier != init_value_variable_name { // Should check identifier is the same as the one which appaered earlier - return preserve_value; - } + } + "div" => { + let mut step = 1; + for args in *exprs { + match args { + Expr::Literal(ExprLiteral::Number( + expr_literal_num, + )) => { + step = expr_literal_num.value.0[3]; + // TODO: make sure this is within u32 constraints + } + Expr::Variable(ExprVariableReference { + identifier, + .. + }) => { + if identifier != init_value_variable_name { + // Should check identifier is the same as the one which appaered earlier + return preserve_value; } - _ => return preserve_value } + _ => return preserve_value, + } + + let mut j = i.clone(); + while j > end_val { + j /= step; + iterations += 1; } } - "div" => { - let step = 1; - let args = *exprs; - } - _ => return preserve_value, } + + _ =>{} } } - _ => return Some(expr) } + _ => return Some(expr), } + } - //TODO: - - // for { let i := 29 } lt(i, exponent) { i := add(i, 1) } - // { - // result := mul(result, base) - // } - - return Some(Expr::Repeat(ExprRepeat { - iterations: 53, - interior_block, - })); - }; + return Some(Expr::Repeat(ExprRepeat { + iterations, + interior_block, + })); + }; Some(expr) } From a5c726f8ec714b37cd8b545276c1442afc6e4432 Mon Sep 17 00:00:00 2001 From: 0xKitsune <0xKitsune@protonmail.com> Date: Fri, 27 Jan 2023 01:43:13 -0500 Subject: [PATCH 4/6] refactored ast optimizations into optimizer mod --- bin/scribe/src/main.rs | 2 +- .../tests/optimization.rs | 2 +- crates/miden-integration-tests/tests/utils.rs | 2 +- crates/papyrus/src/executor.rs | 8 +- crates/papyrus/src/lib.rs | 2 +- crates/papyrus/src/miden_generator.rs | 4 +- .../papyrus/src/optimizer/const_variable.rs | 31 +++ .../for_loop_to_repeat.rs} | 235 +----------------- crates/papyrus/src/optimizer/mod.rs | 153 ++++++++++++ .../src/optimizer/variable_assignment.rs | 59 +++++ 10 files changed, 261 insertions(+), 237 deletions(-) create mode 100644 crates/papyrus/src/optimizer/const_variable.rs rename crates/papyrus/src/{ast_optimization.rs => optimizer/for_loop_to_repeat.rs} (58%) create mode 100644 crates/papyrus/src/optimizer/mod.rs create mode 100644 crates/papyrus/src/optimizer/variable_assignment.rs diff --git a/bin/scribe/src/main.rs b/bin/scribe/src/main.rs index b3ccb75..79aad64 100644 --- a/bin/scribe/src/main.rs +++ b/bin/scribe/src/main.rs @@ -1,5 +1,5 @@ -use papyrus::ast_optimization::optimize_ast; use papyrus::miden_generator; +use papyrus::optimizer::optimize_ast; use papyrus::parser; use papyrus::type_inference::infer_types; diff --git a/crates/miden-integration-tests/tests/optimization.rs b/crates/miden-integration-tests/tests/optimization.rs index ed4c6f2..21a8e02 100644 --- a/crates/miden-integration-tests/tests/optimization.rs +++ b/crates/miden-integration-tests/tests/optimization.rs @@ -1,6 +1,6 @@ use crate::utils::compile_example; use indoc::indoc; -use papyrus::ast_optimization::optimize_ast; +use papyrus::optimizer::optimize_ast; #[test] fn optimization_basic_constant_replacement() { diff --git a/crates/miden-integration-tests/tests/utils.rs b/crates/miden-integration-tests/tests/utils.rs index 6e17acd..d300c19 100644 --- a/crates/miden-integration-tests/tests/utils.rs +++ b/crates/miden-integration-tests/tests/utils.rs @@ -1,9 +1,9 @@ use colored::*; use miden_core::StarkField; -use papyrus::ast_optimization::optimize_ast; use papyrus::executor; use papyrus::miden_generator; use papyrus::miden_generator::CompileOptions; +use papyrus::optimizer::optimize_ast; use papyrus::parser; use papyrus::type_inference::infer_types; use papyrus::types::expressions_to_tree; diff --git a/crates/papyrus/src/executor.rs b/crates/papyrus/src/executor.rs index 6beb9cf..23a1931 100644 --- a/crates/papyrus/src/executor.rs +++ b/crates/papyrus/src/executor.rs @@ -8,8 +8,12 @@ pub fn execute(program: String, _pub_inputs: Vec) -> Result, options: CompileOptions) -> Str //transpiling the body of the miden program for expr in &ast { if let Expr::FunctionDefinition(op) = expr { - transpiler.transpile_function_declaration(op) + transpiler.transpile_function_declaration(&op) } } diff --git a/crates/papyrus/src/optimizer/const_variable.rs b/crates/papyrus/src/optimizer/const_variable.rs new file mode 100644 index 0000000..4673fc8 --- /dev/null +++ b/crates/papyrus/src/optimizer/const_variable.rs @@ -0,0 +1,31 @@ +use std::collections::HashMap; + +use crate::types::{Expr, ExprLiteral}; + +use super::ExpressionVisitor; + +//TODO: Keeps track of constant variables +#[derive(Default)] +struct ConstVariableVisitor { + const_variables: HashMap, +} + +impl ExpressionVisitor for ConstVariableVisitor { + fn visit_expr(&mut self, _expr: Expr) -> Option { + todo!(); + // match &expr { + // Expr::DeclareVariable(ExprDeclareVariable { identifier, rhs: _ }) => { + // if self.const_variables.get(identifier).is_some() { + // return None; + // } + // } + // Expr::Variable(ExprVariableReference { identifier }) => { + // if let Some(value) = self.const_variables.get(identifier) { + // return Some(Expr::Literal(value.clone())); + // } + // } + // _ => {} + // } + // Some(expr) + } +} diff --git a/crates/papyrus/src/ast_optimization.rs b/crates/papyrus/src/optimizer/for_loop_to_repeat.rs similarity index 58% rename from crates/papyrus/src/ast_optimization.rs rename to crates/papyrus/src/optimizer/for_loop_to_repeat.rs index cdaf8e0..774245f 100644 --- a/crates/papyrus/src/ast_optimization.rs +++ b/crates/papyrus/src/optimizer/for_loop_to_repeat.rs @@ -1,71 +1,15 @@ -#![allow(dead_code)] -use std::{collections::HashMap, vec}; - -use crate::types::*; - //TODO: Update this mod and comment the functions -pub fn optimize_ast(ast: Vec) -> Vec { - // let mut assignment_visitor = VariableAssignmentVisitor::default(); - // let ast = walk_ast(ast, &mut assignment_visitor); +use crate::types::{ + Expr, ExprDeclareVariable, ExprForLoop, ExprFunctionCall, ExprLiteral, ExprRepeat, + ExprVariableReference, TypedIdentifier, +}; - // let const_variables = assignment_visitor.get_const_variables(); - // let ast = walk_ast(ast, &mut ConstVariableVisitor { const_variables }); - - // walk_ast(ast, &mut ForLoopToRepeatVisitor {}) - // TODO: fix optimizations - ast -} - -// Walks through each expression in the abstract syntax tree, optimizing the AST where possible. A new, optimized AST is returned -//Which is then passed into the Miden generation logic. -fn walk_ast(ast: Vec, visitor: &mut V) -> Vec { - let mut new_ast = vec![]; - for expr in ast { - if let Some(expr) = walk_expr(expr, visitor) { - new_ast.push(expr); - } - } - new_ast -} - -trait ExpressionVisitor { - fn visit_expr(&mut self, expr: Expr) -> Option; -} - -//TODO: Keeps track of constant variables -#[derive(Default)] -struct ConstVariableVisitor { - const_variables: HashMap, -} +use super::ExpressionVisitor; #[derive(Default)] struct ForLoopToRepeatVisitor {} -//The variable assignment visitor keeps track of variables that are reused through the code and the last assignment. -//Variables that do not change can be optimized by converting them into constants. -#[derive(Default)] -struct VariableAssignmentVisitor { - assignment_counter: HashMap, - last_assignment: HashMap, -} - -impl VariableAssignmentVisitor { - // Checks for variables that are only assigned once and returns a hashmap of the variables to convert into constants. - fn get_const_variables(&self) -> HashMap { - self.assignment_counter - .iter() - .filter(|(_k, v)| **v == 1) - .filter_map(|(k, _)| { - if let Some(value) = self.last_assignment.get(k) { - return Some((k.clone(), value.clone())); - } - None - }) - .collect::>() - } -} - // TODO: unstable for now, as it will incorrectly transform for loops that modify the iterator in // the interior block. To fix this we should have the variable assignment visitor walk the interior // block, for assignments. Also need to make sure the var isn't referenced within the for loop @@ -278,7 +222,7 @@ impl ExpressionVisitor for ForLoopToRepeatVisitor { } } - _ =>{} + _ => {} } } } @@ -295,170 +239,3 @@ impl ExpressionVisitor for ForLoopToRepeatVisitor { Some(expr) } } - -impl ExpressionVisitor for VariableAssignmentVisitor { - fn visit_expr(&mut self, _expr: Expr) -> Option { - todo!(); - // match &expr { - // Expr::DeclareVariable(ExprDeclareVariable { identifier, rhs }) => { - // if let Some(Expr::Literal(literal)) = rhs.clone().map(|r| *r) { - // self.last_assignment.insert(identifier.clone(), literal); - // } - // let count = self - // .assignment_counter - // .entry(identifier.clone()) - // .or_insert(0); - // *count += 1; - // } - // Expr::Assignment(ExprAssignment { - // typed_identifier: identifier, - // rhs: _, - // }) => { - // let count = self - // .assignment_counter - // .entry(identifier.clone()) - // .or_insert(0); - // *count += 1; - // } - // _ => {} - // } - // Some(expr) - } -} - -impl ExpressionVisitor for ConstVariableVisitor { - fn visit_expr(&mut self, _expr: Expr) -> Option { - todo!(); - // match &expr { - // Expr::DeclareVariable(ExprDeclareVariable { identifier, rhs: _ }) => { - // if self.const_variables.get(identifier).is_some() { - // return None; - // } - // } - // Expr::Variable(ExprVariableReference { identifier }) => { - // if let Some(value) = self.const_variables.get(identifier) { - // return Some(Expr::Literal(value.clone())); - // } - // } - // _ => {} - // } - // Some(expr) - } -} - -// TODO: it would be nice if there wasn't so much cloning in here -fn walk_expr(expr: Expr, visitor: &mut V) -> Option { - let expr = visitor.visit_expr(expr); - if let Some(expr) = expr { - return Some(match expr { - //Expr is literal - Expr::Literal(ref _x) => expr, - - //Expr is function call - Expr::FunctionCall(ExprFunctionCall { - function_name, - inferred_return_types, - inferred_param_types, - exprs, - }) => Expr::FunctionCall(ExprFunctionCall { - function_name, - inferred_return_types, - inferred_param_types, - exprs: Box::new(vec![ - walk_expr(exprs[0].clone(), visitor).unwrap(), - walk_expr(exprs[1].clone(), visitor).unwrap(), - ]), - }), - - //Expr is if statement - Expr::IfStatement(ExprIfStatement { - first_expr, - second_expr, - }) => Expr::IfStatement(ExprIfStatement { - first_expr: Box::new(walk_expr(*first_expr, visitor).unwrap()), - second_expr: Box::new(ExprBlock { - exprs: walk_ast(second_expr.exprs, visitor), - }), - }), - - //Expr is assignment - Expr::Assignment(ExprAssignment { - inferred_types, - identifiers, - rhs, - }) => Expr::Assignment(ExprAssignment { - identifiers, - inferred_types, - rhs: Box::new(walk_expr(*rhs, visitor).unwrap()), - }), - - //Expr is declare variable - Expr::DeclareVariable(ExprDeclareVariable { - typed_identifiers, - rhs, - }) => Expr::DeclareVariable(ExprDeclareVariable { - typed_identifiers, - rhs: rhs.map(|rhs| Box::new(walk_expr(*rhs, visitor).unwrap())), - }), - - //TODO: Expr is function definition - Expr::FunctionDefinition(ExprFunctionDefinition { - function_name: _, - params: _, - returns: _, - block: _, - }) => todo!(), - - //TODO: Expr is break - Expr::Break => todo!(), - - //TODO: Expr is continue - Expr::Continue => todo!(), - Expr::Leave => todo!(), - - //Expr is repeat - Expr::Repeat(ExprRepeat { - interior_block, - iterations, - }) => Expr::Repeat(ExprRepeat { - iterations, - interior_block: Box::new(ExprBlock { - exprs: walk_ast(interior_block.exprs, visitor), - }), - }), - - //Expr is for loop - Expr::ForLoop(ExprForLoop { - init_block, - conditional, - after_block, - interior_block, - }) => Expr::ForLoop(ExprForLoop { - init_block: Box::new(ExprBlock { - exprs: walk_ast(init_block.exprs, visitor), - }), - conditional: Box::new(walk_expr(*conditional, visitor).unwrap()), - after_block: Box::new(ExprBlock { - exprs: walk_ast(after_block.exprs, visitor), - }), - interior_block: Box::new(ExprBlock { - exprs: walk_ast(interior_block.exprs, visitor), - }), - }), - - //Expr is block - Expr::Block(ExprBlock { exprs }) => Expr::Block(ExprBlock { - exprs: walk_ast(exprs, visitor), - }), - - //Expr is variable - Expr::Variable(ExprVariableReference { - identifier: _, - inferred_type: _, - }) => expr, - Expr::Case(_) => todo!(), - Expr::Switch(_) => todo!(), - }); - } - None -} diff --git a/crates/papyrus/src/optimizer/mod.rs b/crates/papyrus/src/optimizer/mod.rs new file mode 100644 index 0000000..31913cd --- /dev/null +++ b/crates/papyrus/src/optimizer/mod.rs @@ -0,0 +1,153 @@ +use crate::types::{ + Expr, ExprAssignment, ExprBlock, ExprDeclareVariable, ExprForLoop, ExprFunctionCall, + ExprFunctionDefinition, ExprIfStatement, ExprRepeat, ExprVariableReference, +}; + +pub mod const_variable; +pub mod for_loop_to_repeat; +pub mod variable_assignment; + +pub fn optimize_ast(ast: Vec) -> Vec { + // let mut assignment_visitor = VariableAssignmentVisitor::default(); + // let ast = walk_ast(ast, &mut assignment_visitor); + + // let const_variables = assignment_visitor.get_const_variables(); + // let ast = walk_ast(ast, &mut ConstVariableVisitor { const_variables }); + + // walk_ast(ast, &mut ForLoopToRepeatVisitor {}) + // TODO: fix optimizations + ast +} + +// Walks through each expression in the abstract syntax tree, optimizing the AST where possible. A new, optimized AST is returned +//Which is then passed into the Miden generation logic. +fn walk_ast(ast: Vec, visitor: &mut V) -> Vec { + let mut new_ast = vec![]; + for expr in ast { + if let Some(expr) = walk_expr(expr, visitor) { + new_ast.push(expr); + } + } + new_ast +} + +// TODO: it would be nice if there wasn't so much cloning in here +fn walk_expr(expr: Expr, visitor: &mut V) -> Option { + let expr = visitor.visit_expr(expr); + if let Some(expr) = expr { + return Some(match expr { + //Expr is literal + Expr::Literal(ref _x) => expr, + + //Expr is function call + Expr::FunctionCall(ExprFunctionCall { + function_name, + inferred_return_types, + inferred_param_types, + exprs, + }) => Expr::FunctionCall(ExprFunctionCall { + function_name, + inferred_return_types, + inferred_param_types, + exprs: Box::new(vec![ + walk_expr(exprs[0].clone(), visitor).unwrap(), + walk_expr(exprs[1].clone(), visitor).unwrap(), + ]), + }), + + //Expr is if statement + Expr::IfStatement(ExprIfStatement { + first_expr, + second_expr, + }) => Expr::IfStatement(ExprIfStatement { + first_expr: Box::new(walk_expr(*first_expr, visitor).unwrap()), + second_expr: Box::new(ExprBlock { + exprs: walk_ast(second_expr.exprs, visitor), + }), + }), + + //Expr is assignment + Expr::Assignment(ExprAssignment { + inferred_types, + identifiers, + rhs, + }) => Expr::Assignment(ExprAssignment { + identifiers, + inferred_types, + rhs: Box::new(walk_expr(*rhs, visitor).unwrap()), + }), + + //Expr is declare variable + Expr::DeclareVariable(ExprDeclareVariable { + typed_identifiers, + rhs, + }) => Expr::DeclareVariable(ExprDeclareVariable { + typed_identifiers, + rhs: rhs.map(|rhs| Box::new(walk_expr(*rhs, visitor).unwrap())), + }), + + //TODO: Expr is function definition + Expr::FunctionDefinition(ExprFunctionDefinition { + function_name: _, + params: _, + returns: _, + block: _, + }) => todo!(), + + //TODO: Expr is break + Expr::Break => todo!(), + + //TODO: Expr is continue + Expr::Continue => todo!(), + Expr::Leave => todo!(), + + //Expr is repeat + Expr::Repeat(ExprRepeat { + interior_block, + iterations, + }) => Expr::Repeat(ExprRepeat { + iterations, + interior_block: Box::new(ExprBlock { + exprs: walk_ast(interior_block.exprs, visitor), + }), + }), + + //Expr is for loop + Expr::ForLoop(ExprForLoop { + init_block, + conditional, + after_block, + interior_block, + }) => Expr::ForLoop(ExprForLoop { + init_block: Box::new(ExprBlock { + exprs: walk_ast(init_block.exprs, visitor), + }), + conditional: Box::new(walk_expr(*conditional, visitor).unwrap()), + after_block: Box::new(ExprBlock { + exprs: walk_ast(after_block.exprs, visitor), + }), + interior_block: Box::new(ExprBlock { + exprs: walk_ast(interior_block.exprs, visitor), + }), + }), + + //Expr is block + Expr::Block(ExprBlock { exprs }) => Expr::Block(ExprBlock { + exprs: walk_ast(exprs, visitor), + }), + + //Expr is variable + Expr::Variable(ExprVariableReference { + identifier: _, + inferred_type: _, + }) => expr, + Expr::Case(_) => todo!(), + Expr::Switch(_) => todo!(), + }); + } + None +} + +trait ExpressionVisitor { + fn visit_expr(&mut self, expr: Expr) -> Option; +} diff --git a/crates/papyrus/src/optimizer/variable_assignment.rs b/crates/papyrus/src/optimizer/variable_assignment.rs new file mode 100644 index 0000000..5769804 --- /dev/null +++ b/crates/papyrus/src/optimizer/variable_assignment.rs @@ -0,0 +1,59 @@ +use std::collections::HashMap; + +use crate::types::{Expr, ExprLiteral}; + +use super::ExpressionVisitor; + +//The variable assignment visitor keeps track of variables that are reused through the code and the last assignment. +//Variables that do not change can be optimized by converting them into constants. +#[derive(Default)] +struct VariableAssignmentVisitor { + assignment_counter: HashMap, + last_assignment: HashMap, +} + +impl VariableAssignmentVisitor { + // Checks for variables that are only assigned once and returns a hashmap of the variables to convert into constants. + fn get_const_variables(&self) -> HashMap { + self.assignment_counter + .iter() + .filter(|(_k, v)| **v == 1) + .filter_map(|(k, _)| { + if let Some(value) = self.last_assignment.get(k) { + return Some((k.clone(), value.clone())); + } + None + }) + .collect::>() + } +} + +impl ExpressionVisitor for VariableAssignmentVisitor { + fn visit_expr(&mut self, _expr: Expr) -> Option { + todo!(); + // match &expr { + // Expr::DeclareVariable(ExprDeclareVariable { identifier, rhs }) => { + // if let Some(Expr::Literal(literal)) = rhs.clone().map(|r| *r) { + // self.last_assignment.insert(identifier.clone(), literal); + // } + // let count = self + // .assignment_counter + // .entry(identifier.clone()) + // .or_insert(0); + // *count += 1; + // } + // Expr::Assignment(ExprAssignment { + // typed_identifier: identifier, + // rhs: _, + // }) => { + // let count = self + // .assignment_counter + // .entry(identifier.clone()) + // .or_insert(0); + // *count += 1; + // } + // _ => {} + // } + // Some(expr) + } +} From 1fe29050d022312fcc471d49e3c2c976c90ac010 Mon Sep 17 00:00:00 2001 From: 0xKitsune <0xKitsune@protonmail.com> Date: Fri, 27 Jan 2023 02:25:23 -0500 Subject: [PATCH 5/6] added CostEstimator and MidenOperation, allowing the optimizer to keep track of how expensive a sequence is --- .../papyrus/src/optimizer/cost_estimator.rs | 261 ++++++++++++++++++ crates/papyrus/src/optimizer/mod.rs | 1 + 2 files changed, 262 insertions(+) create mode 100644 crates/papyrus/src/optimizer/cost_estimator.rs diff --git a/crates/papyrus/src/optimizer/cost_estimator.rs b/crates/papyrus/src/optimizer/cost_estimator.rs new file mode 100644 index 0000000..b9e141f --- /dev/null +++ b/crates/papyrus/src/optimizer/cost_estimator.rs @@ -0,0 +1,261 @@ +pub struct CostEstimator { + pub total_cost: u128, +} + +impl CostEstimator { + pub fn new() -> Self { + CostEstimator { total_cost: 0 } + } + + pub fn increment(&mut self, operation: MidenOperation) { + self.total_cost += operation.cost() as u128; + } +} + +pub enum MidenOperation { + //Field Operations + Assert, + AssertEq, + Add, + Sub, + Mul, + + Div, + Neg, + Inv, + CheckedPow2, + UncheckedPow2, + NOT, + AND, + OR, + XOR, + //Comparison Operations + Eq, + Neq, + Lt, + Lte, + Gt, + Gte, + Eqw, + //U32 Conversions and tests + U32Test, + U32TestW, + U32Assert, + U32Assert1, + U32Assert2, + U32AssertW, + U32Cast, + U32Split, + //U32 Arithmetic operations + U32CheckedAdd, + U32OverflowingAdd, + U32OverflowingAdd3, + U32CheckedSuB, + U32CheckedSubB, + U32OverflowingSub, + U32OverflowingSubB, + U32WrappingSub, + U32CheckedMul, + U32WrappingMul, + U32OverflowingMadd, + U32CheckedDiv, + U32UncheckedDiv, + U32CheckedMod, + U32UncheckedMod, + U32CheckedDivMod, + U32UncheckedDivMod, + //U32 Bitwise operations + U32CheckedAND, + U32CheckedOR, + U32CheckedXOR, + U32CheckedNOT, + U32CheckedShl, + U32UncheckedShl, + U32CheckedShr, + U32UncheckedShr, + U32CheckedRotl, + U32UncheckedRotl, + U32CheckedRotr, + U32UncheckedRotr, + //U32 Comparison operations + U32CheckedEq, + U32CheckedNeq, + U32CheckedLt, + U32UncheckedLt, + U32CheckedLte, + U32UncheckedLte, + U32CheckedGt, + U32UnCheckedGt, + U32CheckedGte, + U32UncheckedGte, + U32CheckedMin, + U32UncheckedMin, + U32CheckedMax, + U32UncheckedMax, + //Stack manipulation + Drop, + DropW, + Pad, + Dup, + DupW, + Swap, + SwapW, + Movup, + MovupW, + Movdn, + MovdnW, + CSwap, + CSwapW, + CDrop, + CDropW, + //Input/Output Operations + Push, + PushEnvSdepth, + PushEnvLocaddr, + PushAdv, + LoadWAdv, + PushMem, + PushWMem, + LoadWMem, + PopMem, + PopWMem, + StoreWMem, + PushLocal, + PushWLocal, + LoadWLocal, + PopLocal, + PopWlocal, + StoreWLocal, + //Cryptographic operations + Rpperm, + Rphash, + MtreeGet, + MtreeSet, + MtreeCwm, +} + +impl MidenOperation { + //TODO: we might need to update this to be a larger value, I am not sure what the cost of each op might look like + pub fn cost(&self) -> u8 { + match self { + //Field Operations + Self::Assert => 0, + Self::AssertEq => 0, + Self::Add => 0, + Self::Sub => 0, + Self::Mul => 0, + Self::Div => 0, + Self::Neg => 0, + Self::Inv => 0, + Self::CheckedPow2 => 0, + Self::UncheckedPow2 => 0, + Self::NOT => 0, + Self::AND => 0, + Self::OR => 0, + Self::XOR => 0, + //Comparison Operations + Self::Eq => 0, + Self::Neq => 0, + Self::Lt => 0, + Self::Lte => 0, + Self::Gt => 0, + Self::Gte => 0, + Self::Eqw => 0, + //U32 Conversions and tests + Self::U32Test => 0, + Self::U32TestW => 0, + Self::U32Assert => 0, + Self::U32Assert1 => 0, + Self::U32Assert2 => 0, + Self::U32AssertW => 0, + Self::U32Cast => 0, + Self::U32Split => 0, + //U32 Arithmetic operations + Self::U32CheckedAdd => 0, + Self::U32OverflowingAdd => 0, + Self::U32OverflowingAdd3 => 0, + Self::U32CheckedSuB => 0, + Self::U32CheckedSubB => 0, + Self::U32OverflowingSub => 0, + Self::U32OverflowingSubB => 0, + Self::U32WrappingSub => 0, + Self::U32CheckedMul => 0, + Self::U32WrappingMul => 0, + Self::U32OverflowingMadd => 0, + Self::U32CheckedDiv => 0, + Self::U32UncheckedDiv => 0, + Self::U32CheckedMod => 0, + Self::U32UncheckedMod => 0, + Self::U32CheckedDivMod => 0, + Self::U32UncheckedDivMod => 0, + //U32 Bitwise operations + Self::U32CheckedAND => 0, + Self::U32CheckedOR => 0, + Self::U32CheckedXOR => 0, + Self::U32CheckedNOT => 0, + Self::U32CheckedShl => 0, + Self::U32UncheckedShl => 0, + Self::U32CheckedShr => 0, + Self::U32UncheckedShr => 0, + Self::U32CheckedRotl => 0, + Self::U32UncheckedRotl => 0, + Self::U32CheckedRotr => 0, + Self::U32UncheckedRotr => 0, + //U32 Comparison operations + Self::U32CheckedEq => 0, + Self::U32CheckedNeq => 0, + Self::U32CheckedLt => 0, + Self::U32UncheckedLt => 0, + Self::U32CheckedLte => 0, + Self::U32UncheckedLte => 0, + Self::U32CheckedGt => 0, + Self::U32UnCheckedGt => 0, + Self::U32CheckedGte => 0, + Self::U32UncheckedGte => 0, + Self::U32CheckedMin => 0, + Self::U32UncheckedMin => 0, + Self::U32CheckedMax => 0, + Self::U32UncheckedMax => 0, + //Stack manipulation + Self::Drop => 0, + Self::DropW => 0, + Self::Pad => 0, + Self::Dup => 0, + Self::DupW => 0, + Self::Swap => 0, + Self::SwapW => 0, + Self::Movup => 0, + Self::MovupW => 0, + Self::Movdn => 0, + Self::MovdnW => 0, + Self::CSwap => 0, + Self::CSwapW => 0, + Self::CDrop => 0, + Self::CDropW => 0, + //Input/Output Operations + Self::Push => 0, + Self::PushEnvSdepth => 0, + Self::PushEnvLocaddr => 0, + Self::PushAdv => 0, + Self::LoadWAdv => 0, + Self::PushMem => 0, + Self::PushWMem => 0, + Self::LoadWMem => 0, + Self::PopMem => 0, + Self::PopWMem => 0, + Self::StoreWMem => 0, + Self::PushLocal => 0, + Self::PushWLocal => 0, + Self::LoadWLocal => 0, + Self::PopLocal => 0, + Self::PopWlocal => 0, + Self::StoreWLocal => 0, + //Cryptographic operations + Self::Rpperm => 0, + Self::Rphash => 0, + Self::MtreeGet => 0, + Self::MtreeSet => 0, + Self::MtreeCwm => 0, + } + } +} diff --git a/crates/papyrus/src/optimizer/mod.rs b/crates/papyrus/src/optimizer/mod.rs index 31913cd..5627e47 100644 --- a/crates/papyrus/src/optimizer/mod.rs +++ b/crates/papyrus/src/optimizer/mod.rs @@ -4,6 +4,7 @@ use crate::types::{ }; pub mod const_variable; +pub mod cost_estimator; pub mod for_loop_to_repeat; pub mod variable_assignment; From 4af1b5919c56cb613923726f95b9f6076a2b0fe3 Mon Sep 17 00:00:00 2001 From: 0xKitsune <0xKitsune@protonmail.com> Date: Fri, 27 Jan 2023 02:26:47 -0500 Subject: [PATCH 6/6] added evaluate function --- crates/papyrus/src/optimizer/cost_estimator.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/papyrus/src/optimizer/cost_estimator.rs b/crates/papyrus/src/optimizer/cost_estimator.rs index b9e141f..d4839ed 100644 --- a/crates/papyrus/src/optimizer/cost_estimator.rs +++ b/crates/papyrus/src/optimizer/cost_estimator.rs @@ -7,6 +7,10 @@ impl CostEstimator { CostEstimator { total_cost: 0 } } + //TODO: not sure what the program passed in might look like, but we can this function return the cost of the sequence + pub fn evaluate() {} + + //TODO: we might not need this, might only need the evaluate function pub fn increment(&mut self, operation: MidenOperation) { self.total_cost += operation.cost() as u128; }