rework interaction generation to only generate possible queries + do less allocations

simulator/generation/plan.rs

@@ -8,6 +8,7 @@ use std::{
 };
 
 use indexmap::IndexSet;
+use rand::distr::weighted::WeightedIndex;
 use serde::{Deserialize, Serialize};
 
 use sql_generation::{
@@ -26,7 +27,7 @@ use turso_core::{Connection, Result, StepResult};
 
 use crate::{
     SimulatorEnv,
-    generation::Shadow,
+    generation::{Shadow, property::possiple_properties, query::possible_queries},
     model::Query,
     runner::env::{ShadowTablesMut, SimConnection, SimulationType},
 };
@@ -1077,26 +1078,27 @@ fn random_create<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usiz
     Interactions::new(conn_index, InteractionsType::Query(Query::Create(create)))
 }
 
-fn random_read<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usize) -> Interactions {
-    Interactions::new(
-        conn_index,
-        InteractionsType::Query(Query::Select(Select::arbitrary(
-            rng,
-            &env.connection_context(conn_index),
-        ))),
-    )
+fn random_select<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usize) -> Interactions {
+    if rng.random_bool(0.7) {
+        Interactions::new(
+            conn_index,
+            InteractionsType::Query(Query::Select(Select::arbitrary(
+                rng,
+                &env.connection_context(conn_index),
+            ))),
+        )
+    } else {
+        // Random expression
+        Interactions::new(
+            conn_index,
+            InteractionsType::Query(Query::Select(
+                SelectFree::arbitrary(rng, &env.connection_context(conn_index)).0,
+            )),
+        )
+    }
 }
 
-fn random_expr<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usize) -> Interactions {
-    Interactions::new(
-        conn_index,
-        InteractionsType::Query(Query::Select(
-            SelectFree::arbitrary(rng, &env.connection_context(conn_index)).0,
-        )),
-    )
-}
-
-fn random_write<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usize) -> Interactions {
+fn random_insert<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usize) -> Interactions {
     Interactions::new(
         conn_index,
         InteractionsType::Query(Query::Insert(Insert::arbitrary(
@@ -1164,14 +1166,14 @@ fn random_create_index<R: rand::Rng>(
     ))
 }
 
-fn random_fault<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv) -> Interactions {
+fn random_fault<R: rand::Rng>(rng: &mut R, env: &SimulatorEnv, conn_index: usize) -> Interactions {
     let faults = if env.opts.disable_reopen_database {
         vec![Fault::Disconnect]
     } else {
         vec![Fault::Disconnect, Fault::ReopenDatabase]
     };
     let fault = faults[rng.random_range(0..faults.len())];
-    Interactions::new(env.choose_conn(rng), InteractionsType::Fault(fault))
+    Interactions::new(conn_index, InteractionsType::Fault(fault))
 }
 
 impl ArbitraryFrom<(&SimulatorEnv, InteractionStats, usize)> for Interactions {
@@ -1186,10 +1188,24 @@ impl ArbitraryFrom<(&SimulatorEnv, InteractionStats, usize)> for Interactions {
             &stats,
             env.profile.experimental_mvcc,
         );
+
+        // TODO: find a way to be more efficient and pass the weights and properties down to the ArbitraryFrom functions
+        let queries = possible_queries(conn_ctx.tables());
+        let query_weights =
+            WeightedIndex::new(queries.iter().map(|query| query.weight(&remaining_))).unwrap();
+
+        let properties = possiple_properties(conn_ctx.tables());
+        let property_weights = WeightedIndex::new(
+            properties
+                .iter()
+                .map(|property| property.weight(env, &remaining_, conn_ctx.opts())),
+        )
+        .unwrap();
+
         frequency(
             vec![
                 (
-                    u32::min(remaining_.select, remaining_.insert) + remaining_.create,
+                    property_weights.total_weight(),
                     Box::new(|rng: &mut R| {
                         Interactions::new(
                             conn_index,
@@ -1202,52 +1218,25 @@ impl ArbitraryFrom<(&SimulatorEnv, InteractionStats, usize)> for Interactions {
                     }),
                 ),
                 (
-                    remaining_.select,
-                    Box::new(|rng: &mut R| random_read(rng, env, conn_index)),
-                ),
-                (
-                    remaining_.select / 3,
-                    Box::new(|rng: &mut R| random_expr(rng, env, conn_index)),
-                ),
-                (
-                    remaining_.insert,
-                    Box::new(|rng: &mut R| random_write(rng, env, conn_index)),
-                ),
-                (
-                    remaining_.create,
-                    Box::new(|rng: &mut R| random_create(rng, env, conn_index)),
-                ),
-                (
-                    remaining_.create_index,
+                    query_weights.total_weight(),
                     Box::new(|rng: &mut R| {
-                        if let Some(interaction) = random_create_index(rng, env, conn_index) {
-                            interaction
-                        } else {
-                            // if no tables exist, we can't create an index, so fallback to creating a table
-                            random_create(rng, env, conn_index)
-                        }
+                        Interactions::new(
+                            conn_index,
+                            InteractionsType::Query(Query::arbitrary_from(
+                                rng,
+                                conn_ctx,
+                                &remaining_,
+                            )),
+                        )
                     }),
                 ),
-                (
-                    remaining_.delete,
-                    Box::new(|rng: &mut R| random_delete(rng, env, conn_index)),
-                ),
-                (
-                    remaining_.update,
-                    Box::new(|rng: &mut R| random_update(rng, env, conn_index)),
-                ),
-                (
-                    // remaining_.drop,
-                    0,
-                    Box::new(|rng: &mut R| random_drop(rng, env, conn_index)),
-                ),
                 (
                     remaining_
                         .select
                         .min(remaining_.insert)
                         .min(remaining_.create)
                         .max(1),
-                    Box::new(|rng: &mut R| random_fault(rng, env)),
+                    Box::new(|rng: &mut R| random_fault(rng, env, conn_index)),
                 ),
             ],
             rng,

simulator/generation/property.rs

@@ -1,6 +1,7 @@
+use rand::distr::{Distribution, weighted::WeightedIndex};
 use serde::{Deserialize, Serialize};
 use sql_generation::{
-    generation::{Arbitrary, ArbitraryFrom, GenerationContext, frequency, pick, pick_index},
+    generation::{Arbitrary, ArbitraryFrom, GenerationContext, Opts, pick, pick_index},
     model::{
         query::{
             Create, Delete, Drop, Insert, Select,
@@ -9,16 +10,17 @@ use sql_generation::{
             transaction::{Begin, Commit, Rollback},
             update::Update,
         },
-        table::SimValue,
+        table::{SimValue, Table},
     },
 };
+use strum::IntoEnumIterator;
 use turso_core::{LimboError, types};
 use turso_parser::ast::{self, Distinctness};
 
 use crate::{
     common::print_diff,
-    generation::{Shadow as _, plan::InteractionType},
-    model::Query,
+    generation::{Shadow as _, plan::InteractionType, query::possible_queries},
+    model::{Query, QueryCapabilities, QueryDiscriminants},
     profiles::query::QueryProfile,
     runner::env::SimulatorEnv,
 };
@@ -27,7 +29,8 @@ use super::plan::{Assertion, Interaction, InteractionStats, ResultSet};
 
 /// Properties are representations of executable specifications
 /// about the database behavior.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, strum::EnumDiscriminants)]
+#[strum_discriminants(derive(strum::EnumIter))]
 pub enum Property {
     /// Insert-Select is a property in which the inserted row
     /// must be in the resulting rows of a select query that has a
@@ -1308,7 +1311,9 @@ fn property_insert_values_select<R: rand::Rng>(
 
 fn property_read_your_updates_back<R: rand::Rng>(
     rng: &mut R,
+    _remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // e.g. UPDATE t SET a=1, b=2 WHERE c=1;
     let update = Update::arbitrary(rng, ctx);
@@ -1330,7 +1335,9 @@ fn property_read_your_updates_back<R: rand::Rng>(
 
 fn property_table_has_expected_content<R: rand::Rng>(
     rng: &mut R,
+    _remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
@@ -1339,7 +1346,12 @@ fn property_table_has_expected_content<R: rand::Rng>(
     }
 }
 
-fn property_select_limit<R: rand::Rng>(rng: &mut R, ctx: &impl GenerationContext) -> Property {
+fn property_select_limit<R: rand::Rng>(
+    rng: &mut R,
+    _remaining: &Remaining,
+    ctx: &impl GenerationContext,
+    _mvcc: bool,
+) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
     // Select the table
@@ -1357,6 +1369,7 @@ fn property_double_create_failure<R: rand::Rng>(
     rng: &mut R,
     remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Create the table
     let create_query = Create::arbitrary(rng, ctx);
@@ -1389,6 +1402,7 @@ fn property_delete_select<R: rand::Rng>(
     rng: &mut R,
     remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
@@ -1447,6 +1461,7 @@ fn property_drop_select<R: rand::Rng>(
     rng: &mut R,
     remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
@@ -1480,7 +1495,9 @@ fn property_drop_select<R: rand::Rng>(
 
 fn property_select_select_optimizer<R: rand::Rng>(
     rng: &mut R,
+    _remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
@@ -1501,7 +1518,9 @@ fn property_select_select_optimizer<R: rand::Rng>(
 
 fn property_where_true_false_null<R: rand::Rng>(
     rng: &mut R,
+    _remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
@@ -1520,7 +1539,9 @@ fn property_where_true_false_null<R: rand::Rng>(
 
 fn property_union_all_preserves_cardinality<R: rand::Rng>(
     rng: &mut R,
+    _remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     // Get a random table
     let table = pick(ctx.tables(), rng);
@@ -1547,6 +1568,7 @@ fn property_fsync_no_wait<R: rand::Rng>(
     rng: &mut R,
     remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     Property::FsyncNoWait {
         query: Query::arbitrary_from(rng, ctx, remaining),
@@ -1558,6 +1580,7 @@ fn property_faulty_query<R: rand::Rng>(
     rng: &mut R,
     remaining: &Remaining,
     ctx: &impl GenerationContext,
+    _mvcc: bool,
 ) -> Property {
     Property::FaultyQuery {
         query: Query::arbitrary_from(rng, ctx, remaining),
@@ -1565,6 +1588,161 @@ fn property_faulty_query<R: rand::Rng>(
     }
 }
 
+type PropertyGenFunc<R, G> = fn(&mut R, &Remaining, &G, bool) -> Property;
+
+impl PropertyDiscriminants {
+    pub fn gen_function<R, G>(&self) -> PropertyGenFunc<R, G>
+    where
+        R: rand::Rng,
+        G: GenerationContext,
+    {
+        match self {
+            PropertyDiscriminants::InsertValuesSelect => property_insert_values_select,
+            PropertyDiscriminants::ReadYourUpdatesBack => property_read_your_updates_back,
+            PropertyDiscriminants::TableHasExpectedContent => property_table_has_expected_content,
+            PropertyDiscriminants::DoubleCreateFailure => property_double_create_failure,
+            PropertyDiscriminants::SelectLimit => property_select_limit,
+            PropertyDiscriminants::DeleteSelect => property_delete_select,
+            PropertyDiscriminants::DropSelect => property_drop_select,
+            PropertyDiscriminants::SelectSelectOptimizer => property_select_select_optimizer,
+            PropertyDiscriminants::WhereTrueFalseNull => property_where_true_false_null,
+            PropertyDiscriminants::UNIONAllPreservesCardinality => {
+                property_union_all_preserves_cardinality
+            }
+            PropertyDiscriminants::FsyncNoWait => property_fsync_no_wait,
+            PropertyDiscriminants::FaultyQuery => property_faulty_query,
+            PropertyDiscriminants::Queries => {
+                unreachable!("should not try to generate queries property")
+            }
+        }
+    }
+
+    pub fn weight(&self, env: &SimulatorEnv, remaining: &Remaining, opts: &Opts) -> u32 {
+        match self {
+            PropertyDiscriminants::InsertValuesSelect => {
+                if !env.opts.disable_insert_values_select {
+                    u32::min(remaining.select, remaining.insert).max(1)
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::ReadYourUpdatesBack => {
+                u32::min(remaining.select, remaining.insert).max(1)
+            }
+            PropertyDiscriminants::TableHasExpectedContent => remaining.select.max(1),
+            PropertyDiscriminants::DoubleCreateFailure => {
+                if !env.opts.disable_double_create_failure {
+                    remaining.create / 2
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::SelectLimit => {
+                if !env.opts.disable_select_limit {
+                    remaining.select
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::DeleteSelect => {
+                if !env.opts.disable_delete_select {
+                    u32::min(remaining.select, remaining.insert).min(remaining.delete)
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::DropSelect => {
+                if !env.opts.disable_drop_select {
+                    // remaining.drop
+                    0
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::SelectSelectOptimizer => {
+                if !env.opts.disable_select_optimizer {
+                    remaining.select / 2
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::WhereTrueFalseNull => {
+                if opts.indexes && !env.opts.disable_where_true_false_null {
+                    remaining.select / 2
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::UNIONAllPreservesCardinality => {
+                if opts.indexes && !env.opts.disable_union_all_preserves_cardinality {
+                    remaining.select / 3
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::FsyncNoWait => {
+                if env.profile.io.enable && !env.opts.disable_fsync_no_wait {
+                    50 // Freestyle number
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::FaultyQuery => {
+                if env.profile.io.enable
+                    && env.profile.io.fault.enable
+                    && !env.opts.disable_faulty_query
+                {
+                    20
+                } else {
+                    0
+                }
+            }
+            PropertyDiscriminants::Queries => {
+                unreachable!("queries property should not be generated")
+            }
+        }
+    }
+
+    fn can_generate(queries: &[QueryDiscriminants]) -> Vec<PropertyDiscriminants> {
+        let queries_capabilities = QueryCapabilities::from_list_queries(queries);
+
+        PropertyDiscriminants::iter()
+            .filter(|property| queries_capabilities.contains(property.requirements()))
+            .collect()
+    }
+
+    pub const fn requirements(&self) -> QueryCapabilities {
+        match self {
+            PropertyDiscriminants::InsertValuesSelect => {
+                QueryCapabilities::SELECT.union(QueryCapabilities::INSERT)
+            }
+            PropertyDiscriminants::ReadYourUpdatesBack => {
+                QueryCapabilities::SELECT.union(QueryCapabilities::UPDATE)
+            }
+            PropertyDiscriminants::TableHasExpectedContent => QueryCapabilities::SELECT,
+            PropertyDiscriminants::DoubleCreateFailure => QueryCapabilities::CREATE,
+            PropertyDiscriminants::SelectLimit => QueryCapabilities::SELECT,
+            PropertyDiscriminants::DeleteSelect => {
+                QueryCapabilities::SELECT.union(QueryCapabilities::DELETE)
+            }
+            PropertyDiscriminants::DropSelect => {
+                QueryCapabilities::SELECT.union(QueryCapabilities::DROP)
+            }
+            PropertyDiscriminants::SelectSelectOptimizer => QueryCapabilities::SELECT,
+            PropertyDiscriminants::WhereTrueFalseNull => QueryCapabilities::SELECT,
+            PropertyDiscriminants::UNIONAllPreservesCardinality => QueryCapabilities::SELECT,
+            PropertyDiscriminants::FsyncNoWait => QueryCapabilities::all(),
+            PropertyDiscriminants::FaultyQuery => QueryCapabilities::all(),
+            PropertyDiscriminants::Queries => panic!("queries property should not be generated"),
+        }
+    }
+}
+
+pub fn possiple_properties(tables: &[Table]) -> Vec<PropertyDiscriminants> {
+    let queries = possible_queries(tables);
+    PropertyDiscriminants::can_generate(queries)
+}
+
 impl ArbitraryFrom<(&SimulatorEnv, &InteractionStats)> for Property {
     fn arbitrary_from<R: rand::Rng, C: GenerationContext>(
         rng: &mut R,
@@ -1579,110 +1757,19 @@ impl ArbitraryFrom<(&SimulatorEnv, &InteractionStats)> for Property {
             env.profile.experimental_mvcc,
         );
 
-        #[allow(clippy::type_complexity)]
-        let choices: Vec<(_, Box<dyn Fn(&mut R) -> Property>)> = vec![
-            (
-                if !env.opts.disable_insert_values_select {
-                    u32::min(remaining_.select, remaining_.insert).max(1)
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| {
-                    property_insert_values_select(
-                        rng,
-                        &remaining_,
-                        conn_ctx,
-                        env.profile.experimental_mvcc,
-                    )
-                }),
-            ),
-            (
-                remaining_.select.max(1),
-                Box::new(|rng: &mut R| property_table_has_expected_content(rng, conn_ctx)),
-            ),
-            (
-                u32::min(remaining_.select, remaining_.insert).max(1),
-                Box::new(|rng: &mut R| property_read_your_updates_back(rng, conn_ctx)),
-            ),
-            (
-                if !env.opts.disable_double_create_failure {
-                    remaining_.create / 2
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_double_create_failure(rng, &remaining_, conn_ctx)),
-            ),
-            (
-                if !env.opts.disable_select_limit {
-                    remaining_.select
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_select_limit(rng, conn_ctx)),
-            ),
-            (
-                if !env.opts.disable_delete_select {
-                    u32::min(remaining_.select, remaining_.insert).min(remaining_.delete)
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_delete_select(rng, &remaining_, conn_ctx)),
-            ),
-            (
-                if !env.opts.disable_drop_select {
-                    // remaining_.drop
-                    0
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_drop_select(rng, &remaining_, conn_ctx)),
-            ),
-            (
-                if !env.opts.disable_select_optimizer {
-                    remaining_.select / 2
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_select_select_optimizer(rng, conn_ctx)),
-            ),
-            (
-                if opts.indexes && !env.opts.disable_where_true_false_null {
-                    remaining_.select / 2
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_where_true_false_null(rng, conn_ctx)),
-            ),
-            (
-                if opts.indexes && !env.opts.disable_union_all_preserves_cardinality {
-                    remaining_.select / 3
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_union_all_preserves_cardinality(rng, conn_ctx)),
-            ),
-            (
-                if env.profile.io.enable && !env.opts.disable_fsync_no_wait {
-                    50 // Freestyle number
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_fsync_no_wait(rng, &remaining_, conn_ctx)),
-            ),
-            (
-                if env.profile.io.enable
-                    && env.profile.io.fault.enable
-                    && !env.opts.disable_faulty_query
-                {
-                    20
-                } else {
-                    0
-                },
-                Box::new(|rng: &mut R| property_faulty_query(rng, &remaining_, conn_ctx)),
-            ),
-        ];
+        let properties = possiple_properties(conn_ctx.tables());
+        let weights = WeightedIndex::new(
+            properties
+                .iter()
+                .map(|property| property.weight(env, &remaining_, opts)),
+        )
+        .unwrap();
 
-        frequency(choices, rng)
+        let idx = weights.sample(rng);
+        let property_fn = properties[idx].gen_function();
+        let property = (property_fn)(rng, &remaining_, conn_ctx, env.profile.experimental_mvcc);
+
+        property
     }
 }
 

simulator/generation/query.rs

@@ -76,18 +76,12 @@ fn random_create_index<R: rand::Rng>(rng: &mut R, conn_ctx: &impl GenerationCont
 /// Possible queries that can be generated given the table state
 ///
 /// Does not take into account transactional statements
-pub fn possible_queries(tables: &[Table]) -> Vec<QueryDiscriminants> {
-    let mut queries = vec![QueryDiscriminants::Select, QueryDiscriminants::Create];
-    if !tables.is_empty() {
-        queries.extend([
-            QueryDiscriminants::Insert,
-            QueryDiscriminants::Update,
-            QueryDiscriminants::Delete,
-            QueryDiscriminants::Drop,
-            QueryDiscriminants::CreateIndex,
-        ]);
+pub const fn possible_queries(tables: &[Table]) -> &'static [QueryDiscriminants] {
+    if tables.is_empty() {
+        &[QueryDiscriminants::Select, QueryDiscriminants::Create]
+    } else {
+        QueryDiscriminants::ALL_NO_TRANSACTION
     }
-    queries
 }
 
 type QueryGenFunc<R, G> = fn(&mut R, &G) -> Query;