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::prelude::IndexedRandom;
20use rand::seq::{IteratorRandom, SliceRandom};
21use rand::Rng;
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 fn random_capitalize_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
69    let mut v = s.value.chars().collect::<Vec<_>>();
70
71    let str_len = s.value.len();
72    let select = rng.random_range(0..str_len);
73    for idx in (0..str_len).choose_multiple(rng, select) {
74        v[idx] = v[idx].to_uppercase().next().unwrap();
75    }
76
77    Ident {
78        quote_style: s.quote_style,
79        value: v.into_iter().collect::<String>(),
80    }
81}
82
83lazy_static! {
84    static ref KEYWORDS_SET: HashSet<&'static str> = sqlparser::keywords::ALL_KEYWORDS
85        .iter()
86        .cloned()
87        .collect::<HashSet<_>>();
88}
89
90/// Returns true if it's a keyword.
91pub fn is_keyword(word: impl AsRef<str>) -> bool {
92    KEYWORDS_SET.contains(word.as_ref())
93}
94
95/// Returns true if it contains uppercase char.
96pub fn contain_uppercase_char(s: &str) -> bool {
97    s.chars().any(|c| c.is_uppercase())
98}
99
100/// Returns true if it's a keyword or contains uppercase char.
101pub fn is_keyword_or_contain_uppercase(s: &str) -> bool {
102    is_keyword(s.to_uppercase()) || contain_uppercase_char(s)
103}
104
105pub fn make_backtick_map<R: Rng + 'static, F: Fn(&str) -> bool>(
106    f: F,
107) -> impl Fn(&mut R, Ident) -> Ident {
108    move |_rng, s| -> Ident {
109        let need = f(&s.value);
110
111        if need {
112            Ident {
113                value: s.value,
114                quote_style: Some('`'),
115            }
116        } else {
117            s
118        }
119    }
120}
121
122pub fn make_quote_map<R: Rng + 'static, F: Fn(&str) -> bool>(
123    f: F,
124) -> impl Fn(&mut R, Ident) -> Ident {
125    move |_rng, s| -> Ident {
126        let need = f(&s.value);
127
128        if need {
129            Ident {
130                value: s.value,
131                quote_style: Some('"'),
132            }
133        } else {
134            s
135        }
136    }
137}
138
139/// Adds backticks if it contains uppercase chars.
140pub fn auto_backtick_map<R: Rng + 'static>(_rng: &mut R, s: Ident) -> Ident {
141    let need = s.value.chars().any(|c| c.is_uppercase());
142
143    if need {
144        Ident {
145            value: s.value,
146            quote_style: Some('`'),
147        }
148    } else {
149        s
150    }
151}
152
153/// Adds backticks if it contains uppercase chars.
154pub fn uppercase_and_keyword_backtick_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
155    make_backtick_map(is_keyword_or_contain_uppercase)(rng, s)
156}
157
158/// Adds quotes if it contains uppercase chars.
159pub fn auto_quote_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
160    make_quote_map(contain_uppercase_char)(rng, s)
161}
162
163/// Adds quotes if it contains uppercase chars.
164pub fn uppercase_and_keyword_quote_map<R: Rng + 'static>(rng: &mut R, s: Ident) -> Ident {
165    make_quote_map(is_keyword_or_contain_uppercase)(rng, s)
166}
167
168pub fn merge_two_word_map_fn<R: Rng>(
169    f1: impl Fn(&mut R, Ident) -> Ident,
170    f2: impl Fn(&mut R, Ident) -> Ident,
171) -> impl Fn(&mut R, Ident) -> Ident {
172    move |rng, s| -> Ident {
173        let s = f1(rng, s);
174        f2(rng, s)
175    }
176}
177
178impl<T, F, R, V> MappedGenerator<T, F, R, V>
179where
180    T: Random<V, R>,
181    F: Fn(&mut R, V) -> V,
182    R: Rng,
183{
184    pub fn new(base: T, map: F) -> Self {
185        Self {
186            base,
187            map,
188            _r: Default::default(),
189            _v: Default::default(),
190        }
191    }
192}
193
194impl<T, F, R, V> Random<V, R> for MappedGenerator<T, F, R, V>
195where
196    T: Random<V, R>,
197    F: Fn(&mut R, V) -> V,
198    R: Rng,
199{
200    fn choose(&self, rng: &mut R, amount: usize) -> Vec<V> {
201        self.base
202            .choose(rng, amount)
203            .into_iter()
204            .map(|s| (self.map)(rng, s))
205            .collect()
206    }
207}
tests_fuzz/fake.rs

tests_fuzz/
fake.rs
