[add]: composable runtime pipeline

3 files changed, 594 insertions, 35 deletions

diff --git a/libpsi-core/src/core/circuit.rs b/libpsi-core/src/core/circuit.rs
index e5d7847..d0e4450 100644
--- a/libpsi-core/src/core/circuit.rs
+++ b/libpsi-core/src/core/circuit.rs
@@ -1,4 +1,4 @@
-use super::{CustomGate, QuantumState, Runtime};
+use super::{CustomGate, QuantumState, Runtime, RuntimeConfig};
 use crate::{format_amplitude, format_probability, Vector};
 use core::fmt;
 use std::sync::Arc;
@@ -193,6 +193,15 @@ impl QuantumCircuit {
         self.computed_state.as_ref().unwrap()
     }
 
+    pub fn compute_with_config(&mut self, config: RuntimeConfig) -> &QuantumState {
+        if self.computed_state.is_some() {
+            return self.computed_state.as_ref().unwrap();
+        }
+
+        self.computed_state = Some(config.compute(self.num_qubits, &self.operations));
+        self.computed_state.as_ref().unwrap()
+    }
+
     pub fn state(&mut self) -> &QuantumState {
         self.compute()
     }
@@ -201,6 +210,10 @@ impl QuantumCircuit {
         self.compute_with(runtime)
     }
 
+    pub fn state_with_config(&mut self, config: RuntimeConfig) -> &QuantumState {
+        self.compute_with_config(config)
+    }
+
     pub fn h(&mut self, target: usize) -> &mut Self {
         self.operations.push(GateOp::H(target));
         self.computed_state = None;
diff --git a/libpsi-core/src/core/kernel.rs b/libpsi-core/src/core/kernel.rs
index 9f65d09..7b66eea 100644
--- a/libpsi-core/src/core/kernel.rs
+++ b/libpsi-core/src/core/kernel.rs
@@ -3,27 +3,88 @@ use crate::maths::simd::{
 };
 use crate::{complex, Complex, Matrix};
 use rayon::prelude::*;
+use std::collections::HashSet;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum GateType {
+    Diagonal,
+    NonDiagonal,
+    Controlled,
+}
 
 #[derive(Clone)]
 pub struct Kernel {
     pub matrix: Matrix<Complex<f64>>,
     pub targets: Vec<usize>,
     pub name: String,
+    pub gate_type: GateType,
 }
 
 impl Kernel {
     pub fn new(name: &str, matrix: Matrix<Complex<f64>>, targets: Vec<usize>) -> Self {
+        let gate_type = Self::detect_gate_type(name, &matrix);
         Self {
             matrix,
             targets,
             name: name.to_string(),
+            gate_type,
+        }
+    }
+
+    fn detect_gate_type(name: &str, matrix: &Matrix<Complex<f64>>) -> GateType {
+        let diagonal_gates = [
+            "Z", "S", "T", "Sdg", "Tdg", "Rz", "P", "U1", "CZ", "CP", "CRz",
+        ];
+        if diagonal_gates.iter().any(|&g| name.starts_with(g)) {
+            return GateType::Diagonal;
         }
+
+        let controlled_gates = [
+            "CNOT", "CZ", "SWAP", "CRx", "CRy", "CRz", "CP", "CCNOT", "CSWAP",
+        ];
+        if controlled_gates.iter().any(|&g| name.starts_with(g)) {
+            return GateType::Controlled;
+        }
+
+        if matrix.rows == 2 && matrix.cols == 2 {
+            let is_diag = matrix.data[1].real.abs() < 1e-10
+                && matrix.data[1].imaginary.abs() < 1e-10
+                && matrix.data[2].real.abs() < 1e-10
+                && matrix.data[2].imaginary.abs() < 1e-10;
+            if is_diag {
+                return GateType::Diagonal;
+            }
+        }
+
+        GateType::NonDiagonal
     }
 
     pub fn num_qubits(&self) -> usize {
         self.targets.len()
     }
 
+    pub fn target_set(&self) -> HashSet<usize> {
+        self.targets.iter().cloned().collect()
+    }
+
+    pub fn shares_qubits(&self, other: &Kernel) -> bool {
+        self.targets.iter().any(|t| other.targets.contains(t))
+    }
+
+    pub fn commutes_with(&self, other: &Kernel) -> bool {
+        if !self.shares_qubits(other) {
+            return true;
+        }
+
+        if self.gate_type == GateType::Diagonal && other.gate_type == GateType::Diagonal {
+            if self.targets == other.targets {
+                return true;
+            }
+        }
+
+        false
+    }
+
     pub fn can_fuse_with(&self, other: &Kernel) -> bool {
         if self.targets.len() != 1 || other.targets.len() != 1 {
             return false;
@@ -36,10 +97,17 @@ impl Kernel {
             return None;
         }
         let fused_matrix = other.matrix.dot(&self.matrix)?;
+        let new_type =
+            if self.gate_type == GateType::Diagonal && other.gate_type == GateType::Diagonal {
+                GateType::Diagonal
+            } else {
+                GateType::NonDiagonal
+            };
         Some(Kernel {
             matrix: fused_matrix,
             targets: self.targets.clone(),
             name: format!("{}+{}", self.name, other.name),
+            gate_type: new_type,
         })
     }
 }
@@ -264,3 +332,291 @@ impl KernelBuilder {
         self.num_qubits
     }
 }
+
+#[derive(Clone)]
+pub struct ExecutionLayer {
+    pub kernels: Vec<Kernel>,
+}
+
+impl ExecutionLayer {
+    pub fn new() -> Self {
+        Self {
+            kernels: Vec::new(),
+        }
+    }
+
+    pub fn can_add(&self, kernel: &Kernel) -> bool {
+        !self.kernels.iter().any(|k| k.shares_qubits(kernel))
+    }
+
+    pub fn add(&mut self, kernel: Kernel) {
+        self.kernels.push(kernel);
+    }
+
+    pub fn affected_qubits(&self) -> HashSet<usize> {
+        self.kernels
+            .iter()
+            .flat_map(|k| k.targets.iter().cloned())
+            .collect()
+    }
+}
+
+impl Default for ExecutionLayer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pub struct StructureAwareKernelBatch {
+    kernels: Vec<Kernel>,
+    layers: Vec<ExecutionLayer>,
+    num_qubits: usize,
+    optimised: bool,
+}
+
+impl StructureAwareKernelBatch {
+    pub fn new(num_qubits: usize) -> Self {
+        Self {
+            kernels: Vec::new(),
+            layers: Vec::new(),
+            num_qubits,
+            optimised: false,
+        }
+    }
+
+    pub fn add(&mut self, kernel: Kernel) {
+        self.kernels.push(kernel);
+        self.optimised = false;
+    }
+
+    pub fn len(&self) -> usize {
+        self.kernels.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.kernels.is_empty()
+    }
+
+    pub fn kernels(&self) -> &[Kernel] {
+        &self.kernels
+    }
+
+    pub fn layers(&self) -> &[ExecutionLayer] {
+        &self.layers
+    }
+
+    pub fn num_layers(&self) -> usize {
+        self.layers.len()
+    }
+
+    pub fn optimise(&mut self) {
+        if self.optimised || self.kernels.len() < 2 {
+            return;
+        }
+
+        self.reorder_commuting_gates();
+        self.multi_pass_fusion();
+        self.build_execution_layers();
+        self.optimised = true;
+    }
+
+    fn reorder_commuting_gates(&mut self) {
+        let mut changed = true;
+        let mut iterations = 0;
+        const MAX_ITERATIONS: usize = 100;
+
+        while changed && iterations < MAX_ITERATIONS {
+            changed = false;
+            iterations += 1;
+
+            for i in 0..self.kernels.len().saturating_sub(1) {
+                let current = &self.kernels[i];
+                let next = &self.kernels[i + 1];
+
+                if current.targets.len() == 1
+                    && next.targets.len() == 1
+                    && current.targets[0] != next.targets[0]
+                    && current.commutes_with(next)
+                {
+                    for j in (i + 2)..self.kernels.len() {
+                        let candidate = &self.kernels[j];
+
+                        if candidate.targets.len() == 1
+                            && candidate.targets[0] == current.targets[0]
+                        {
+                            let can_move = (i + 1..j).all(|k| {
+                                let between = &self.kernels[k];
+                                !between.shares_qubits(current) || current.commutes_with(between)
+                            });
+
+                            if can_move && current.can_fuse_with(candidate) {
+                                let kernel_to_move = self.kernels.remove(j);
+                                self.kernels.insert(i + 1, kernel_to_move);
+                                changed = true;
+                                break;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    fn multi_pass_fusion(&mut self) {
+        let mut changed = true;
+        let mut iterations = 0;
+        const MAX_ITERATIONS: usize = 50;
+
+        while changed && iterations < MAX_ITERATIONS {
+            changed = false;
+            iterations += 1;
+
+            let mut new_kernels: Vec<Kernel> = Vec::with_capacity(self.kernels.len());
+            let mut i = 0;
+
+            while i < self.kernels.len() {
+                if i + 1 < self.kernels.len() {
+                    let current = &self.kernels[i];
+                    let next = &self.kernels[i + 1];
+
+                    if let Some(fused) = current.fuse(next) {
+                        new_kernels.push(fused);
+                        i += 2;
+                        changed = true;
+                        continue;
+                    }
+                }
+
+                new_kernels.push(self.kernels[i].clone());
+                i += 1;
+            }
+
+            self.kernels = new_kernels;
+        }
+    }
+
+    fn build_execution_layers(&mut self) {
+        self.layers.clear();
+
+        for kernel in &self.kernels {
+            let mut placed = false;
+
+            for layer in &mut self.layers {
+                if layer.can_add(kernel) {
+                    layer.add(kernel.clone());
+                    placed = true;
+                    break;
+                }
+            }
+
+            if !placed {
+                let mut new_layer = ExecutionLayer::new();
+                new_layer.add(kernel.clone());
+                self.layers.push(new_layer);
+            }
+        }
+    }
+
+    pub fn execute(&self, state: &mut Vec<Complex<f64>>) {
+        for kernel in &self.kernels {
+            *state = apply_kernel(state, kernel, self.num_qubits);
+        }
+    }
+
+    pub fn execute_parallel(&self, state: &mut Vec<Complex<f64>>) {
+        for kernel in &self.kernels {
+            *state = apply_kernel_parallel(state, kernel, self.num_qubits);
+        }
+    }
+
+    pub fn execute_layered(&self, state: &mut Vec<Complex<f64>>) {
+        for layer in &self.layers {
+            for kernel in &layer.kernels {
+                *state = apply_kernel(state, kernel, self.num_qubits);
+            }
+        }
+    }
+
+    pub fn execute_layered_parallel(&self, state: &mut Vec<Complex<f64>>) {
+        for layer in &self.layers {
+            for kernel in &layer.kernels {
+                *state = apply_kernel_parallel(state, kernel, self.num_qubits);
+            }
+        }
+    }
+
+    pub fn execute_simd(&self, state: &mut Vec<Complex<f64>>) {
+        for kernel in &self.kernels {
+            if kernel.targets.len() == 1 {
+                let gate = matrix_to_2x2(&kernel.matrix);
+                apply_single_qubit_gate_simd(state, &gate, kernel.targets[0], self.num_qubits);
+            } else {
+                *state = apply_kernel(state, kernel, self.num_qubits);
+            }
+        }
+    }
+
+    pub fn execute_simd_parallel(&self, state: &mut Vec<Complex<f64>>) {
+        for kernel in &self.kernels {
+            if kernel.targets.len() == 1 && self.num_qubits >= 10 {
+                let gate = matrix_to_2x2(&kernel.matrix);
+                apply_single_qubit_gate_simd_parallel(
+                    state,
+                    &gate,
+                    kernel.targets[0],
+                    self.num_qubits,
+                );
+            } else if kernel.targets.len() == 1 {
+                let gate = matrix_to_2x2(&kernel.matrix);
+                apply_single_qubit_gate_simd(state, &gate, kernel.targets[0], self.num_qubits);
+            } else {
+                *state = apply_kernel_parallel(state, kernel, self.num_qubits);
+            }
+        }
+    }
+
+    pub fn stats(&self) -> KernelStats {
+        let single_qubit = self.kernels.iter().filter(|k| k.targets.len() == 1).count();
+        let two_qubit = self.kernels.iter().filter(|k| k.targets.len() == 2).count();
+        let multi_qubit = self.kernels.iter().filter(|k| k.targets.len() > 2).count();
+        let diagonal = self
+            .kernels
+            .iter()
+            .filter(|k| k.gate_type == GateType::Diagonal)
+            .count();
+
+        KernelStats {
+            total_kernels: self.kernels.len(),
+            single_qubit,
+            two_qubit,
+            multi_qubit,
+            diagonal,
+            execution_layers: self.layers.len(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct KernelStats {
+    pub total_kernels: usize,
+    pub single_qubit: usize,
+    pub two_qubit: usize,
+    pub multi_qubit: usize,
+    pub diagonal: usize,
+    pub execution_layers: usize,
+}
+
+impl std::fmt::Display for KernelStats {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "Kernels: {} (1q: {}, 2q: {}, 3q+: {}, diag: {}), Layers: {}",
+            self.total_kernels,
+            self.single_qubit,
+            self.two_qubit,
+            self.multi_qubit,
+            self.diagonal,
+            self.execution_layers
+        )
+    }
+}
diff --git a/libpsi-core/src/core/runtime.rs b/libpsi-core/src/core/runtime.rs
index a75435e..382fa2f 100644
--- a/libpsi-core/src/core/runtime.rs
+++ b/libpsi-core/src/core/runtime.rs
@@ -1,9 +1,13 @@
-use super::{GateOp, Kernel, KernelBatch, QuantumGate, QuantumRegister, QuantumState};
+use super::{
+    GateOp, Kernel, KernelBatch, QuantumGate, QuantumRegister, QuantumState,
+    StructureAwareKernelBatch,
+};
 use crate::gates::{
     cp_matrix, crx_matrix, cry_matrix, crz_matrix, p_matrix, rx_matrix, ry_matrix, rz_matrix,
     u1_matrix, u2_matrix, u3_matrix, CNOT, CZ, FREDKIN, HADAMARD, PAULI_X, PAULI_Y, PAULI_Z,
     SDG_GATE, SWAP, SXDG_GATE, SX_GATE, S_GATE, TDG_GATE, TOFFOLI, T_GATE,
 };
+use crate::maths::simd::{apply_single_qubit_gate_simd, apply_single_qubit_gate_simd_parallel};
 use crate::maths::vector::Vector;
 use crate::{complex, Complex, Matrix};
 use rayon::prelude::*;
@@ -11,6 +15,129 @@ use rayon::prelude::*;
 const PARALLEL_THRESHOLD: usize = 8;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub struct RuntimeConfig {
+    pub parallel: bool,
+    pub simd: bool,
+    pub batched: bool,
+    pub structure_aware: bool,
+    pub parallel_threshold: usize,
+}
+
+impl RuntimeConfig {
+    pub fn new() -> Self {
+        Self {
+            parallel: false,
+            simd: false,
+            batched: false,
+            structure_aware: false,
+            parallel_threshold: PARALLEL_THRESHOLD,
+        }
+    }
+
+    pub fn parallel(mut self) -> Self {
+        self.parallel = true;
+        self
+    }
+
+    pub fn simd(mut self) -> Self {
+        self.simd = true;
+        self
+    }
+
+    pub fn batched(mut self) -> Self {
+        self.batched = true;
+        self
+    }
+
+    pub fn structure_aware(mut self) -> Self {
+        self.structure_aware = true;
+        self
+    }
+
+    pub fn with_threshold(mut self, threshold: usize) -> Self {
+        self.parallel_threshold = threshold;
+        self
+    }
+
+    pub fn optimal() -> Self {
+        Self::new().structure_aware().simd().parallel()
+    }
+
+    pub fn compute(&self, num_qubits: usize, operations: &[GateOp]) -> QuantumState {
+        let dim = 1 << num_qubits;
+        let mut state: Vec<Complex<f64>> = vec![complex!(0.0, 0.0); dim];
+        state[0] = complex!(1.0, 0.0);
+
+        let use_parallel = self.parallel && num_qubits >= self.parallel_threshold;
+
+        if self.structure_aware {
+            let mut batch = Runtime::build_structure_aware_batch(num_qubits, operations);
+            batch.optimise();
+            self.execute_kernels(&mut state, batch.kernels(), num_qubits, use_parallel);
+        } else if self.batched {
+            let mut batch = Runtime::build_kernel_batch(num_qubits, operations);
+            batch.optimize();
+            self.execute_kernels(&mut state, batch.kernels(), num_qubits, use_parallel);
+        } else {
+            let batch = Runtime::build_kernel_batch(num_qubits, operations);
+            self.execute_kernels(&mut state, batch.kernels(), num_qubits, use_parallel);
+        }
+
+        QuantumState::new(state)
+    }
+
+    fn execute_kernels(
+        &self,
+        state: &mut Vec<Complex<f64>>,
+        kernels: &[Kernel],
+        num_qubits: usize,
+        use_parallel: bool,
+    ) {
+        for kernel in kernels {
+            if self.simd && kernel.targets.len() == 1 {
+                let gate = matrix_to_2x2(&kernel.matrix);
+                if use_parallel {
+                    apply_single_qubit_gate_simd_parallel(
+                        state,
+                        &gate,
+                        kernel.targets[0],
+                        num_qubits,
+                    );
+                } else {
+                    apply_single_qubit_gate_simd(state, &gate, kernel.targets[0], num_qubits);
+                }
+            } else if use_parallel {
+                *state = apply_gate_parallel(state, &kernel.matrix, &kernel.targets, num_qubits);
+            } else {
+                *state = apply_kernel_direct(state, kernel, num_qubits);
+            }
+        }
+    }
+}
+
+impl std::fmt::Display for RuntimeConfig {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut features = Vec::new();
+        if self.structure_aware {
+            features.push("structure-aware");
+        }
+        if self.batched && !self.structure_aware {
+            features.push("batched");
+        }
+        if self.simd {
+            features.push("SIMD");
+        }
+        if self.parallel {
+            features.push("parallel");
+        }
+        if features.is_empty() {
+            features.push("basic");
+        }
+        write!(f, "Runtime[{}]", features.join("+"))
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 pub enum Runtime {
     #[default]
     BasicRT,
@@ -19,20 +146,43 @@ pub enum Runtime {
     BatchedRTMT,
     SimdRT,
     SimdRTMT,
+    StructureAwareRT,
+    StructureAwareMT,
     WFEvolution,
     WFEvolutionMT,
     GPUAccelerated,
+    Custom(RuntimeConfig),
 }
 
 impl Runtime {
+    pub fn custom() -> RuntimeConfig {
+        RuntimeConfig::new()
+    }
+
+    pub fn optimal() -> RuntimeConfig {
+        RuntimeConfig::optimal()
+    }
+
+    pub fn to_config(&self) -> RuntimeConfig {
+        match self {
+            Runtime::BasicRT => RuntimeConfig::new(),
+            Runtime::BasicRTMT => RuntimeConfig::new().parallel(),
+            Runtime::BatchedRT => RuntimeConfig::new().batched(),
+            Runtime::BatchedRTMT => RuntimeConfig::new().batched().parallel(),
+            Runtime::SimdRT => RuntimeConfig::new().batched().simd(),
+            Runtime::SimdRTMT => RuntimeConfig::new().batched().simd().parallel(),
+            Runtime::StructureAwareRT => RuntimeConfig::new().structure_aware().simd(),
+            Runtime::StructureAwareMT => RuntimeConfig::new().structure_aware().simd().parallel(),
+            Runtime::Custom(config) => *config,
+            _ => RuntimeConfig::new(),
+        }
+    }
+
     pub fn compute(&self, num_qubits: usize, operations: &[GateOp]) -> QuantumState {
         match self {
             Runtime::BasicRT => Self::compute_basic(num_qubits, operations),
             Runtime::BasicRTMT => Self::compute_basic_mt(num_qubits, operations),
-            Runtime::BatchedRT => Self::compute_batched(num_qubits, operations, false),
-            Runtime::BatchedRTMT => Self::compute_batched(num_qubits, operations, true),
-            Runtime::SimdRT => Self::compute_simd(num_qubits, operations, false),
-            Runtime::SimdRTMT => Self::compute_simd(num_qubits, operations, true),
+            Runtime::Custom(config) => config.compute(num_qubits, operations),
             Runtime::WFEvolution => {
                 unimplemented!("WFEvolution (Schrödinger equation) runtime not yet implemented")
             }
@@ -44,6 +194,7 @@ impl Runtime {
             Runtime::GPUAccelerated => {
                 unimplemented!("GPUAccelerated runtime not yet implemented")
             }
+            _ => self.to_config().compute(num_qubits, operations),
         }
     }
 
@@ -97,38 +248,19 @@ impl Runtime {
         Some(Kernel::new(name, matrix, targets))
     }
 
-    fn compute_batched(num_qubits: usize, operations: &[GateOp], parallel: bool) -> QuantumState {
-        let dim = 1 << num_qubits;
-        let mut state: Vec<Complex<f64>> = vec![complex!(0.0, 0.0); dim];
-        state[0] = complex!(1.0, 0.0);
-
-        let mut batch = Self::build_kernel_batch(num_qubits, operations);
-        batch.optimize();
-
-        if parallel && num_qubits >= PARALLEL_THRESHOLD {
-            batch.execute_parallel(&mut state);
-        } else {
-            batch.execute(&mut state);
-        }
-
-        QuantumState::new(state)
-    }
-
-    fn compute_simd(num_qubits: usize, operations: &[GateOp], parallel: bool) -> QuantumState {
-        let dim = 1 << num_qubits;
-        let mut state: Vec<Complex<f64>> = vec![complex!(0.0, 0.0); dim];
-        state[0] = complex!(1.0, 0.0);
-
-        let mut batch = Self::build_kernel_batch(num_qubits, operations);
-        batch.optimize();
+    pub fn build_structure_aware_batch(
+        num_qubits: usize,
+        operations: &[GateOp],
+    ) -> StructureAwareKernelBatch {
+        let mut batch = StructureAwareKernelBatch::new(num_qubits);
 
-        if parallel && num_qubits >= PARALLEL_THRESHOLD {
-            batch.execute_simd_parallel(&mut state);
-        } else {
-            batch.execute_simd(&mut state);
+        for op in operations {
+            if let Some(kernel) = Self::op_to_kernel(op) {
+                batch.add(kernel);
+            }
         }
 
-        QuantumState::new(state)
+        batch
     }
 
     fn compute_basic(num_qubits: usize, operations: &[GateOp]) -> QuantumState {
@@ -393,3 +525,61 @@ fn apply_gate_parallel(
 
     new_state
 }
+
+fn matrix_to_2x2(matrix: &Matrix<Complex<f64>>) -> [[Complex<f64>; 2]; 2] {
+    [
+        [matrix.data[0], matrix.data[1]],
+        [matrix.data[2], matrix.data[3]],
+    ]
+}
+
+fn apply_kernel_direct(
+    state: &[Complex<f64>],
+    kernel: &Kernel,
+    num_qubits: usize,
+) -> Vec<Complex<f64>> {
+    let dim = 1 << num_qubits;
+    let g = kernel.targets.len();
+    let gate_dim = 1 << g;
+
+    let target_bits: Vec<usize> = kernel.targets.iter().map(|&t| num_qubits - 1 - t).collect();
+
+    let mut non_target_mask: usize = (1 << num_qubits) - 1;
+    for &pos in &target_bits {
+        non_target_mask &= !(1 << pos);
+    }
+
+    let mut new_state = vec![complex!(0.0, 0.0); dim];
+
+    for i in 0..dim {
+        let mut target_idx = 0usize;
+        for (k, &pos) in target_bits.iter().enumerate() {
+            if (i >> pos) & 1 == 1 {
+                target_idx |= 1 << (g - 1 - k);
+            }
+        }
+
+        let mut sum = complex!(0.0, 0.0);
+
+        for j in 0..gate_dim {
+            let gate_elem = kernel.matrix.data[target_idx * gate_dim + j];
+
+            if gate_elem.real.abs() < 1e-15 && gate_elem.imaginary.abs() < 1e-15 {
+                continue;
+            }
+
+            let mut source_idx = i & non_target_mask;
+            for (k, &pos) in target_bits.iter().enumerate() {
+                if (j >> (g - 1 - k)) & 1 == 1 {
+                    source_idx |= 1 << pos;
+                }
+            }
+
+            sum = sum + gate_elem * state[source_idx];
+        }
+
+        new_state[i] = sum;
+    }
+
+    new_state
+}