support under_new_indices

qip/src/circuit.rs

@@ -18,14 +18,16 @@ impl<CB: CircuitBuilder, const N: usize> std::fmt::Debug for Circuit<CB, N> {
     }
 }
 
-impl<CB: CircuitBuilder + 'static, const N: usize> Circuit<CB, N> {
+impl<CB: CircuitBuilder, const N: usize> Default for Circuit<CB, N> {
     /// Init a empty circuit
-    pub fn new() -> Self {
+    fn default() -> Self {
         Self {
             func: Rc::new(|_, rs| Ok(rs)),
         }
     }
+}
 
+impl<CB: CircuitBuilder + 'static, const N: usize> Circuit<CB, N> {
     /// Apply a function to part of this circuit
     pub fn apply<F, const L: usize>(self, f: F, indices: [usize; L]) -> Self
     where
@@ -36,7 +38,7 @@ impl<CB: CircuitBuilder + 'static, const N: usize> Circuit<CB, N> {
             func: Rc::new(move |cb, rs| {
                 let out = (*func)(cb, rs)?;
                 let mut out = out.map(Some);
-                let f_input = indices.map(|idx| mem::take(&mut out[idx]).unwrap());
+                let f_input = indices.map(|idx| out[idx].take().unwrap());
                 let f_out = f(cb, f_input)?;
 
                 iter::zip(indices, f_out).for_each(|(idx, r)| out[idx] = Some(r));
@@ -46,6 +48,64 @@ impl<CB: CircuitBuilder + 'static, const N: usize> Circuit<CB, N> {
         }
     }
 
+    /// Apply a sub circuit for specific qubits under some new indices combine
+    pub fn under_new_indices<const L: usize>(
+        self,
+        new_indices: [&[(usize, usize)]; L],
+        sub_circuit: Circuit<CB, L>,
+    ) -> Self {
+        let func = self.func.clone();
+        let new_indices = new_indices.map(|idx_pairs| idx_pairs.to_vec());
+
+        Self {
+            func: Rc::new(move |cb, rs| {
+                let out = (*func)(cb, rs)?;
+                let mut out = out.map(RegisterRepr::Origin);
+
+                // combine to new  registers
+                let f_input = new_indices.clone().map(|idx_pairs| {
+                    let rs = idx_pairs
+                        .iter()
+                        .map(|&(reg_idx, idx)| {
+                            if matches!(out[reg_idx], RegisterRepr::Origin(_)) {
+                                let r =
+                                    mem::replace(out.get_mut(reg_idx).unwrap(), RegisterRepr::Tmp);
+                                let RegisterRepr::Origin(r) = r else {
+                                    unreachable!()
+                                };
+                                let qubits = cb.split_all_register(r);
+                                out[reg_idx] =
+                                    RegisterRepr::Splited(qubits.into_iter().map(Some).collect());
+                            }
+
+                            let RegisterRepr::Splited(rs) = out.get_mut(reg_idx).unwrap() else {
+                                unreachable!()
+                            };
+                            rs[idx].take().unwrap()
+                        })
+                        .collect::<Vec<_>>();
+
+                    cb.merge_registers(rs).unwrap()
+                });
+
+                let f_output = (*sub_circuit.func)(cb, f_input)?;
+                let f_output_qubits = f_output.map(|r| cb.split_all_register(r));
+
+                // restore
+                iter::zip(new_indices.clone(), f_output_qubits).for_each(|(idx_pairs, qubits)| {
+                    iter::zip(idx_pairs, qubits).for_each(|((reg_idx, idx), qubit)| {
+                        let RegisterRepr::Splited(rs) = out.get_mut(reg_idx).unwrap() else {
+                            unreachable!()
+                        };
+                        rs[idx] = Some(qubit);
+                    });
+                });
+
+                Ok(out.map(|rr| rr.into_origin(cb)))
+            }),
+        }
+    }
+
     /// Set input for circuit
     pub fn input(self, input: [CB::Register; N]) -> CircuitWithInput<CB, N> {
         CircuitWithInput {
@@ -55,6 +115,24 @@ impl<CB: CircuitBuilder + 'static, const N: usize> Circuit<CB, N> {
     }
 }
 
+enum RegisterRepr<CB: CircuitBuilder> {
+    Origin(CB::Register),
+    Splited(Vec<Option<CB::Register>>),
+    Tmp,
+}
+
+impl<CB: CircuitBuilder> RegisterRepr<CB> {
+    fn into_origin(self, cb: &mut CB) -> CB::Register {
+        match self {
+            Self::Origin(r) => r,
+            Self::Splited(qubits) => cb
+                .merge_registers(qubits.into_iter().map(|r| r.unwrap()))
+                .unwrap(),
+            Self::Tmp => unreachable!(),
+        }
+    }
+}
+
 /// Circuit with input
 #[derive(Debug)]
 pub struct CircuitWithInput<CB: CircuitBuilder, const N: usize> {
@@ -90,8 +168,14 @@ mod tests {
         let ra = b.try_register(3).unwrap();
         let rb = b.try_register(3).unwrap();
 
-        let [ra, rb] = Circuit::new()
+        let [ra, rb] = Circuit::default()
+            // Applies gamma to |ra>|rb>
             .apply(gamma, [0, 1])
+            // Applies gamma to |ra[0] ra[1]>|ra[2]>
+            .under_new_indices(
+                [&[(0, 0), (0, 1)], &[(0, 2)]],
+                Circuit::default().apply(gamma, [0, 1]),
+            )
             .input([ra, rb])
             .run(&mut b)?;