1use std::collections::HashSet;
71
72use rucc_ir::{Block, BlockCall, Builder, Extra, Flags, Func, Imm, Inst, Opcode, Type, Value};
73
74use crate::cfg::Cfg;
75use crate::{Analyses, Fuel, Pass, Preserved, Stats};
76
77const CONVERTED: &str = "switch replaced by a range check and the arithmetic its arms were doing";
79
80const NO_FUEL: &str = "switch left alone, the pass ran out of fuel";
82
83const TOO_FEW: &str = "switch left alone, it has too few labels for arithmetic to be cheaper";
85
86const NOT_CONSECUTIVE: &str = "switch left alone, its labels are not consecutive";
88
89const ARM_IS_SHARED: &str = "switch left alone, an arm is reached from somewhere other than it";
91
92const ARM_DOES_WORK: &str = "switch left alone, an arm does more than work out a constant";
94
95const ARMS_DIFFER: &str = "switch left alone, its arms do not all hand on the same thing";
97
98const NOT_AFFINE: &str = "switch left alone, its answers are not a fixed multiple of the label \
100 plus a constant";
101
102const WIDTHS_DIFFER: &str = "switch left alone, its answers are not as wide as its labels";
104
105const LABELS: usize = 3;
107
108#[derive(Debug)]
110pub struct SwitchConv;
111
112impl Pass for SwitchConv {
113 fn name(&self) -> &'static str {
114 "switch-conv"
115 }
116
117 fn describe(&self) -> &'static str {
118 "a switch whose arms are a fixed multiple of the label becomes a range check and arithmetic"
119 }
120
121 fn preserves(&self) -> Preserved {
122 Preserved::NONE
124 }
125
126 fn run(&self, func: &mut Func, an: &mut Analyses, fuel: &mut Fuel) -> Stats {
127 let mut stats = Stats::new();
128 if func.entry().is_none() {
129 return stats;
130 }
131 let cfg = an.cfg(func);
132 let found: Vec<Inst> = func
133 .blocks()
134 .filter_map(|block| func.terminator(block))
135 .filter(|&inst| func[inst].opcode == Opcode::Switch)
136 .collect();
137
138 let mut plans = Vec::new();
139 for inst in found {
140 match plan(func, cfg, inst) {
141 Ok(plan) => plans.push(plan),
142 Err(why) => stats.missed(why),
143 }
144 }
145
146 let mut changed = false;
147 for plan in plans {
148 if !fuel.take() {
149 stats.missed(NO_FUEL);
150 continue;
151 }
152 apply(func, &plan);
153 stats.optimized(CONVERTED);
154 changed = true;
155 }
156 if changed {
157 an.clear();
158 }
159 stats
160 }
161}
162
163#[derive(Clone, Copy, Debug, PartialEq, Eq)]
165enum Hands {
166 On(Block),
168 Back,
170}
171
172#[derive(Debug)]
174struct Plan {
175 inst: Inst,
177 value: Value,
179 ty: Type,
181 hands: Hands,
183 args: Vec<Value>,
186 answer: usize,
188 scale: i128,
190 offset: i128,
192 arms: Vec<Block>,
194}
195
196fn plan(func: &Func, cfg: &Cfg, inst: Inst) -> Result<Plan, &'static str> {
198 let Extra::Switch(info) = func[inst].extra else { return Err(ARMS_DIFFER) };
199 let info = func[info];
200 let Some(&value) = func[func[inst].args].first() else { return Err(ARMS_DIFFER) };
201 let ty = func[value].ty;
202 if !ty.is_int() {
203 return Err(WIDTHS_DIFFER);
204 }
205 let calls: Vec<BlockCall> = func[info.targets].to_vec();
206 let labels: Vec<i128> = func[info.cases].iter().map(|imm| imm.signed(ty)).collect();
207 let Some((&default, arms)) = calls.split_first() else { return Err(ARMS_DIFFER) };
208 if arms.len() != labels.len() || arms.len() < LABELS {
209 return Err(TOO_FEW);
210 }
211 if arms.iter().any(|call| call.block == default.block) {
215 return Err(ARM_IS_SHARED);
216 }
217
218 for pair in labels.windows(2) {
223 if pair[1].checked_sub(pair[0]) != Some(1) {
224 return Err(NOT_CONSECUTIVE);
225 }
226 }
227
228 let mut hands = None;
231 let mut shared: Option<Vec<Value>> = None;
232 let mut answer = None;
233 let mut answers = Vec::new();
234 for call in arms {
235 if !call.args.is_empty() {
236 return Err(ARM_DOES_WORK);
237 }
238 if cfg.predecessors(call.block).len() != 1 {
239 return Err(ARM_IS_SHARED);
240 }
241 let (way, args) = tail(func, call.block)?;
242 if *hands.get_or_insert(way) != way {
243 return Err(ARMS_DIFFER);
244 }
245 let previous = shared.get_or_insert_with(|| args.clone());
246 if previous.len() != args.len() {
247 return Err(ARMS_DIFFER);
248 }
249 for (index, (&mine, &theirs)) in previous.iter().zip(&args).enumerate() {
252 if mine == theirs {
253 continue;
254 }
255 if *answer.get_or_insert(index) != index {
256 return Err(ARMS_DIFFER);
257 }
258 }
259 let at = answer.unwrap_or(0);
260 let Some(&handed) = args.get(at) else { return Err(ARMS_DIFFER) };
261 if func[handed].ty != ty {
262 return Err(WIDTHS_DIFFER);
263 }
264 let Some(number) = constant(func, handed) else { return Err(NOT_AFFINE) };
265 answers.push(number);
266 }
267 let (Some(hands), Some(args)) = (hands, shared) else { return Err(ARMS_DIFFER) };
268 let answer = answer.ok_or(NOT_AFFINE)?;
269
270 let (scale, offset) = line(&labels, &answers, ty).ok_or(NOT_AFFINE)?;
271 Ok(Plan {
272 inst,
273 value,
274 ty,
275 hands,
276 args,
277 answer,
278 scale,
279 offset,
280 arms: arms.iter().map(|call| call.block).collect(),
281 })
282}
283
284fn tail(func: &Func, block: Block) -> Result<(Hands, Vec<Value>), &'static str> {
290 let Some(last) = func.terminator(block) else { return Err(ARM_DOES_WORK) };
291 for inst in func.insts(block) {
292 if inst != last && func[inst].opcode != Opcode::IConst {
293 return Err(ARM_DOES_WORK);
294 }
295 }
296 let args: Vec<Value> = match func[last].opcode {
297 Opcode::Jump => {
298 let Some(call) = func.successors(last).next() else { return Err(ARM_DOES_WORK) };
299 let args = func[call.args].to_vec();
300 return Ok((Hands::On(call.block), args));
301 }
302 Opcode::Return => func[func[last].args].to_vec(),
303 _ => return Err(ARM_DOES_WORK),
304 };
305 Ok((Hands::Back, args))
306}
307
308fn constant(func: &Func, value: Value) -> Option<i128> {
310 crate::discharge::constant(func, value)
311}
312
313fn line(labels: &[i128], answers: &[i128], ty: Type) -> Option<(i128, i128)> {
320 let [first, second, ..] = *labels else { return None };
321 let [low, high, ..] = *answers else { return None };
322 debug_assert_eq!(second - first, 1, "the labels were checked to be consecutive");
323 let scale = high.checked_sub(low)?;
324 let offset = low.checked_sub(scale.checked_mul(first)?)?;
325 for (&label, &answer) in labels.iter().zip(answers) {
326 let want = scale.checked_mul(label)?.checked_add(offset)?;
327 if wrap(want, ty) != answer {
328 return None;
329 }
330 }
331 Some((scale, offset))
332}
333
334fn wrap(value: i128, ty: Type) -> i128 {
339 Imm::int(value, ty).signed(ty)
340}
341
342fn apply(func: &mut Func, plan: &Plan) {
344 let span = func.span(plan.inst);
345 let hit = func.create_block();
346 let mut builder = Builder::new(func, hit).at(span);
347 let scaled = match plan.scale {
348 0 => builder.iconst(plan.ty, plan.offset),
349 1 => plan.value,
350 scale => {
351 let by = builder.iconst(plan.ty, scale);
352 builder.binary(Opcode::Mul, plan.value, by, Flags::NONE)
353 }
354 };
355 let answer = if plan.offset == 0 || plan.scale == 0 {
356 scaled
357 } else {
358 let by = builder.iconst(plan.ty, plan.offset);
359 builder.binary(Opcode::Add, scaled, by, Flags::NONE)
360 };
361 let mut args = plan.args.clone();
362 args[plan.answer] = answer;
363 match plan.hands {
364 Hands::On(block) => builder.jump(block, &args),
365 Hands::Back => builder.ret(&args),
366 };
367
368 let Extra::Switch(info) = func[plan.inst].extra else { return };
371 let empty = func.push_values(&[]);
372 let mut calls: Vec<BlockCall> = func[func[info].targets].to_vec();
373 for call in &mut calls[1..] {
374 *call = BlockCall::new(hit, empty);
377 }
378 let targets = func.push_block_calls(&calls);
379 let cases = func[info].cases;
380 let info = func.add_switch(rucc_ir::SwitchInfo { targets, cases });
381 func[plan.inst].extra = Extra::Switch(info);
382
383 let mut gone = HashSet::new();
386 for &arm in &plan.arms {
387 if gone.insert(arm) {
388 func.remove_block(arm);
389 }
390 }
391}
392
393#[cfg(test)]
394mod tests {
395 use std::collections::HashMap;
396
397 use rucc_base::Interner;
398 use rucc_ir::{Block, Builder, Func, Opcode, Signature, Type, Value};
399
400 use super::SwitchConv;
401 use crate::stats::Kind;
402 use crate::{Fuel, Pass, Stats};
403
404 fn i32() -> Type {
406 Type::int(32)
407 }
408
409 fn convert(func: &mut Func) -> Stats {
411 SwitchConv.run(func, &mut crate::machine::fixtures::analyses(), &mut Fuel::unlimited())
412 }
413
414 fn returning(ty: Type, labels: &[i128], answers: &[i128]) -> Func {
419 let mut names = Interner::new();
420 let mut func = Func::new(names.intern("f"), Signature::new());
421 let head = func.create_block();
422 let value = func.append_param(head, ty);
423 let default = func.create_block();
424 let arms: Vec<Block> = answers.iter().map(|_| func.create_block()).collect();
425 for (&arm, &answer) in arms.iter().zip(answers) {
426 let mut build = Builder::new(&mut func, arm);
427 let it = build.iconst(ty, answer);
428 build.ret(&[it]);
429 }
430 let mut build = Builder::new(&mut func, default);
431 let it = build.iconst(ty, 999);
432 build.ret(&[it]);
433 let cases: Vec<(i128, Block)> = labels.iter().copied().zip(arms.iter().copied()).collect();
434 Builder::new(&mut func, head).switch(value, default, &cases);
435 func
436 }
437
438 fn cases(func: &Func) -> Vec<usize> {
440 let head = func.entry().expect("a function with blocks in it");
441 let term = func.terminator(head).expect("a head block has one");
442 func.successors(term).skip(1).map(|call| call.block.index()).collect()
443 }
444
445 fn arm(func: &Func) -> Block {
447 let blocks = cases(func);
448 let first = blocks[0];
449 assert!(blocks.iter().all(|&block| block == first), "the case edges did not all move");
450 Block::from_usize(first)
451 }
452
453 fn opcodes(func: &Func, block: Block) -> Vec<Opcode> {
455 func.insts(block).map(|inst| func[inst].opcode).collect()
456 }
457
458 fn answer(func: &Func, block: Block, label: i128) -> i128 {
464 let head = func.entry().expect("a function with blocks in it");
465 let mut values: HashMap<Value, i128> = HashMap::new();
466 values.insert(func[head].params[0], label);
467 for inst in func.insts(block) {
468 let data = func[inst];
469 let Some(result) = data.first_result else {
470 let args = func[data.args].to_vec();
471 let handed = match data.opcode {
472 Opcode::Return => args[0],
473 Opcode::Jump => {
474 func[func.successors(inst).next().expect("a jump goes").args][0]
475 }
476 other => panic!("a block this pass wrote ends in {other:?}"),
477 };
478 return values[&handed];
479 };
480 let args: Vec<i128> = func[data.args].iter().map(|arg| values[arg]).collect();
481 let it = match data.opcode {
482 Opcode::IConst => {
483 let (imm, ty) = crate::fold::constant(func, result).expect("a constant is one");
484 imm.signed(ty)
485 }
486 Opcode::Mul => args[0].wrapping_mul(args[1]),
487 Opcode::Add => args[0].wrapping_add(args[1]),
488 other => panic!("this pass does not write {other:?}"),
489 };
490 values.insert(result, super::wrap(it, func[result].ty));
491 }
492 panic!("a block with no terminator");
493 }
494
495 fn fired(stats: &Stats) -> bool {
497 stats.total(Kind::Optimized) > 0
498 }
499
500 #[test]
501 fn labels_that_run_with_their_answers_become_one_addition() {
502 let mut func = returning(i32(), &[0, 1, 2, 3], &[1, 2, 3, 4]);
503 assert!(fired(&convert(&mut func)));
504 let arm = arm(&func);
505 assert_eq!(opcodes(&func, arm), [Opcode::IConst, Opcode::Add, Opcode::Return]);
506 for label in 0..4 {
507 assert_eq!(answer(&func, arm, label), label + 1);
508 }
509 }
510
511 #[test]
512 fn answers_that_are_a_multiple_of_the_label_become_a_multiplication() {
513 let mut func = returning(i32(), &[3, 4, 5, 6], &[30, 40, 50, 60]);
514 assert!(fired(&convert(&mut func)));
515 let arm = arm(&func);
516 assert_eq!(opcodes(&func, arm), [Opcode::IConst, Opcode::Mul, Opcode::Return]);
517 for label in 3..7 {
518 assert_eq!(answer(&func, arm, label), label * 10);
519 }
520 }
521
522 #[test]
523 fn answers_that_are_all_the_same_become_the_constant_they_all_were() {
524 let mut func = returning(i32(), &[7, 8, 9, 10], &[9, 9, 9, 9]);
525 assert!(fired(&convert(&mut func)));
526 let arm = arm(&func);
527 assert_eq!(opcodes(&func, arm), [Opcode::IConst, Opcode::Return]);
528 assert_eq!(answer(&func, arm, 8), 9);
529 }
530
531 #[test]
532 fn labels_that_run_below_zero_are_a_run_like_any_other() {
533 let mut func = returning(i32(), &[-2, -1, 0, 1], &[-4, -2, 0, 2]);
534 assert!(fired(&convert(&mut func)));
535 let arm = arm(&func);
536 for label in -2..2 {
537 assert_eq!(answer(&func, arm, label), label * 2);
538 }
539 }
540
541 #[test]
548 fn a_line_that_only_holds_by_wrapping_still_holds() {
549 let ty = Type::int(8);
550 let mut func = returning(ty, &[0, 1, 2], &[0, 100, -56]);
551 assert!(fired(&convert(&mut func)));
552 let arm = arm(&func);
553 assert_eq!(answer(&func, arm, 2), -56);
554 }
555
556 #[test]
557 fn labels_with_a_hole_in_them_are_left_alone() {
558 let mut func = returning(i32(), &[0, 1, 3], &[1, 2, 4]);
559 assert!(!fired(&convert(&mut func)));
560 assert_eq!(cases(&func).len(), 3);
561 }
562
563 #[test]
564 fn answers_that_are_not_a_line_are_left_alone() {
565 let mut func = returning(i32(), &[0, 1, 2], &[5, 9, 2]);
566 assert!(!fired(&convert(&mut func)));
567 }
568
569 #[test]
570 fn two_labels_are_not_enough_to_pay_for_the_arithmetic() {
571 let mut func = returning(i32(), &[0, 1], &[1, 2]);
572 assert!(!fired(&convert(&mut func)));
573 }
574
575 #[test]
576 fn an_answer_wider_than_its_label_is_left_alone() {
577 let mut names = Interner::new();
578 let mut func = Func::new(names.intern("f"), Signature::new());
579 let head = func.create_block();
580 let value = func.append_param(head, i32());
581 let default = func.create_block();
582 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
583 for (index, &arm) in arms.iter().enumerate() {
584 let mut build = Builder::new(&mut func, arm);
585 let it = build.iconst(Type::int(64), index as i128 + 1);
586 build.ret(&[it]);
587 }
588 let mut build = Builder::new(&mut func, default);
589 let it = build.iconst(Type::int(64), 0);
590 build.ret(&[it]);
591 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
592 Builder::new(&mut func, head).switch(value, default, &cases);
593 assert!(!fired(&convert(&mut func)));
594 }
595
596 #[test]
597 fn an_arm_something_else_reaches_is_left_alone() {
598 let mut func = returning(i32(), &[0, 1, 2], &[1, 2, 3]);
599 let default = Block::from_usize(1);
602 let arm = Block::from_usize(2);
603 let term = func.terminator(default).expect("the default returns");
604 func.remove_inst(term);
605 Builder::new(&mut func, default).jump(arm, &[]);
606 assert!(!fired(&convert(&mut func)));
607 }
608
609 #[test]
610 fn an_arm_that_is_also_the_default_is_left_alone() {
611 let mut names = Interner::new();
612 let mut func = Func::new(names.intern("f"), Signature::new());
613 let head = func.create_block();
614 let value = func.append_param(head, i32());
615 let shared = func.create_block();
616 let mut build = Builder::new(&mut func, shared);
617 let it = build.iconst(i32(), 1);
618 build.ret(&[it]);
619 let others: Vec<Block> = (0..2).map(|_| func.create_block()).collect();
620 for (index, &arm) in others.iter().enumerate() {
621 let mut build = Builder::new(&mut func, arm);
622 let it = build.iconst(i32(), index as i128 + 2);
623 build.ret(&[it]);
624 }
625 let cases = [(0, shared), (1, others[0]), (2, others[1])];
626 Builder::new(&mut func, head).switch(value, shared, &cases);
627 assert!(!fired(&convert(&mut func)));
628 }
629
630 #[test]
631 fn arms_that_join_keep_what_they_pass_beside_the_answer() {
632 let mut names = Interner::new();
633 let mut func = Func::new(names.intern("f"), Signature::new());
634 let head = func.create_block();
635 let value = func.append_param(head, i32());
636 let alongside = func.append_param(head, i32());
637 let join = func.create_block();
638 let handed = func.append_param(join, i32());
639 let carried = func.append_param(join, i32());
640 Builder::new(&mut func, join).ret(&[handed, carried]);
641 let default = func.create_block();
642 let mut build = Builder::new(&mut func, default);
643 let it = build.iconst(i32(), 999);
644 build.jump(join, &[it, alongside]);
645 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
646 for (index, &arm) in arms.iter().enumerate() {
647 let mut build = Builder::new(&mut func, arm);
648 let it = build.iconst(i32(), index as i128 + 1);
649 build.jump(join, &[it, alongside]);
650 }
651 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
652 Builder::new(&mut func, head).switch(value, default, &cases);
653 assert!(fired(&convert(&mut func)));
654
655 let arm = arm(&func);
656 assert_eq!(answer(&func, arm, 2), 3);
657 let term = func.terminator(arm).expect("the block ends in a jump");
659 let call = func.successors(term).next().expect("a jump goes somewhere");
660 assert_eq!(func[call.args][1], alongside);
661 }
662
663 #[test]
664 fn arms_that_hand_on_two_different_things_are_left_alone() {
665 let mut names = Interner::new();
666 let mut func = Func::new(names.intern("f"), Signature::new());
667 let head = func.create_block();
668 let value = func.append_param(head, i32());
669 let join = func.create_block();
670 let first = func.append_param(join, i32());
671 let second = func.append_param(join, i32());
672 Builder::new(&mut func, join).ret(&[first, second]);
673 let default = func.create_block();
674 let mut build = Builder::new(&mut func, default);
675 let it = build.iconst(i32(), 999);
676 build.jump(join, &[it, it]);
677 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
678 for (index, &arm) in arms.iter().enumerate() {
679 let mut build = Builder::new(&mut func, arm);
680 let one = build.iconst(i32(), index as i128 + 1);
681 let two = build.iconst(i32(), index as i128 + 10);
682 build.jump(join, &[one, two]);
683 }
684 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
685 Builder::new(&mut func, head).switch(value, default, &cases);
686 assert!(!fired(&convert(&mut func)));
687 }
688
689 #[test]
690 fn an_arm_that_does_something_is_left_alone() {
691 let mut names = Interner::new();
692 let mut func = Func::new(names.intern("f"), Signature::new());
693 let head = func.create_block();
694 let value = func.append_param(head, i32());
695 let default = func.create_block();
696 let mut build = Builder::new(&mut func, default);
697 let it = build.iconst(i32(), 999);
698 build.ret(&[it]);
699 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
700 for (index, &arm) in arms.iter().enumerate() {
701 let mut build = Builder::new(&mut func, arm);
702 let it = build.iconst(i32(), index as i128 + 1);
703 let sum = build.binary(Opcode::Add, it, value, rucc_ir::Flags::NONE);
705 build.ret(&[sum]);
706 }
707 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
708 Builder::new(&mut func, head).switch(value, default, &cases);
709 assert!(!fired(&convert(&mut func)));
710 }
711
712 #[test]
713 fn the_default_goes_where_it_went() {
714 let mut func = returning(i32(), &[0, 1, 2, 3], &[1, 2, 3, 4]);
715 let head = func.entry().expect("a function with blocks in it");
716 let before = func.terminator(head).expect("a head block has one");
717 let was = func.successors(before).next().expect("a switch has a default").block;
718 assert!(fired(&convert(&mut func)));
719 let after = func.terminator(head).expect("a head block has one");
720 let now = func.successors(after).next().expect("a switch has a default").block;
721 assert_eq!(was, now, "the default moved");
722 }
723}