refactor shadowing code to take into account snapshot isolation

simulator/generation/mod.rs

@@ -1,6 +1,6 @@
 use sql_generation::generation::GenerationContext;
 
-use crate::runner::env::SimulatorTables;
+use crate::runner::env::ShadowTablesMut;
 
 pub mod plan;
 pub mod property;
@@ -17,7 +17,7 @@ pub mod query;
 /// might return a vector of rows that were inserted into the table.
 pub(crate) trait Shadow {
     type Result;
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result;
+    fn shadow(&self, tables: &mut ShadowTablesMut<'_>) -> Self::Result;
 }
 
 /// Generation context that will always panic when called

simulator/generation/plan.rs

@@ -27,7 +27,7 @@ use crate::{
     SimulatorEnv,
     generation::{PanicGenerationContext, Shadow},
     model::Query,
-    runner::env::{SimConnection, SimulationType, SimulatorTables},
+    runner::env::{ShadowTablesMut, SimConnection, SimulationType},
 };
 
 use super::property::{Property, remaining};
@@ -228,7 +228,7 @@ impl InteractionPlan {
             tracing::debug!("Generating interaction {}/{}", plan.len(), num_interactions);
             let interactions =
                 Interactions::arbitrary_from(rng, &PanicGenerationContext, (env, plan.stats()));
-            interactions.shadow(env.get_conn_tables_mut(interactions.connection_index));
+            interactions.shadow(&mut env.get_conn_tables_mut(interactions.connection_index));
             plan.push(interactions);
         }
 
@@ -311,7 +311,7 @@ pub enum InteractionsType {
 impl Shadow for Interactions {
     type Result = ();
 
-    fn shadow(&self, tables: &mut SimulatorTables) {
+    fn shadow(&self, tables: &mut ShadowTablesMut) {
         match &self.interactions {
             InteractionsType::Property(property) => {
                 let initial_tables = tables.clone();
@@ -319,7 +319,7 @@ impl Shadow for Interactions {
                     let res = interaction.shadow(tables);
                     if res.is_err() {
                         // If any interaction fails, we reset the tables to the initial state
-                        *tables = initial_tables.clone();
+                        **tables = initial_tables.clone();
                         break;
                     }
                 }
@@ -576,7 +576,7 @@ impl Display for InteractionType {
 
 impl Shadow for InteractionType {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
-    fn shadow(&self, env: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, env: &mut ShadowTablesMut) -> Self::Result {
         match self {
             Self::Query(query) => query.shadow(env),
             Self::FsyncQuery(query) => {

simulator/generation/property.rs

@@ -796,8 +796,8 @@ impl Property {
                         let last = stack.last().unwrap();
                         match last {
                             Ok(_) => {
-                                let _ =
-                                    query_clone.shadow(env.get_conn_tables_mut(connection_index));
+                                let _ = query_clone
+                                    .shadow(&mut env.get_conn_tables_mut(connection_index));
                                 Ok(Ok(()))
                             }
                             Err(err) => {
@@ -1040,7 +1040,8 @@ fn assert_all_table_values(
         let assertion = InteractionType::Assertion(Assertion::new(format!("table {table} should contain all of its expected values"), {
                 let table = table.clone();
                 move |stack: &Vec<ResultSet>, env: &mut SimulatorEnv| {
-                    let table = env.get_conn_tables(connection_index).iter().find(|t| t.name == table).ok_or_else(|| {
+                    let conn_ctx = env.get_conn_tables(connection_index);
+                    let table = conn_ctx.iter().find(|t| t.name == table).ok_or_else(|| {
                         LimboError::InternalError(format!(
                             "table {table} should exist in simulator env"
                         ))

simulator/model/mod.rs

@@ -15,7 +15,7 @@ use sql_generation::model::{
 };
 use turso_parser::ast::Distinctness;
 
-use crate::{generation::Shadow, runner::env::SimulatorTables};
+use crate::{generation::Shadow, runner::env::ShadowTablesMut};
 
 // This type represents the potential queries on the database.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -83,7 +83,7 @@ impl Display for Query {
 impl Shadow for Query {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, env: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, env: &mut ShadowTablesMut) -> Self::Result {
         match self {
             Query::Create(create) => create.shadow(env),
             Query::Insert(insert) => insert.shadow(env),
@@ -102,7 +102,7 @@ impl Shadow for Query {
 impl Shadow for Create {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
         if !tables.iter().any(|t| t.name == self.table.name) {
             tables.push(self.table.clone());
             Ok(vec![])
@@ -117,9 +117,8 @@ impl Shadow for Create {
 
 impl Shadow for CreateIndex {
     type Result = Vec<Vec<SimValue>>;
-    fn shadow(&self, env: &mut SimulatorTables) -> Vec<Vec<SimValue>> {
-        env.tables
-            .iter_mut()
+    fn shadow(&self, env: &mut ShadowTablesMut) -> Vec<Vec<SimValue>> {
+        env.iter_mut()
             .find(|t| t.name == self.table_name)
             .unwrap()
             .indexes
@@ -131,8 +130,8 @@ impl Shadow for CreateIndex {
 impl Shadow for Delete {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
-        let table = tables.tables.iter_mut().find(|t| t.name == self.table);
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
+        let table = tables.iter_mut().find(|t| t.name == self.table);
 
         if let Some(table) = table {
             // If the table exists, we can delete from it
@@ -153,7 +152,7 @@ impl Shadow for Delete {
 impl Shadow for Drop {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
         if !tables.iter().any(|t| t.name == self.table) {
             // If the table does not exist, we return an error
             return Err(anyhow::anyhow!(
@@ -162,7 +161,7 @@ impl Shadow for Drop {
             ));
         }
 
-        tables.tables.retain(|t| t.name != self.table);
+        tables.retain(|t| t.name != self.table);
 
         Ok(vec![])
     }
@@ -171,10 +170,10 @@ impl Shadow for Drop {
 impl Shadow for Insert {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
         match self {
             Insert::Values { table, values } => {
-                if let Some(t) = tables.tables.iter_mut().find(|t| &t.name == table) {
+                if let Some(t) = tables.iter_mut().find(|t| &t.name == table) {
                     t.rows.extend(values.clone());
                 } else {
                     return Err(anyhow::anyhow!(
@@ -185,7 +184,7 @@ impl Shadow for Insert {
             }
             Insert::Select { table, select } => {
                 let rows = select.shadow(tables)?;
-                if let Some(t) = tables.tables.iter_mut().find(|t| &t.name == table) {
+                if let Some(t) = tables.iter_mut().find(|t| &t.name == table) {
                     t.rows.extend(rows);
                 } else {
                     return Err(anyhow::anyhow!(
@@ -202,9 +201,7 @@ impl Shadow for Insert {
 
 impl Shadow for FromClause {
     type Result = anyhow::Result<JoinTable>;
-    fn shadow(&self, env: &mut SimulatorTables) -> Self::Result {
-        let tables = &mut env.tables;
-
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
         let first_table = tables
             .iter()
             .find(|t| t.name == self.table)
@@ -259,7 +256,7 @@ impl Shadow for FromClause {
 impl Shadow for SelectInner {
     type Result = anyhow::Result<JoinTable>;
 
-    fn shadow(&self, env: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, env: &mut ShadowTablesMut) -> Self::Result {
         if let Some(from) = &self.from {
             let mut join_table = from.shadow(env)?;
             let col_count = join_table.columns().count();
@@ -327,7 +324,7 @@ impl Shadow for SelectInner {
 impl Shadow for Select {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, env: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, env: &mut ShadowTablesMut) -> Self::Result {
         let first_result = self.body.select.shadow(env)?;
 
         let mut rows = first_result.rows;
@@ -357,28 +354,26 @@ impl Shadow for Select {
 
 impl Shadow for Begin {
     type Result = Vec<Vec<SimValue>>;
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
         // FIXME: currently the snapshot is taken eagerly
         // this is wrong for Deffered transactions
-        tables.snapshot = Some(tables.tables.clone());
+        tables.create_snapshot();
         vec![]
     }
 }
 
 impl Shadow for Commit {
     type Result = Vec<Vec<SimValue>>;
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
-        tables.snapshot = None;
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
+        tables.apply_snapshot();
         vec![]
     }
 }
 
 impl Shadow for Rollback {
     type Result = Vec<Vec<SimValue>>;
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
-        if let Some(tables_) = tables.snapshot.take() {
-            tables.tables = tables_;
-        }
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
+        tables.delete_snapshot();
         vec![]
     }
 }
@@ -386,8 +381,8 @@ impl Shadow for Rollback {
 impl Shadow for Update {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
-    fn shadow(&self, tables: &mut SimulatorTables) -> Self::Result {
-        let table = tables.tables.iter_mut().find(|t| t.name == self.table);
+    fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
+        let table = tables.iter_mut().find(|t| t.name == self.table);
 
         let table = if let Some(table) = table {
             table

simulator/runner/env.rs

@@ -1,6 +1,6 @@
 use std::fmt::Display;
 use std::mem;
-use std::ops::Deref;
+use std::ops::{Deref, DerefMut};
 use std::panic::UnwindSafe;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -32,6 +32,79 @@ pub(crate) enum SimulationPhase {
     Shrink,
 }
 
+#[derive(Debug)]
+pub struct ShadowTables<'a> {
+    commited_tables: &'a Vec<Table>,
+    transaction_tables: Option<&'a Vec<Table>>,
+}
+
+#[derive(Debug)]
+pub struct ShadowTablesMut<'a> {
+    commited_tables: &'a mut Vec<Table>,
+    transaction_tables: &'a mut Option<Vec<Table>>,
+}
+
+impl<'a> ShadowTables<'a> {
+    fn tables(&self) -> &'a Vec<Table> {
+        self.transaction_tables.map_or(self.commited_tables, |v| v)
+    }
+}
+
+impl<'a> Deref for ShadowTables<'a> {
+    type Target = Vec<Table>;
+
+    fn deref(&self) -> &Self::Target {
+        self.tables()
+    }
+}
+
+impl<'a, 'b> ShadowTablesMut<'a>
+where
+    'a: 'b,
+{
+    fn tables(&'a self) -> &'a Vec<Table> {
+        self.transaction_tables
+            .as_ref()
+            .unwrap_or(self.commited_tables)
+    }
+
+    fn tables_mut(&'b mut self) -> &'b mut Vec<Table> {
+        self.transaction_tables
+            .as_mut()
+            .unwrap_or(self.commited_tables)
+    }
+
+    pub fn create_snapshot(&mut self) {
+        *self.transaction_tables = Some(self.commited_tables.clone());
+    }
+
+    pub fn apply_snapshot(&mut self) {
+        // TODO: as we do not have concurrent tranasactions yet in the simulator
+        // there is no conflict we are ignoring conflict problems right now
+        if let Some(transation_tables) = self.transaction_tables.take() {
+            *self.commited_tables = transation_tables
+        }
+    }
+
+    pub fn delete_snapshot(&mut self) {
+        *self.transaction_tables = None;
+    }
+}
+
+impl<'a> Deref for ShadowTablesMut<'a> {
+    type Target = Vec<Table>;
+
+    fn deref(&self) -> &Self::Target {
+        self.tables()
+    }
+}
+
+impl<'a> DerefMut for ShadowTablesMut<'a> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.tables_mut()
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct SimulatorTables {
     pub(crate) tables: Vec<Table>,
@@ -75,9 +148,9 @@ pub(crate) struct SimulatorEnv {
     pub memory_io: bool,
 
     /// If connection state is None, means we are not in a transaction
-    pub connection_tables: Vec<Option<SimulatorTables>>,
+    pub connection_tables: Vec<Option<Vec<Table>>>,
     // Table data that is committed into the database or wal
-    pub committed_tables: SimulatorTables,
+    pub committed_tables: Vec<Table>,
 }
 
 impl UnwindSafe for SimulatorEnv {}
@@ -300,7 +373,7 @@ impl SimulatorEnv {
             phase: SimulationPhase::Test,
             memory_io: cli_opts.memory_io,
             profile: profile.clone(),
-            committed_tables: SimulatorTables::new(),
+            committed_tables: Vec::new(),
             connection_tables: vec![None; profile.max_connections],
         }
     }
@@ -363,7 +436,7 @@ impl SimulatorEnv {
             }
         }
 
-        let tables = &self.get_conn_tables(conn_index).tables;
+        let tables = self.get_conn_tables(conn_index).tables();
 
         ConnectionGenContext {
             opts: &self.profile.query.gen_opts,
@@ -371,20 +444,18 @@ impl SimulatorEnv {
         }
     }
 
-    pub fn get_conn_tables(&self, conn_index: usize) -> &SimulatorTables {
-        self.connection_tables
-            .get(conn_index)
-            .unwrap()
-            .as_ref()
-            .unwrap_or(&self.committed_tables)
+    pub fn get_conn_tables<'a>(&'a self, conn_index: usize) -> ShadowTables<'a> {
+        ShadowTables {
+            transaction_tables: self.connection_tables.get(conn_index).unwrap().as_ref(),
+            commited_tables: &self.committed_tables,
+        }
     }
 
-    pub fn get_conn_tables_mut(&mut self, conn_index: usize) -> &mut SimulatorTables {
-        self.connection_tables
-            .get_mut(conn_index)
-            .unwrap()
-            .as_mut()
-            .unwrap_or(&mut self.committed_tables)
+    pub fn get_conn_tables_mut<'a>(&'a mut self, conn_index: usize) -> ShadowTablesMut<'a> {
+        ShadowTablesMut {
+            transaction_tables: self.connection_tables.get_mut(conn_index).unwrap(),
+            commited_tables: &mut self.committed_tables,
+        }
     }
 }
 

simulator/runner/execution.rs

@@ -236,7 +236,7 @@ pub fn execute_interaction_turso(
             }
         }
     }
-    let _ = interaction.shadow(env.get_conn_tables_mut(interaction.connection_index));
+    let _ = interaction.shadow(&mut env.get_conn_tables_mut(interaction.connection_index));
     Ok(ExecutionContinuation::NextInteraction)
 }
 
@@ -329,7 +329,7 @@ fn execute_interaction_rusqlite(
         }
     }
 
-    let _ = interaction.shadow(env.get_conn_tables_mut(interaction.connection_index));
+    let _ = interaction.shadow(&mut env.get_conn_tables_mut(interaction.connection_index));
     Ok(ExecutionContinuation::NextInteraction)
 }