Record vreg classes explicitly during liverange pass. (#35)

This resolves an issue seen when the source program uses multiple
regclasses (Int and Float): in some cases, the logic that grabs the
vregs and retains them (with class) in `vreg_regs` missed a register and
we had a class mismatch. This occurred because data structures were
initialized assuming `Int` regclass at first.

This PR instead removes the `vreg_regs` array, stores the class
explicitly as an `Option<RegClass>` in the `VRegData`, and provides a
`Env::vreg()` method that reconstitutes a `VReg` given its index and its
observed class. We "observe" the class of every vreg seen during the
liveness pass (and we assert that every occurrence of the vreg index has
the same class). In this way, we still have a single source-of-truth for
the vreg class (the mention of the vreg itself) and we explicitly
represent the "not observed yet" state (and panic on attempting to use
such a vreg) rather than implicitly taking the wrong class.

src/ion/data_structures.rs

@@ -259,6 +259,8 @@ pub struct VRegData {
     pub ranges: LiveRangeList,
     pub blockparam: Block,
     pub is_ref: bool,
+    // We don't initially know the RegClass until we observe a use of the VReg.
+    pub class: Option<RegClass>,
 }
 
 #[derive(Clone, Debug)]
@@ -340,7 +342,6 @@ pub struct Env<'a, F: Function> {
     pub bundles: Vec<LiveBundle>,
     pub spillsets: Vec<SpillSet>,
     pub vregs: Vec<VRegData>,
-    pub vreg_regs: Vec<VReg>,
     pub pregs: Vec<PRegData>,
     pub allocation_queue: PrioQueue,
     pub safepoints: Vec<Inst>, // Sorted list of safepoint insts.
@@ -397,6 +398,31 @@ pub struct Env<'a, F: Function> {
     pub annotations_enabled: bool,
 }
 
+impl<'a, F: Function> Env<'a, F> {
+    /// Get the VReg (with bundled RegClass) from a vreg index.
+    #[inline]
+    pub fn vreg(&self, index: VRegIndex) -> VReg {
+        let class = self.vregs[index.index()]
+            .class
+            .expect("trying to get a VReg before observing its class");
+        VReg::new(index.index(), class)
+    }
+
+    /// Record the class of a VReg. We learn this only when we observe
+    /// the VRegs in use.
+    pub fn observe_vreg_class(&mut self, vreg: VReg) {
+        let old_class = self.vregs[vreg.vreg()].class.replace(vreg.class());
+        // We should never observe two different classes for two
+        // mentions of a VReg in the source program.
+        debug_assert!(old_class == None || old_class == Some(vreg.class()));
+    }
+
+    /// Is this vreg actually used in the source program?
+    pub fn is_vreg_used(&self, index: VRegIndex) -> bool {
+        self.vregs[index.index()].class.is_some()
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct SpillSlotData {
     pub ranges: LiveRangeSet,

src/ion/liveranges.rs

@@ -119,6 +119,8 @@ impl<'a, F: Function> Env<'a, F> {
                     ranges: smallvec![],
                     blockparam: Block::invalid(),
                     is_ref: false,
+                    // We'll learn the RegClass as we scan the code.
+                    class: None,
                 },
             );
         }
@@ -137,8 +139,8 @@ impl<'a, F: Function> Env<'a, F> {
 
     pub fn add_vreg(&mut self, reg: VReg, data: VRegData) -> VRegIndex {
         let idx = self.vregs.len();
+        debug_assert_eq!(reg.vreg(), idx);
         self.vregs.push(data);
-        self.vreg_regs.push(reg);
         VRegIndex::new(idx)
     }
 
@@ -322,8 +324,9 @@ impl<'a, F: Function> Env<'a, F> {
             // Include outgoing blockparams in the initial live set.
             if self.func.is_branch(insns.last()) {
                 for i in 0..self.func.block_succs(block).len() {
-                    for param in self.func.branch_blockparams(block, insns.last(), i) {
+                    for &param in self.func.branch_blockparams(block, insns.last(), i) {
                         live.set(param.vreg(), true);
+                        self.observe_vreg_class(param);
                     }
                 }
             }
@@ -332,6 +335,8 @@ impl<'a, F: Function> Env<'a, F> {
                 if let Some((src, dst)) = self.func.is_move(inst) {
                     live.set(dst.vreg().vreg(), false);
                     live.set(src.vreg().vreg(), true);
+                    self.observe_vreg_class(src.vreg());
+                    self.observe_vreg_class(dst.vreg());
                 }
 
                 for pos in &[OperandPos::Late, OperandPos::Early] {
@@ -347,12 +352,14 @@ impl<'a, F: Function> Env<'a, F> {
                                     live.set(op.vreg().vreg(), false);
                                 }
                             }
+                            self.observe_vreg_class(op.vreg());
                         }
                     }
                 }
             }
             for &blockparam in self.func.block_params(block) {
                 live.set(blockparam.vreg(), false);
+                self.observe_vreg_class(blockparam);
             }
 
             for &pred in self.func.block_preds(block) {
@@ -372,7 +379,7 @@ impl<'a, F: Function> Env<'a, F> {
         // for pinned vregs. (The client should create a virtual
         // instruction that defines any other liveins if necessary.)
         for livein in self.liveins[self.func.entry_block().index()].iter() {
-            let livein = self.vreg_regs[livein];
+            let livein = self.vreg(VRegIndex::new(livein));
             if self.func.is_pinned_vreg(livein).is_none() {
                 trace!("non-pinned-vreg livein to entry block: {}", livein);
                 return Err(RegAllocError::EntryLivein);
@@ -451,8 +458,7 @@ impl<'a, F: Function> Env<'a, F> {
             }
 
             // Create vreg data for blockparams.
-            for param in self.func.block_params(block) {
+            for &param in self.func.block_params(block) {
-                self.vreg_regs[param.vreg()] = *param;
                 self.vregs[param.vreg()].blockparam = block;
             }
 
@@ -841,7 +847,6 @@ impl<'a, F: Function> Env<'a, F> {
                                 self.ranges[dst_lr.index()].set_flag(LiveRangeFlag::StartsAtDef);
                                 live.set(dst.vreg().vreg(), false);
                                 vreg_ranges[dst.vreg().vreg()] = LiveRangeIndex::invalid();
-                                self.vreg_regs[dst.vreg().vreg()] = dst.vreg();
 
                                 // Handle the use w.r.t. liveranges: make it live
                                 // and create an initial LR back to the start of
@@ -930,9 +935,6 @@ impl<'a, F: Function> Env<'a, F> {
                             OperandKind::Def | OperandKind::Mod => {
                                 trace!("Def of {} at {:?}", operand.vreg(), pos);
 
-                                // Fill in vreg's actual data.
-                                self.vreg_regs[operand.vreg().vreg()] = operand.vreg();
-
                                 // Get or create the LiveRange.
                                 let mut lr = vreg_ranges[operand.vreg().vreg()];
                                 trace!(" -> has existing LR {:?}", lr);
@@ -1107,7 +1109,7 @@ impl<'a, F: Function> Env<'a, F> {
                     // Create a virtual use.
                     let pos = ProgPoint::before(self.safepoints[safepoint_idx]);
                     let operand = Operand::new(
-                        self.vreg_regs[vreg.index()],
+                        self.vreg(vreg),
                         OperandConstraint::Stack,
                         OperandKind::Use,
                         OperandPos::Early,

src/ion/merge.rs

@@ -52,11 +52,11 @@ impl<'a, F: Function> Env<'a, F> {
             // Sanity check: both bundles should contain only ranges with appropriate VReg classes.
             for entry in &self.bundles[from.index()].ranges {
                 let vreg = self.ranges[entry.index.index()].vreg;
-                debug_assert_eq!(from_rc, self.vreg_regs[vreg.index()].class());
+                debug_assert_eq!(from_rc, self.vreg(vreg).class());
             }
             for entry in &self.bundles[to.index()].ranges {
                 let vreg = self.ranges[entry.index.index()].vreg;
-                debug_assert_eq!(to_rc, self.vreg_regs[vreg.index()].class());
+                debug_assert_eq!(to_rc, self.vreg(vreg).class());
             }
         }
 
@@ -234,7 +234,7 @@ impl<'a, F: Function> Env<'a, F> {
 
             // If this is a pinned vreg, go ahead and add it to the
             // commitment map, and avoid creating a bundle entirely.
-            if let Some(preg) = self.func.is_pinned_vreg(self.vreg_regs[vreg.index()]) {
+            if let Some(preg) = self.func.is_pinned_vreg(self.vreg(vreg)) {
                 for entry in &self.vregs[vreg.index()].ranges {
                     let key = LiveRangeKey::from_range(&entry.range);
                     self.pregs[preg.index()]
@@ -281,7 +281,7 @@ impl<'a, F: Function> Env<'a, F> {
 
             // Create a spillslot for this bundle.
             let ssidx = SpillSetIndex::new(self.spillsets.len());
-            let reg = self.vreg_regs[vreg.index()];
+            let reg = self.vreg(vreg);
             let size = self.func.spillslot_size(reg.class()) as u8;
             self.spillsets.push(SpillSet {
                 vregs: smallvec![vreg],
@@ -361,8 +361,8 @@ impl<'a, F: Function> Env<'a, F> {
                 dst
             );
 
-            let dst_vreg = self.vreg_regs[self.ranges[dst.index()].vreg.index()];
+            let dst_vreg = self.vreg(self.ranges[dst.index()].vreg);
-            let src_vreg = self.vreg_regs[self.ranges[src.index()].vreg.index()];
+            let src_vreg = self.vreg(self.ranges[src.index()].vreg);
             if self.func.is_pinned_vreg(src_vreg).is_some()
                 && self.func.is_pinned_vreg(dst_vreg).is_some()
             {

src/ion/mod.rs

@@ -57,7 +57,6 @@ impl<'a, F: Function> Env<'a, F> {
             ranges: Vec::with_capacity(4 * n),
             spillsets: Vec::with_capacity(n),
             vregs: Vec::with_capacity(n),
-            vreg_regs: Vec::with_capacity(n),
             pregs: vec![],
             allocation_queue: PrioQueue::new(),
             safepoints: vec![],

src/ion/moves.rs

@@ -188,8 +188,11 @@ impl<'a, F: Function> Env<'a, F> {
         let mut prog_move_dst_idx = 0;
         for vreg in 0..self.vregs.len() {
             let vreg = VRegIndex::new(vreg);
+            if !self.is_vreg_used(vreg) {
+                continue;
+            }
 
-            let pinned_alloc = self.func.is_pinned_vreg(self.vreg_regs[vreg.index()]);
+            let pinned_alloc = self.func.is_pinned_vreg(self.vreg(vreg));
 
             // For each range in each vreg, insert moves or
             // half-moves.  We also scan over `blockparam_ins` and
@@ -283,7 +286,7 @@ impl<'a, F: Function> Env<'a, F> {
                             InsertMovePrio::Regular,
                             prev_alloc,
                             alloc,
-                            Some(self.vreg_regs[vreg.index()]),
+                            Some(self.vreg(vreg)),
                         );
                     }
                 }
@@ -725,7 +728,7 @@ impl<'a, F: Function> Env<'a, F> {
                     prio,
                     src.alloc,
                     dest.alloc,
-                    Some(self.vreg_regs[dest.to_vreg().index()]),
+                    Some(self.vreg(dest.to_vreg())),
                 );
                 last = Some(dest.alloc);
             }
@@ -747,7 +750,7 @@ impl<'a, F: Function> Env<'a, F> {
                 InsertMovePrio::MultiFixedReg,
                 from_alloc,
                 to_alloc,
-                Some(self.vreg_regs[fixup.vreg.index()]),
+                Some(self.vreg(fixup.vreg)),
             );
             self.set_alloc(
                 fixup.pos.inst(),
@@ -877,7 +880,7 @@ impl<'a, F: Function> Env<'a, F> {
                 InsertMovePrio::Regular,
                 from_alloc,
                 to_alloc,
-                Some(self.vreg_regs[to_vreg.index()]),
+                Some(self.vreg(to_vreg)),
             );
         }