branching? yay?

src/mir.rs

@@ -5,6 +5,7 @@ use std::collections::btree_map::Entry;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::fmt::Display;
 use std::hash::{Hash, Hasher};
+use std::iter::FusedIterator;
 use std::u32;
 
 use itertools::Itertools;
@@ -729,22 +730,22 @@ impl NodeRef {
         self.0 as usize
     }
 
-    fn into_reference_range(self) -> std::ops::Range<(NodeRef, NodeRef)> {
+    const fn into_reference_range(self) -> std::ops::Range<(NodeRef, NodeRef)> {
-        (self, Self::MIN)..(self, Self::MAX)
+        (self.inclusive_start())..(self.inclusive_end())
     }
     #[allow(dead_code)]
-    fn inclusive_start(self) -> (Self, Self) {
+    const fn inclusive_start(self) -> (Self, Self) {
         (self, Self::MIN)
     }
-    fn exclusive_start(self) -> (Self, Self) {
+    const fn exclusive_start(self) -> (Self, Self) {
         (self, Self::MAX)
     }
     #[allow(dead_code)]
-    fn inclusive_end(self) -> (Self, Self) {
+    const fn inclusive_end(self) -> (Self, Self) {
         (self, Self::MAX)
     }
     #[allow(dead_code)]
-    fn exclusive_end(self) -> (Self, Self) {
+    const fn exclusive_end(self) -> (Self, Self) {
         (self, Self::MIN)
     }
 
@@ -772,7 +773,7 @@ impl Mir {
         self.data[node.index()] = data;
     }
 
-    fn indices(&self) -> impl Iterator<Item = NodeRef> {
+    fn indices(&self) -> impl Iterator<Item = NodeRef> + DoubleEndedIterator + FusedIterator {
         (0..self.nodes.len() as u32).map(|n| NodeRef::from(n))
     }
 }
@@ -1369,7 +1370,394 @@ pub mod liveness {
     use super::*;
     pub struct Liveness {
         register_map: BTreeMap<NodeRef, Register>,
-        inference_graph: petgraph::graph::UnGraph<(), ()>,
+        branches: Branches,
+    }
+
+    pub struct Branches {
+        // branches are sequential, so a range, inclusive
+        branches: Vec<(MirRange, InterferenceGraph)>,
+        intervals: Vec<(NodeRef, NodeRef)>,
+        sorted_branches: Vec<MirRange>,
+        branch_graph: BranchGraph,
+    }
+
+    type BranchGraph = petgraph::graph::DiGraph<(), ()>;
+
+    impl Branches {
+        fn dirty_registers_at_node(
+            &self,
+            register_map: &BTreeMap<NodeRef, Register>,
+            node: NodeRef,
+        ) -> Vec<Register> {
+            let others = self
+                .intervals
+                .iter()
+                .filter(|(from, to)| node >= *from && node <= *to)
+                .filter(|interval| {
+                    Self::interval_intersects_with_node(
+                        &self.sorted_branches,
+                        &self.branch_graph,
+                        **interval,
+                        node,
+                    )
+                })
+                .filter_map(|(from, _)| register_map.get(from).cloned());
+
+            let dirty = self
+                .interference_graph_for_node(node)
+                .map(|graph| {
+                    graph
+                        .neighbors(node.0.into())
+                        .filter_map(|n| register_map.get(&NodeRef(n.index() as u32)).cloned())
+                })
+                .into_iter()
+                .flatten()
+                .chain(others)
+                .collect::<Vec<_>>();
+
+            dirty
+        }
+        fn interval_intersects_with_node(
+            sorted_branches: &Vec<MirRange>,
+            branch_graph: &BranchGraph,
+            interval: (NodeRef, NodeRef),
+            node: NodeRef,
+        ) -> bool {
+            let (from, to) = interval;
+            if node < from || node > to {
+                return false;
+            }
+
+            let from_index = sorted_branches
+                .binary_search_by(|range| range.partial_cmp(&from).unwrap())
+                .unwrap();
+            let to_index = sorted_branches
+                .binary_search_by(|range| range.partial_cmp(&to).unwrap())
+                .unwrap();
+
+            let ours_index = sorted_branches
+                .binary_search_by(|range| range.partial_cmp(&node).unwrap())
+                .unwrap();
+
+            let from_branch = sorted_branches[from_index].start();
+            let to_branch = sorted_branches[to_index].start();
+            let our_branch = sorted_branches[ours_index].start();
+
+            if from_branch == our_branch || to_branch == our_branch {
+                return true;
+            }
+
+            let mut space = petgraph::algo::DfsSpace::new(&branch_graph);
+            if petgraph::algo::has_path_connecting(
+                &branch_graph,
+                from_branch.0.into(),
+                our_branch.0.into(),
+                Some(&mut space),
+            ) && petgraph::algo::has_path_connecting(
+                &branch_graph,
+                our_branch.0.into(),
+                to_branch.0.into(),
+                Some(&mut space),
+            ) {
+                true
+            } else {
+                false
+            }
+        }
+
+        pub fn new(mir: &Mir, references: &BTreeSet<(NodeRef, NodeRef)>) -> Self {
+            let mut sorted_branches = Vec::new();
+            let mut branch_graph_edges = Vec::new();
+
+            let mut start = NodeRef::MIN;
+            let mut end = start;
+            for node in mir.indices() {
+                let (inst, data) = mir.get_node(node);
+
+                match inst {
+                    Inst::Label => {
+                        sorted_branches.push(MirRange::new(start, end));
+                        start = node;
+                    }
+                    Inst::Jump => {
+                        let to = data.as_node();
+                        branch_graph_edges.push((start.0, to));
+                    }
+                    Inst::Branch(_) => {
+                        let (a, b) = data.as_binary();
+                        branch_graph_edges.push((start.0, a));
+                        branch_graph_edges.push((start.0, b));
+                    }
+                    _ => {}
+                }
+                end = node;
+            }
+            sorted_branches.push(MirRange::new(start, end));
+            sorted_branches.sort();
+
+            let branch_graph = BranchGraph::from_edges(branch_graph_edges);
+
+            eprintln!("branch graph: {branch_graph:?}");
+            eprintln!("references: {references:?}");
+
+            let mut intervals = mir
+                .indices()
+                .filter_map(|node| {
+                    references
+                        .range(node.into_reference_range())
+                        .map(|&(_, to)| to)
+                        .reduce(|acc, to| acc.max(to))
+                        .map(|max| (node, max))
+                })
+                // .filter(|(from, _)| mir.is_register(from.0))
+                .collect::<Vec<_>>();
+            intervals.sort();
+
+            let edges = intervals
+                .iter()
+                .flat_map(|(from, to)| {
+                    references
+                        .range(from.exclusive_start()..to.exclusive_end())
+                        .filter(|(node, _)| {
+                            Self::interval_intersects_with_node(
+                                &sorted_branches,
+                                &branch_graph,
+                                (*from, *to),
+                                *node,
+                            )
+                        })
+                        .map(move |(other, _)| (*from, *other))
+                })
+                .collect::<Vec<_>>();
+
+            eprintln!("intervals: {intervals:#?}");
+            eprintln!("edges: {edges:#?}");
+
+            // build interference graph with mutliple branches.
+            // BTreeMap<MirRange, Interferencegraph>
+            type Edge = (NodeRef, NodeRef);
+            let mut per_branch_edges = BTreeMap::<MirRange, Vec<Edge>>::new();
+
+            for edge in edges.iter() {
+                let (from, to) = edge;
+                // find starting label of branch of from and to
+                let from_branch = sorted_branches
+                    .binary_search_by(|range| range.partial_cmp(from).unwrap())
+                    .unwrap();
+                let to_branch = sorted_branches
+                    .binary_search_by(|range| range.partial_cmp(to).unwrap())
+                    .unwrap();
+
+                match per_branch_edges.entry(sorted_branches[from_branch]) {
+                    Entry::Vacant(v) => {
+                        v.insert(vec![*edge]);
+                    }
+                    Entry::Occupied(mut o) => o.get_mut().push(*edge),
+                }
+                match per_branch_edges.entry(sorted_branches[to_branch]) {
+                    Entry::Vacant(v) => {
+                        v.insert(vec![*edge]);
+                    }
+                    Entry::Occupied(mut o) => o.get_mut().push(*edge),
+                }
+            }
+
+            eprintln!("per_branch_edges: {per_branch_edges:?}");
+
+            let branches = per_branch_edges
+                .into_iter()
+                .map(|(range, edges)| {
+                    (
+                        range,
+                        InterferenceGraph::from_edges(edges.into_iter().map(|(a, b)| (a.0, b.0))),
+                    )
+                })
+                .collect::<Vec<_>>();
+            eprintln!("branches: {branches:#?}");
+
+            Self {
+                branches,
+                intervals,
+                sorted_branches,
+                branch_graph,
+            }
+        }
+
+        fn does_interfere(&self, node: NodeRef, other: NodeRef) -> bool {
+            if let Some(index) = self.branch_index_for_node(node) {
+                if self.branches[index].0.contains(&other) {
+                    return self.branches[index]
+                        .1
+                        .find_edge(node.0.into(), other.0.into())
+                        .is_some();
+                }
+            }
+            false
+        }
+
+        fn branch_index_for_node(&self, node: NodeRef) -> Option<usize> {
+            self.branches
+                .binary_search_by(|(range, _)| range.partial_cmp(&node).unwrap())
+                .ok()
+        }
+
+        pub fn interference_graph_for_node(&self, node: NodeRef) -> Option<&InterferenceGraph> {
+            if let Some(index) = self.branch_index_for_node(node) {
+                Some(&self.branches[index].1)
+            } else {
+                None
+            }
+        }
+    }
+
+    #[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord)]
+    struct MirRange {
+        start: NodeRef,
+        inclusive_end: NodeRef,
+    }
+
+    impl PartialEq<NodeRef> for MirRange {
+        fn eq(&self, other: &NodeRef) -> bool {
+            self.contains(other)
+        }
+    }
+
+    impl PartialOrd<NodeRef> for MirRange {
+        fn partial_cmp(&self, other: &NodeRef) -> Option<Ordering> {
+            if self.contains(other) {
+                Some(Ordering::Equal)
+            } else if &self.inclusive_end < other {
+                Some(Ordering::Less)
+            } else {
+                assert!(&self.start > other);
+                Some(Ordering::Greater)
+            }
+        }
+    }
+
+    impl MirRange {
+        fn start(&self) -> NodeRef {
+            self.start
+        }
+
+        fn new(start: NodeRef, inclusive_end: NodeRef) -> MirRange {
+            Self {
+                start,
+                inclusive_end,
+            }
+        }
+        fn contains(&self, other: &NodeRef) -> bool {
+            &self.start <= other && &self.inclusive_end >= other
+        }
+    }
+
+    #[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord)]
+    struct LivenessReference {
+        node: NodeRef,
+        referenced_by: NodeRef,
+    }
+
+    pub struct BranchedLivenessBuilder<'a> {
+        mir: &'a Mir,
+        // tree of (node, referenced_by) pairs.
+        // a full range for each node is (node, NodeRef::MIN)..(node, NodeRef::MAX)
+        // QUESTION: do I want to treat every interval in this tree a separate node to color?
+        //  or just the (node, NodeRef::MIN)..(node, NodeRef::MAX) range?
+        // references: BTreeSet<(NodeRef, NodeRef)>,
+        branches: Branches,
+        // preferred_colors: BTreeMap<NodeRef, Register>,
+    }
+
+    type InterferenceGraph = petgraph::graph::UnGraph<(), ()>;
+
+    impl<'a> BranchedLivenessBuilder<'a> {
+        pub fn new(mir: &'a Mir) -> Self {
+            let references = build_reference_tree(mir);
+            let branches = Branches::new(mir, &references);
+            Self {
+                mir,
+                //references,
+                branches,
+            }
+        }
+
+        fn pre_color_colors(&self) -> BTreeMap<NodeRef, Register> {
+            /* do i want to find colors from the bottom up?
+             * - on one hand, consumers with in/outs like mul/div are below and want their inputs in some specific register
+             * - on the other hand, producers like mul/div are above and want to dictate their output registers
+             * i think the answer is coloring mul/div first, specifically, then collapsing from there.
+             */
+            let mut want_colors = BTreeMap::<NodeRef, Register>::new();
+
+            use Register::*;
+            let mut in_colors = [r9, r8, rcx, rdx, rdi, rsi].to_vec();
+            let mut in_colors_sse = [xmm7, xmm6, xmm5, xmm4, xmm3, xmm2, xmm1, xmm0].to_vec();
+
+            for node in self.mir.indices() {
+                let want_color = match self.mir.get_node(node).0 {
+                    Inst::Parameter(ty) => {
+                        if ty.is_floating() {
+                            in_colors_sse.pop()
+                        } else {
+                            in_colors.pop()
+                        }
+                    }
+                    Inst::ShiftLeft(_)
+                    | Inst::ShiftRightSigned(_)
+                    | Inst::ShiftRightUnsigned(_) => Some(rcx),
+                    Inst::Mul(_) | Inst::Div(_) | Inst::DivSigned(_) => Some(rax),
+                    Inst::Rem(_) | Inst::RemSigned(_) => Some(rdx),
+                    Inst::ReturnValue(ty) => {
+                        if ty.is_floating() {
+                            Some(xmm0)
+                        } else {
+                            Some(rax)
+                        }
+                    }
+                    _ => {
+                        if let Some(dst) = self.mir.dst_node(node) {
+                            // check if there is interference
+                            if !self.branches.does_interfere(node, dst) {
+                                // want the same color as our dst node to avoid copying
+                                want_colors.get(&dst).cloned()
+                            } else {
+                                None
+                            }
+                        } else {
+                            None
+                        }
+                    }
+                };
+
+                if let Some(color) = want_color {
+                    want_colors.insert(node, color);
+                }
+            }
+
+            want_colors
+        }
+
+        pub fn build(self) -> Liveness {
+            let preferred = self.pre_color_colors();
+
+            let mut colorizer = Colorizer {
+                mir: &self.mir,
+                preferred,
+                graph: &self.branches,
+                colors: BTreeMap::new(),
+            };
+
+            // // prepass: assign preferred colours for in/out values and specific
+            // // instructions like mul/div which require one operand to be in rax
+            colorizer.prepass();
+
+            let register_map = colorizer.colorise();
+
+            Liveness {
+                register_map,
+                branches: self.branches,
+            }
+        }
     }
 
     impl Liveness {
@@ -1381,23 +1769,19 @@ pub mod liveness {
         }
         pub fn get_scratch_register_at_node(&self, node: NodeRef) -> Option<Register> {
             let dirty = self
-                .inference_graph
+                .branches
-                .neighbors(node.0.into())
+                .dirty_registers_at_node(&self.register_map, node);
-                .filter_map(|n| self.register_map.get(&NodeRef(n.index() as u32)))
-                .cloned()
-                .collect::<Vec<_>>();
 
             Register::SYSV_SCRATCH_GPRS
                 .into_iter()
                 .filter(|c| !dirty.contains(c))
                 .next()
         }
+
         pub fn is_register_in_use_at_node(&self, node: NodeRef, reg: Register) -> bool {
-            self.inference_graph
+            self.branches
-                .neighbors(node.0.into())
+                .dirty_registers_at_node(&self.register_map, node)
-                .filter_map(|n| self.register_map.get(&NodeRef(n.index() as u32)))
+                .contains(&reg)
-                .find(|&&r| r.parent_reg() == reg.parent_reg())
-                .is_some()
         }
     }
 
@@ -1416,26 +1800,23 @@ pub mod liveness {
         // preferred_color: BTreeMap<NodeRef, Register>,
     }
 
-    #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
     enum Color {
-        #[default]
-        Unassigned,
         Tentative(Register),
         Final(Register),
     }
 
     impl Color {
-        fn color(self) -> Option<Register> {
+        fn color(self) -> Register {
             match self {
-                Color::Unassigned => None,
+                Color::Tentative(color) | Color::Final(color) => color,
-                Color::Tentative(color) | Color::Final(color) => Some(color),
             }
         }
     }
 
     impl<'a> LivenessBuilder<'a> {
         pub fn new(mir: &'a Mir) -> Self {
-            let references = Self::build_references(mir);
+            let references = build_reference_tree(mir);
 
             let intervals = mir.indices().filter_map(|node| {
                 let interval = references
@@ -1468,42 +1849,6 @@ pub mod liveness {
             }
         }
 
-        pub fn build(self) -> Liveness {
-            let preferred = self.pre_color_colors();
-            let Self {
-                mir,
-                inference_graph,
-                ..
-            } = self;
-
-            let mut colorizer = Colorizer {
-                mir,
-                preferred,
-                graph: inference_graph,
-                colors: BTreeMap::new(),
-            };
-
-            // // prepass: assign preferred colours for in/out values and specific
-            // // instructions like mul/div which require one operand to be in rax
-            colorizer.prepass();
-
-            // for &node in references.keys().rev() {
-            //     if !self.nodes[node as usize].result_is_register() {
-            //         continue;
-            //     }
-            //     colorizer.color_node(node.into());
-            // }
-
-            let register_map = colorizer.colorise();
-
-            let Colorizer { graph, .. } = colorizer;
-
-            Liveness {
-                register_map,
-                inference_graph: graph,
-            }
-        }
-
         fn pre_color_colors(&self) -> BTreeMap<NodeRef, Register> {
             /* do i want to find colors from the bottom up?
              * - on one hand, consumers with in/outs like mul/div are below and want their inputs in some specific register
@@ -1564,22 +1909,173 @@ pub mod liveness {
 
             want_colors
         }
+    }
 
-        fn build_references(mir: &Mir) -> BTreeSet<(NodeRef, NodeRef)> {
+    fn build_reference_tree(mir: &Mir) -> BTreeSet<(NodeRef, NodeRef)> {
         let mut references = BTreeSet::new();
         for node in mir.indices() {
-                Self::reference_node_operands(mir, &mut references, node);
+            mir.reference_node_operands(node, &mut references);
         }
 
         references
     }
 
+    struct Colorizer<'a> {
+        mir: &'a Mir,
+        graph: &'a Branches,
+        colors: BTreeMap<NodeRef, Color>,
+        preferred: BTreeMap<NodeRef, Register>,
+    }
+
+    impl<'a> Colorizer<'a> {
+        fn node_colors(&self, node: petgraph::graph::NodeIndex) -> &[Register] {
+            let colors = if self.mir.nodes[node.index()]
+                .value_type()
+                .map(|t| t.is_floating())
+                == Some(true)
+            {
+                &Register::SSE[..]
+            } else {
+                &Register::GPR[..]
+            };
+            colors
+        }
+
+        fn prepass(&mut self) {
+            // parameters are first in line
+            let keys = self.preferred.keys().cloned().collect::<Vec<_>>();
+            for node in keys.into_iter() {
+                self.precolor_node(node.0.into());
+            }
+        }
+
+        fn precolor_node(&mut self, node: petgraph::graph::NodeIndex) {
+            // prepass: assign preferred colours for in/out values and specific
+            // instructions like mul/div which require one operand to be in rax
+            let node_u32 = node.index() as u32;
+            let node_ref = NodeRef(node_u32);
+            // only apply color here if we have a preference
+            if let Some(preferred_color) = self.preferred.remove(&node_ref) {
+                let color_conflict = self
+                    .graph
+                    .interference_graph_for_node(node_ref)
+                    .map(|graph| {
+                        graph
+                            .neighbors(node)
+                            .filter_map(|n| self.colors.get(&NodeRef(n.index() as u32)).cloned())
+                            .find(|color| color.color() == preferred_color)
+                    })
+                    .flatten();
+
+                if color_conflict.is_none() {
+                    match self.colors.entry(node_ref) {
+                        Entry::Vacant(v) => {
+                            v.insert(Color::Tentative(preferred_color));
+                        }
+                        _ => {} // Entry::Occupied(mut o) => match o.get_mut() {
+                                //     Color::Tentative(color) => {
+                                //         *color = preferred_color;
+                                //     }
+                                //     Color::Final(_) => {}
+                                // },
+                    }
+                }
+            };
+        }
+
+        fn color_node(&mut self, node: petgraph::graph::NodeIndex) {
+            // final pass:
+            // look at clique colors and prefer to steal colors from
+            // tentatively colored nodes. this results in preferential
+            // coloring depending on the order of the prepass.
+            let node_u32 = node.index() as u32;
+            let node_ref = NodeRef(node_u32);
+            let clique_colors = self
+                .graph
+                .interference_graph_for_node(node_ref)
+                .map(|graph| {
+                    graph
+                        .neighbors(node)
+                        .filter_map(|n| self.colors.get(&NodeRef(n.index() as u32)).cloned())
+                        .collect::<BTreeSet<_>>()
+                })
+                .unwrap_or_default();
+
+            let colors = self
+                .node_colors(node)
+                .into_iter()
+                .filter(|&&r| !clique_colors.contains(&Color::Final(r)))
+                .cloned()
+                .collect::<BTreeSet<_>>();
+
+            eprintln!("coloring %{node_u32}:");
+            eprintln!("\twants: {:?}", self.colors.get(&node_ref));
+            eprintln!("\tclique: {clique_colors:?}");
+            eprintln!("\tcandidates: {colors:?}");
+
+            let color = match self.colors.entry(node_ref) {
+                Entry::Vacant(v) => {
+                    // here we want to first check clique_colors with tentative coloring.
+                    let color = colors
+                        .into_iter()
+                        .find_or_first(|&c| !clique_colors.contains(&Color::Tentative(c)))
+                        .expect("ran out of registers :(");
+                    v.insert(Color::Final(color));
+                    color
+                }
+                Entry::Occupied(mut e) => {
+                    // we prefer to steal
+                    variant!(e.get() => &Color::Tentative(reg));
+                    let color = colors
+                        .into_iter()
+                        .find_or_first(|&c| c == reg)
+                        .expect("ran out of registers :(");
+                    e.insert(Color::Final(color));
+                    color
+                }
+            };
+
+            // if this node has a dst node (an operand that is both source and
+            // destination), give it our tentative color
+            if let Some(dst) = self.mir.dst_node(node_ref) {
+                _ = self.colors.try_insert(dst, Color::Tentative(color));
+            }
+
+            // for any Phi(y_1,y_2, y_n) give y_i the color of Phi
+            // reasonably certain that this will never fail to color all phi nodes the same color.
+            if let Some(inputs) = self.mir.get_phi_inputs(node_ref) {
+                eprintln!("coloring {inputs:?} {color}");
+                for node in inputs {
+                    _ = self.colors.insert(node, Color::Tentative(color));
+                }
+            }
+        }
+
+        fn colorise(&mut self) -> BTreeMap<NodeRef, Register> {
+            for node in self.mir.indices().rev() {
+                if !self.mir.is_register(node.0) {
+                    continue;
+                }
+                self.color_node(node.0.into());
+            }
+            self.colors
+                .iter()
+                .filter_map(|(&n, &c)| match c {
+                    Color::Final(reg) => Some((n, reg)),
+                    _ => None,
+                })
+                .collect()
+        }
+    }
+}
+
+impl Mir {
     fn reference_node_operands(
-            mir: &Mir,
+        &self,
-            references: &mut BTreeSet<(NodeRef, NodeRef)>,
         node: NodeRef,
+        references: &mut BTreeSet<(NodeRef, NodeRef)>,
     ) {
-            let (inst, data) = mir.get_node(node);
+        let (inst, data) = self.get_node(node);
 
         match inst {
             Inst::ReturnValue(_)
@@ -1647,142 +2143,7 @@ pub mod liveness {
             | Inst::Alloca => {}
         }
     }
-    }
 
-    struct Colorizer<'a> {
-        mir: &'a Mir,
-        graph: petgraph::graph::UnGraph<(), ()>,
-        colors: BTreeMap<NodeRef, Color>,
-        preferred: BTreeMap<NodeRef, Register>,
-    }
-
-    impl<'a> Colorizer<'a> {
-        fn node_colors(&self, node: petgraph::graph::NodeIndex) -> &[Register] {
-            let colors = if self.mir.nodes[node.index()]
-                .value_type()
-                .map(|t| t.is_floating())
-                == Some(true)
-            {
-                &Register::SSE[..]
-            } else {
-                &Register::GPR[..]
-            };
-            colors
-        }
-
-        fn prepass(&mut self) {
-            // parameters are first in line
-            let keys = self.preferred.keys().cloned().collect::<Vec<_>>();
-            for node in keys {
-                self.precolor_node(node.0.into());
-            }
-        }
-
-        fn precolor_node(&mut self, node: petgraph::graph::NodeIndex) {
-            // prepass: assign preferred colours for in/out values and specific
-            // instructions like mul/div which require one operand to be in rax
-            let node_u32 = node.index() as u32;
-            let noderef = NodeRef(node_u32);
-            // only apply color here if we have a preference
-            if let Some(preferred_color) = self.preferred.remove(&noderef) {
-                let mut clique_colors = self
-                    .graph
-                    .neighbors(node)
-                    .filter_map(|n| self.colors.get(&NodeRef(n.index() as u32)).cloned());
-
-                if clique_colors
-                    .find(|color| color.color() == Some(preferred_color))
-                    .is_none()
-                {
-                    self.colors
-                        .insert(noderef, Color::Tentative(preferred_color));
-                }
-            };
-            // .chain(self.node_colors(node).into_iter().cloned());
-        }
-
-        fn color_node(&mut self, node: petgraph::graph::NodeIndex) {
-            // final pass:
-            // look at clique colors and prefer to steal colors from
-            // tentatively colored nodes. this results in preferential
-            // coloring depending on the order of the prepass.
-            let node_u32 = node.index() as u32;
-            let noderef = NodeRef(node_u32);
-            let clique_colors = self
-                .graph
-                .neighbors(node)
-                .filter_map(|n| self.colors.get(&NodeRef(n.index() as u32)).cloned())
-                .collect::<BTreeSet<_>>();
-
-            let colors = self
-                .node_colors(node)
-                .into_iter()
-                .filter(|&&r| !clique_colors.contains(&Color::Final(r)))
-                .cloned()
-                .collect::<BTreeSet<_>>();
-
-            // eprintln!("coloring %{node_u32}:");
-            // eprintln!("\twants: {:?}", self.colors.get(&node_u32));
-            // eprintln!("\tclique: {clique_colors:?}");
-            // eprintln!("\tcandidates: {colors:?}");
-
-            let color = match self.colors.entry(noderef) {
-                Entry::Vacant(v) => {
-                    // here we want to first check clique_colors with tentative coloring.
-                    let color = colors
-                        .into_iter()
-                        .find_or_first(|&c| !clique_colors.contains(&Color::Tentative(c)))
-                        .expect("ran out of registers :(");
-                    v.insert(Color::Final(color));
-                    color
-                }
-                Entry::Occupied(mut e) => {
-                    // we prefer to steal
-                    variant!(e.get() => &Color::Tentative(reg));
-                    let color = colors
-                        .into_iter()
-                        .find_or_first(|&c| c == reg)
-                        .expect("ran out of registers :(");
-                    e.insert(Color::Final(color));
-                    color
-                }
-            };
-
-            // if this node has a dst node (an operand that is both source and
-            // destination), give it our tentative color
-            if let Some(dst) = self.mir.dst_node(noderef) {
-                _ = self.colors.try_insert(dst, Color::Tentative(color));
-            }
-
-            // for any Phi(y_1,y_2, y_n) give y_i the color of Phi
-            // reasonably certain that this will never fail to color all phi nodes the same color.
-            if let Some(inputs) = self.mir.get_phi_inputs(noderef) {
-                eprintln!("coloring {inputs:?} {color}");
-                for node in inputs {
-                    _ = self.colors.insert(node, Color::Tentative(color));
-                }
-            }
-        }
-
-        fn colorise(&mut self) -> BTreeMap<NodeRef, Register> {
-            for node in self.mir.indices() {
-                if !self.mir.is_register(node.0) {
-                    continue;
-                }
-                self.color_node(node.0.into());
-            }
-            self.colors
-                .iter()
-                .filter_map(|(&node, &c)| match c {
-                    Color::Final(reg) => Some((node, reg)),
-                    _ => None,
-                })
-                .collect()
-        }
-    }
-}
-
-impl Mir {
     fn get_phi_inputs(&self, node: NodeRef) -> Option<Vec<NodeRef>> {
         let (inst, data) = self.get_node(node);
         match inst {
@@ -1858,7 +2219,7 @@ impl Mir {
     }
 
     pub fn build_liveness(&self) -> Liveness {
-        LivenessBuilder::new(self).build()
+        liveness::BranchedLivenessBuilder::new(self).build()
     }
 }
 
@@ -2082,13 +2443,11 @@ impl Function {
         write!(w, "{}", self.branches.remove(&NodeRef::MIN).unwrap())?;
 
         for (branch, content) in &self.branches {
-            if name != "main" {
+            writeln!(w, ".{name}__L{}:", branch.0)?;
-                writeln!(w, "{name}_L{}:", branch.0)?;
             write!(w, "{content}")?;
         }
-        }
 
-        writeln!(w, "{name}__epilogue:")?;
+        writeln!(w, ".{name}__epilogue:")?;
         writeln!(w, "mov rsp, rbp")?;
         writeln!(w, "pop rbp")?;
 
@@ -2207,8 +2566,6 @@ impl Mir {
             let inst = self.nodes[i];
             let data = self.data[i];
 
-            self.render_node(func.current_branch(), strings, &liveness, node);
-
             match inst {
                 Inst::Label => {
                     func.create_new_branch(NodeRef(node));
@@ -3160,15 +3517,15 @@ impl Mir {
                             )?;
                         }
                     }
-                    writeln!(func.current_branch(), "jmp {name}__epilogue")?;
+                    writeln!(func.current_branch(), "jmp .{name}__epilogue")?;
                 }
                 Inst::Return => {
-                    writeln!(func.current_branch(), "jmp {name}__epilogue")?;
+                    writeln!(func.current_branch(), "jmp .{name}__epilogue")?;
                 }
                 Inst::Jump => {
                     let lhs = data.as_node();
                     if lhs != node + 1 {
-                        writeln!(func.current_branch(), "jmp {name}__L{lhs}")?;
+                        writeln!(func.current_branch(), "jmp .{name}__L{lhs}")?;
                     }
                 }
                 Inst::Branch(condition) => {
@@ -3179,14 +3536,14 @@ impl Mir {
 
                     match (lhs, rhs) {
                         _ if lhs == node + 1 => {
-                            writeln!(func.current_branch(), "jz {name}__L{rhs}")?;
+                            writeln!(func.current_branch(), "jz .{name}__L{rhs}")?;
                         }
                         _ if rhs == node + 1 => {
-                            writeln!(func.current_branch(), "jnz {name}__L{lhs}")?;
+                            writeln!(func.current_branch(), "jnz .{name}__L{lhs}")?;
                         }
                         _ => {
-                            writeln!(func.current_branch(), "jnz {name}__L{lhs}")?;
+                            writeln!(func.current_branch(), "jnz .{name}__L{lhs}")?;
-                            writeln!(func.current_branch(), "jz {name}__L{rhs}")?;
+                            writeln!(func.current_branch(), "jz .{name}__L{rhs}")?;
                         }
                     }
                 }

tests/legal/if.sea

@@ -1,6 +1,6 @@
 
 fn square_of_greater(a: i32, b: i32) -> i32 {
-   if (a < b)
+   if (a > b)
       a * a
    else
       b * b