diff --git a/src/gates/gmimc.rs b/src/gates/gmimc.rs
index 1cb823bf..682bd1fa 100644
--- a/src/gates/gmimc.rs
+++ b/src/gates/gmimc.rs
@@ -364,7 +364,7 @@ mod tests {
             );
         }
 
-        let mut partition_witness = PartitionWitness::new(gate.num_wires(), gate.num_wires(), 1);
+        let mut partition_witness = PartitionWitness::new(gate.num_wires(), gate.num_wires(), 1, 0);
         for input in 0..gate.num_wires() {
             partition_witness.add(Target::Wire(Wire { gate: 0, input }));
         }
diff --git a/src/iop/witness.rs b/src/iop/witness.rs
index 26b15a24..f1aca8a1 100644
--- a/src/iop/witness.rs
+++ b/src/iop/witness.rs
@@ -196,7 +196,8 @@ pub struct PartitionWitness<F: Field> {
 
 impl<F: Field> Witness<F> for PartitionWitness<F> {
     fn try_get_target(&self, target: Target) -> Option<F> {
-        self.forest[self.forest[self.target_index(target)].parent].value
+        let parent_index = self.forest[self.target_index(target)].parent;
+        self.forest[parent_index].value
     }
 
     fn set_target(&mut self, target: Target, value: F) {
@@ -215,7 +216,7 @@ impl<F: Field> Witness<F> for PartitionWitness<F> {
 }
 
 impl<F: Field> PartitionWitness<F> {
-    pub const fn target_index(&self, target: Target) -> usize {
+    pub fn target_index(&self, target: Target) -> usize {
         match target {
             Target::Wire(Wire { gate, input }) => gate * self.num_wires + input,
             Target::VirtualTarget { index } => self.degree * self.num_wires + index,
diff --git a/src/lib.rs b/src/lib.rs
index 56fd179f..1249e9db 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,4 @@
 #![feature(destructuring_assignment)]
-#![feature(const_fn_trait_bound)]
 
 pub mod field;
 pub mod fri;
diff --git a/src/plonk/circuit_builder.rs b/src/plonk/circuit_builder.rs
index ac6196bf..56bee0a5 100644
--- a/src/plonk/circuit_builder.rs
+++ b/src/plonk/circuit_builder.rs
@@ -506,8 +506,12 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     ) -> (Vec<PolynomialValues<F>>, PartitionWitness<F>) {
         let degree = self.gate_instances.len();
         let degree_log = log2_strict(degree);
-        let mut partition_witness =
-            PartitionWitness::new(self.config.num_wires, self.config.num_routed_wires, degree);
+        let mut partition_witness = PartitionWitness::new(
+            self.config.num_wires,
+            self.config.num_routed_wires,
+            degree,
+            self.virtual_target_index,
+        );
 
         for gate in 0..degree {
             for input in 0..self.config.num_wires {
diff --git a/src/plonk/permutation_argument.rs b/src/plonk/permutation_argument.rs
index dbb6bee9..384ac9e1 100644
--- a/src/plonk/permutation_argument.rs
+++ b/src/plonk/permutation_argument.rs
@@ -21,9 +21,14 @@ pub struct ForestNode<T: Debug + Copy + Eq + PartialEq, V: Field> {
 
 /// Disjoint Set Forest data-structure following https://en.wikipedia.org/wiki/Disjoint-set_data_structure.
 impl<F: Field> PartitionWitness<F> {
-    pub fn new(num_wires: usize, num_routed_wires: usize, degree: usize) -> Self {
+    pub fn new(
+        num_wires: usize,
+        num_routed_wires: usize,
+        degree: usize,
+        num_virtual_targets: usize,
+    ) -> Self {
         Self {
-            forest: vec![],
+            forest: Vec::with_capacity(degree * num_wires + num_virtual_targets),
             num_wires,
             num_routed_wires,
             degree,