Add type-safe VecMap for slices indexed by newtypes

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};
@@ -450,8 +451,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)>,
@@ -479,7 +480,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
@@ -563,11 +564,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);
                 }
             }
         }
@@ -625,25 +625,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;
@@ -703,27 +702,12 @@ impl State {
         // If we apply a single swap it could be that we route 2 nodes; that is a setup like
         //  A - B - A - B
         // and we swap the middle two qubits. This cannot happen if we apply 2 or more swaps.
-        if current_swaps.len() > 1 {
-            smallvec![closest_node]
-        } else {
-            // check if the closest node has neighbors that are now routable -- for that we get
-            // the other physical qubit that was swapped and check whether the node on it
-            // is now routable
-            let mut possible_other_qubit = current_swaps[0]
+        match current_swaps.as_slice() {
+            [swap] => swap
                 .iter()
-                // check if other nodes are in the front layer that are connected by this swap
-                .filter_map(|&swap_qubit| self.front_layer.qubits()[swap_qubit.index()])
-                // remove the closest_node, which we know we already routed
-                .filter(|(node_index, _other_qubit)| *node_index != closest_node)
-                .map(|(_node_index, other_qubit)| other_qubit);
-
-            // if there is indeed another candidate, check if that gate is routable
-            if let Some(other_qubit) = possible_other_qubit.next() {
-                if let Some(also_routed) = self.routable_node_on_qubit(problem, other_qubit) {
-                    return smallvec![closest_node, also_routed];
-                }
-            }
-            smallvec![closest_node]
+                .filter_map(|q| self.routable_node_on_qubit(problem, *q))
+                .collect(),
+            _ => smallvec![closest_node],
         }
     }
 
@@ -781,8 +765,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]]);
             }
         }
 
@@ -875,14 +858,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),
@@ -918,8 +901,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(),
+        }
+    }
+}