1use std::collections::{HashMap, HashSet};
58
59use rucc_base::Interner;
60use rucc_mir as mir;
61use rucc_target::{FrameInsts, Role};
62
63pub fn addresses(
72 func: &mut mir::Func,
73 insts: &FrameInsts,
74 names: &mut Interner,
75 waiting: &HashSet<mir::Inst>,
76) -> usize {
77 let lea = mir::Opcode::new(names.intern(&format!("{}{}", insts.prefix, insts.lea)));
78 let reads = reads(func);
79 let mut folded = 0;
80 for block in func.blocks().collect::<Vec<_>>() {
81 let mut open: HashMap<mir::Reg, mir::Inst> = HashMap::new();
84 for inst in func.insts(block).collect::<Vec<_>>() {
85 if let Some(folding) = candidate(func, &open, inst) {
86 let operands = func.push_operands(&folding.operands);
87 let mem = func.add_amode(folding.amode);
88 func[inst].operands = operands;
89 func[inst].mem = Some(mem);
90 open.remove(&folding.base);
91 func.remove_inst(folding.from);
92 folded += 1;
93 }
94 for written in written(func, inst) {
95 open.retain(|reg, &mut held| *reg != written && !touches(func, held, written));
96 }
97 if func[inst].opcode == lea && !waiting.contains(&inst) {
98 if let Some(reg) = written_once(func, &reads, inst) {
99 open.insert(reg, inst);
100 }
101 }
102 }
103 }
104 folded
105}
106
107pub(crate) fn reads(func: &mir::Func) -> HashMap<mir::Reg, usize> {
118 let mut counts = HashMap::new();
119 for block in func.blocks() {
120 for inst in func.insts(block) {
121 for operand in &func[func[inst].operands] {
122 if operand.role == Role::Use {
123 *counts.entry(operand.reg).or_insert(0) += 1;
124 }
125 }
126 }
127 for call in &func[block].succs {
128 for &arg in &call.args {
129 *counts.entry(arg).or_insert(0) += 1;
130 }
131 }
132 }
133 counts
134}
135
136fn written_once(
139 func: &mir::Func,
140 reads: &HashMap<mir::Reg, usize>,
141 inst: mir::Inst,
142) -> Option<mir::Reg> {
143 let operands = &func[func[inst].operands];
144 let mut defs = operands.iter().filter(|operand| operand.role != Role::Use);
145 let def = defs.next()?;
146 if defs.next().is_some() || !def.reg.is_virtual() || reads.get(&def.reg) != Some(&1) {
147 return None;
148 }
149 Some(def.reg)
150}
151
152fn written(func: &mir::Func, inst: mir::Inst) -> Vec<mir::Reg> {
154 func[func[inst].operands]
155 .iter()
156 .filter(|operand| operand.role != Role::Use)
157 .map(|operand| operand.reg)
158 .collect()
159}
160
161fn touches(func: &mir::Func, inst: mir::Inst, reg: mir::Reg) -> bool {
164 let Some(mem) = func[inst].mem else { return false };
165 let amode = func[mem];
166 let operands = &func[func[inst].operands];
167 [amode.base, amode.index]
168 .into_iter()
169 .flatten()
170 .filter_map(|at| operands.get(usize::from(at)))
171 .any(|operand| operand.reg == reg)
172}
173
174fn base_reg(func: &mir::Func, inst: mir::Inst) -> Option<mir::Reg> {
180 let amode = func[func[inst].mem?];
181 if amode.index.is_some() || amode.symbol.is_some() || amode.got {
182 return None;
183 }
184 Some(func[func[inst].operands].get(usize::from(amode.base?))?.reg)
185}
186
187struct Folding {
193 from: mir::Inst,
195 base: mir::Reg,
197 operands: Vec<mir::Operand>,
199 amode: mir::Amode,
201}
202
203fn candidate(
212 func: &mir::Func,
213 open: &HashMap<mir::Reg, mir::Inst>,
214 inst: mir::Inst,
215) -> Option<Folding> {
216 let base = base_reg(func, inst)?;
217 let from = *open.get(&base)?;
218 let address = func[func[from].mem?];
219 let disp = i64::from(address.disp) + i64::from(func[func[inst].mem?].disp);
224 let mut amode = mir::Amode { disp: i32::try_from(disp).ok()?, ..address };
225
226 let taken = &func[func[from].operands];
227 let reader = &func[func[inst].operands];
228 let mut operands = reader.get(..reader.len().checked_sub(1)?)?.to_vec();
229 for (at, into) in [(address.base, &mut amode.base), (address.index, &mut amode.index)] {
230 let Some(at) = at else { continue };
231 operands.push(*taken.get(usize::from(at))?);
232 *into = Some(u8::try_from(operands.len() - 1).ok()?);
233 }
234 Some(Folding { from, base, operands, amode })
235}
236
237#[cfg(test)]
238mod tests {
239 use rucc_target::x86_64::{FRAME, GPR, RDI};
240
241 use super::*;
242
243 fn empty() -> (Interner, mir::Func, mir::Block) {
245 let mut names = Interner::new();
246 let mut func = mir::Func::new(names.intern("f"));
247 let block = func.create_block();
248 (names, func, block)
249 }
250
251 fn op(names: &mut Interner, name: &str) -> mir::Opcode {
253 mir::Opcode::new(names.intern(&format!("{}{name}", FRAME.prefix)))
254 }
255
256 fn shape(func: &mir::Func, names: &Interner, block: mir::Block) -> Vec<(String, mir::Amode)> {
258 func.insts(block)
259 .map(|inst| {
260 let amode = func[inst].mem.map_or(mir::Amode::NOTHING, |mem| func[mem]);
261 (names.resolve(func[inst].opcode.name()).to_owned(), amode)
262 })
263 .collect()
264 }
265
266 fn address_regs(func: &mir::Func, inst: mir::Inst) -> Vec<mir::Reg> {
268 let amode = func[func[inst].mem.expect("a memory operand")];
269 let operands = &func[func[inst].operands];
270 [amode.base, amode.index]
271 .into_iter()
272 .flatten()
273 .map(|at| operands[usize::from(at)].reg)
274 .collect()
275 }
276
277 #[test]
280 fn an_address_a_load_reads_once_becomes_the_load_s_own_addressing_mode() {
281 let (mut names, mut func, block) = empty();
282 let array = func.new_vreg(GPR);
283 let index = func.new_vreg(GPR);
284 let address = func.new_vreg(GPR);
285 let value = func.new_vreg(GPR);
286 let lea = op(&mut names, FRAME.lea);
287 let load = op(&mut names, "mov_rm_32");
288 func.build(block, lea)
289 .def(address, GPR)
290 .mem(
291 mir::Mem::at(mir::Operand::read(array, GPR))
292 .indexed(mir::Operand::read(index, GPR), 4),
293 )
294 .finish();
295 func.build(block, load)
296 .def(value, GPR)
297 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
298 .finish();
299
300 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 1);
301
302 let left = shape(&func, &names, block);
303 assert_eq!(left.len(), 1, "the address is worked out twice: {left:?}");
304 assert_eq!(left[0].0, format!("{}mov_rm_32", FRAME.prefix));
305 assert_eq!(left[0].1.scale, 4);
306 assert_eq!(left[0].1.disp, 0);
307 let inst = func.insts(block).next().expect("the load is still there");
308 assert_eq!(address_regs(&func, inst), vec![array, index], "the load reads the wrong pair");
309 }
310
311 #[test]
314 fn the_displacements_of_the_two_addresses_are_added() {
315 let (mut names, mut func, block) = empty();
316 let array = func.new_vreg(GPR);
317 let address = func.new_vreg(GPR);
318 let value = func.new_vreg(GPR);
319 let lea = op(&mut names, FRAME.lea);
320 let load = op(&mut names, "mov_rm_32");
321 func.build(block, lea)
322 .def(address, GPR)
323 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(16))
324 .finish();
325 func.build(block, load)
326 .def(value, GPR)
327 .mem(mir::Mem::at(mir::Operand::read(address, GPR)).plus(8))
328 .finish();
329
330 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 1);
331
332 let left = shape(&func, &names, block);
333 assert_eq!(left.len(), 1);
334 assert_eq!(left[0].1.disp, 24, "the field is at the sum of the two offsets or nowhere");
335 }
336
337 #[test]
340 fn a_store_keeps_the_value_it_is_storing() {
341 let (mut names, mut func, block) = empty();
342 let array = func.new_vreg(GPR);
343 let index = func.new_vreg(GPR);
344 let address = func.new_vreg(GPR);
345 let value = func.new_vreg(GPR);
346 let lea = op(&mut names, FRAME.lea);
347 let store = op(&mut names, "mov_mr_32");
348 func.build(block, lea)
349 .def(address, GPR)
350 .mem(
351 mir::Mem::at(mir::Operand::read(array, GPR))
352 .indexed(mir::Operand::read(index, GPR), 8),
353 )
354 .finish();
355 func.build(block, store)
356 .uses(value, GPR)
357 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
358 .finish();
359
360 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 1);
361
362 let inst = func.insts(block).next().expect("the store is still there");
363 let regs: Vec<mir::Reg> = func[func[inst].operands].iter().map(|op| op.reg).collect();
364 assert_eq!(regs, vec![value, array, index], "the value the store writes went missing");
365 assert_eq!(func[func[inst].mem.expect("a memory operand")].scale, 8);
366 }
367
368 #[test]
371 fn an_address_two_instructions_read_is_left_where_it_is() {
372 let (mut names, mut func, block) = empty();
373 let array = func.new_vreg(GPR);
374 let address = func.new_vreg(GPR);
375 let lea = op(&mut names, FRAME.lea);
376 let load = op(&mut names, "mov_rm_32");
377 func.build(block, lea)
378 .def(address, GPR)
379 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(16))
380 .finish();
381 for _ in 0..2 {
382 let value = func.new_vreg(GPR);
383 func.build(block, load)
384 .def(value, GPR)
385 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
386 .finish();
387 }
388
389 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 0);
390 assert_eq!(shape(&func, &names, block).len(), 3);
391 }
392
393 #[test]
396 fn a_reader_that_already_has_an_index_is_left_alone() {
397 let (mut names, mut func, block) = empty();
398 let array = func.new_vreg(GPR);
399 let index = func.new_vreg(GPR);
400 let address = func.new_vreg(GPR);
401 let value = func.new_vreg(GPR);
402 let lea = op(&mut names, FRAME.lea);
403 let load = op(&mut names, "mov_rm_32");
404 func.build(block, lea)
405 .def(address, GPR)
406 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(16))
407 .finish();
408 func.build(block, load)
409 .def(value, GPR)
410 .mem(
411 mir::Mem::at(mir::Operand::read(address, GPR))
412 .indexed(mir::Operand::read(index, GPR), 4),
413 )
414 .finish();
415
416 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 0);
417 assert_eq!(shape(&func, &names, block).len(), 2);
418 }
419
420 #[test]
424 fn two_displacements_that_do_not_fit_together_are_not_put_together() {
425 let (mut names, mut func, block) = empty();
426 let array = func.new_vreg(GPR);
427 let address = func.new_vreg(GPR);
428 let value = func.new_vreg(GPR);
429 let lea = op(&mut names, FRAME.lea);
430 let load = op(&mut names, "mov_rm_32");
431 func.build(block, lea)
432 .def(address, GPR)
433 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(i32::MAX))
434 .finish();
435 func.build(block, load)
436 .def(value, GPR)
437 .mem(mir::Mem::at(mir::Operand::read(address, GPR)).plus(1))
438 .finish();
439
440 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 0);
441 assert_eq!(shape(&func, &names, block).len(), 2);
442 }
443
444 #[test]
447 fn a_register_the_address_reads_being_written_in_between_ends_the_chance() {
448 let (mut names, mut func, block) = empty();
449 let array = mir::Reg::physical(RDI);
450 let address = func.new_vreg(GPR);
451 let value = func.new_vreg(GPR);
452 let lea = op(&mut names, FRAME.lea);
453 let load = op(&mut names, "mov_rm_32");
454 let put = op(&mut names, "mov_ri_64");
455 func.build(block, lea)
456 .def(address, GPR)
457 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(16))
458 .finish();
459 func.build(block, put).def(array, GPR).imm(7).finish();
460 func.build(block, load)
461 .def(value, GPR)
462 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
463 .finish();
464
465 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 0);
466 assert_eq!(shape(&func, &names, block).len(), 3);
467 }
468
469 #[test]
472 fn a_reader_in_another_block_is_not_one_this_folds_into() {
473 let (mut names, mut func, block) = empty();
474 let next = func.create_block();
475 let array = func.new_vreg(GPR);
476 let address = func.new_vreg(GPR);
477 let value = func.new_vreg(GPR);
478 let lea = op(&mut names, FRAME.lea);
479 let load = op(&mut names, "mov_rm_32");
480 func.build(block, lea)
481 .def(address, GPR)
482 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(16))
483 .finish();
484 *func.succs_mut(block) = vec![mir::BlockCall::to(next)];
485 func.build(next, load)
486 .def(value, GPR)
487 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
488 .finish();
489
490 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 0);
491 }
492
493 #[test]
497 fn a_chain_of_two_addresses_is_folded_the_whole_way_in_one_pass() {
498 let (mut names, mut func, block) = empty();
499 let array = func.new_vreg(GPR);
500 let index = func.new_vreg(GPR);
501 let element = func.new_vreg(GPR);
502 let field = func.new_vreg(GPR);
503 let value = func.new_vreg(GPR);
504 let lea = op(&mut names, FRAME.lea);
505 let load = op(&mut names, "mov_rm_32");
506 func.build(block, lea)
507 .def(element, GPR)
508 .mem(
509 mir::Mem::at(mir::Operand::read(array, GPR))
510 .indexed(mir::Operand::read(index, GPR), 8),
511 )
512 .finish();
513 func.build(block, lea)
514 .def(field, GPR)
515 .mem(mir::Mem::at(mir::Operand::read(element, GPR)).plus(4))
516 .finish();
517 func.build(block, load)
518 .def(value, GPR)
519 .mem(mir::Mem::at(mir::Operand::read(field, GPR)))
520 .finish();
521
522 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 2);
523
524 let left = shape(&func, &names, block);
525 assert_eq!(left.len(), 1, "one of the two addresses is still its own instruction");
526 assert_eq!(left[0].1.scale, 8);
527 assert_eq!(left[0].1.disp, 4);
528 let inst = func.insts(block).next().expect("the load is still there");
529 assert_eq!(address_regs(&func, inst), vec![array, index]);
530 }
531
532 #[test]
536 fn an_address_of_a_global_folds_into_the_reader_symbol_and_all() {
537 let (mut names, mut func, block) = empty();
538 let global = names.intern("counters");
539 let address = func.new_vreg(GPR);
540 let value = func.new_vreg(GPR);
541 let lea = op(&mut names, FRAME.lea);
542 let load = op(&mut names, "mov_rm_32");
543 func.build(block, lea).def(address, GPR).mem(mir::Mem::of(global)).finish();
544 func.build(block, load)
545 .def(value, GPR)
546 .mem(mir::Mem::at(mir::Operand::read(address, GPR)).plus(12))
547 .finish();
548
549 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 1);
550
551 let left = shape(&func, &names, block);
552 assert_eq!(left.len(), 1);
553 assert_eq!(left[0].1.symbol, Some(global));
554 assert_eq!(left[0].1.disp, 12);
555 }
556
557 #[test]
561 fn an_address_whose_displacement_is_still_to_be_written_is_left_where_it_is() {
562 let (mut names, mut func, block) = empty();
563 let sp = mir::Reg::physical(RDI);
564 let address = func.new_vreg(GPR);
565 let value = func.new_vreg(GPR);
566 let lea = op(&mut names, FRAME.lea);
567 let load = op(&mut names, "mov_rm_32");
568 let local = func
569 .build(block, lea)
570 .def(address, GPR)
571 .mem(mir::Mem::at(mir::Operand::read(sp, GPR)))
572 .finish();
573 func.build(block, load)
574 .def(value, GPR)
575 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
576 .finish();
577
578 let waiting = HashSet::from([local]);
579 assert_eq!(addresses(&mut func, &FRAME, &mut names, &waiting), 0);
580 assert_eq!(shape(&func, &names, block).len(), 2);
581
582 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 1);
585 }
586
587 #[test]
591 fn only_the_target_s_address_instruction_is_one_this_folds() {
592 let (mut names, mut func, block) = empty();
593 let array = func.new_vreg(GPR);
594 let address = func.new_vreg(GPR);
595 let value = func.new_vreg(GPR);
596 let load = op(&mut names, "mov_rm_64");
597 let read = op(&mut names, "mov_rm_32");
598 func.build(block, load)
599 .def(address, GPR)
600 .mem(mir::Mem::at(mir::Operand::read(array, GPR)).plus(16))
601 .finish();
602 func.build(block, read)
603 .def(value, GPR)
604 .mem(mir::Mem::at(mir::Operand::read(address, GPR)))
605 .finish();
606
607 assert_eq!(addresses(&mut func, &FRAME, &mut names, &HashSet::new()), 0);
608 assert_eq!(shape(&func, &names, block).len(), 2);
609 }
610}