feat: add source mapping for calibration expansion

quil-py/quil/program/__init__.pyi

@@ -1,4 +1,4 @@
-from typing import Callable, Dict, FrozenSet, List, Optional, Sequence, Set, Tuple, Union, final
+from typing import Callable, Dict, FrozenSet, List, Optional, Sequence, Set, Union, final
 
 import numpy as np
 from numpy.typing import NDArray
@@ -270,19 +270,38 @@ class BasicBlock:
 class CalibrationExpansion:
     def calibration_used(self) -> CalibrationSource: ...
     """The calibration which was used to expand the instruction."""
-    def range(self) -> Tuple[int, int]: ...
+    def range(self) -> range: ...
     """The range of instructions in the expanded list which were generated by this expansion."""
     def expansions(self) -> CalibrationExpansionSourceMap: ...
     """The source map describing the nested expansions made."""
 
 @final
 class CalibrationExpansionSourceMap:
     def entries(self) -> List[CalibrationExpansionSourceMapEntry]: ...
+    def list_sources_for_target_index(self, target_index: int) -> List[int]:
+        """Return the locations in the source which were expanded to generate that instruction.
+
+        This is `O(n)` where `n` is the number of first-level calibration expansions performed.
+        """
+        ...
+
+    def list_sources_for_calibration_used(self, calibration_used: CalibrationSource) -> List[int]:
+        """Return the locations in the source program which were expanded using a calibration.
+
+        This is `O(n)` where `n` is the number of first-level calibration expansions performed.
+        """
+        ...
+
+    def list_targets_for_source_index(self, source_index: int) -> List[CalibrationExpansion]:
+        """Given a source index, return information about its expansion.
+
+        This is `O(n)` where `n` is the number of first-level calibration expansions performed.
+        """
+        ...
 
 @final
 class CalibrationExpansionSourceMapEntry:
-    """
-    A description of the expansion of one instruction into other instructions.
+    """A description of the expansion of one instruction into other instructions.
 
     If present, the instruction located at `source_location` was expanded using calibrations
     into the instructions located at `target_location`.
@@ -322,9 +341,9 @@ class CalibrationExpansionSourceMapEntry:
 
 @final
 class CalibrationSource:
-    """
-    The source of a calibration expansion, which can be either a calibration (`DEFCAL`)
-    or a measure calibration (`DEFCAL MEASURE`).
+    """The source of a calibration expansion.
+
+    Can be either a calibration (`DEFCAL`) or a measure calibration (`DEFCAL MEASURE`).
     """
     def as_calibration(self) -> CalibrationIdentifier: ...
     def as_measure_calibration(self) -> MeasureCalibrationIdentifier: ...
@@ -400,8 +419,9 @@ class CalibrationSet:
 
 @final
 class MaybeCalibrationExpansion:
-    """
-    The result of having expanded a certain instruction within a program. Has two variants:
+    """The result of having expanded a certain instruction within a program.
+
+    Has two variants:
 
     - `expanded`: The instruction was expanded into other instructions, described by a `CalibrationExpansion`.
     - `int`: The instruction was not expanded and is described by an integer, the index of the instruction
@@ -524,12 +544,32 @@ class ProgramCalibrationExpansion:
 @final
 class ProgramCalibrationExpansionSourceMap:
     def entries(self) -> List[ProgramCalibrationExpansionSourceMapEntry]: ...
+    def list_sources_for_target_index(self, target_index: int) -> List[int]:
+        """Return the locations in the source which were expanded to generate that instruction.
+
+        This is `O(n)` where `n` is the number of source instructions.
+        """
+        ...
+
+    def list_sources_for_calibration_used(self, calibration_used: CalibrationSource) -> List[int]:
+        """Return the locations in the source program which were expanded using a calibration.
+
+        This is `O(n)` where `n` is the number of source instructions.
+        """
+        ...
+
+    def list_targets_for_source_index(self, source_index: int) -> List[MaybeCalibrationExpansion]:
+        """Given a source index, return information about its expansion.
+
+        This is `O(n)` where `n` is the number of source instructions.
+        """
+        ...
 
 @final
 class ProgramCalibrationExpansionSourceMapEntry:
-    """
-    A description of the possible expansion of one instruction into other instructions
-    within the scope of a program's calibrations.
+    """A description of the possible expansion of one instruction into other instructions.
+
+    Valid within the scope of a program's calibrations.
     """
     def source_location(self) -> int: ...
     """The instruction index within the source program's body instructions."""

quil-py/src/program/source_map.rs

@@ -1,8 +1,8 @@
 use std::ops::Range;
 
 use pyo3::{
-    types::{PyInt, PyTuple},
-    Py, PyResult, Python,
+    types::{PyInt, PyModule, PyTuple},
+    Py, PyAny, PyResult, Python,
 };
 use quil_rs::program::{
     CalibrationExpansion, CalibrationSource, MaybeCalibrationExpansion,
@@ -35,14 +35,11 @@ impl PyCalibrationExpansion {
         self.as_inner().calibration_used().into()
     }
 
-    // Reviewer: is there a better return type?
-    pub fn range(&self) -> PyResult<(usize, usize)> {
-        Python::with_gil(|py| {
-            let range = py.import("builtins")?.get_item("range")?;
-            let Range { start, end } = self.as_inner().range();
-            let args = PyTuple::new(py, [start, end]);
-            range.call1(args)?.extract()
-        })
+    pub fn range<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        let range = PyModule::import(py, "builtins")?.getattr("range")?;
+        let Range { start, end } = self.as_inner().range();
+        let tuple = PyTuple::new(py, [start, end]);
+        range.call1(tuple)?.extract()
     }
 
     pub fn expansions(&self) -> PyCalibrationExpansionSourceMap {
@@ -67,6 +64,47 @@ impl PyCalibrationExpansionSourceMap {
             .map(|entry| entry.into())
             .collect()
     }
+
+    /// Given an instruction index within the resulting expansion, return the locations in the source
+    /// which were expanded to generate that instruction.
+    ///
+    /// This is `O(n)` where `n` is the number of first-level calibration expansions performed.
+    pub fn list_sources_for_target_index(&self, target_index: usize) -> Vec<usize> {
+        self.as_inner()
+            .list_sources(&target_index)
+            .into_iter()
+            .copied()
+            .collect()
+    }
+
+    /// Given a particular calibration (`DEFCAL` or `DEFCAL MEASURE`), return the locations in the source
+    /// program which were expanded using that calibration.
+    ///
+    /// This is `O(n)` where `n` is the number of first-level calibration expansions performed.
+    pub fn list_sources_for_calibration_used(
+        &self,
+        calibration_used: PyCalibrationSource,
+    ) -> Vec<usize> {
+        self.as_inner()
+            .list_sources(calibration_used.as_inner())
+            .into_iter()
+            .copied()
+            .collect()
+    }
+
+    /// Given a source index, return information about its expansion.
+    ///
+    /// This is `O(n)` where `n` is the number of first-level calibration expansions performed.
+    pub fn list_targets_for_source_index(
+        &self,
+        source_index: usize,
+    ) -> Vec<PyCalibrationExpansion> {
+        self.as_inner()
+            .list_targets(&source_index)
+            .into_iter()
+            .map(|expansion| expansion.into())
+            .collect()
+    }
 }
 
 py_wrap_type! {
@@ -146,6 +184,47 @@ impl PyProgramCalibrationExpansionSourceMap {
             .map(|entry| entry.into())
             .collect()
     }
+
+    /// Given an instruction index within the resulting expansion, return the locations in the source
+    /// which were expanded to generate that instruction.
+    ///
+    /// This is `O(n)` where `n` is the number of source instructions.
+    pub fn list_sources_for_target_index(&self, target_index: usize) -> Vec<usize> {
+        self.as_inner()
+            .list_sources(&target_index)
+            .into_iter()
+            .copied()
+            .collect()
+    }
+
+    /// Given a particular calibration (`DEFCAL` or `DEFCAL MEASURE`), return the locations in the source
+    /// program which were expanded using that calibration.
+    ///
+    /// This is `O(n)` where `n` is the number of source instructions.
+    pub fn list_sources_for_calibration_used(
+        &self,
+        calibration_used: PyCalibrationSource,
+    ) -> Vec<usize> {
+        self.as_inner()
+            .list_sources(calibration_used.as_inner())
+            .into_iter()
+            .copied()
+            .collect()
+    }
+
+    /// Given a source index, return information about its expansion.
+    ///
+    /// This is `O(n)` where `n` is the number of source instructions.
+    pub fn list_targets_for_source_index(
+        &self,
+        source_index: usize,
+    ) -> Vec<PyMaybeCalibrationExpansion> {
+        self.as_inner()
+            .list_targets(&source_index)
+            .into_iter()
+            .map(|expansion| expansion.into())
+            .collect()
+    }
 }
 
 py_wrap_type! {

quil-py/test/program/test_program.py

@@ -186,3 +186,49 @@ def test_filter_instructions(snapshot: SnapshotAssertion):
     program = Program.parse(input)
     program_without_quil_t = program.filter_instructions(lambda instruction: not instruction.is_quil_t())
     assert program_without_quil_t.to_quil() == snapshot
+
+
+def test_calibration_expansion():
+    program = Program.parse(
+        """DEFCAL X 0:
+    Y 0
+
+DEFCAL Y 0:
+    Z 0
+
+X 0
+Y 0
+"""
+    )
+    expansion = program.expand_calibrations_with_source_map()
+    source_map = expansion.source_map()
+
+    assert (
+        expansion.program().to_quil()
+        == Program.parse("""DEFCAL X 0:
+    Y 0
+
+DEFCAL Y 0:
+    Z 0
+
+Z 0
+Z 0
+""").to_quil()
+    )
+
+    # The X at index 0 should have been replaced with a Z at index 0
+    targets = source_map.list_targets_for_source_index(0)
+    assert len(targets) == 1
+    expanded = targets[0].as_expanded()
+    assert expanded.range() == range(0, 1)
+    assert source_map.list_sources_for_target_index(0) == [0]
+
+    # The Y at index 1 should have been replaced with a Z at index 1
+    targets = source_map.list_targets_for_source_index(1)
+    assert len(targets) == 1
+    expanded = targets[0].as_expanded()
+    assert expanded.range() == range(1, 2)
+    assert source_map.list_sources_for_target_index(1) == [1]
+
+    # There is no source index 2 and so there should be no mapping
+    assert source_map.list_targets_for_source_index(2) == []

quil-rs/src/instruction/calibration.rs

@@ -153,7 +153,7 @@ impl CalibrationIdentifier {
                         (calib, gate) => calib == gate,
                     }
                 });
-        !fixed_parameters_match
+        fixed_parameters_match
     }
 }
 
@@ -229,7 +229,7 @@ impl Quil for MeasureCalibrationDefinition {
     }
 }
 
-// For review: how would we feel about making this a subfield of the `MeasureCalibrationDefinition` itself?
+/// A unique identifier for a measurement calibration definition within a program
 #[derive(Clone, Debug, Default, PartialEq)]
 pub struct MeasureCalibrationIdentifier {
     /// The qubit which is the target of measurement, if any