From c5152d81f8548335f4424dea44f2889e09e9b7b8 Mon Sep 17 00:00:00 2001
From: Raynel Sanchez <raynelfss@hotmail.com>
Date: Sat, 17 Aug 2024 21:33:15 -0400
Subject: [PATCH] Initial: implement fixed capacity constructor for
 `DAGCircuit` - Implement `DAGCircuit` with `with_capacity` to create an
 initially allocated instance, for rust only. - Implement `with_capacity` for
 `BitData` and `IndexInterner` that creates an initally allocated instance of
 each type. - Other small tweaks and fixes.

---
 crates/circuit/src/bit_data.rs    |  9 ++++++
 crates/circuit/src/dag_circuit.rs | 52 +++++++++++++++++++++++++++++++
 crates/circuit/src/interner.rs    |  7 +++++
 3 files changed, 68 insertions(+)

diff --git a/crates/circuit/src/bit_data.rs b/crates/circuit/src/bit_data.rs
index 0c0b20a02522..6a8e4af6b920 100644
--- a/crates/circuit/src/bit_data.rs
+++ b/crates/circuit/src/bit_data.rs
@@ -95,6 +95,15 @@ where
         }
     }
 
+    pub fn with_capacity(py: Python<'_>, description: String, capacity: usize) -> Self {
+        BitData {
+            description,
+            bits: Vec::with_capacity(capacity),
+            indices: HashMap::with_capacity(capacity),
+            cached: PyList::empty_bound(py).unbind(),
+        }
+    }
+
     /// Gets the number of bits.
     pub fn len(&self) -> usize {
         self.bits.len()
diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs
index 4e72da5735ec..987d647f20fe 100644
--- a/crates/circuit/src/dag_circuit.rs
+++ b/crates/circuit/src/dag_circuit.rs
@@ -6112,6 +6112,58 @@ impl DAGCircuit {
         }
         Ok(())
     }
+
+    /// Alternative constructor, builds a DAGCircuit with a fixed capacity.
+    pub fn with_capacity(
+        py: Python,
+        num_qubits: usize,
+        num_clbits: usize,
+        num_ops: Option<usize>,
+        num_vars: Option<usize>,
+    ) -> PyResult<Self> {
+        let num_ops = num_ops.unwrap_or_default();
+        let num_vars = num_vars.unwrap_or_default();
+
+        let dag_nodes = num_qubits * 2 + // One input + One output node per qubit
+            num_clbits * 2 +    // One input + One output node per clbit
+            num_vars * 2 +  // One input + output node per variable
+            num_ops;
+
+        let dag_edges = num_qubits +    // 1 edge between the input node and the output node or 1st op node.
+            num_clbits +    // 1 edge between the input node and the output node or 1st op node.
+            num_vars +      // 1 edge between the input node and the output node or 1st op node.
+            num_ops * 3; // In Average there will be 3 edges (2 qubits and 1 clbit, or 3 qubits) per op_node.
+
+        Ok(Self {
+            name: None,
+            metadata: Some(PyDict::new_bound(py).unbind().into()),
+            calibrations: HashMap::with_capacity(0),
+            dag: StableDiGraph::with_capacity(dag_nodes, dag_edges),
+            qregs: PyDict::new_bound(py).unbind(),
+            cregs: PyDict::new_bound(py).unbind(),
+            qargs_cache: IndexedInterner::with_capacity(num_qubits),
+            cargs_cache: IndexedInterner::with_capacity(num_clbits),
+            qubits: BitData::with_capacity(py, "qubits".to_string(), num_qubits),
+            clbits: BitData::with_capacity(py, "clbits".to_string(), num_clbits),
+            global_phase: Param::Float(0.),
+            duration: None,
+            unit: "dt".to_string(),
+            qubit_locations: PyDict::new_bound(py).unbind(),
+            clbit_locations: PyDict::new_bound(py).unbind(),
+            qubit_io_map: IndexMap::with_capacity_and_hasher(num_qubits, RandomState::new()),
+            clbit_io_map: IndexMap::with_capacity_and_hasher(num_clbits, RandomState::new()),
+            var_input_map: _VarIndexMap::new(py),
+            var_output_map: _VarIndexMap::new(py),
+            op_names: IndexMap::with_capacity_and_hasher(num_ops, RandomState::new()), // Based on worst case scenario, could be improved
+            control_flow_module: PyControlFlowModule::new(py)?,
+            vars_info: HashMap::with_capacity(num_vars),
+            vars_by_type: [
+                PySet::empty_bound(py)?.unbind(),
+                PySet::empty_bound(py)?.unbind(),
+                PySet::empty_bound(py)?.unbind(),
+            ],
+        })
+    }
 }
 
 /// Add to global phase. Global phase can only be Float or ParameterExpression so this
diff --git a/crates/circuit/src/interner.rs b/crates/circuit/src/interner.rs
index e19f56e87a7d..3aaa9ef9e0a8 100644
--- a/crates/circuit/src/interner.rs
+++ b/crates/circuit/src/interner.rs
@@ -79,6 +79,13 @@ impl<T: Eq + Hash> IndexedInterner<T> {
             index_lookup: HashMap::new(),
         }
     }
+
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            entries: Vec::with_capacity(capacity),
+            index_lookup: HashMap::with_capacity(capacity),
+        }
+    }
 }
 
 impl<T: Eq + Hash> Default for IndexedInterner<T> {