use std::{iter::Sum, ops::SubAssign}; use anarchist_readable_name_generator_lib::readable_name_custom; use rand::{distr::uniform::SampleUniform, Rng}; use crate::model::table::Table; pub mod expr; pub mod predicate; pub mod query; pub mod table; #[derive(Debug, Clone, Copy)] pub struct Opts { /// Indexes enabled pub indexes: bool, } /// Trait used to provide context to generation functions pub trait GenerationContext { fn tables(&self) -> &Vec; fn opts(&self) -> Opts; } type ArbitraryFromFunc<'a, R, T> = Box T + 'a>; type Choice<'a, R, T> = (usize, Box Option + 'a>); /// Arbitrary trait for generating random values /// An implementation of arbitrary is assumed to be a uniform sampling of /// the possible values of the type, with a bias towards smaller values for /// practicality. pub trait Arbitrary { fn arbitrary(rng: &mut R) -> Self; } /// ArbitrarySized trait for generating random values of a specific size /// An implementation of arbitrary_sized is assumed to be a uniform sampling of /// the possible values of the type, with a bias towards smaller values for /// practicality, but with the additional constraint that the generated value /// must fit in the given size. This is useful for generating values that are /// constrained by a specific size, such as integers or strings. pub trait ArbitrarySized { fn arbitrary_sized(rng: &mut R, size: usize) -> Self; } /// ArbitraryFrom trait for generating random values from a given value /// ArbitraryFrom allows for constructing relations, where the generated /// value is dependent on the given value. These relations could be constraints /// such as generating an integer within an interval, or a value that fits in a table, /// or a predicate satisfying a given table row. pub trait ArbitraryFrom { fn arbitrary_from(rng: &mut R, t: T) -> Self; } /// ArbitrarySizedFrom trait for generating random values from a given value /// ArbitrarySizedFrom allows for constructing relations, where the generated /// value is dependent on the given value and a size constraint. These relations /// could be constraints such as generating an integer within an interval, /// or a value that fits in a table, or a predicate satisfying a given table row, /// but with the additional constraint that the generated value must fit in the given size. /// This is useful for generating values that are constrained by a specific size, /// such as integers or strings, while still being dependent on the given value. pub trait ArbitrarySizedFrom { fn arbitrary_sized_from(rng: &mut R, t: T, size: usize) -> Self; } /// ArbitraryFromMaybe trait for fallibally generating random values from a given value pub trait ArbitraryFromMaybe { fn arbitrary_from_maybe(rng: &mut R, t: T) -> Option where Self: Sized; } /// Frequency is a helper function for composing different generators with different frequency /// of occurrences. /// The type signature for the `N` parameter is a bit complex, but it /// roughly corresponds to a type that can be summed, compared, subtracted and sampled, which are /// the operations we require for the implementation. // todo: switch to a simpler type signature that can accommodate all integer and float types, which // should be enough for our purposes. pub fn frequency( choices: Vec<(N, ArbitraryFromFunc)>, rng: &mut R, ) -> T { let total = choices.iter().map(|(weight, _)| *weight).sum::(); let mut choice = rng.random_range(N::default()..total); for (weight, f) in choices { if choice < weight { return f(rng); } choice -= weight; } unreachable!() } /// one_of is a helper function for composing different generators with equal probability of occurrence. pub fn one_of(choices: Vec>, rng: &mut R) -> T { let index = rng.random_range(0..choices.len()); choices[index](rng) } /// backtrack is a helper function for composing different "failable" generators. /// The function takes a list of functions that return an Option, along with number of retries /// to make before giving up. pub fn backtrack(mut choices: Vec>, rng: &mut R) -> Option { loop { // If there are no more choices left, we give up let choices_ = choices .iter() .enumerate() .filter(|(_, (retries, _))| *retries > 0) .collect::>(); if choices_.is_empty() { tracing::trace!("backtrack: no more choices left"); return None; } // Run a one_of on the remaining choices let (choice_index, choice) = pick(&choices_, rng); let choice_index = *choice_index; // If the choice returns None, we decrement the number of retries and try again let result = choice.1(rng); if result.is_some() { return result; } else { choices[choice_index].0 -= 1; } } } /// pick is a helper function for uniformly picking a random element from a slice pub fn pick<'a, T, R: Rng>(choices: &'a [T], rng: &mut R) -> &'a T { let index = rng.random_range(0..choices.len()); &choices[index] } /// pick_index is typically used for picking an index from a slice to later refer to the element /// at that index. pub fn pick_index(choices: usize, rng: &mut R) -> usize { rng.random_range(0..choices) } /// pick_n_unique is a helper function for uniformly picking N unique elements from a range. /// The elements themselves are usize, typically representing indices. pub fn pick_n_unique(range: std::ops::Range, n: usize, rng: &mut R) -> Vec { use rand::seq::SliceRandom; let mut items: Vec = range.collect(); items.shuffle(rng); items.into_iter().take(n).collect() } /// gen_random_text uses `anarchist_readable_name_generator_lib` to generate random /// readable names for tables, columns, text values etc. pub fn gen_random_text(rng: &mut T) -> String { let big_text = rng.random_ratio(1, 1000); if big_text { // let max_size: u64 = 2 * 1024 * 1024 * 1024; let max_size: u64 = 2 * 1024; let size = rng.random_range(1024..max_size); let mut name = String::with_capacity(size as usize); for i in 0..size { name.push(((i % 26) as u8 + b'A') as char); } name } else { let name = readable_name_custom("_", rng); name.replace("-", "_") } } pub fn pick_unique( items: &[T], count: usize, rng: &mut impl rand::Rng, ) -> Vec where ::Owned: PartialEq, { let mut picked: Vec = Vec::new(); while picked.len() < count { let item = pick(items, rng); if !picked.contains(&item.to_owned()) { picked.push(item.to_owned()); } } picked }