use crate::random::RandomProvider;
use crate::random::ResponseError;
use apollo_compiler::executable::Field;
use apollo_compiler::Name;
use apollo_compiler::Node;
use indexmap::IndexMap;
use serde_json_bytes::serde_json::Number;
use serde_json_bytes::Value;
use std::collections::HashMap;
pub trait Generator<R: RandomProvider> {
fn generate(
&mut self,
rng: &mut R,
generators: &mut Generators<R>,
fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Result<Value, ResponseError>;
fn boxed(self) -> Box<dyn Generator<R>>
where
Self: Sized + 'static,
{
Box::new(self)
}
}
pub struct Generators<R: RandomProvider> {
map: HashMap<Name, Box<dyn Generator<R>>>,
}
impl<R: RandomProvider> Generators<R> {
pub(crate) fn insert(&mut self, name: Name, generator: Box<dyn Generator<R>>) {
self.map.insert(name, generator);
}
pub fn try_generate(
&mut self,
type_name: &Name,
rng: &mut R,
fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Option<Result<Value, ResponseError>> {
let mut generator = self.map.remove(type_name)?;
let result = generator.generate(rng, self, fields);
self.map.insert(type_name.clone(), generator);
Some(result)
}
pub fn generate_scalar(
&mut self,
type_name: &Name,
rng: &mut R,
) -> Result<Value, ResponseError> {
let empty = IndexMap::new();
if let Some(result) = self.try_generate(type_name, rng, &empty) {
return result;
}
let mut fallback = StringGenerator {
min_len: 1,
max_len: 10,
};
fallback.generate(rng, self, &empty)
}
}
#[derive(Debug, Default, Clone)]
pub struct BooleanGenerator;
impl<R: RandomProvider> Generator<R> for BooleanGenerator {
fn generate(
&mut self,
rng: &mut R,
_generators: &mut Generators<R>,
_fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Result<Value, ResponseError> {
Ok(Value::Bool(rng.gen_bool()?))
}
}
#[derive(Debug, Clone)]
pub struct IntGenerator {
pub min: i32,
pub max: i32,
}
impl<R: RandomProvider> Generator<R> for IntGenerator {
fn generate(
&mut self,
rng: &mut R,
_generators: &mut Generators<R>,
_fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Result<Value, ResponseError> {
Ok(Value::Number(rng.gen_i32_range(self.min, self.max)?.into()))
}
}
impl Default for IntGenerator {
fn default() -> Self {
Self { min: 0, max: 100 }
}
}
#[derive(Debug, Clone)]
pub struct FloatGenerator {
pub min: f64,
pub max: f64,
}
impl<R: RandomProvider> Generator<R> for FloatGenerator {
fn generate(
&mut self,
rng: &mut R,
_generators: &mut Generators<R>,
_fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Result<Value, ResponseError> {
let f = rng.gen_f64_range(self.min, self.max)?;
let num = Number::from_f64(f)
.ok_or_else(|| ResponseError::InvalidFormat("generated non-finite float".into()))?;
Ok(Value::Number(num))
}
}
impl Default for FloatGenerator {
fn default() -> Self {
Self {
min: -1.0,
max: 1.0,
}
}
}
#[derive(Debug, Clone)]
pub struct StringGenerator {
pub min_len: usize,
pub max_len: usize,
}
impl<R: RandomProvider> Generator<R> for StringGenerator {
fn generate(
&mut self,
rng: &mut R,
_generators: &mut Generators<R>,
_fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Result<Value, ResponseError> {
let len = rng.gen_usize_range(self.min_len, self.max_len)?;
let s: Result<std::string::String, _> =
(0..len).map(|_| rng.gen_alphanumeric_char()).collect();
Ok(Value::String(s?.into()))
}
}
impl Default for StringGenerator {
fn default() -> Self {
Self {
min_len: 1,
max_len: 10,
}
}
}
#[derive(Debug, Clone)]
pub struct IdGenerator {
pub min: i32,
pub max: i32,
}
impl<R: RandomProvider> Generator<R> for IdGenerator {
fn generate(
&mut self,
rng: &mut R,
_generators: &mut Generators<R>,
_fields: &IndexMap<String, Vec<Node<Field>>>,
) -> Result<Value, ResponseError> {
Ok(Value::String(
rng.gen_i32_range(self.min, self.max)?.to_string().into(),
))
}
}
impl Default for IdGenerator {
fn default() -> Self {
Self { min: 0, max: 100 }
}
}
impl<R: RandomProvider> Default for Generators<R> {
fn default() -> Self {
let map: HashMap<Name, Box<dyn Generator<R>>> = [
(Name::new_unchecked("Boolean"), BooleanGenerator.boxed()),
(Name::new_unchecked("Int"), IntGenerator::default().boxed()),
(Name::new_unchecked("ID"), IdGenerator::default().boxed()),
(
Name::new_unchecked("Float"),
FloatGenerator::default().boxed(),
),
(
Name::new_unchecked("String"),
StringGenerator::default().boxed(),
),
]
.into_iter()
.collect();
Self { map }
}
}