[Stretch] Support stretch and duration expressions for Delay instruction.

crates/circuit/src/converters.rs

@@ -32,6 +32,8 @@ pub struct QuantumCircuitData<'py> {
     pub input_vars: Vec<Bound<'py, PyAny>>,
     pub captured_vars: Vec<Bound<'py, PyAny>>,
     pub declared_vars: Vec<Bound<'py, PyAny>>,
+    pub captured_stretches: Vec<Bound<'py, PyAny>>,
+    pub declared_stretches: Vec<Bound<'py, PyAny>>,
 }
 
 impl<'py> FromPyObject<'py> for QuantumCircuitData<'py> {
@@ -63,6 +65,14 @@ impl<'py> FromPyObject<'py> for QuantumCircuitData<'py> {
                 .call_method0(intern!(py, "iter_declared_vars"))?
                 .try_iter()?
                 .collect::<PyResult<Vec<_>>>()?,
+            captured_stretches: ob
+                .call_method0(intern!(py, "iter_captured_stretches"))?
+                .try_iter()?
+                .collect::<PyResult<Vec<_>>>()?,
+            declared_stretches: ob
+                .call_method0(intern!(py, "iter_declared_stretches"))?
+                .try_iter()?
+                .collect::<PyResult<Vec<_>>>()?,
         })
     }
 }

crates/circuit/src/dag_circuit.rs

@@ -237,6 +237,9 @@ pub struct DAGCircuit {
     control_flow_module: PyControlFlowModule,
     vars_info: HashMap<String, DAGVarInfo>,
     vars_by_type: [Py<PySet>; 3],
+
+    captured_stretches: IndexMap<String, Py<PyAny>, RandomState>,
+    declared_stretches: IndexMap<String, Py<PyAny>, RandomState>,
 }
 
 #[derive(Clone, Debug)]
@@ -391,6 +394,8 @@ impl DAGCircuit {
                 PySet::empty(py)?.unbind(),
                 PySet::empty(py)?.unbind(),
             ],
+            captured_stretches: IndexMap::default(),
+            declared_stretches: IndexMap::default(),
         })
     }
 
@@ -1497,6 +1502,12 @@ impl DAGCircuit {
             {
                 target_dag.add_var(py, &var, DAGVarType::Declare)?;
             }
+            for stretch in self.captured_stretches.values() {
+                target_dag.add_captured_stretch(py, stretch.bind(py))?;
+            }
+            for stretch in self.declared_stretches.values() {
+                target_dag.add_declared_stretch(stretch.bind(py))?;
+            }
         } else if vars_mode == "captures" {
             for var in self.vars_by_type[DAGVarType::Input as usize]
                 .bind(py)
@@ -1516,6 +1527,12 @@ impl DAGCircuit {
             {
                 target_dag.add_var(py, &var, DAGVarType::Capture)?;
             }
+            for stretch in self.captured_stretches.values() {
+                target_dag.add_captured_stretch(py, stretch.bind(py))?;
+            }
+            for stretch in self.declared_stretches.values() {
+                target_dag.add_captured_stretch(py, stretch.bind(py))?;
+            }
         } else if vars_mode != "drop" {
             return Err(PyValueError::new_err(format!(
                 "unknown vars_mode: '{}'",
@@ -1813,20 +1830,26 @@ impl DAGCircuit {
             dag.add_input_var(py, &var?)?;
         }
         if inline_captures {
-            for var in other.iter_captured_vars(py)?.bind(py) {
+            for var in other.iter_captures(py)?.bind(py) {
                 let var = var?;
-                if !dag.has_var(&var)? {
+                if !dag.has_identifier(&var)? {
                     return Err(DAGCircuitError::new_err(format!("Variable '{}' to be inlined is not in the base DAG. If you wanted it to be automatically added, use `inline_captures=False`.", var)));
                 }
             }
         } else {
             for var in other.iter_captured_vars(py)?.bind(py) {
                 dag.add_captured_var(py, &var?)?;
             }
+            for stretch in other.iter_captured_stretches(py)?.bind(py) {
+                dag.add_captured_stretch(py, &stretch?)?;
+            }
         }
         for var in other.iter_declared_vars(py)?.bind(py) {
             dag.add_declared_var(py, &var?)?;
         }
+        for var in other.iter_declared_stretches(py)?.bind(py) {
+            dag.add_declared_stretch(&var?)?;
+        }
 
         let variable_mapper = PyVariableMapper::new(
             py,
@@ -2244,6 +2267,32 @@ impl DAGCircuit {
             }
         }
 
+        if self.captured_stretches.len() != other.captured_stretches.len()
+            || self.declared_stretches.len() != other.declared_stretches.len()
+        {
+            return Ok(false);
+        }
+
+        for (our_stretch, their_stretch) in self
+            .captured_stretches
+            .values()
+            .zip(other.captured_stretches.values())
+        {
+            if !our_stretch.bind(py).eq(their_stretch)? {
+                return Ok(false);
+            }
+        }
+
+        for (our_stretch, their_stretch) in self
+            .declared_stretches
+            .values()
+            .zip(other.declared_stretches.values())
+        {
+            if !our_stretch.bind(py).eq(their_stretch)? {
+                return Ok(false);
+            }
+        }
+
         let self_bit_indices = {
             let indices = self
                 .qubits
@@ -4265,6 +4314,47 @@ impl DAGCircuit {
         Ok(())
     }
 
+    /// Add a captured stretch variable to the circuit.
+    ///
+    /// Args:
+    ///     var: the stretch variable to add.
+    fn add_captured_stretch(&mut self, py: Python, var: &Bound<PyAny>) -> PyResult<()> {
+        if !self.vars_by_type[DAGVarType::Input as usize]
+            .bind(py)
+            .is_empty()
+        {
+            return Err(DAGCircuitError::new_err(
+                "cannot add captures to a circuit with inputs",
+            ));
+        }
+        let var_name: String = var.getattr("name")?.extract::<String>()?;
+        if self.vars_info.contains_key(&var_name) {
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing var",
+            ));
+        }
+        if let Some(previous) = self.declared_stretches.get(&var_name) {
+            if var.eq(previous)? {
+                return Err(DAGCircuitError::new_err("already present in the circuit"));
+            }
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing var",
+            ));
+        }
+        if let Some(previous) = self.captured_stretches.get(&var_name) {
+            if var.eq(previous)? {
+                return Err(DAGCircuitError::new_err("already present in the circuit"));
+            }
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing var",
+            ));
+        }
+
+        self.captured_stretches
+            .insert(var_name, var.clone().unbind());
+        Ok(())
+    }
+
     /// Add a declared local variable to the circuit.
     ///
     /// Args:
@@ -4274,6 +4364,39 @@ impl DAGCircuit {
         Ok(())
     }
 
+    /// Add a declared stretch variable to the circuit.
+    ///
+    /// Args:
+    ///     var: the stretch variable to add.
+    fn add_declared_stretch(&mut self, var: &Bound<PyAny>) -> PyResult<()> {
+        let var_name: String = var.getattr("name")?.extract::<String>()?;
+        if self.vars_info.contains_key(&var_name) {
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing var",
+            ));
+        }
+        if let Some(previous) = self.declared_stretches.get(&var_name) {
+            if var.eq(previous)? {
+                return Err(DAGCircuitError::new_err("already present in the circuit"));
+            }
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing var",
+            ));
+        }
+        if let Some(previous) = self.captured_stretches.get(&var_name) {
+            if var.eq(previous)? {
+                return Err(DAGCircuitError::new_err("already present in the circuit"));
+            }
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing var",
+            ));
+        }
+
+        self.declared_stretches
+            .insert(var_name, var.clone().unbind());
+        Ok(())
+    }
+
     /// Total number of classical variables tracked by the circuit.
     #[getter]
     fn num_vars(&self) -> usize {
@@ -4320,6 +4443,36 @@ impl DAGCircuit {
         }
     }
 
+    /// Is this stretch variable in the DAG?
+    ///
+    /// Args:
+    ///     var: the stretch variable or name to check.
+    fn has_stretch(&self, var: &Bound<PyAny>) -> PyResult<bool> {
+        match var.extract::<String>() {
+            Ok(name) => Ok(self.declared_stretches.contains_key(&name)
+                || self.captured_stretches.contains_key(&name)),
+            Err(_) => {
+                let raw_name = var.getattr("name")?;
+                let var_name: String = raw_name.extract()?;
+                if let Some(stretch) = self.declared_stretches.get(&var_name) {
+                    return Ok(stretch.is(var));
+                }
+                if let Some(stretch) = self.captured_stretches.get(&var_name) {
+                    return Ok(stretch.is(var));
+                }
+                Ok(false)
+            }
+        }
+    }
+
+    /// Is this identifier in the DAG?
+    ///
+    /// Args:
+    ///     var: the identifier or name to check.
+    fn has_identifier(&self, var: &Bound<PyAny>) -> PyResult<bool> {
+        Ok(self.has_var(var)? || self.has_stretch(var)?)
+    }
+
     /// Iterable over the input classical variables tracked by the circuit.
     fn iter_input_vars(&self, py: Python) -> PyResult<Py<PyIterator>> {
         Ok(self.vars_by_type[DAGVarType::Input as usize]
@@ -4340,6 +4493,29 @@ impl DAGCircuit {
             .unbind())
     }
 
+    /// Iterable over the captured stretch variables tracked by the circuit.
+    fn iter_captured_stretches(&self, py: Python) -> PyResult<Py<PyIterator>> {
+        Ok(PyList::new(py, self.captured_stretches.values())?
+            .into_any()
+            .try_iter()?
+            .unbind())
+    }
+
+    /// Iterable over all captured identifiers tracked by the circuit.
+    fn iter_captures(&self, py: Python) -> PyResult<Py<PyIterator>> {
+        let out_set = PySet::empty(py)?;
+        for var in self.vars_by_type[DAGVarType::Capture as usize]
+            .bind(py)
+            .iter()
+        {
+            out_set.add(var)?;
+        }
+        for stretch in self.captured_stretches.values() {
+            out_set.add(stretch)?;
+        }
+        Ok(out_set.into_any().try_iter()?.unbind())
+    }
+
     /// Iterable over the declared classical variables tracked by the circuit.
     fn iter_declared_vars(&self, py: Python) -> PyResult<Py<PyIterator>> {
         Ok(self.vars_by_type[DAGVarType::Declare as usize]
@@ -4350,6 +4526,14 @@ impl DAGCircuit {
             .unbind())
     }
 
+    /// Iterable over the declared stretch variables tracked by the circuit.
+    fn iter_declared_stretches(&self, py: Python) -> PyResult<Py<PyIterator>> {
+        Ok(PyList::new(py, self.declared_stretches.values())?
+            .into_any()
+            .try_iter()?
+            .unbind())
+    }
+
     /// Iterable over all the classical variables tracked by the circuit.
     fn iter_vars(&self, py: Python) -> PyResult<Py<PyIterator>> {
         let out_set = PySet::empty(py)?;
@@ -4361,6 +4545,18 @@ impl DAGCircuit {
         Ok(out_set.into_any().try_iter()?.unbind())
     }
 
+    /// Iterable over all the stretch variables tracked by the circuit.
+    fn iter_stretches(&self, py: Python) -> PyResult<Py<PyIterator>> {
+        let out_set = PySet::empty(py)?;
+        for s in self.captured_stretches.values() {
+            out_set.add(s)?;
+        }
+        for s in self.declared_stretches.values() {
+            out_set.add(s)?;
+        }
+        Ok(out_set.into_any().try_iter()?.unbind())
+    }
+
     fn _has_edge(&self, source: usize, target: usize) -> bool {
         self.dag
             .contains_edge(NodeIndex::new(source), NodeIndex::new(target))
@@ -5792,6 +5988,13 @@ impl DAGCircuit {
                 "cannot add var as its name shadows an existing var",
             ));
         }
+        if self.declared_stretches.contains_key(&var_name)
+            || self.captured_stretches.contains_key(&var_name)
+        {
+            return Err(DAGCircuitError::new_err(
+                "cannot add var as its name shadows an existing stretch",
+            ));
+        }
 
         let var_idx = self.vars.add(py, var, true)?;
         let (in_index, out_index) = self.add_wire(Wire::Var(var_idx))?;
@@ -5910,6 +6113,8 @@ impl DAGCircuit {
                 PySet::empty(py)?.unbind(),
                 PySet::empty(py)?.unbind(),
             ],
+            captured_stretches: IndexMap::default(),
+            declared_stretches: IndexMap::default(),
         })
     }
 
@@ -6406,6 +6611,24 @@ impl DAGCircuit {
             new_dag.add_var(py, var, DAGVarType::Capture)?;
         }
 
+        new_dag
+            .captured_stretches
+            .reserve(qc.captured_stretches.len());
+
+        for var in qc.captured_stretches {
+            let name: String = var.getattr("name")?.extract::<String>()?;
+            new_dag.captured_stretches.insert(name, var.unbind());
+        }
+
+        new_dag
+            .declared_stretches
+            .reserve(qc.declared_stretches.len());
+
+        for var in qc.declared_stretches {
+            let name: String = var.getattr("name")?.extract::<String>()?;
+            new_dag.declared_stretches.insert(name, var.unbind());
+        }
+
         // Add all the registers
         if let Some(qregs) = qc.qregs {
             for qreg in qregs.iter() {
@@ -6455,6 +6678,8 @@ impl DAGCircuit {
             input_vars: Vec::new(),
             captured_vars: Vec::new(),
             declared_vars: Vec::new(),
+            captured_stretches: Vec::new(),
+            declared_stretches: Vec::new(),
         };
         Self::from_circuit(py, circ, copy_op, None, None)
     }