use std::collections::{BTreeMap};
use std::collections::{HashMap, VecDeque};
#[derive(Debug, Clone, PartialEq)]
pub struct LoopUnrollingDecision {
pub should_unroll: bool,
pub unroll_factor: usize,
pub benefit_score: f64,
pub estimated_iterations: usize,
pub reason: String,
}
#[derive(Debug, Clone)]
pub struct LoopInfo {
pub loop_type: LoopType,
pub start_line: usize,
pub end_line: usize,
pub variable: String,
pub bounds: LoopBounds,
pub body: Vec<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum LoopType {
For,
While,
DoWhile,
}
#[derive(Debug, Clone)]
pub struct LoopBounds {
pub start_value: Option<i64>,
pub end_value: Option<i64>,
pub increment: Option<i64>,
pub is_constant: bool,
}
pub struct LoopUnrollingOptimizer {
default_unroll_factor: usize,
max_unroll_factor: usize,
min_iterations: usize,
analysis_history: HashMap<String, UnrollingStats>,
}
#[derive(Debug, Clone)]
struct UnrollingStats {
pub total_unrolled: usize,
pub total_benefit: f64,
pub avg_benefit: f64,
pub last_analysis: Instant,
}
impl LoopUnrollingOptimizer {
pub fn new() -> Self {
Self {
default_unroll_factor: 4,
max_unroll_factor: 16,
min_iterations: 4,
analysis_history: HashMap::new(),
}
}
pub fn with_params(default_factor: usize, max_factor: usize, min_iter: usize) -> Self {
Self {
default_unroll_factor: default_factor,
max_unroll_factor: max_factor,
min_iterations: min_iter,
analysis_history: HashMap::new(),
}
}
pub fn analyze_loops(&self, code: &str) -> Vec<LoopInfo> {
let lines: Vec<&str> = code.lines().collect();
let mut loops = Vec::new();
let mut i = 0;
while i < lines.len() {
let line: _ = lines[i].trim();
if let Some(loop_info) = self.extract_for_loop(&lines, i) {
loops.push(loop_info);
i = loop_info.end_line + 1;
continue;
}
if let Some(loop_info) = self.extract_while_loop(&lines, i) {
loops.push(loop_info);
i = loop_info.end_line + 1;
continue;
}
i += 1;
}
loops
}
pub fn make_unrolling_decision(&self, loop_info: &LoopInfo) -> LoopUnrollingDecision {
let estimated_iterations: _ = self.estimate_iterations(loop_info);
let should_unroll: _ = if loop_info.bounds.is_constant {
estimated_iterations >= self.min_iterations
} else {
estimated_iterations >= self.min_iterations / 2
};
let unroll_factor: _ = if should_unroll {
self.calculate_unroll_factor(loop_info, estimated_iterations)
} else {
1
};
let benefit_score: _ = if should_unroll {
self.calculate_benefit_score(loop_info, estimated_iterations, unroll_factor)
} else {
0.0
};
let reason: _ = if should_unroll {
format!(
"Unroll {}x: estimated {} iterations, complexity score {:.2}",
unroll_factor,
estimated_iterations,
self.calculate_complexity_score(loop_info)
)
} else {
format!(
"No unrolling: estimated {} iterations, minimum {} required",
estimated_iterations,
self.min_iterations
)
};
LoopUnrollingDecision {
should_unroll,
unroll_factor,
benefit_score,
estimated_iterations,
reason,
}
}
fn extract_for_loop(&self, lines: &[&str], start_idx: usize) -> Option<LoopInfo> {
let line: _ = lines[start_idx].trim();
if !line.starts_with("for ") {
return None;
}
if let Some((init, cond, incr)) = self.parse_for_loop(line) {
let (variable, bounds) = self.parse_for_init(init);
let body: _ = self.extract_loop_body(lines, start_idx);
let mut brace_count = 0;
let mut end_line = start_idx;
for i in start_idx..lines.len() {
let line: _ = lines[i];
brace_count += line.matches('{').count();
brace_count -= line.matches('}').count();
if brace_count == 0 && i > start_idx {
end_line = i;
break;
}
}
return Some(LoopInfo {
loop_type: LoopType::For,
start_line: start_idx,
end_line,
variable,
bounds,
body,
});
}
None
}
fn extract_while_loop(&self, lines: &[&str], start_idx: usize) -> Option<LoopInfo> {
let line: _ = lines[start_idx].trim();
if !line.starts_with("while ") {
return None;
}
let condition: _ = self.extract_while_condition(line);
let body: _ = self.extract_loop_body(lines, start_idx);
let mut brace_count = 0;
let mut end_line = start_idx;
for i in start_idx..lines.len() {
let line: _ = lines[i];
brace_count += line.matches('{').count();
brace_count -= line.matches('}').count();
if brace_count == 0 && i > start_idx {
end_line = i;
break;
}
}
let bounds: _ = LoopBounds {
start_value: None,
end_value: None,
increment: None,
is_constant: false,
};
Some(LoopInfo {
loop_type: LoopType::While,
start_line: start_idx,
end_line,
variable: "i".to_string(), bounds,
body,
})
}
fn parse_for_loop(&self, line: &str) -> Option<(String, String, String)> {
if let Some(start) = line.find('(') {
if let Some(end) = line[start..].find(')') {
let inside: _ = &line[start + 1..start + end];
let parts: Vec<&str> = inside.split(';').collect();
if parts.len() == 3 {
return Some((
parts[0].trim().to_string(),
parts[1].trim().to_string(),
parts[2].trim().to_string(),
));
}
}
}
None
}
fn parse_for_init(&self, init: &str) -> (String, LoopBounds) {
let mut variable = "i".to_string();
let mut start_value = Some(0i64);
let mut is_constant = true;
if init.starts_with("let ") || init.starts_with("const ") || init.starts_with("var ") {
let after_keyword: _ = if init.starts_with("let ") {
&init[4..]
} else if init.starts_with("const ") {
&init[6..]
} else {
&init[4..]
};
if let Some(eq_pos) = after_keyword.find('=') {
let var_part: _ = after_keyword[..eq_pos].trim();
let val_part: _ = after_keyword[eq_pos + 1..].trim();
variable = var_part.to_string();
start_value = val_part.parse::<i64>().ok();
if start_value.is_none() {
is_constant = false;
}
}
}
(
variable,
LoopBounds {
start_value,
end_value: None,
increment: None,
is_constant,
},
)
}
fn extract_while_condition(&self, line: &str) -> String {
if let Some(start) = line.find('(') {
if let Some(end) = line[start..].find(')') {
return line[start + 1..start + end].trim().to_string();
}
}
String::new()
}
fn extract_loop_body(&self, lines: &[&str], start_idx: usize) -> Vec<String> {
let mut body = Vec::new();
let mut brace_count = 0;
let mut found_start = false;
for i in start_idx..lines.len() {
let line: _ = lines[i];
if !found_start {
if line.contains('{') {
found_start = true;
brace_count += line.matches('{').count();
brace_count -= line.matches('}').count();
if let Some(pos) = line.find('{') {
let after_brace: _ = &line[pos + 1..];
if !after_brace.trim().is_empty() {
body.push(after_brace.trim().to_string());
}
}
}
} else {
brace_count += line.matches('{').count();
brace_count -= line.matches('}').count();
if brace_count == 0 {
break;
}
body.push(line.to_string());
}
}
body
}
fn estimate_iterations(&self, loop_info: &LoopInfo) -> usize {
if !loop_info.bounds.is_constant {
return 10; }
if let (Some(start), Some(end)) = (loop_info.bounds.start_value, loop_info.bounds.end_value) {
if end > start {
return (end - start) as usize;
}
}
1 }
fn calculate_unroll_factor(&self, loop_info: &LoopInfo, iterations: usize) -> usize {
let mut factor = self.default_unroll_factor;
if iterations > 100 {
factor = self.max_unroll_factor;
} else if iterations > 50 {
factor = 8;
} else if iterations > 20 {
factor = 6;
}
std::cmp::min(factor, iterations)
}
fn calculate_benefit_score(&self, loop_info: &LoopInfo, iterations: usize, unroll_factor: usize) -> f64 {
let complexity_score: _ = self.calculate_complexity_score(loop_info);
let iteration_factor: _ = iterations as f64 / self.min_iterations as f64;
let unroll_benefit: _ = unroll_factor as f64;
iteration_factor * unroll_benefit * (complexity_score / 10.0)
}
fn calculate_complexity_score(&self, loop_info: &LoopInfo) -> f64 {
let mut score = 10.0; score += loop_info.body.len() as f64 * 2.0;
match loop_info.loop_type {
LoopType::For => score += 5.0,
LoopType::While => score += 8.0,
LoopType::DoWhile => score += 8.0,
}
let nested_loops: _ = self.count_nested_loops(&loop_info.body);
score += nested_loops as f64 * 15.0;
score
}
fn count_nested_loops(&self, body: &[String]) -> usize {
let mut count = 0;
for line in body {
if line.contains("for ") || line.contains("while ") {
count += 1;
}
}
count
}
pub fn record_unrolling(&mut self, loop_hash: &str, benefit: f64) {
let stats: _ = self.analysis_history.entry(loop_hash.to_string()).or_insert(UnrollingStats {
total_unrolled: 0,
total_benefit: 0.0,
avg_benefit: 0.0,
last_analysis: Instant::now(),
});
stats.total_unrolled += 1;
stats.total_benefit += benefit;
stats.avg_benefit = stats.total_benefit / stats.total_unrolled as f64;
stats.last_analysis = Instant::now();
}
pub fn get_unrolling_stats(&self, loop_hash: &str) -> Option<&UnrollingStats> {
self.analysis_history.get(loop_hash)
}
pub fn reset_stats(&mut self) {
self.analysis_history.clear();
}
}
impl Default for LoopUnrollingOptimizer {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::{Duration, Instant};
#[test]
fn test_simple_for_loop_analysis() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let code: _ = r#"
for (let i: _ = 0; i < 10; i++) {
console.log(i);
}
"#;
let loops: _ = optimizer.analyze_loops(code);
assert_eq!(loops.len(), 1);
assert_eq!(loops[0].loop_type, LoopType::For);
assert_eq!(loops[0].variable, "i");
assert!(loops[0].bounds.is_constant);
}
#[test]
fn test_for_loop_bounds_extraction() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let code: _ = r#"
for (let i: _ = 1; i < 100; i++) {
sum += i;
}
"#;
let loops: _ = optimizer.analyze_loops(code);
assert_eq!(loops.len(), 1);
let bounds: _ = &loops[0].bounds;
assert_eq!(bounds.start_value, Some(1));
assert!(bounds.is_constant);
}
#[test]
fn test_while_loop_analysis() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let code: _ = r#"
while (count < 10) {
count++;
}
"#;
let loops: _ = optimizer.analyze_loops(code);
assert_eq!(loops.len(), 1);
assert_eq!(loops[0].loop_type, LoopType::While);
}
#[test]
fn test_nested_loop_analysis() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let code: _ = r#"
for (let i: _ = 0; i < 10; i++) {
for (let j: _ = 0; j < 10; j++) {
sum += i * j;
}
}
"#;
let loops: _ = optimizer.analyze_loops(code);
assert_eq!(loops.len(), 2);
assert!(loops.iter().all(|l| l.loop_type == LoopType::For));
}
#[test]
fn test_unrolling_decision_simple_loop() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let code: _ = r#"
for (let i: _ = 0; i < 10; i++) {
console.log(i);
}
"#;
let loops: _ = optimizer.analyze_loops(code);
let decision: _ = optimizer.make_unrolling_decision(&loops[0]);
assert!(decision.should_unroll);
assert!(decision.unroll_factor > 1);
assert!(decision.benefit_score > 0.0);
}
#[test]
fn test_unrolling_decision_small_loop() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let code: _ = r#"
for (let i: _ = 0; i < 2; i++) {
console.log(i);
}
"#;
let loops: _ = optimizer.analyze_loops(code);
let decision: _ = optimizer.make_unrolling_decision(&loops[0]);
assert!(!decision.should_unroll || decision.estimated_iterations < optimizer.min_iterations);
}
#[test]
fn test_unroll_factor_calculation() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let loop_info: _ = LoopInfo {
loop_type: LoopType::For,
start_line: 0,
end_line: 5,
variable: "i".to_string(),
bounds: LoopBounds {
start_value: Some(0),
end_value: Some(1000),
increment: None,
is_constant: true,
},
body: vec!["console.log(i);".to_string()],
};
let iterations: _ = 1000;
let factor: _ = optimizer.calculate_unroll_factor(&loop_info, iterations);
assert_eq!(factor, optimizer.max_unroll_factor);
}
#[test]
fn test_complexity_score() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let simple_loop: _ = LoopInfo {
loop_type: LoopType::For,
start_line: 0,
end_line: 3,
variable: "i".to_string(),
bounds: LoopBounds {
start_value: Some(0),
end_value: Some(10),
increment: None,
is_constant: true,
},
body: vec!["sum += i;".to_string()],
};
let nested_loop: _ = LoopInfo {
loop_type: LoopType::For,
start_line: 0,
end_line: 10,
variable: "i".to_string(),
bounds: LoopBounds {
start_value: Some(0),
end_value: Some(10),
increment: None,
is_constant: true,
},
body: vec![
"for (let j: _ = 0; j < 10; j++)".to_string(),
"sum += i * j;".to_string(),
],
};
let simple_score: _ = optimizer.calculate_complexity_score(&simple_loop);
let nested_score: _ = optimizer.calculate_complexity_score(&nested_loop);
assert!(nested_score > simple_score, "Nested loop should have higher complexity");
}
#[test]
fn test_iteration_estimation() {
let optimizer: _ = LoopUnrollingOptimizer::new();
let loop_info: _ = LoopInfo {
loop_type: LoopType::For,
start_line: 0,
end_line: 5,
variable: "i".to_string(),
bounds: LoopBounds {
start_value: Some(0),
end_value: Some(100),
increment: None,
is_constant: true,
},
body: Vec::new(),
};
let iterations: _ = optimizer.estimate_iterations(&loop_info);
assert_eq!(iterations, 100);
}
#[test]
fn test_unrolling_history() {
let mut optimizer = LoopUnrollingOptimizer::new();
let loop_hash: _ = "test_loop";
optimizer.record_unrolling(loop_hash, 100.0);
optimizer.record_unrolling(loop_hash, 200.0);
let stats: _ = optimizer.get_unrolling_stats(loop_hash).unwrap();
assert_eq!(stats.total_unrolled, 2);
assert_eq!(stats.total_benefit, 300.0);
assert_eq!(stats.avg_benefit, 150.0);
optimizer.reset_stats();
assert!(optimizer.get_unrolling_stats(loop_hash).is_none());
}
#[test]
fn test_custom_parameters() {
let optimizer: _ = LoopUnrollingOptimizer::with_params(2, 8, 10);
let code: _ = r#"
for (let i: _ = 0; i < 20; i++) {
console.log(i);
}
"#;
let loops: _ = optimizer.analyze_loops(code);
let decision: _ = optimizer.make_unrolling_decision(&loops[0]);
assert!(decision.should_unroll || decision.estimated_iterations >= 10);
}
}