Add type-safe VecMap for slices indexed by newtypes

Sabre uses several objects that are logically maps from an index-like
newtype (like `NodeIndex` or `PhysicalQubit`) to some value, and are
implemented as fixed-slice `Vec`s for lookup efficiency.    The newtype
provides type safety while it's in use, but we have to cast it away to
index, which makes it easy to index slices with the wrong object.

This introduces a `VecMap` object, which provides a (minimal) slice-like
interface, but indexes using the relevant newtype.

The current implementation of Sabre does not use this _too_ much, but a
refactoring of the layer structures will have them store one slice
indexed by `PhysicalQubit` and one by `VirtualQubit`, which are
trivially easy to get switched (a frequent mistake that is the base
reason those new types were introduced in the first place).

Author: jakelishman

crates/transpiler/src/passes/sabre/layer.rs

@@ -16,6 +16,8 @@ use rustworkx_core::petgraph::prelude::*;
 
 use qiskit_circuit::PhysicalQubit;
 
+use super::vec_map::VecMap;
+
 /// A container for the current non-routable parts of the front layer.  This only ever holds
 /// two-qubit gates; the only reason a 0q- or 1q operation can be unroutable is because it has an
 /// unsatisfied 2q predecessor, which disqualifies it from being in the front layer.
@@ -31,7 +33,7 @@ pub struct FrontLayer {
     nodes: IndexMap<NodeIndex, [PhysicalQubit; 2], ::ahash::RandomState>,
     /// Map of each qubit to the node that acts on it and the other qubit that node acts on, if this
     /// qubit is active (otherwise `None`).
-    qubits: Vec<Option<(NodeIndex, PhysicalQubit)>>,
+    qubits: VecMap<PhysicalQubit, Option<(NodeIndex, PhysicalQubit)>>,
 }
 
 impl FrontLayer {
@@ -43,7 +45,7 @@ impl FrontLayer {
                 num_qubits as usize / 2,
                 ::ahash::RandomState::default(),
             ),
-            qubits: vec![None; num_qubits as usize],
+            qubits: vec![None; num_qubits as usize].into(),
         }
     }
 
@@ -54,15 +56,15 @@ impl FrontLayer {
 
     /// View onto the mapping between qubits and their `(node, other_qubit)` pair.  Index `i`
     /// corresponds to physical qubit `i`.
-    pub fn qubits(&self) -> &[Option<(NodeIndex, PhysicalQubit)>] {
+    pub fn qubits(&self) -> &VecMap<PhysicalQubit, Option<(NodeIndex, PhysicalQubit)>> {
         &self.qubits
     }
 
     /// Add a node into the front layer, with the two qubits it operates on.
     pub fn insert(&mut self, index: NodeIndex, qubits: [PhysicalQubit; 2]) {
         let [a, b] = qubits;
-        self.qubits[a.index()] = Some((index, b));
-        self.qubits[b.index()] = Some((index, a));
+        self.qubits[a] = Some((index, b));
+        self.qubits[b] = Some((index, a));
         self.nodes.insert(index, qubits);
     }
 
@@ -74,14 +76,14 @@ impl FrontLayer {
             .nodes
             .swap_remove(index)
             .expect("Tried removing index that does not exist.");
-        self.qubits[a.index()] = None;
-        self.qubits[b.index()] = None;
+        self.qubits[a] = None;
+        self.qubits[b] = None;
     }
 
     /// Query whether a qubit has an active node.
     #[inline]
     pub fn is_active(&self, qubit: PhysicalQubit) -> bool {
-        self.qubits[qubit.index()].is_some()
+        self.qubits[qubit].is_some()
     }
 
     /// Calculate the score _difference_ caused by this swap, compared to not making the swap.
@@ -95,10 +97,10 @@ impl FrontLayer {
         // equal anyway, so not affect the score.
         let [a, b] = swap;
         let mut total = 0.0;
-        if let Some((_, c)) = self.qubits[a.index()] {
+        if let Some((_, c)) = self.qubits[a] {
             total += dist[[b.index(), c.index()]] - dist[[a.index(), c.index()]]
         }
-        if let Some((_, c)) = self.qubits[b.index()] {
+        if let Some((_, c)) = self.qubits[b] {
             total += dist[[a.index(), c.index()]] - dist[[b.index(), c.index()]]
         }
         total
@@ -114,25 +116,25 @@ impl FrontLayer {
     /// Apply a physical swap to the current layout data structure.
     pub fn apply_swap(&mut self, swap: [PhysicalQubit; 2]) {
         let [a, b] = swap;
-        match (self.qubits[a.index()], self.qubits[b.index()]) {
+        match (self.qubits[a], self.qubits[b]) {
             (Some((index1, _)), Some((index2, _))) if index1 == index2 => {
                 let entry = self.nodes.get_mut(&index1).unwrap();
                 *entry = [entry[1], entry[0]];
                 return;
             }
             _ => {}
         }
-        if let Some((index, c)) = self.qubits[a.index()] {
-            self.qubits[c.index()] = Some((index, b));
+        if let Some((index, c)) = self.qubits[a] {
+            self.qubits[c] = Some((index, b));
             let entry = self.nodes.get_mut(&index).unwrap();
             *entry = if *entry == [a, c] { [b, c] } else { [c, b] };
         }
-        if let Some((index, c)) = self.qubits[b.index()] {
-            self.qubits[c.index()] = Some((index, a));
+        if let Some((index, c)) = self.qubits[b] {
+            self.qubits[c] = Some((index, a));
             let entry = self.nodes.get_mut(&index).unwrap();
             *entry = if *entry == [b, c] { [a, c] } else { [c, a] };
         }
-        self.qubits.swap(a.index(), b.index());
+        self.qubits.swap(a, b);
     }
 
     /// True if there are no nodes in the current layer.

crates/transpiler/src/passes/sabre/mod.rs

@@ -15,6 +15,7 @@ pub mod heuristic;
 mod layer;
 mod layout;
 pub(crate) mod route;
+mod vec_map;
 
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;

crates/transpiler/src/passes/sabre/route.rs

@@ -37,6 +37,7 @@ use smallvec::{SmallVec, smallvec};
 use super::dag::{InteractionKind, SabreDAG};
 use super::heuristic::{BasicHeuristic, DecayHeuristic, Heuristic, LookaheadHeuristic, SetScaling};
 use super::layer::{ExtendedSet, FrontLayer};
+use super::vec_map::VecMap;
 use crate::TranspilerError;
 use crate::neighbors::Neighbors;
 use crate::target::{Target, TargetCouplingError};
@@ -439,8 +440,8 @@ struct State {
     extended_set: ExtendedSet,
     /// How many predecessors still need to be satisfied for each node index before it is at the
     /// front of the topological iteration through the nodes as they're routed.
-    required_predecessors: Vec<u32>,
-    decay: Vec<f64>,
+    required_predecessors: VecMap<NodeIndex, u32>,
+    decay: VecMap<PhysicalQubit, f64>,
     /// Reusable allocated storage space for accumulating and scoring swaps.  This is owned as part
     /// of the general state to avoid reallocation costs.
     swap_scores: Vec<([PhysicalQubit; 2], f64)>,
@@ -468,7 +469,7 @@ impl State {
         problem: RoutingProblem,
         qubit: PhysicalQubit,
     ) -> Option<NodeIndex> {
-        self.front_layer.qubits()[qubit.index()].and_then(|(node, other)| {
+        self.front_layer.qubits()[qubit].and_then(|(node, other)| {
             problem
                 .target
                 .neighbors
@@ -552,11 +553,10 @@ impl State {
                 .dag
                 .edges_directed(node_id, Direction::Outgoing)
             {
-                let successor_node = edge.target();
-                let successor_index = successor_node.index();
-                self.required_predecessors[successor_index] -= 1;
-                if self.required_predecessors[successor_index] == 0 {
-                    to_visit.push_back(successor_node);
+                let successor = edge.target();
+                self.required_predecessors[successor] -= 1;
+                if self.required_predecessors[successor] == 0 {
+                    to_visit.push_back(successor);
                 }
             }
         }
@@ -614,25 +614,24 @@ impl State {
                 return;
             };
         let mut to_visit = self.front_layer.iter_nodes().copied().collect::<Vec<_>>();
-        let mut decremented: IndexMap<usize, u32, ahash::RandomState> =
+        let mut decremented: IndexMap<NodeIndex, u32, ahash::RandomState> =
             IndexMap::with_hasher(ahash::RandomState::default());
         let mut i = 0;
         while i < to_visit.len() && self.extended_set.len() < extended_set_size {
             let node = to_visit[i];
             for edge in problem.sabre.dag.edges_directed(node, Direction::Outgoing) {
-                let successor_node = edge.target();
-                let successor_index = successor_node.index();
-                *decremented.entry(successor_index).or_insert(0) += 1;
-                self.required_predecessors[successor_index] -= 1;
-                if self.required_predecessors[successor_index] == 0 {
+                let successor = edge.target();
+                *decremented.entry(successor).or_insert(0) += 1;
+                self.required_predecessors[successor] -= 1;
+                if self.required_predecessors[successor] == 0 {
                     // TODO: this looks "through" control-flow ops without seeing them, but we
                     // actually eagerly route control-flow blocks as soon as they're eligible, so
                     // they should be reflected in the extended set.
-                    if let InteractionKind::TwoQ([a, b]) = &problem.sabre.dag[successor_node].kind {
+                    if let InteractionKind::TwoQ([a, b]) = &problem.sabre.dag[successor].kind {
                         self.extended_set
                             .push([a.to_phys(&self.layout), b.to_phys(&self.layout)]);
                     }
-                    to_visit.push(successor_node);
+                    to_visit.push(successor);
                 }
             }
             i += 1;
@@ -755,8 +754,7 @@ impl State {
 
         if let Some(DecayHeuristic { .. }) = problem.heuristic.decay {
             for (swap, score) in self.swap_scores.iter_mut() {
-                *score = (absolute_score + *score)
-                    * self.decay[swap[0].index()].max(self.decay[swap[1].index()]);
+                *score = (absolute_score + *score) * self.decay[swap[0]].max(self.decay[swap[1]]);
             }
         }
 
@@ -849,14 +847,14 @@ pub fn swap_map_trial<'a>(
     let mut order = Order::for_problem(problem);
 
     let num_qubits: u32 = problem.target.num_qubits().try_into().unwrap();
-    let mut required_predecessors = vec![0; problem.sabre.dag.node_count()];
+    let mut required_predecessors = VecMap::from(vec![0; problem.sabre.dag.node_count()]);
     for edge in problem.sabre.dag.edge_references() {
-        required_predecessors[edge.target().index()] += 1;
+        required_predecessors[edge.target()] += 1;
     }
     let mut state = State {
         front_layer: FrontLayer::new(num_qubits),
         extended_set: ExtendedSet::new(num_qubits),
-        decay: vec![1.; num_qubits as usize],
+        decay: vec![1.; num_qubits as usize].into(),
         required_predecessors,
         layout: initial_layout.clone(),
         swap_scores: Vec::with_capacity(problem.target.neighbors.edge_count() / 2),
@@ -892,8 +890,8 @@ pub fn swap_map_trial<'a>(
                     state.decay.fill(1.);
                     num_search_steps = 0;
                 } else {
-                    state.decay[best_swap[0].index()] += increment;
-                    state.decay[best_swap[1].index()] += increment;
+                    state.decay[best_swap[0]] += increment;
+                    state.decay[best_swap[1]] += increment;
                 }
             }
         }

crates/transpiler/src/passes/sabre/vec_map.rs

@@ -0,0 +1,64 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2026
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at https://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use std::marker::PhantomData;
+use std::ops;
+
+use rustworkx_core::petgraph::graph::IndexType;
+
+/// Internal helper struct that represents a `Box<[T]>`, but is indexed by a petgraph-like indexer.
+///
+/// The layer structures involve several flat arrays, each representing a map from an index to a
+/// value, but the index types are often different, so it's less error-prone if we _enforce_ that
+/// you index using the object, rather than calling its `.index` method and erasing the type.
+#[derive(Clone, Debug)]
+pub struct VecMap<Idx: IndexType, T> {
+    phantom: PhantomData<Idx>,
+    data: Box<[T]>,
+}
+impl<Idx: IndexType, T> VecMap<Idx, T> {
+    /// Swap the values of two indices.
+    #[inline]
+    pub fn swap(&mut self, a: Idx, b: Idx) {
+        self.data.swap(a.index(), b.index())
+    }
+
+    /// Fill all entries in the slice with a given value.
+    #[inline]
+    pub fn fill(&mut self, val: T)
+    where
+        T: Clone,
+    {
+        self.data.fill(val)
+    }
+}
+
+impl<Idx: IndexType, T> ops::Index<Idx> for VecMap<Idx, T> {
+    type Output = <[T] as ops::Index<usize>>::Output;
+    fn index(&self, index: Idx) -> &Self::Output {
+        &self.data[index.index()]
+    }
+}
+impl<Idx: IndexType, T> ops::IndexMut<Idx> for VecMap<Idx, T> {
+    fn index_mut(&mut self, index: Idx) -> &mut Self::Output {
+        &mut self.data[index.index()]
+    }
+}
+
+impl<Idx: IndexType, T> From<Vec<T>> for VecMap<Idx, T> {
+    fn from(value: Vec<T>) -> Self {
+        Self {
+            phantom: PhantomData,
+            data: value.into_boxed_slice(),
+        }
+    }
+}