1use crate::bytecode::{Register, Value};
2use crate::jit;
3use crate::jit::trace::{Trace, TraceOp};
4use std::collections::HashSet;
5pub struct TraceOptimizer {
6 hoisted_constants: Vec<(Register, Value)>,
7}
8
9impl TraceOptimizer {
10 pub fn new() -> Self {
11 Self {
12 hoisted_constants: Vec::new(),
13 }
14 }
15
16 pub fn optimize(&mut self, trace: &mut Trace) -> Vec<(Register, Value)> {
17 jit::log(|| "🔧 JIT Optimizer: Starting optimization...".to_string());
18 let original_ops = trace.ops.len();
19 self.hoist_constants(trace);
20 self.unroll_loop(trace, crate::jit::UNROLL_FACTOR);
21 self.eliminate_arithmetic_moves(trace);
22 self.coalesce_registers(trace);
23 let optimized_ops = trace.ops.len();
24 let hoisted = self.hoisted_constants.len();
25 jit::log(|| {
26 format!(
27 "✨ JIT Optimizer: Optimized {} ops → {} ops, hoisted {} constants",
28 original_ops, optimized_ops, hoisted
29 )
30 });
31 self.hoisted_constants.clone()
32 }
33
34 fn hoist_constants(&mut self, trace: &mut Trace) {
35 let mut non_hoistable: HashSet<Register> = HashSet::new();
36 for op in &trace.ops {
37 if let TraceOp::CallNative { callee, .. } = op {
38 non_hoistable.insert(*callee);
39 }
40 }
41
42 let mut new_ops = Vec::new();
43 let mut already_hoisted: HashSet<Register> = HashSet::new();
44 let mut assigned: HashSet<Register> = HashSet::new();
45
46 let dest_of = |op: &TraceOp| -> Option<Register> {
47 match op {
48 TraceOp::Move { dest, .. }
49 | TraceOp::Add { dest, .. }
50 | TraceOp::Sub { dest, .. }
51 | TraceOp::Mul { dest, .. }
52 | TraceOp::Div { dest, .. }
53 | TraceOp::Mod { dest, .. }
54 | TraceOp::Neg { dest, .. }
55 | TraceOp::Eq { dest, .. }
56 | TraceOp::Ne { dest, .. }
57 | TraceOp::Lt { dest, .. }
58 | TraceOp::Le { dest, .. }
59 | TraceOp::Gt { dest, .. }
60 | TraceOp::Ge { dest, .. }
61 | TraceOp::And { dest, .. }
62 | TraceOp::Or { dest, .. }
63 | TraceOp::Not { dest, .. }
64 | TraceOp::Concat { dest, .. }
65 | TraceOp::GetIndex { dest, .. }
66 | TraceOp::ArrayLen { dest, .. }
67 | TraceOp::CallMethod { dest, .. }
68 | TraceOp::GetField { dest, .. }
69 | TraceOp::NewStruct { dest, .. }
70 | TraceOp::CallNative { dest, .. } => Some(*dest),
71 _ => None,
72 }
73 };
74
75 for op in trace.ops.drain(..) {
76 match op {
77 TraceOp::LoadConst { dest, value } => {
78 if non_hoistable.contains(&dest)
79 || already_hoisted.contains(&dest)
80 || assigned.contains(&dest)
81 {
82 new_ops.push(TraceOp::LoadConst { dest, value });
83 } else {
84 self.hoisted_constants.push((dest, value));
85 already_hoisted.insert(dest);
86 }
87 assigned.insert(dest);
88 }
89
90 other => {
91 if let Some(dest) = dest_of(&other) {
92 assigned.insert(dest);
93 }
94 new_ops.push(other);
95 }
96 }
97 }
98
99 trace.ops = new_ops;
100 }
101
102 fn eliminate_arithmetic_moves(&mut self, trace: &mut Trace) {
103 let mut new_ops = Vec::new();
104 let mut i = 0;
105 while i < trace.ops.len() {
106 if i + 1 < trace.ops.len() {
107 let current = &trace.ops[i];
108 let next = &trace.ops[i + 1];
109 if let Some((arith_dest, final_dest)) = self.match_arithmetic_move(current, next) {
110 let mut rewritten = current.clone();
111 self.rewrite_arithmetic_dest(&mut rewritten, final_dest);
112 new_ops.push(rewritten);
113 i += 2;
114 continue;
115 }
116 }
117
118 new_ops.push(trace.ops[i].clone());
119 i += 1;
120 }
121
122 trace.ops = new_ops;
123 }
124
125 fn match_arithmetic_move(&self, op1: &TraceOp, op2: &TraceOp) -> Option<(Register, Register)> {
126 let arith_dest = match op1 {
127 TraceOp::Add { dest, .. }
128 | TraceOp::Sub { dest, .. }
129 | TraceOp::Mul { dest, .. }
130 | TraceOp::Div { dest, .. }
131 | TraceOp::Mod { dest, .. } => *dest,
132 _ => return None,
133 };
134 if let TraceOp::Move {
135 dest: move_dest,
136 src,
137 } = op2
138 {
139 if *src == arith_dest {
140 return Some((arith_dest, *move_dest));
141 }
142 }
143
144 None
145 }
146
147 fn rewrite_arithmetic_dest(&self, op: &mut TraceOp, new_dest: Register) {
148 match op {
149 TraceOp::Add { dest, .. }
150 | TraceOp::Sub { dest, .. }
151 | TraceOp::Mul { dest, .. }
152 | TraceOp::Div { dest, .. }
153 | TraceOp::Mod { dest, .. } => {
154 *dest = new_dest;
155 }
156
157 _ => {}
158 }
159 }
160
161 fn unroll_loop(&mut self, trace: &mut Trace, factor: usize) {
162 if factor <= 1 || trace.ops.is_empty() {
163 return;
164 }
165
166 let loop_condition_op = trace.ops.iter().find_map(|op| match op {
167 TraceOp::Le { dest, lhs, rhs } => Some((op.clone(), *dest)),
168 TraceOp::Lt { dest, lhs, rhs } => Some((op.clone(), *dest)),
169 TraceOp::Ge { dest, lhs, rhs } => Some((op.clone(), *dest)),
170 TraceOp::Gt { dest, lhs, rhs } => Some((op.clone(), *dest)),
171 _ => None,
172 });
173 if loop_condition_op.is_none() {
174 return;
175 }
176
177 let (loop_cmp_op, cond_reg) = loop_condition_op.unwrap();
178 let original_ops = trace.ops.clone();
179 let mut new_ops = Vec::new();
180 new_ops.extend(original_ops.iter().cloned());
181 for iteration in 1..factor {
182 new_ops.push(loop_cmp_op.clone());
183 new_ops.push(TraceOp::GuardLoopContinue {
184 condition_register: cond_reg,
185 bailout_ip: trace.start_ip,
186 });
187 for op in &original_ops {
188 if !matches!(
189 op,
190 TraceOp::Le { .. }
191 | TraceOp::Lt { .. }
192 | TraceOp::Ge { .. }
193 | TraceOp::Gt { .. }
194 ) {
195 new_ops.push(op.clone());
196 }
197 }
198 }
199
200 trace.ops = new_ops;
201 }
202
203 fn coalesce_registers(&mut self, _trace: &mut Trace) {}
204}
205
206impl Default for TraceOptimizer {
207 fn default() -> Self {
208 Self::new()
209 }
210}