tests_fuzz/
generator.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod alter_expr;
16pub mod create_expr;
17pub mod insert_expr;
18pub mod select_expr;
19
20use std::fmt;
21
22use datatypes::data_type::ConcreteDataType;
23use datatypes::types::TimestampType;
24use datatypes::value::Value;
25use rand::Rng;
26
27use crate::error::Error;
28use crate::ir::create_expr::ColumnOption;
29use crate::ir::{AlterTableExpr, CreateTableExpr, Ident};
30
31pub type CreateTableExprGenerator<R> =
32    Box<dyn Generator<CreateTableExpr, R, Error = Error> + Sync + Send>;
33
34pub type AlterTableExprGenerator<R> =
35    Box<dyn Generator<AlterTableExpr, R, Error = Error> + Sync + Send>;
36
37pub type ColumnOptionGenerator<R> = Box<dyn Fn(&mut R, &ConcreteDataType) -> Vec<ColumnOption>>;
38
39pub type ConcreteDataTypeGenerator<R> = Box<dyn Random<ConcreteDataType, R>>;
40
41pub type ValueGenerator<R> =
42    Box<dyn Fn(&mut R, &ConcreteDataType, Option<&dyn Random<Ident, R>>) -> Value>;
43
44pub type TsValueGenerator<R> = Box<dyn Fn(&mut R, TimestampType) -> Value>;
45
46pub trait Generator<T, R: Rng> {
47    type Error: Sync + Send + fmt::Debug;
48
49    fn generate(&self, rng: &mut R) -> Result<T, Self::Error>;
50}
51
52pub trait Random<T, R: Rng> {
53    /// Generates a random element.
54    fn gen(&self, rng: &mut R) -> T {
55        self.choose(rng, 1).remove(0)
56    }
57
58    /// Uniformly sample `amount` distinct elements.
59    fn choose(&self, rng: &mut R, amount: usize) -> Vec<T>;
60}
61
62#[macro_export]
63macro_rules! impl_random {
64    ($type: ident, $value:ident, $values: ident) => {
65        impl<R: Rng> Random<$type, R> for $value {
66            fn choose(&self, rng: &mut R, amount: usize) -> Vec<$type> {
67                // Collects the elements in deterministic order first.
68                let mut result = std::collections::BTreeSet::new();
69                while result.len() != amount {
70                    result.insert($values.choose(rng).unwrap().clone());
71                }
72                let mut result = result.into_iter().map(Into::into).collect::<Vec<_>>();
73                // Shuffles the result slice.
74                result.shuffle(rng);
75                result
76            }
77        }
78    };
79}