From 6e2b0c901eeec4046d7a56393a49933c4e34a279 Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Sun, 21 Sep 2025 14:02:23 -0300
Subject: [PATCH] remove PanicGenerationContext and instead just pass the
 connection context directly

---
 simulator/generation/mod.rs      | 16 ----------------
 simulator/generation/plan.rs     | 21 ++++++++++++---------
 simulator/generation/property.rs |  7 +++----
 3 files changed, 15 insertions(+), 29 deletions(-)

diff --git a/simulator/generation/mod.rs b/simulator/generation/mod.rs
index d2ccf4a00..80a2d0cff 100644
--- a/simulator/generation/mod.rs
+++ b/simulator/generation/mod.rs
@@ -1,5 +1,3 @@
-use sql_generation::generation::GenerationContext;
-
 use crate::runner::env::ShadowTablesMut;
 
 pub mod plan;
@@ -19,17 +17,3 @@ pub(crate) trait Shadow {
     type Result;
     fn shadow(&self, tables: &mut ShadowTablesMut<'_>) -> Self::Result;
 }
-
-/// Generation context that will always panic when called
-/// This is meant to be used when want to ensure that no downstream arbitrary fn will use this context
-pub struct PanicGenerationContext;
-
-impl GenerationContext for PanicGenerationContext {
-    fn tables(&self) -> &Vec<sql_generation::model::table::Table> {
-        unimplemented!("you are not supposed to use this context")
-    }
-
-    fn opts(&self) -> &sql_generation::generation::Opts {
-        unimplemented!("you are not supposed to use this context")
-    }
-}
diff --git a/simulator/generation/plan.rs b/simulator/generation/plan.rs
index 8b8f97aa8..64a6574b8 100644
--- a/simulator/generation/plan.rs
+++ b/simulator/generation/plan.rs
@@ -25,7 +25,7 @@ use turso_core::{Connection, Result, StepResult};
 
 use crate::{
     SimulatorEnv,
-    generation::{PanicGenerationContext, Shadow},
+    generation::Shadow,
     model::Query,
     runner::env::{ShadowTablesMut, SimConnection, SimulationType},
 };
@@ -226,8 +226,12 @@ impl InteractionPlan {
 
         while plan.len() < num_interactions {
             tracing::debug!("Generating interaction {}/{}", plan.len(), num_interactions);
-            let interactions =
-                Interactions::arbitrary_from(rng, &PanicGenerationContext, (env, plan.stats()));
+            let interactions = {
+                let conn_index = env.choose_conn(rng);
+                let conn_ctx = &env.connection_context(conn_index);
+                Interactions::arbitrary_from(rng, conn_ctx, (env, plan.stats(), conn_index))
+            };
+
             interactions.shadow(&mut env.get_conn_tables_mut(interactions.connection_index));
             plan.push(interactions);
         }
@@ -1000,11 +1004,11 @@ fn random_fault<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv) -> Interactions {
     Interactions::new(env.choose_conn(rng), InteractionsType::Fault(fault))
 }
 
-impl ArbitraryFrom<(&SimulatorEnv, InteractionStats)> for Interactions {
+impl ArbitraryFrom<(&SimulatorEnv, InteractionStats, usize)> for Interactions {
     fn arbitrary_from<R: rand::Rng, C: GenerationContext>(
         rng: &mut R,
-        _context: &C,
-        (env, stats): (&SimulatorEnv, InteractionStats),
+        conn_ctx: &C,
+        (env, stats, conn_index): (&SimulatorEnv, InteractionStats, usize),
     ) -> Self {
         let remaining_ = remaining(
             env.opts.max_interactions,
@@ -1012,7 +1016,6 @@ impl ArbitraryFrom<(&SimulatorEnv, InteractionStats)> for Interactions {
             &stats,
             env.profile.experimental_mvcc,
         );
-        let conn_index = env.choose_conn(rng);
         frequency(
             vec![
                 (
@@ -1022,8 +1025,8 @@ impl ArbitraryFrom<(&SimulatorEnv, InteractionStats)> for Interactions {
                             conn_index,
                             InteractionsType::Property(Property::arbitrary_from(
                                 rng,
-                                &PanicGenerationContext,
-                                (env, &stats, conn_index),
+                                conn_ctx,
+                                (env, &stats),
                             )),
                         )
                     }),
diff --git a/simulator/generation/property.rs b/simulator/generation/property.rs
index bb4678872..847c30593 100644
--- a/simulator/generation/property.rs
+++ b/simulator/generation/property.rs
@@ -1514,13 +1514,12 @@ fn property_faulty_query<R: rand::Rng>(
     }
 }
 
-impl ArbitraryFrom<(&SimulatorEnv, &InteractionStats, usize)> for Property {
+impl ArbitraryFrom<(&SimulatorEnv, &InteractionStats)> for Property {
     fn arbitrary_from<R: rand::Rng, C: GenerationContext>(
         rng: &mut R,
-        _context: &C,
-        (env, stats, conn_index): (&SimulatorEnv, &InteractionStats, usize),
+        conn_ctx: &C,
+        (env, stats): (&SimulatorEnv, &InteractionStats),
     ) -> Self {
-        let conn_ctx = &env.connection_context(conn_index);
         let opts = conn_ctx.opts();
         let remaining_ = remaining(
             env.opts.max_interactions,