use crate::json_node::{JsonNode, KeyValuePair};
use crate::types::{Diagnostic, Span};
use std::collections::HashMap;
#[derive(Default)]
struct GeneratorState {
variables: HashMap<String, f64>,
}
pub struct GeneratorOptions {
pub target_size_bytes: Option<usize>,
pub target_count: Option<usize>,
}
pub fn generate_from_schema(
schema: &JsonNode,
options: &GeneratorOptions,
) -> Result<JsonNode, Vec<Diagnostic>> {
let mut diagnostics = Vec::new();
let mut state = GeneratorState::default();
if let Some(target_size) = options.target_size_bytes {
let mut items = Vec::new();
let mut current_size = 0;
while current_size < target_size {
let item = evaluate_node(schema, &mut state, &mut diagnostics);
current_size += estimate_size(&item);
items.push(item);
if items.len() > 1000000 {
break;
}
}
if !diagnostics.is_empty() {
Err(diagnostics)
} else {
Ok(JsonNode::Array(items, Span::default()))
}
} else if let Some(target_count) = options.target_count {
let mut items = Vec::with_capacity(target_count);
for _ in 0..target_count {
items.push(evaluate_node(schema, &mut state, &mut diagnostics));
}
if !diagnostics.is_empty() {
Err(diagnostics)
} else {
Ok(JsonNode::Array(items, Span::default()))
}
} else {
let result = evaluate_node(schema, &mut state, &mut diagnostics);
if !diagnostics.is_empty() {
Err(diagnostics)
} else {
Ok(result)
}
}
}
fn evaluate_node(
node: &JsonNode,
state: &mut GeneratorState,
diagnostics: &mut Vec<Diagnostic>,
) -> JsonNode {
match node {
JsonNode::Object(pairs, span) => {
if let Some(type_pair) = pairs.iter().find(|p| p.key == "@type") {
#[allow(clippy::collapsible_if)]
if let JsonNode::String(t, _) = &type_pair.value {
return evaluate_instruction(t, pairs, state, diagnostics, span.clone());
}
}
let mut new_pairs = Vec::with_capacity(pairs.len());
for pair in pairs {
let key = if pair.key.starts_with("@@") {
pair.key[1..].to_string()
} else {
pair.key.clone()
};
new_pairs.push(KeyValuePair {
key,
value: evaluate_node(&pair.value, state, diagnostics),
});
}
JsonNode::Object(new_pairs, span.clone())
}
JsonNode::Array(items, span) => {
let new_items = items
.iter()
.map(|item| evaluate_node(item, state, diagnostics))
.collect();
JsonNode::Array(new_items, span.clone())
}
_ => node.clone(),
}
}
fn evaluate_instruction(
instruction_type: &str,
pairs: &[KeyValuePair],
state: &mut GeneratorState,
diagnostics: &mut Vec<Diagnostic>,
span: Span,
) -> JsonNode {
match instruction_type {
"uuid" => {
let lcg = state.variables.entry("lcg".to_string()).or_insert(1.0);
*lcg = (*lcg * 1664525.0 + 1013904223.0) % 4294967296.0;
let val = *lcg as u32;
JsonNode::String(
format!("uuid-{:08x}-1234-5678-abcd-123456789012", val),
span,
)
}
"integer" => {
let mut start = 0.0;
let mut step = 1.0;
for pair in pairs {
#[allow(clippy::collapsible_if)]
if pair.key == "@start" {
if let JsonNode::Number(n, _, _) = pair.value {
start = n;
}
} else if pair.key == "@step" {
if let JsonNode::Number(n, _, _) = pair.value {
step = n;
}
}
}
let state_key = format!("int_{}", span.start);
let current = state.variables.entry(state_key).or_insert(start);
let val = *current;
*current += step;
JsonNode::Number(val, val.to_string(), span)
}
"float" => {
let mut min = 0.0;
let mut max = 1.0;
for pair in pairs {
#[allow(clippy::collapsible_if)]
if pair.key == "@min" {
if let JsonNode::Number(n, _, _) = pair.value {
min = n;
}
} else if pair.key == "@max" {
if let JsonNode::Number(n, _, _) = pair.value {
max = n;
}
}
}
let lcg = state.variables.entry("lcg".to_string()).or_insert(1.0);
*lcg = (*lcg * 1664525.0 + 1013904223.0) % 4294967296.0;
let val = min + (*lcg / 4294967296.0) * (max - min);
JsonNode::Number(val, format!("{:.4}", val), span)
}
"bool" => {
let lcg = state.variables.entry("lcg".to_string()).or_insert(1.0);
*lcg = (*lcg * 1664525.0 + 1013904223.0) % 4294967296.0;
JsonNode::Bool(*lcg > 2147483648.0, span)
}
"string" => {
let mut template = "".to_string();
let mut pool = Vec::new();
let mut vars_schema = None;
for pair in pairs {
#[allow(clippy::collapsible_if)]
if pair.key == "@template" {
if let JsonNode::String(s, _) = &pair.value {
template = s.clone();
}
} else if pair.key == "@pool" {
if let JsonNode::Array(items, _) = &pair.value {
for item in items {
if let JsonNode::String(s, _) = item {
pool.push(s.clone());
}
}
}
} else if pair.key == "@vars" {
vars_schema = Some(&pair.value);
}
}
if !pool.is_empty() {
let lcg = state.variables.entry("lcg".to_string()).or_insert(1.0);
*lcg = (*lcg * 1664525.0 + 1013904223.0) % 4294967296.0;
let idx = (*lcg as usize) % pool.len();
JsonNode::String(pool[idx].clone(), span)
} else if !template.is_empty() {
let mut output = template;
if let Some(JsonNode::Object(var_pairs, _)) = vars_schema {
for var_pair in var_pairs {
let var_val = evaluate_node(&var_pair.value, state, diagnostics);
let var_str = match var_val {
JsonNode::String(s, _) => s,
JsonNode::Number(_, raw, _) => raw,
JsonNode::Bool(b, _) => b.to_string(),
_ => "".to_string(),
};
let placeholder = format!("{{{}}}", var_pair.key);
output = output.replace(&placeholder, &var_str);
}
}
JsonNode::String(output, span)
} else {
JsonNode::String("dummy".to_string(), span)
}
}
"array" => {
let mut count = 1;
let mut item_schema = &JsonNode::Null(span.clone());
for pair in pairs {
if pair.key == "@count" {
if let JsonNode::Number(n, _, _) = pair.value {
count = n as usize;
}
} else if pair.key == "@item" {
item_schema = &pair.value;
}
}
let mut items = Vec::with_capacity(count);
for _ in 0..count {
items.push(evaluate_node(item_schema, state, diagnostics));
}
JsonNode::Array(items, span)
}
_ => {
diagnostics.push(Diagnostic {
span: span.clone(),
message: format!("Unknown generator instruction: {}", instruction_type),
});
JsonNode::Null(span)
}
}
}
fn estimate_size(node: &JsonNode) -> usize {
match node {
JsonNode::Null(_) => 4,
JsonNode::Bool(_, _) => 4,
JsonNode::Number(_, raw, _) => raw.len(),
JsonNode::String(s, _) => s.len() + 2,
JsonNode::Array(items, _) => {
items.iter().map(estimate_size).sum::<usize>() + items.len() + 2
}
JsonNode::Object(pairs, _) => {
pairs
.iter()
.map(|p| p.key.len() + 4 + estimate_size(&p.value))
.sum::<usize>()
+ 2
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::parser::parse;
#[test]
fn test_generator_evaluates_integer_and_uuid() {
let schema_str = r#"{
"id": { "@type": "uuid" },
"index": { "@type": "integer", "@start": 5, "@step": 2 }
}"#;
let schema = parse(schema_str).unwrap();
let opts = GeneratorOptions {
target_size_bytes: None,
target_count: Some(2),
};
let result = generate_from_schema(&schema, &opts).unwrap();
if let JsonNode::Array(items, _) = result {
assert_eq!(items.len(), 2);
if let JsonNode::Object(pairs, _) = &items[0] {
let index_node = pairs.iter().find(|p| p.key == "index").unwrap();
if let JsonNode::Number(val, _, _) = index_node.value {
assert_eq!(val, 5.0);
} else {
panic!("Expected Number");
}
}
if let JsonNode::Object(pairs, _) = &items[1] {
let index_node = pairs.iter().find(|p| p.key == "index").unwrap();
if let JsonNode::Number(val, _, _) = index_node.value {
assert_eq!(val, 7.0);
} else {
panic!("Expected Number");
}
}
} else {
panic!("Expected array");
}
}
#[test]
fn test_generator_unknown_instruction_returns_diagnostic() {
let schema_str = r#"{ "field": { "@type": "fake_type" } }"#;
let schema = parse(schema_str).unwrap();
let opts = GeneratorOptions {
target_size_bytes: None,
target_count: Some(1),
};
let result = generate_from_schema(&schema, &opts);
assert!(result.is_err());
let diags = result.unwrap_err();
assert_eq!(diags.len(), 1);
assert!(diags[0].message.contains("Unknown generator instruction"));
}
#[test]
fn test_generator_escaped_keys() {
let schema_str = r#"{ "@@type": "User" }"#;
let schema = parse(schema_str).unwrap();
let opts = GeneratorOptions {
target_size_bytes: None,
target_count: None,
};
let result = generate_from_schema(&schema, &opts).unwrap();
if let JsonNode::Object(pairs, _) = result {
assert_eq!(pairs[0].key, "@type");
assert!(matches!(pairs[0].value, JsonNode::String(_, _)));
} else {
panic!("Expected object");
}
}
}