rucc-opt 0.11.12

//! What `__builtin_expect` said, moved onto the branch it was said about.
//!
//! Design: `spec/optimizer/11-profile.md` section 11.2, and tamnd/rucc#364.
//!
//! ```c
//! if (__builtin_expect(error, 0)) { report(); }
//! ```
//!
//! The front end builds an `expect` instruction holding the value and what the program says it will
//! be. The value is what the branch is on, and this pass is what turns the pair into a statement
//! about the branch: the arm the hint names gets ninety parts in a hundred, the other gets ten, and
//! the instruction comes out. On a `switch` the arm the expected value goes to gets the ninety and
//! the other arms share the ten. Everything downstream reads the number off the arm rather than
//! chasing the condition back to a node, which is what makes a profile and a hint the same thing to
//! everything that consumes either.
//!
//! # Why it runs first and at every level
//!
//! An `expect` sits on the branch condition, and a wrapper on a branch condition is a wrapper on
//! whatever the peephole, the folder and the comparison simplifier were about to match. Left
//! standing through a pipeline it would cost code quality on exactly the programs that took the
//! trouble to say which way their branches go, which is the wrong way round. So it comes off in the
//! first pass, before anything has had a chance to fail to match through it.
//!
//! At `-O0` too, where nothing else runs, for the same reason a stray node is a cost there as well:
//! `-O0` emits what it is given, and what it is given would otherwise have an instruction in it per
//! `__builtin_expect` the program wrote. gcc has no such instruction at any level.
//!
//! # What it does not do
//!
//! It does not predict anything. `crate::predict` is where the prediction is, this only records
//! what the program claimed, and the difference matters: a hint is a fact about what somebody wrote
//! and a prediction is a guess, so the ten static predictors write no hints at all. A hint that a
//! heuristic could have written would be indistinguishable afterwards from one the program did.
//!
//! It does not touch a branch whose condition it cannot follow back to an `expect`, and it removes
//! the instruction either way. A hint nothing can place is a hint nothing can use, and leaving the
//! instruction standing for a later pass to find would be leaving the wrapper in the way of
//! everything, which is the thing this exists to avoid.

use std::collections::HashMap;

use rucc_cost::heuristics::PREDICT_EXPECT;
use rucc_ir::{Block, Def, Extra, Func, Hint, Imm, Inst, IntPred, Opcode, Type, Value};

use crate::fold::{constant, convert};
use crate::{Analyses, Analysis, Fuel, Pass, Preserved, Stats, uses};

/// A branch now says which way the program expects it to go.
const PLACED: &str = "branch weight written from a __builtin_expect on its condition";

/// The value was expected and nothing branches on it.
const NO_BRANCH: &str = "__builtin_expect dropped, no branch in this function is on its value";

/// Ran out.
const NO_FUEL: &str = "__builtin_expect kept, the pass ran out of fuel";

/// The pass.
#[derive(Debug)]
pub struct Expect;

impl Pass for Expect {
    fn name(&self) -> &'static str {
        "expect"
    }

    fn describe(&self) -> &'static str {
        "what __builtin_expect said moves onto the arms of the branch it was said about"
    }

    fn preserves(&self) -> Preserved {
        // No block moves and no edge moves, so the graph and everything built on it stand. What
        // changes is who reads which value, which liveness is a statement about, and what the
        // branches say about themselves, which is what the frequencies are worked out from.
        Preserved::ALL.without(Analysis::Liveness).without(Analysis::Frequencies)
    }

    fn required(&self) -> bool {
        // This is the only thing that removes an `Opcode::Expect`, and the back end has no rule
        // that lowers one, so `-fno-expect` would not be a compile without the hints. It would be
        // a compile that stops on a construct the program never wrote.
        true
    }

    fn run(&self, func: &mut Func, _an: &mut Analyses, fuel: &mut Fuel) -> Stats {
        let mut stats = Stats::new();
        let mut hints: Vec<Inst> = Vec::new();
        for block in func.blocks().collect::<Vec<Block>>() {
            for inst in func.insts(block) {
                if func[inst].opcode == Opcode::Expect {
                    hints.push(inst);
                }
            }
        }
        // The whole of almost every function, since almost no program says anything about its
        // branches, and the walk above is the only cost this pass has on one that does not.
        if hints.is_empty() {
            return stats;
        }

        let mut placed = 0;
        for block in func.blocks().collect::<Vec<Block>>() {
            let Some(term) = func.terminator(block) else { continue };
            let Some(hints) = claims(func, term) else { continue };
            if !fuel.take() {
                stats.missed(NO_FUEL);
                break;
            }
            write(func, term, &hints);
            stats.optimized(PLACED);
            placed += 1;
        }

        // Every one of them, and not only the ones a branch was found for. The instruction has done
        // whatever it is going to do by this point, and what it would do from here is sit in front
        // of the folder.
        let mut forward: HashMap<Value, Value> = HashMap::new();
        for &inst in &hints {
            let args = &func[func[inst].args];
            let (Some(&result), Some(&value)) = (func[inst].first_result.as_ref(), args.first())
            else {
                continue;
            };
            forward.insert(result, value);
        }
        uses::substitute(func, &forward);
        for &inst in &hints {
            func.remove_inst(inst);
        }
        for _ in placed..hints.len() {
            stats.note(NO_BRANCH);
        }
        stats
    }
}

/// What the hint behind this terminator says about each of its arms, in the order the arms are in.
fn claims(func: &Func, term: Inst) -> Option<Vec<Hint>> {
    let &value = func[func[term].args].first()?;
    match func[term].opcode {
        Opcode::BrIf => {
            let (inst, sense) = through(func, value)?;
            let hint = Hint::parts(claim(func, inst, sense)?);
            Some(vec![hint, hint.complement()])
        }
        Opcode::Switch => switched(func, term, value),
        _ => None,
    }
}

/// The hints for the arms of a `switch` on an expected value, the default first as the arms are.
///
/// The arm the expected value goes to gets what the program claimed, which is the case of that
/// value or the default when no case has it, and the other arms share the rest evenly. That is what
/// gcc does with `__builtin_expect` on a switch, and it is all the hint says: which value is likely,
/// and nothing about the ones that are not. What evenly leaves over goes to the likely arm, so the
/// arms add up to certainty.
fn switched(func: &Func, term: Inst, value: Value) -> Option<Vec<Hint>> {
    let Extra::Switch(info) = func[term].extra else { return None };
    let info = func[info];
    let cases = &func[info.cases];
    let arms = u32::try_from(cases.len()).ok()?;
    if arms == 0 {
        return None;
    }
    let (inst, steps) = expected(func, value)?;
    let args = &func[func[inst].args];
    let ty = func[func[inst].first_result?].ty;
    let mut wanted = Imm::int(literal(func, *args.get(1)?)?, ty);
    for &(opcode, from, to) in steps.iter().rev() {
        wanted = convert(opcode, wanted, from, to);
    }
    let ty = func[value].ty;
    let hot =
        cases.iter().position(|case| case.signed(ty) == wanted.signed(ty)).map_or(0, |at| at + 1);
    let parts = probability(func, args.get(2))?;
    let rest = (Hint::SCALE - parts) / arms;
    let most = Hint::SCALE - rest * arms;
    let hints = (0..=cases.len()).map(|at| Hint::parts(if at == hot { most } else { rest }));
    Some(hints.collect())
}

/// A conversion between an `expect` and the `switch` on it: the opcode, and the types it goes from
/// and to.
type Step = (Opcode, Type, Type);

/// The `expect` a `switch` operand comes from, and the conversions between the two, outermost
/// first.
///
/// A `switch` on an `int` is on the `long` the prototype returned, narrowed back, so there is a
/// truncation in between, and the argument's widening is under the `expect` rather than over it.
/// A truncation is safe to walk through here, where it is not under a branch, because the question
/// is which case the operand lands on rather than whether it is zero, and the expected value goes
/// through the same conversions and lands on the case the operand would.
fn expected(func: &Func, value: Value) -> Option<(Inst, Vec<Step>)> {
    let mut value = value;
    let mut steps = Vec::new();
    loop {
        let Def::Result { inst, .. } = func[value].def else { return None };
        let data = &func[inst];
        match data.opcode {
            Opcode::Expect => return Some((inst, steps)),
            Opcode::Trunc | Opcode::ZExt | Opcode::SExt => {
                let from = *func[data.args].first()?;
                steps.push((data.opcode, func[from].ty, func[value].ty));
                value = from;
            }
            _ => return None,
        }
    }
}

/// The `expect` a branch condition comes from, and whether the condition is the expectation being
/// met or its opposite.
///
/// A condition is one bit and an `expect` is as wide as the `long` the prototype converted its
/// arguments to, so there is always something in between. What the lowering walk writes is a
/// comparison against zero, and what the peephole writes for `!x` is the same comparison the other
/// way round, so the chain is comparisons against zero and widenings that do not change whether a
/// value is zero.
fn through(func: &Func, cond: Value) -> Option<(Inst, bool)> {
    let mut value = cond;
    let mut sense = true;
    // Bounded, because a chain this walks is a chain of instructions the function has and each step
    // moves to the operand of the one before it.
    loop {
        let Def::Result { inst, .. } = func[value].def else { return None };
        let data = &func[inst];
        match data.opcode {
            Opcode::Expect => return Some((inst, sense)),
            // A widening keeps a value zero and keeps it non-zero, whichever bit pattern it makes.
            Opcode::ZExt | Opcode::SExt => value = *func[data.args].first()?,
            Opcode::ICmp => {
                let Extra::IntPred(pred) = data.extra else { return None };
                let args = &func[data.args];
                let lhs = *args.first()?;
                let rhs = *args.get(1)?;
                if literal(func, rhs)? != 0 {
                    return None;
                }
                match pred {
                    IntPred::Ne => {}
                    IntPred::Eq => sense = !sense,
                    _ => return None,
                }
                value = lhs;
            }
            _ => return None,
        }
    }
}

/// How often the first arm of the branch is taken, in parts of [`Hint::SCALE`], given what the
/// `expect` says and which way round the condition is.
///
/// The probability is the program's where it wrote one and ninety percent otherwise, which is
/// GCC's `param_builtin_expect_probability` and is section 11.2's number. It is the probability
/// that the value turns out to be the hint, so a hint of zero is the same claim about the other
/// arm and the complement is what the branch gets.
fn claim(func: &Func, inst: Inst, sense: bool) -> Option<u32> {
    let args = &func[func[inst].args];
    let value = literal(func, *args.get(1)?)?;
    let parts = probability(func, args.get(2))?;
    let met = (value != 0) == sense;
    Some(if met { parts } else { Hint::SCALE - parts })
}

/// How often the expected value is the one that turns up, in parts of [`Hint::SCALE`], given the
/// third operand where `__builtin_expect_with_probability` wrote one.
fn probability(func: &Func, given: Option<&Value>) -> Option<u32> {
    Some(match given {
        Some(&given) => u32::try_from(literal(func, given)?).ok()?.min(Hint::SCALE),
        None => PREDICT_EXPECT * Hint::SCALE / 100,
    })
}

/// The constant a value is, looking through the widenings a converted argument arrives behind.
///
/// `__builtin_expect(x, 0)` writes its second argument as an `int` and the prototype converts it to
/// `long`, so what the lowering walk leaves is a sign extension of a constant and not a constant.
/// This pass runs before anything that would fold one away, which is the point of it running first,
/// so the walk is here rather than left to the folder.
///
/// A widening only, because it is the one conversion whose answer is the value it was given. A
/// truncation is not: it can turn a value that is not zero into one that is, which would be the pass
/// reading a hint the program did not write.
fn literal(func: &Func, value: Value) -> Option<i128> {
    let mut value = value;
    loop {
        if let Some((bits, ty)) = constant(func, value) {
            return Some(bits.signed(ty));
        }
        let Def::Result { inst, .. } = func[value].def else { return None };
        let data = &func[inst];
        match data.opcode {
            Opcode::ZExt | Opcode::SExt => value = *func[data.args].first()?,
            _ => return None,
        }
    }
}

/// Writes the claim onto the arms of a branch or a `switch`, one hint per arm.
fn write(func: &mut Func, term: Inst, hints: &[Hint]) {
    for (at, &hint) in func.target_list(term).iter().zip(hints) {
        let call = func[at];
        func.set_block_call(at, rucc_ir::BlockCall { hint, ..call });
    }
}

#[cfg(test)]
mod tests {
    use rucc_base::Interner;
    use rucc_ir::{Builder, InstData, Signature, Type};

    use super::*;

    /// A function with `blocks` empty blocks in it and nothing else.
    fn blank(blocks: usize) -> (Interner, Func, Vec<Block>) {
        let mut names = Interner::new();
        let mut func = Func::new(names.intern("f"), Signature::new());
        let list = (0..blocks).map(|_| func.create_block()).collect();
        (names, func, list)
    }

    /// A function whose entry branches on `__builtin_expect(x, hint)`, with `x` a parameter.
    ///
    /// `parts` is the third operand where there is one, which is what
    /// `__builtin_expect_with_probability` wrote.
    fn shaped(hint: i128, parts: Option<i128>) -> (Interner, Func, Vec<Block>) {
        let (names, mut func, at) = blank(3);
        let i64_ = Type::int(64);
        let value = func.append_param(at[0], i64_);
        let mut build = Builder::new(&mut func, at[0]);
        let hint = build.iconst(i64_, hint);
        let mut operands = vec![value, hint];
        if let Some(parts) = parts {
            let parts = build.iconst(i64_, parts);
            operands.push(parts);
        }
        let args = build.func().push_values(&operands);
        let wrapped = build.value(InstData { args, ..InstData::new(Opcode::Expect) }, i64_);
        let zero = build.iconst(i64_, 0);
        let cond = build.icmp(IntPred::Ne, wrapped, zero);
        build.br_if(cond, at[1], &[], at[2], &[]);
        for block in [at[1], at[2]] {
            let mut build = Builder::new(&mut func, block);
            let answer = build.iconst(Type::int(32), 0);
            build.ret(&[answer]);
        }
        (names, func, at)
    }

    /// What the two arms of the entry's branch say about themselves.
    fn arms(func: &Func, block: Block) -> Vec<Option<u32>> {
        let term = func.terminator(block).expect("a branch");
        func.target_list(term).iter().map(|at| func[at].hint.taken()).collect()
    }

    /// Runs the pass over a function, and says what it did.
    fn run(func: &mut Func) -> Stats {
        Expect.run(func, &mut crate::machine::fixtures::analyses(), &mut Fuel::unlimited())
    }

    /// A function whose entry is `switch ((int)__builtin_expect(x, hint))` with cases 1, 3 and 5,
    /// with `x` an `int` parameter, the default the first block after the entry and one block for
    /// each case after that.
    fn switched_on(hint: i128) -> (Interner, Func, Vec<Block>) {
        let (names, mut func, at) = blank(5);
        let (int, long) = (Type::int(32), Type::int(64));
        let x = func.append_param(at[0], int);
        let mut build = Builder::new(&mut func, at[0]);
        let widened = build.unary(Opcode::SExt, x, long);
        let hint = build.iconst(long, hint);
        let args = build.func().push_values(&[widened, hint]);
        let wrapped = build.value(InstData { args, ..InstData::new(Opcode::Expect) }, long);
        let narrowed = build.unary(Opcode::Trunc, wrapped, int);
        build.switch(narrowed, at[1], &[(1, at[2]), (3, at[3]), (5, at[4])]);
        for &block in &at[1..] {
            let mut build = Builder::new(&mut func, block);
            let answer = build.iconst(int, 0);
            build.ret(&[answer]);
        }
        (names, func, at)
    }

    #[test]
    fn the_case_a_switch_expects_gets_the_hint_and_the_other_arms_share_the_rest() {
        let (_, mut func, at) = switched_on(3);
        let stats = run(&mut func);
        assert_eq!(stats.count(crate::stats::Kind::Optimized, PLACED), 1);
        assert_eq!(arms(&func, at[0]), [Some(333), Some(333), Some(9_001), Some(333)]);
        assert!(!func.blocks().any(|block| func.insts(block).any(|inst| is_expect(&func)(&inst))));
    }

    #[test]
    fn an_expected_value_no_case_has_is_the_default() {
        let (_, mut func, at) = switched_on(4);
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(9_001), Some(333), Some(333), Some(333)]);
    }

    /// The operand is the expected `long` narrowed to an `int`, so a value past the top of an `int`
    /// lands on the case its low bits name, the way the operand would.
    #[test]
    fn an_expected_value_is_narrowed_the_way_the_operand_is() {
        let (_, mut func, at) = switched_on((1 << 32) + 5);
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(333), Some(333), Some(333), Some(9_001)]);
    }

    #[test]
    fn a_hint_of_one_names_the_arm_taken_when_the_condition_holds() {
        let (_, mut func, at) = shaped(1, None);
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(9_000), Some(1_000)]);
    }

    #[test]
    fn a_hint_of_zero_names_the_other_arm() {
        let (_, mut func, at) = shaped(0, None);
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(1_000), Some(9_000)]);
    }

    /// The shape every program has, since `__builtin_expect(x, 0)` writes an `int` where the
    /// prototype asks for a `long` and the conversion is an instruction of its own at this point.
    #[test]
    fn a_hint_behind_the_conversion_the_prototype_asked_for_is_still_a_hint() {
        let (_, mut func, at) = blank(3);
        let i64_ = Type::int(64);
        let value = func.append_param(at[0], i64_);
        let mut build = Builder::new(&mut func, at[0]);
        let narrow = build.iconst(Type::int(32), 0);
        let hint = build.unary(Opcode::SExt, narrow, i64_);
        let args = build.func().push_values(&[value, hint]);
        let wrapped = build.value(InstData { args, ..InstData::new(Opcode::Expect) }, i64_);
        let zero = build.iconst(i64_, 0);
        let cond = build.icmp(IntPred::Ne, wrapped, zero);
        build.br_if(cond, at[1], &[], at[2], &[]);
        for block in [at[1], at[2]] {
            let mut build = Builder::new(&mut func, block);
            let answer = build.iconst(Type::int(32), 0);
            build.ret(&[answer]);
        }
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(1_000), Some(9_000)]);
    }

    #[test]
    fn a_probability_the_program_wrote_is_the_one_the_branch_gets() {
        let (_, mut func, at) = shaped(1, Some(7_500));
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(7_500), Some(2_500)]);
    }

    /// The one with three operands says how often the value is the hint, so a hint of zero and a
    /// probability of three quarters is three quarters for the arm taken when it is zero.
    #[test]
    fn a_probability_with_a_hint_of_zero_is_about_the_other_arm() {
        let (_, mut func, at) = shaped(0, Some(7_500));
        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(2_500), Some(7_500)]);
    }

    #[test]
    fn the_instruction_goes_and_its_readers_read_what_it_was_given() {
        let (_, mut func, at) = shaped(1, None);
        run(&mut func);
        let left: Vec<Inst> =
            func.blocks().flat_map(|block| func.insts(block)).filter(is_expect(&func)).collect();
        assert!(left.is_empty(), "the wrapper is gone");
        // The comparison now reads the parameter itself, which is what the wrapper answered with.
        let term = func.terminator(at[0]).expect("a branch");
        let cond = *func[func[term].args].first().expect("a condition");
        let Def::Result { inst, .. } = func[cond].def else { panic!("a comparison") };
        let read = *func[func[inst].args].first().expect("a left hand side");
        assert!(matches!(func[read].def, Def::Param { .. }), "it reads the parameter");
    }

    /// Whether an instruction is the wrapper, as a closure so that the filter above reads.
    fn is_expect(func: &Func) -> impl Fn(&Inst) -> bool + use<'_> {
        move |&inst| func[inst].opcode == Opcode::Expect
    }

    #[test]
    fn a_function_with_no_hint_in_it_is_left_alone() {
        let (_, mut func, at) = blank(3);
        let i64_ = Type::int(64);
        let value = func.append_param(at[0], i64_);
        let mut build = Builder::new(&mut func, at[0]);
        let zero = build.iconst(i64_, 0);
        let cond = build.icmp(IntPred::Ne, value, zero);
        build.br_if(cond, at[1], &[], at[2], &[]);
        for block in [at[1], at[2]] {
            let mut build = Builder::new(&mut func, block);
            let answer = build.iconst(Type::int(32), 0);
            build.ret(&[answer]);
        }

        let stats = run(&mut func);
        assert!(!stats.changed(), "nothing to do");
        assert_eq!(arms(&func, at[0]), [None, None]);
    }

    /// A condition written the other way round is the same claim about the other arm, which is
    /// what the sense in [`through`] is for.
    #[test]
    fn a_condition_that_is_a_comparison_against_zero_the_other_way_flips_the_arms() {
        let (_, mut func, at) = blank(3);
        let i64_ = Type::int(64);
        let value = func.append_param(at[0], i64_);
        let mut build = Builder::new(&mut func, at[0]);
        let hint = build.iconst(i64_, 1);
        let args = build.func().push_values(&[value, hint]);
        let wrapped = build.value(InstData { args, ..InstData::new(Opcode::Expect) }, i64_);
        let zero = build.iconst(i64_, 0);
        let cond = build.icmp(IntPred::Eq, wrapped, zero);
        build.br_if(cond, at[1], &[], at[2], &[]);
        for block in [at[1], at[2]] {
            let mut build = Builder::new(&mut func, block);
            let answer = build.iconst(Type::int(32), 0);
            build.ret(&[answer]);
        }

        run(&mut func);
        assert_eq!(arms(&func, at[0]), [Some(1_000), Some(9_000)]);
    }

    /// A hint about a value nothing branches on is dropped, and the instruction still goes.
    #[test]
    fn a_hint_on_a_value_no_branch_reads_leaves_nothing_behind() {
        let (_, mut func, at) = blank(1);
        let i64_ = Type::int(64);
        let value = func.append_param(at[0], i64_);
        let mut build = Builder::new(&mut func, at[0]);
        let hint = build.iconst(i64_, 1);
        let args = build.func().push_values(&[value, hint]);
        let wrapped = build.value(InstData { args, ..InstData::new(Opcode::Expect) }, i64_);
        build.ret(&[wrapped]);

        run(&mut func);
        let left: Vec<Inst> =
            func.blocks().flat_map(|block| func.insts(block)).filter(is_expect(&func)).collect();
        assert!(left.is_empty(), "the wrapper is gone");
        let term = func.terminator(at[0]).expect("a return");
        let answer = *func[func[term].args].first().expect("a returned value");
        assert!(matches!(func[answer].def, Def::Param { .. }), "it returns the parameter");
    }
}