rucc_codegen/

select.rs

//! Matching a target's lowering rules against a term.
//!
//! Design: `spec/10-backend.md` section 10.2. The rules themselves are in `rules/`, one file per
//! target, and the automaton they compile into is generated by `rucc-rules` when this crate is
//! built. What is here is the walk over that automaton, which is the same walk for every target
//! and is written once.
//!
//! # What a subject is
//!
//! A rule matches a term, and the compiler does not have terms: it has a function full of
//! instructions, and what a pattern is about is one of them and whatever it was computed from.
//! So the walk is written against [`Subject`], which is the three questions the automaton asks
//! of whatever it is matching, and the selector answers them out of the IR without building a
//! term to be thrown away. A test can answer them out of anything at all, which is what the
//! tests at the bottom of this file do.
//!
//! # What a match gives back
//!
//! The rule that fired and what its pattern bound, in the order the pattern binds it. The
//! bindings are positions rather than names because that is what the walk has, and the rule
//! carries the names for anything that has to say what it did. Building the replacement out of
//! [`Piece`] is the selector's job rather than this file's, because what a machine term becomes
//! is a machine instruction, and this module is about matching.
//!
//! # Order
//!
//! At every node the concrete tests are tried before the branch that takes anything, so a rule
//! naming an operand is tried before a rule taking whatever is there. That is the maximal munch
//! `spec/10-backend.md` asks for, and it falls out of the shape of the trie rather than being
//! sorted for. Among rules that are equally specific the first one written wins.
//!
//! A guard is part of deciding whether a rule fires, so a rule whose guard is false is a rule
//! that did not match, and the walk carries on looking rather than giving up. What that costs is
//! the search from where the guard failed, which is the price of a guard being allowed to be
//! about the values rather than only about the shape.

pub mod x86_64;

/// The bits of a term the automaton asks about.
///
/// A node is whatever the thing doing the matching calls one of its terms: an IR value, an index
/// into an arena, a pointer. It has to be cheap to copy because the walk keeps a stack of them.
pub trait Subject {
    /// What this subject calls one of its terms.
    type Node: Copy;

    /// The head of a term and how many arguments it has, or nothing if the term is not an
    /// application. An IR instruction answers with its opcode and its width, spelled the way the
    /// rule file spells it.
    fn head(&self, node: Self::Node) -> Option<(&str, usize)>;

    /// One argument of a term, counted from zero. Only ever asked for an argument the answer to
    /// [`Subject::head`] said was there.
    fn arg(&self, node: Self::Node, index: usize) -> Self::Node;

    /// The value of a term that is a constant, or nothing if it is not one. This is what a
    /// pattern matching a literal is asking, and what a guard reads.
    fn int(&self, node: Self::Node) -> Option<i128>;
}

/// One test on one subterm.
#[derive(Debug)]
pub enum Test {
    /// The subterm has to be this head applied to this many arguments.
    App {
        /// The name in head position.
        head: &'static str,
        /// How many arguments it takes.
        arity: usize,
    },
    /// The subterm has to be this constant.
    Int(i128),
}

/// One node of the trie over the patterns.
#[derive(Debug)]
pub struct Node {
    /// The tests to try, in the order the rules were written, before the wildcard.
    pub tests: &'static [(Test, u32)],
    /// The branch that takes anything, and the name the first rule to reach it gave that hole.
    pub wildcard: Option<(&'static str, u32)>,
    /// The rule that ends here, if one does.
    pub accept: Option<u32>,
}

/// One piece of a replacement, in the pre-order that builds it.
#[derive(Debug)]
pub enum Piece {
    /// Whatever the pattern bound at this position.
    Var {
        /// The name the rule gave it, for anything that has to say what it did.
        name: &'static str,
        /// Which binding of the match it is.
        index: usize,
    },
    /// A constant written in the rule.
    Int(i128),
    /// A machine term, which is an instruction once the selector has built it.
    App {
        /// The name in head position.
        head: &'static str,
        /// How many arguments it takes.
        arity: usize,
    },
}

/// A condition on the constants a pattern matched.
///
/// It is handed one entry per binding, holding the value of that binding when it has one. A
/// guard about a binding that is not a constant is false, which is how a rule about a number
/// declines an operand that is a register.
pub type Guard = fn(&[Option<i128>]) -> bool;

/// One lowering rule, as much of it as matching needs.
#[derive(Debug)]
pub struct Rule {
    /// The pattern as it is written in the rule file, for diagnostics and for tests.
    pub pattern: &'static str,
    /// What to put in the matched term's place, flattened into pre-order.
    pub replacement: &'static [Piece],
    /// The condition on the match, if the rule has one.
    pub guard: Option<Guard>,
    /// The line of the rule file this rule starts on.
    pub line: u32,
}

impl Rule {
    /// The head of the replacement, which is the instruction this rule selects.
    #[must_use]
    pub fn head(&self) -> Option<&'static str> {
        match self.replacement.first() {
            Some(Piece::App { head, .. }) => Some(head),
            _ => None,
        }
    }
}

/// A target's lowering rules, as an automaton over their patterns.
#[derive(Debug)]
pub struct Table {
    /// The rule file this was built from, so that anything said about a rule can name a file
    /// somebody can open.
    pub source: &'static str,
    /// The trie. Node zero is the root.
    pub nodes: &'static [Node],
    /// The rules, in the order the file writes them.
    pub rules: &'static [Rule],
}

/// What a successful match found.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Match<N> {
    /// Which rule of the table fired.
    pub rule: usize,
    /// What the pattern bound, in the order it binds it.
    pub bindings: Vec<N>,
}

impl Table {
    /// The rule that fires on this term, and what it bound.
    ///
    /// The term is matched as a whole. Finding the terms in a function worth matching is the
    /// selector's job and not this one's.
    #[must_use]
    pub fn find<S: Subject>(&self, subject: &S, term: S::Node) -> Option<Match<S::Node>> {
        let mut bindings = Vec::new();
        let rule = self.run(subject, 0, vec![term], &mut bindings)?;
        Some(Match { rule, bindings })
    }

    /// The rule a match found, which is the one thing every caller wants out of it.
    #[must_use]
    pub fn rule<N>(&self, found: &Match<N>) -> &Rule {
        &self.rules[found.rule]
    }

    /// Walk the trie and the subject together.
    ///
    /// `left` is the subterms still to be matched, innermost last, so that popping gives the
    /// pre-order the patterns were flattened in.
    fn run<S: Subject>(
        &self,
        subject: &S,
        at: usize,
        mut left: Vec<S::Node>,
        bindings: &mut Vec<S::Node>,
    ) -> Option<usize> {
        let Some(term) = left.pop() else {
            return self.accept(subject, at, bindings);
        };
        let node = &self.nodes[at];
        let head = subject.head(term);

        for (test, next) in node.tests {
            let matched = match test {
                Test::Int(want) => subject.int(term) == Some(*want),
                Test::App { head: want, arity } => {
                    head.is_some_and(|(have, count)| have == *want && count == *arity)
                }
            };
            if !matched {
                continue;
            }
            let mut deeper = left.clone();
            if let Some((_, arity)) = head {
                for index in (0..arity).rev() {
                    deeper.push(subject.arg(term, index));
                }
            }
            let depth = bindings.len();
            if let Some(rule) = self.run(subject, *next as usize, deeper, bindings) {
                return Some(rule);
            }
            bindings.truncate(depth);
        }

        // The wildcard is last, which is the whole of what specificity order means here.
        let (_, next) = node.wildcard.as_ref()?;
        let depth = bindings.len();
        bindings.push(term);
        if let Some(rule) = self.run(subject, *next as usize, left, bindings) {
            return Some(rule);
        }
        bindings.truncate(depth);
        None
    }

    /// The rule that ends at this node, if one does and if its guard holds.
    fn accept<S: Subject>(&self, subject: &S, at: usize, bindings: &[S::Node]) -> Option<usize> {
        let rule = self.nodes[at].accept? as usize;
        if let Some(guard) = self.rules[rule].guard {
            // The values are collected here rather than as the bindings are made, because most
            // rules have no guard and would pay for it every time.
            let values: Vec<Option<i128>> =
                bindings.iter().map(|&node| subject.int(node)).collect();
            if !guard(&values) {
                return None;
            }
        }
        Some(rule)
    }
}

#[cfg(test)]
mod tests {
    use super::x86_64::TABLE;
    use super::{Piece, Subject};

    /// A term, in the only shape a test needs: a flat arena, because that is the shape the IR
    /// has and answering the questions out of one is what the selector will be doing.
    #[derive(Debug)]
    enum Node {
        Int(i128),
        App(String, Vec<usize>),
    }

    #[derive(Debug, Default)]
    struct Terms {
        nodes: Vec<Node>,
    }

    impl Terms {
        fn constant(&mut self, value: i128) -> usize {
            self.nodes.push(Node::Int(value));
            self.nodes.len() - 1
        }

        fn app(&mut self, head: &str, args: &[usize]) -> usize {
            self.nodes.push(Node::App(head.to_owned(), args.to_vec()));
            self.nodes.len() - 1
        }

        /// A register operand, which is a term with a head the rules write and nothing under it.
        fn value(&mut self, width: u32, name: &str) -> usize {
            let inner = self.app(name, &[]);
            self.app(&format!("value.i{width}"), &[inner])
        }
    }

    impl Subject for Terms {
        type Node = usize;

        fn head(&self, node: usize) -> Option<(&str, usize)> {
            match &self.nodes[node] {
                Node::App(head, args) => Some((head.as_str(), args.len())),
                Node::Int(_) => None,
            }
        }

        fn arg(&self, node: usize, index: usize) -> usize {
            match &self.nodes[node] {
                Node::App(_, args) => args[index],
                Node::Int(_) => unreachable!("a constant has no arguments"),
            }
        }

        fn int(&self, node: usize) -> Option<i128> {
            match self.nodes[node] {
                Node::Int(value) => Some(value),
                Node::App(..) => None,
            }
        }
    }

    /// What the head of the rule that fired selects, which is the answer every one of these
    /// tests is really about.
    fn selects(terms: &Terms, term: usize) -> Option<&'static str> {
        let found = TABLE.find(terms, term)?;
        TABLE.rule(&found).head()
    }

    #[test]
    fn the_table_holds_every_rule_the_file_writes() {
        let text = include_str!("../rules/x86-64.rules");
        let written = text.lines().filter(|line| line.starts_with("(rule ")).count();
        assert_eq!(TABLE.rules.len(), written, "the table and the rule file disagree");
        assert_eq!(TABLE.source, "rules/x86-64.rules");
    }

    #[test]
    fn an_addition_of_two_registers_is_the_register_form() {
        let mut terms = Terms::default();
        let x = terms.value(64, "v0");
        let y = terms.value(64, "v1");
        let add = terms.app("add.i64", &[x, y]);
        assert_eq!(selects(&terms, add), Some("x64.add_rr_64"));
    }

    /// The bindings are the operands in the order the pattern names them, and the replacement
    /// says which of them goes where. This is the whole of what the selector will read.
    ///
    /// What a name is bound to is what the pattern put it under, so `(value.i32 x)` binds the
    /// register and not the term saying it is one. That is the difference between the operand of
    /// the instruction this becomes and a wrapper that exists to say how wide it is.
    #[test]
    fn a_match_gives_back_the_operands_the_pattern_named() {
        let mut terms = Terms::default();
        let first = terms.app("v0", &[]);
        let second = terms.app("v1", &[]);
        let x = terms.app("value.i32", &[first]);
        let y = terms.app("value.i32", &[second]);
        let sub = terms.app("sub.i32", &[x, y]);
        let found = TABLE.find(&terms, sub).expect("a rule fires");
        let rule = TABLE.rule(&found);
        assert_eq!(rule.pattern, "(sub.i32 (value.i32 x) (value.i32 y))");
        assert_eq!(found.bindings, vec![first, second]);
        let names: Vec<&str> = rule
            .replacement
            .iter()
            .filter_map(|piece| match piece {
                Piece::Var { name, index } => {
                    assert_eq!(found.bindings[*index], if *index == 0 { first } else { second });
                    Some(*name)
                }
                _ => None,
            })
            .collect();
        assert_eq!(names, ["x", "y"]);
    }

    /// An immediate the instruction has room for takes the immediate form. The rule for it is
    /// guarded, so this is also the test that a guard which holds does not stop a rule firing.
    #[test]
    fn an_addition_of_an_immediate_that_fits_is_the_immediate_form() {
        let mut terms = Terms::default();
        let x = terms.value(64, "v0");
        let k = terms.constant(4);
        let k = terms.app("iconst.i64", &[k]);
        let add = terms.app("add.i64", &[x, k]);
        assert_eq!(selects(&terms, add), Some("x64.add_ri_64"));
    }

    /// An immediate too wide for the encoding is what the guard is there to refuse. Nothing else
    /// matches such a term, and that is the right answer: the constant has to be put in a
    /// register first, which is a decision for the selector and not for the table.
    #[test]
    fn an_addition_of_an_immediate_too_wide_for_the_form_matches_nothing() {
        let mut terms = Terms::default();
        let x = terms.value(64, "v0");
        let k = terms.constant(1 << 40);
        let k = terms.app("iconst.i64", &[k]);
        let add = terms.app("add.i64", &[x, k]);
        assert_eq!(selects(&terms, add), None);
    }

    /// The other shape of guard, which is a shift count the width allows.
    #[test]
    fn a_shift_by_a_count_the_width_allows_is_the_immediate_form() {
        let mut terms = Terms::default();
        let x = terms.value(64, "v0");
        let k = terms.constant(3);
        let k = terms.app("iconst.i64", &[k]);
        let shl = terms.app("shl.i64", &[x, k]);
        assert_eq!(selects(&terms, shl), Some("x64.shl_ri_64"));
    }

    #[test]
    fn a_shift_by_a_count_the_width_does_not_allow_matches_nothing() {
        let mut terms = Terms::default();
        let x = terms.value(64, "v0");
        let k = terms.constant(64);
        let k = terms.app("iconst.i64", &[k]);
        let shl = terms.app("shl.i64", &[x, k]);
        assert_eq!(selects(&terms, shl), None);
    }

    /// A term the rule set says nothing about is nothing rather than a wrong answer, which is
    /// what the completeness check in `spec/10-backend.md` will be for.
    #[test]
    fn a_term_no_rule_covers_finds_no_rule() {
        let mut terms = Terms::default();
        let x = terms.value(64, "v0");
        let y = terms.value(64, "v1");
        let odd = terms.app("no.such.opcode", &[x, y]);
        assert_eq!(selects(&terms, odd), None);
    }
}