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).clone();
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 { block: hit, args: empty };
375 }
376 let targets = func.push_block_calls(&calls);
377 let cases = func[info].cases;
378 let info = func.add_switch(rucc_ir::SwitchInfo { targets, cases });
379 func[plan.inst].extra = Extra::Switch(info);
380
381 let mut gone = HashSet::new();
384 for &arm in &plan.arms {
385 if gone.insert(arm) {
386 func.remove_block(arm);
387 }
388 }
389}
390
391#[cfg(test)]
392mod tests {
393 use std::collections::HashMap;
394
395 use rucc_base::Interner;
396 use rucc_ir::{Block, Builder, Func, Opcode, Signature, Type, Value};
397
398 use super::SwitchConv;
399 use crate::stats::Kind;
400 use crate::{Fuel, Pass, Stats};
401
402 fn i32() -> Type {
404 Type::int(32)
405 }
406
407 fn convert(func: &mut Func) -> Stats {
409 SwitchConv.run(func, &mut crate::machine::fixtures::analyses(), &mut Fuel::unlimited())
410 }
411
412 fn returning(ty: Type, labels: &[i128], answers: &[i128]) -> Func {
417 let mut names = Interner::new();
418 let mut func = Func::new(names.intern("f"), Signature::new());
419 let head = func.create_block();
420 let value = func.append_param(head, ty);
421 let default = func.create_block();
422 let arms: Vec<Block> = answers.iter().map(|_| func.create_block()).collect();
423 for (&arm, &answer) in arms.iter().zip(answers) {
424 let mut build = Builder::new(&mut func, arm);
425 let it = build.iconst(ty, answer);
426 build.ret(&[it]);
427 }
428 let mut build = Builder::new(&mut func, default);
429 let it = build.iconst(ty, 999);
430 build.ret(&[it]);
431 let cases: Vec<(i128, Block)> = labels.iter().copied().zip(arms.iter().copied()).collect();
432 Builder::new(&mut func, head).switch(value, default, &cases);
433 func
434 }
435
436 fn cases(func: &Func) -> Vec<usize> {
438 let head = func.entry().expect("a function with blocks in it");
439 let term = func.terminator(head).expect("a head block has one");
440 func.successors(term).skip(1).map(|call| call.block.index()).collect()
441 }
442
443 fn arm(func: &Func) -> Block {
445 let blocks = cases(func);
446 let first = blocks[0];
447 assert!(blocks.iter().all(|&block| block == first), "the case edges did not all move");
448 Block::from_usize(first)
449 }
450
451 fn opcodes(func: &Func, block: Block) -> Vec<Opcode> {
453 func.insts(block).map(|inst| func[inst].opcode).collect()
454 }
455
456 fn answer(func: &Func, block: Block, label: i128) -> i128 {
462 let head = func.entry().expect("a function with blocks in it");
463 let mut values: HashMap<Value, i128> = HashMap::new();
464 values.insert(func[head].params[0], label);
465 for inst in func.insts(block) {
466 let data = func[inst];
467 let Some(result) = data.first_result else {
468 let args = func[data.args].to_vec();
469 let handed = match data.opcode {
470 Opcode::Return => args[0],
471 Opcode::Jump => {
472 func[func.successors(inst).next().expect("a jump goes").args][0]
473 }
474 other => panic!("a block this pass wrote ends in {other:?}"),
475 };
476 return values[&handed];
477 };
478 let args: Vec<i128> = func[data.args].iter().map(|arg| values[arg]).collect();
479 let it = match data.opcode {
480 Opcode::IConst => {
481 let (imm, ty) = crate::fold::constant(func, result).expect("a constant is one");
482 imm.signed(ty)
483 }
484 Opcode::Mul => args[0].wrapping_mul(args[1]),
485 Opcode::Add => args[0].wrapping_add(args[1]),
486 other => panic!("this pass does not write {other:?}"),
487 };
488 values.insert(result, super::wrap(it, func[result].ty));
489 }
490 panic!("a block with no terminator");
491 }
492
493 fn fired(stats: &Stats) -> bool {
495 stats.total(Kind::Optimized) > 0
496 }
497
498 #[test]
499 fn labels_that_run_with_their_answers_become_one_addition() {
500 let mut func = returning(i32(), &[0, 1, 2, 3], &[1, 2, 3, 4]);
501 assert!(fired(&convert(&mut func)));
502 let arm = arm(&func);
503 assert_eq!(opcodes(&func, arm), [Opcode::IConst, Opcode::Add, Opcode::Return]);
504 for label in 0..4 {
505 assert_eq!(answer(&func, arm, label), label + 1);
506 }
507 }
508
509 #[test]
510 fn answers_that_are_a_multiple_of_the_label_become_a_multiplication() {
511 let mut func = returning(i32(), &[3, 4, 5, 6], &[30, 40, 50, 60]);
512 assert!(fired(&convert(&mut func)));
513 let arm = arm(&func);
514 assert_eq!(opcodes(&func, arm), [Opcode::IConst, Opcode::Mul, Opcode::Return]);
515 for label in 3..7 {
516 assert_eq!(answer(&func, arm, label), label * 10);
517 }
518 }
519
520 #[test]
521 fn answers_that_are_all_the_same_become_the_constant_they_all_were() {
522 let mut func = returning(i32(), &[7, 8, 9, 10], &[9, 9, 9, 9]);
523 assert!(fired(&convert(&mut func)));
524 let arm = arm(&func);
525 assert_eq!(opcodes(&func, arm), [Opcode::IConst, Opcode::Return]);
526 assert_eq!(answer(&func, arm, 8), 9);
527 }
528
529 #[test]
530 fn labels_that_run_below_zero_are_a_run_like_any_other() {
531 let mut func = returning(i32(), &[-2, -1, 0, 1], &[-4, -2, 0, 2]);
532 assert!(fired(&convert(&mut func)));
533 let arm = arm(&func);
534 for label in -2..2 {
535 assert_eq!(answer(&func, arm, label), label * 2);
536 }
537 }
538
539 #[test]
546 fn a_line_that_only_holds_by_wrapping_still_holds() {
547 let ty = Type::int(8);
548 let mut func = returning(ty, &[0, 1, 2], &[0, 100, -56]);
549 assert!(fired(&convert(&mut func)));
550 let arm = arm(&func);
551 assert_eq!(answer(&func, arm, 2), -56);
552 }
553
554 #[test]
555 fn labels_with_a_hole_in_them_are_left_alone() {
556 let mut func = returning(i32(), &[0, 1, 3], &[1, 2, 4]);
557 assert!(!fired(&convert(&mut func)));
558 assert_eq!(cases(&func).len(), 3);
559 }
560
561 #[test]
562 fn answers_that_are_not_a_line_are_left_alone() {
563 let mut func = returning(i32(), &[0, 1, 2], &[5, 9, 2]);
564 assert!(!fired(&convert(&mut func)));
565 }
566
567 #[test]
568 fn two_labels_are_not_enough_to_pay_for_the_arithmetic() {
569 let mut func = returning(i32(), &[0, 1], &[1, 2]);
570 assert!(!fired(&convert(&mut func)));
571 }
572
573 #[test]
574 fn an_answer_wider_than_its_label_is_left_alone() {
575 let mut names = Interner::new();
576 let mut func = Func::new(names.intern("f"), Signature::new());
577 let head = func.create_block();
578 let value = func.append_param(head, i32());
579 let default = func.create_block();
580 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
581 for (index, &arm) in arms.iter().enumerate() {
582 let mut build = Builder::new(&mut func, arm);
583 let it = build.iconst(Type::int(64), index as i128 + 1);
584 build.ret(&[it]);
585 }
586 let mut build = Builder::new(&mut func, default);
587 let it = build.iconst(Type::int(64), 0);
588 build.ret(&[it]);
589 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
590 Builder::new(&mut func, head).switch(value, default, &cases);
591 assert!(!fired(&convert(&mut func)));
592 }
593
594 #[test]
595 fn an_arm_something_else_reaches_is_left_alone() {
596 let mut func = returning(i32(), &[0, 1, 2], &[1, 2, 3]);
597 let default = Block::from_usize(1);
600 let arm = Block::from_usize(2);
601 let term = func.terminator(default).expect("the default returns");
602 func.remove_inst(term);
603 Builder::new(&mut func, default).jump(arm, &[]);
604 assert!(!fired(&convert(&mut func)));
605 }
606
607 #[test]
608 fn an_arm_that_is_also_the_default_is_left_alone() {
609 let mut names = Interner::new();
610 let mut func = Func::new(names.intern("f"), Signature::new());
611 let head = func.create_block();
612 let value = func.append_param(head, i32());
613 let shared = func.create_block();
614 let mut build = Builder::new(&mut func, shared);
615 let it = build.iconst(i32(), 1);
616 build.ret(&[it]);
617 let others: Vec<Block> = (0..2).map(|_| func.create_block()).collect();
618 for (index, &arm) in others.iter().enumerate() {
619 let mut build = Builder::new(&mut func, arm);
620 let it = build.iconst(i32(), index as i128 + 2);
621 build.ret(&[it]);
622 }
623 let cases = [(0, shared), (1, others[0]), (2, others[1])];
624 Builder::new(&mut func, head).switch(value, shared, &cases);
625 assert!(!fired(&convert(&mut func)));
626 }
627
628 #[test]
629 fn arms_that_join_keep_what_they_pass_beside_the_answer() {
630 let mut names = Interner::new();
631 let mut func = Func::new(names.intern("f"), Signature::new());
632 let head = func.create_block();
633 let value = func.append_param(head, i32());
634 let alongside = func.append_param(head, i32());
635 let join = func.create_block();
636 let handed = func.append_param(join, i32());
637 let carried = func.append_param(join, i32());
638 Builder::new(&mut func, join).ret(&[handed, carried]);
639 let default = func.create_block();
640 let mut build = Builder::new(&mut func, default);
641 let it = build.iconst(i32(), 999);
642 build.jump(join, &[it, alongside]);
643 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
644 for (index, &arm) in arms.iter().enumerate() {
645 let mut build = Builder::new(&mut func, arm);
646 let it = build.iconst(i32(), index as i128 + 1);
647 build.jump(join, &[it, alongside]);
648 }
649 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
650 Builder::new(&mut func, head).switch(value, default, &cases);
651 assert!(fired(&convert(&mut func)));
652
653 let arm = arm(&func);
654 assert_eq!(answer(&func, arm, 2), 3);
655 let term = func.terminator(arm).expect("the block ends in a jump");
657 let call = func.successors(term).next().expect("a jump goes somewhere");
658 assert_eq!(func[call.args][1], alongside);
659 }
660
661 #[test]
662 fn arms_that_hand_on_two_different_things_are_left_alone() {
663 let mut names = Interner::new();
664 let mut func = Func::new(names.intern("f"), Signature::new());
665 let head = func.create_block();
666 let value = func.append_param(head, i32());
667 let join = func.create_block();
668 let first = func.append_param(join, i32());
669 let second = func.append_param(join, i32());
670 Builder::new(&mut func, join).ret(&[first, second]);
671 let default = func.create_block();
672 let mut build = Builder::new(&mut func, default);
673 let it = build.iconst(i32(), 999);
674 build.jump(join, &[it, it]);
675 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
676 for (index, &arm) in arms.iter().enumerate() {
677 let mut build = Builder::new(&mut func, arm);
678 let one = build.iconst(i32(), index as i128 + 1);
679 let two = build.iconst(i32(), index as i128 + 10);
680 build.jump(join, &[one, two]);
681 }
682 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
683 Builder::new(&mut func, head).switch(value, default, &cases);
684 assert!(!fired(&convert(&mut func)));
685 }
686
687 #[test]
688 fn an_arm_that_does_something_is_left_alone() {
689 let mut names = Interner::new();
690 let mut func = Func::new(names.intern("f"), Signature::new());
691 let head = func.create_block();
692 let value = func.append_param(head, i32());
693 let default = func.create_block();
694 let mut build = Builder::new(&mut func, default);
695 let it = build.iconst(i32(), 999);
696 build.ret(&[it]);
697 let arms: Vec<Block> = (0..3).map(|_| func.create_block()).collect();
698 for (index, &arm) in arms.iter().enumerate() {
699 let mut build = Builder::new(&mut func, arm);
700 let it = build.iconst(i32(), index as i128 + 1);
701 let sum = build.binary(Opcode::Add, it, value, rucc_ir::Flags::NONE);
703 build.ret(&[sum]);
704 }
705 let cases: Vec<(i128, Block)> = (0..3).zip(arms.iter().copied()).collect();
706 Builder::new(&mut func, head).switch(value, default, &cases);
707 assert!(!fired(&convert(&mut func)));
708 }
709
710 #[test]
711 fn the_default_goes_where_it_went() {
712 let mut func = returning(i32(), &[0, 1, 2, 3], &[1, 2, 3, 4]);
713 let head = func.entry().expect("a function with blocks in it");
714 let before = func.terminator(head).expect("a head block has one");
715 let was = func.successors(before).next().expect("a switch has a default").block;
716 assert!(fired(&convert(&mut func)));
717 let after = func.terminator(head).expect("a head block has one");
718 let now = func.successors(after).next().expect("a switch has a default").block;
719 assert_eq!(was, now, "the default moved");
720 }
721}