diff --git a/plonky2/examples/square_root.rs b/plonky2/examples/square_root.rs
index 3fce9030..07f1d230 100644
--- a/plonky2/examples/square_root.rs
+++ b/plonky2/examples/square_root.rs
@@ -1,9 +1,36 @@
+use std::marker::PhantomData;
+
 use anyhow::Result;
 use plonky2::field::types::Field;
-use plonky2::iop::witness::PartialWitness;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::generator::{SimpleGenerator, GeneratedValues};
+use plonky2::iop::target::Target;
+use plonky2::iop::witness::{PartialWitness, PartitionWitness, Witness};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::circuit_data::CircuitConfig;
 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+use plonky2_field::extension::Extendable;
+
+#[derive(Debug)]
+struct SquareGenerator<F: RichField + Extendable<D>, const D: usize> {
+    x: Target,
+    x_squared: Target,
+    _phantom: PhantomData<F>,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F> for SquareGenerator<F, D>
+{
+    fn dependencies(&self) -> Vec<Target> {
+        vec![self.x]
+    }
+
+    fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
+        let x = witness.get_target(self.x);
+        let x_squared = x * x;
+
+        out_buffer.set_target(self.x_squared, x_squared);
+    }
+}
 
 fn main() -> Result<()> {
     const D: usize = 2;
@@ -12,23 +39,29 @@ fn main() -> Result<()> {
 
     let config = CircuitConfig::standard_recursion_config();
 
-    let pw = PartialWitness::new();
     let mut builder = CircuitBuilder::<F, D>::new(config);
 
-    let x = F::rand();
-    let x_squared = x * x;
-    let x_target = builder.constant(x);
-    let x_squared_target = builder.constant(x_squared);
+    let x = builder.add_virtual_target();
+    let x_squared = builder.mul(x, x);
 
-    let x_squared_computed = builder.mul(x_target, x_target);
-    builder.connect(x_squared_target, x_squared_computed);
+    builder.register_public_input(x);
+    builder.register_public_input(x_squared);
 
-    builder.register_public_input(x_target);
+    builder.add_simple_generator(SquareGenerator::<F, D> {
+        x,
+        x_squared,
+        _phantom: PhantomData,
+    });
+
+    let x_value = F::rand();
+
+    let mut pw = PartialWitness::new();
+    pw.set_target(x, x_value);
 
     let data = builder.build::<C>();
     let proof = data.prove(pw)?;
 
-    println!("Random field element: {}", x_squared);
+    println!("Random field element: {}", proof.public_inputs[1]);
     println!("Its square root: {}", proof.public_inputs[0]);
 
     data.verify(proof)