nak: Fix BasicBlock::phi*() for OpAnnotate

They assumed that phis would be at the start of the block which isn't
true if there are annotations in the way.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/29591>

src/nouveau/compiler/nak/ir.rs

@@ -5857,47 +5857,55 @@ impl BasicBlock {
         self.instrs = instrs;
     }
 
-    pub fn phi_dsts(&self) -> Option<&OpPhiDsts> {
-        if let Op::PhiDsts(phi) = &self.instrs.first()?.op {
-            return Some(phi);
+    pub fn phi_dsts_ip(&self) -> Option<usize> {
+        for (ip, instr) in self.instrs.iter().enumerate() {
+            match &instr.op {
+                Op::Annotate(_) => (),
+                Op::PhiDsts(_) => return Some(ip),
+                _ => break,
+            }
         }
         None
     }
 
+    pub fn phi_dsts(&self) -> Option<&OpPhiDsts> {
+        self.phi_dsts_ip().map(|ip| match &self.instrs[ip].op {
+            Op::PhiDsts(phi) => phi,
+            _ => panic!("Expected to find the phi"),
+        })
+    }
+
     #[allow(dead_code)]
     pub fn phi_dsts_mut(&mut self) -> Option<&mut OpPhiDsts> {
-        if let Op::PhiDsts(phi) = &mut self.instrs.first_mut()?.op {
-            return Some(phi);
-        }
-        None
+        self.phi_dsts_ip().map(|ip| match &mut self.instrs[ip].op {
+            Op::PhiDsts(phi) => phi,
+            _ => panic!("Expected to find the phi"),
+        })
     }
 
-    pub fn phi_srcs(&self) -> Option<&OpPhiSrcs> {
-        for instr in self.instrs.iter().rev() {
-            if instr.is_branch() {
-                continue;
-            }
-
+    pub fn phi_srcs_ip(&self) -> Option<usize> {
+        for (ip, instr) in self.instrs.iter().enumerate().rev() {
             match &instr.op {
-                Op::PhiSrcs(phi) => return Some(phi),
+                Op::Annotate(_) => (),
+                Op::PhiSrcs(_) => return Some(ip),
+                _ if instr.is_branch() => (),
                 _ => break,
             }
         }
         None
     }
+    pub fn phi_srcs(&self) -> Option<&OpPhiSrcs> {
+        self.phi_srcs_ip().map(|ip| match &self.instrs[ip].op {
+            Op::PhiSrcs(phi) => phi,
+            _ => panic!("Expected to find the phi"),
+        })
+    }
 
     pub fn phi_srcs_mut(&mut self) -> Option<&mut OpPhiSrcs> {
-        for instr in self.instrs.iter_mut().rev() {
-            if instr.is_branch() {
-                continue;
-            }
-
-            match &mut instr.op {
-                Op::PhiSrcs(phi) => return Some(phi),
-                _ => break,
-            }
-        }
-        None
+        self.phi_srcs_ip().map(|ip| match &mut self.instrs[ip].op {
+            Op::PhiSrcs(phi) => phi,
+            _ => panic!("Expected to find the phi"),
+        })
     }
 
     pub fn branch(&self) -> Option<&Instr> {
@@ -5912,6 +5920,18 @@ impl BasicBlock {
         }
     }
 
+    pub fn branch_ip(&self) -> Option<usize> {
+        if let Some(i) = self.instrs.last() {
+            if i.is_branch() {
+                Some(self.instrs.len() - 1)
+            } else {
+                None
+            }
+        } else {
+            None
+        }
+    }
+
     #[allow(dead_code)]
     pub fn branch_mut(&mut self) -> Option<&mut Instr> {
         if let Some(i) = self.instrs.last_mut() {

src/nouveau/compiler/nak/repair_ssa.rs

@@ -82,38 +82,28 @@ fn get_ssa_or_phi(
 }
 
 fn get_or_insert_phi_dsts(bb: &mut BasicBlock) -> &mut OpPhiDsts {
-    let has_phi = matches!(&bb.instrs[0].op, Op::PhiDsts(_));
-    if !has_phi {
+    let ip = if let Some(ip) = bb.phi_dsts_ip() {
+        ip
+    } else {
         bb.instrs.insert(0, Instr::new_boxed(OpPhiDsts::new()));
-    }
-    match &mut bb.instrs[0].op {
+        0
+    };
+    match &mut bb.instrs[ip].op {
         Op::PhiDsts(phi) => phi,
         _ => panic!("Expected to find the phi we just inserted"),
     }
 }
 
 fn get_or_insert_phi_srcs(bb: &mut BasicBlock) -> &mut OpPhiSrcs {
-    let mut has_phi = false;
-    let mut ip = bb.instrs.len();
-    for (i, instr) in bb.instrs.iter_mut().enumerate().rev() {
-        match &mut instr.op {
-            Op::PhiSrcs(_) => {
-                ip = i;
-                has_phi = true;
-                break;
-            }
-            _ => {
-                if instr.is_branch() {
-                    ip = i;
-                } else {
-                    break;
-                }
-            }
-        }
-    }
-    if !has_phi {
+    let ip = if let Some(ip) = bb.phi_srcs_ip() {
+        ip
+    } else if let Some(ip) = bb.branch_ip() {
         bb.instrs.insert(ip, Instr::new_boxed(OpPhiSrcs::new()));
-    }
+        ip
+    } else {
+        bb.instrs.push(Instr::new_boxed(OpPhiSrcs::new()));
+        bb.instrs.len() - 1
+    };
     match &mut bb.instrs[ip].op {
         Op::PhiSrcs(phi) => phi,
         _ => panic!("Expected to find the phi we just inserted"),

src/nouveau/compiler/nak/spill_values.rs

@@ -841,17 +841,10 @@ fn spill_values<S: Spill>(
         }
 
         // Insert spills and fills right after the phi (if any)
-        let mut ip = pb.instrs.len();
-        while ip > 0 {
-            let instr = &pb.instrs[ip - 1];
-            if !instr.is_branch() {
-                match instr.op {
-                    Op::PhiSrcs(_) => (),
-                    _ => break,
-                }
-            }
-            ip -= 1;
-        }
+        let ip = pb
+            .phi_srcs_ip()
+            .or_else(|| pb.branch_ip())
+            .unwrap_or_else(|| pb.instrs.len());
         pb.instrs.splice(ip..ip, instrs.into_iter());
     }
 }