From 0cb5cb68fcfe74de2da4733675c23d91193f597e Mon Sep 17 00:00:00 2001 From: webbrain-one <295484252+webbrain-one@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:04:41 +0300 Subject: [PATCH] Add one-time stack reordering for repeat loops Add a one-time stack reordering step in AST optimization for repeat loops. Closes #51 --- crates/papyrus/src/ast_optimization.rs | 197 ++++++++++++++++--------- 1 file changed, 125 insertions(+), 72 deletions(-) diff --git a/crates/papyrus/src/ast_optimization.rs b/crates/papyrus/src/ast_optimization.rs index c94d7ba..9fcccdc 100644 --- a/crates/papyrus/src/ast_optimization.rs +++ b/crates/papyrus/src/ast_optimization.rs @@ -1,5 +1,6 @@ #![allow(dead_code)] use std::{collections::HashMap, vec}; +use primitive_types::U256; use crate::types::*; @@ -11,7 +12,7 @@ pub fn optimize_ast(ast: Vec) -> Vec { // let const_variables = assignment_visitor.get_const_variables(); // let ast = walk_ast(ast, &mut ConstVariableVisitor { const_variables }); - // walk_ast(ast, &mut ForLoopToRepeatVisitor {}) + let ast = walk_ast(ast, &mut ForLoopToRepeatVisitor {}); // TODO: fix optimizations ast } @@ -73,77 +74,129 @@ 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 { + match &expr { + Expr::ForLoop(ExprForLoop { + init_block, + conditional, + after_block, + interior_block, + }) => { + // The init block must contain exactly one variable declaration with a literal value + if init_block.exprs.len() != 1 { + return Some(expr); + } + + let (iterator_name, start_value) = { + if let Expr::DeclareVariable(ExprDeclareVariable { + typed_identifiers, + rhs, + }) = &init_block.exprs[0] + { + if typed_identifiers.len() != 1 { + return Some(expr); + } + let name = typed_identifiers[0].identifier.clone(); + if let Some(rhs) = rhs { + if let Expr::Literal(ExprLiteral::Number(lit)) = &**rhs { + (name, lit.value) + } else { + return Some(expr); + } + } else { + return Some(expr); + } + } else { + return Some(expr); + } + }; + + // The after block must be: iterator := add(iterator, 1) + if after_block.exprs.len() != 1 { + return Some(expr); + } + + if let Expr::Assignment(ExprAssignment { + identifiers, + rhs, + .. + }) = &after_block.exprs[0] + { + if identifiers.len() != 1 || identifiers[0] != iterator_name { + return Some(expr); + } + if let Expr::FunctionCall(ExprFunctionCall { + function_name, + exprs, + .. + }) = &**rhs + { + if function_name != "add" || exprs.len() != 2 { + return Some(expr); + } + if let Expr::Variable(ExprVariableReference { + identifier, .. + }) = &exprs[0] + { + if identifier != &iterator_name { + return Some(expr); + } + } else { + return Some(expr); + } + if let Expr::Literal(ExprLiteral::Number(lit)) = &exprs[1] { + if lit.value != U256::from(1) { + return Some(expr); + } + } else { + return Some(expr); + } + } else { + return Some(expr); + } + } else { + return Some(expr); + } + + // The conditional must be: lt(iterator, N) where N is a literal + if let Expr::FunctionCall(ExprFunctionCall { + function_name, + exprs, + .. + }) = &**conditional + { + if function_name != "lt" || exprs.len() != 2 { + return Some(expr); + } + if let Expr::Variable(ExprVariableReference { + identifier, .. + }) = &exprs[0] + { + if identifier != &iterator_name { + return Some(expr); + } + } else { + return Some(expr); + } + if let Expr::Literal(ExprLiteral::Number(lit)) = &exprs[1] { + let end_value = lit.value; + if end_value <= start_value { + return Some(expr); + } + let iterations = (end_value - start_value).as_u64(); + return Some(Expr::Repeat(ExprRepeat { + interior_block: interior_block.clone(), + iterations, + })); + } else { + return Some(expr); + } + } else { + return Some(expr); + } + } + _ => Some(expr), + } } }