tests_fuzz/
fake.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
15use std::collections::HashSet;
16use std::marker::PhantomData;
17
18use lazy_static::lazy_static;
19use rand::Rng;
20use rand::prelude::IndexedRandom;
21use rand::seq::{IteratorRandom, SliceRandom};
22
23use crate::generator::Random;
24use crate::impl_random;
25use crate::ir::Ident;
26
27lazy_static! {
28    pub static ref LOREM_WORDS: Vec<String> = include_str!("data/lorem_words")
29        .lines()
30        .map(String::from)
31        .collect();
32}
33
34/// Modified from https://github.com/ucarion/faker_rand/blob/ea70c660e1ecd7320156eddb31d2830a511f8842/src/lib.rs
35macro_rules! faker_impl_from_values {
36    ($name: ident, $values: expr) => {
37        impl rand::distr::Distribution<$name> for rand::distr::StandardUniform {
38            fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> $name {
39                $name($values[rng.random_range(0..$values.len())].clone())
40            }
41        }
42
43        impl std::fmt::Display for $name {
44            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45                write!(f, "{}", self.0)
46            }
47        }
48    };
49}
50
51pub struct Word(String);
52faker_impl_from_values!(Word, LOREM_WORDS);
53pub struct WordGenerator;
54impl_random!(Ident, WordGenerator, LOREM_WORDS);
55
56pub struct MappedGenerator<T, F, R, V>
57where
58    T: Random<V, R>,
59    F: Fn(&mut R, V) -> V,
60    R: Rng,
61{
62    base: T,
63    map: F,
64    _r: PhantomData<R>,
65    _v: PhantomData<V>,
66}
67
68pub struct ConstGenerator<V> {
69    value: V,
70}
71
72impl<V> ConstGenerator<V> {
73    pub fn new(value: V) -> Self {
74        Self { value }
75    }
76}
77
78impl<R, V> Random<V, R> for ConstGenerator<V>
79where
80    R: Rng,
81    V: Clone,
82{
83    fn choose(&self, _rng: &mut R, amount: usize) -> Vec<V> {
84        vec![self.value.clone(); amount]
85    }
86}
87
88pub fn random_capitalize_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
89    let mut v = s.value.chars().collect::<Vec<_>>();
90
91    let str_len = s.value.len();
92    let select = rng.random_range(0..str_len);
93    for idx in (0..str_len).choose_multiple(rng, select) {
94        v[idx] = v[idx].to_uppercase().next().unwrap();
95    }
96
97    Ident {
98        quote_style: s.quote_style,
99        value: v.into_iter().collect::<String>(),
100    }
101}
102
103lazy_static! {
104    static ref KEYWORDS_SET: HashSet<&'static str> = sqlparser::keywords::ALL_KEYWORDS
105        .iter()
106        .cloned()
107        .collect::<HashSet<_>>();
108}
109
110/// Returns true if it's a keyword.
111pub fn is_keyword(word: impl AsRef<str>) -> bool {
112    KEYWORDS_SET.contains(word.as_ref())
113}
114
115/// Returns true if it contains uppercase char.
116pub fn contain_uppercase_char(s: &str) -> bool {
117    s.chars().any(|c| c.is_uppercase())
118}
119
120/// Returns true if it's a keyword or contains uppercase char.
121pub fn is_keyword_or_contain_uppercase(s: &str) -> bool {
122    is_keyword(s.to_uppercase()) || contain_uppercase_char(s)
123}
124
125pub fn make_backtick_map<R: Rng + 'static, F: Fn(&str) -> bool>(
126    f: F,
127) -> impl Fn(&mut R, Ident) -> Ident {
128    move |_rng, s| -> Ident {
129        let need = f(&s.value);
130
131        if need {
132            Ident {
133                value: s.value,
134                quote_style: Some('`'),
135            }
136        } else {
137            s
138        }
139    }
140}
141
142pub fn make_quote_map<R: Rng + 'static, F: Fn(&str) -> bool>(
143    f: F,
144) -> impl Fn(&mut R, Ident) -> Ident {
145    move |_rng, s| -> Ident {
146        let need = f(&s.value);
147
148        if need {
149            Ident {
150                value: s.value,
151                quote_style: Some('"'),
152            }
153        } else {
154            s
155        }
156    }
157}
158
159/// Adds backticks if it contains uppercase chars.
160pub fn auto_backtick_map<R: Rng + 'static>(_rng: &mut R, s: Ident) -> Ident {
161    let need = s.value.chars().any(|c| c.is_uppercase());
162
163    if need {
164        Ident {
165            value: s.value,
166            quote_style: Some('`'),
167        }
168    } else {
169        s
170    }
171}
172
173/// Adds backticks if it contains uppercase chars.
174pub fn uppercase_and_keyword_backtick_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
175    make_backtick_map(is_keyword_or_contain_uppercase)(rng, s)
176}
177
178/// Adds quotes if it contains uppercase chars.
179pub fn auto_quote_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
180    make_quote_map(contain_uppercase_char)(rng, s)
181}
182
183/// Adds quotes if it contains uppercase chars.
184pub fn uppercase_and_keyword_quote_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
185    make_quote_map(is_keyword_or_contain_uppercase)(rng, s)
186}
187
188pub fn merge_two_word_map_fn<R: Rng>(
189    f1: impl Fn(&mut R, Ident) -> Ident,
190    f2: impl Fn(&mut R, Ident) -> Ident,
191) -> impl Fn(&mut R, Ident) -> Ident {
192    move |rng, s| -> Ident {
193        let s = f1(rng, s);
194        f2(rng, s)
195    }
196}
197
198impl<T, F, R, V> MappedGenerator<T, F, R, V>
199where
200    T: Random<V, R>,
201    F: Fn(&mut R, V) -> V,
202    R: Rng,
203{
204    pub fn new(base: T, map: F) -> Self {
205        Self {
206            base,
207            map,
208            _r: Default::default(),
209            _v: Default::default(),
210        }
211    }
212}
213
214impl<T, F, R, V> Random<V, R> for MappedGenerator<T, F, R, V>
215where
216    T: Random<V, R>,
217    F: Fn(&mut R, V) -> V,
218    R: Rng,
219{
220    fn choose(&self, rng: &mut R, amount: usize) -> Vec<V> {
221        self.base
222            .choose(rng, amount)
223            .into_iter()
224            .map(|s| (self.map)(rng, s))
225            .collect()
226    }
227}
tests_fuzz/fake.rs

tests_fuzz/
fake.rs
