use graph_api_lib::Graph;
use graph_api_test::{Edge, Language, Project, Refs, Vertex, populate_graph};
use rand::prelude::*;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use std::collections::HashSet;
use uuid::Uuid;
#[derive(Clone, Copy)]
pub enum GraphSize {
Small, Medium, Large, Huge, }
impl GraphSize {
pub fn vertex_count(&self) -> usize {
match self {
GraphSize::Small => 100,
GraphSize::Medium => 1_000,
GraphSize::Large => 10_000,
GraphSize::Huge => 100_000,
}
}
pub fn edge_multiplier(&self) -> usize {
3 }
}
pub fn generate_test_graph<G>(graph: &mut G) -> Refs<G>
where
G: Graph<Vertex = Vertex, Edge = Edge>,
{
populate_graph(graph)
}
pub fn generate_random_graph<G>(graph: &mut G, size: GraphSize, seed: u64) -> Vec<G::VertexId>
where
G: Graph<Vertex = Vertex, Edge = Edge>,
{
let mut rng = StdRng::seed_from_u64(seed);
let vertex_count = size.vertex_count();
let edge_count = vertex_count * size.edge_multiplier();
let mut vertex_ids = Vec::with_capacity(vertex_count);
for i in 0..vertex_count {
let vertex = match rng.random_range(0..3) {
0 => Vertex::Person {
name: format!("Person-{}", i),
age: rng.random_range(18..80),
unique_id: Uuid::new_v4(),
username: format!("user_{}", i),
biography: format!("Bio for person {}: {}", i, random_biography(&mut rng)),
},
1 => Vertex::Project(Project {
name: format!("Project-{}", i),
}),
_ => Vertex::Rust,
};
vertex_ids.push(graph.add_vertex(vertex));
}
let mut added_edges = HashSet::new();
for _ in 0..edge_count {
let src_idx = rng.random_range(0..vertex_count);
let tgt_idx = rng.random_range(0..vertex_count);
let src = vertex_ids[src_idx];
let tgt = vertex_ids[tgt_idx];
if src == tgt || added_edges.contains(&(src, tgt)) {
continue;
}
let edge = match rng.random_range(0..3) {
0 => Edge::Knows {
since: rng.random_range(1980..2023),
},
1 => Edge::Created,
_ => Edge::Language(Language {
name: match rng.random_range(0..4) {
0 => "Rust".to_string(),
1 => "Java".to_string(),
2 => "Python".to_string(),
_ => "JavaScript".to_string(),
},
}),
};
graph.add_edge(src, tgt, edge);
added_edges.insert((src, tgt));
}
vertex_ids
}
pub fn generate_social_graph<G>(graph: &mut G, size: GraphSize, seed: u64) -> Vec<G::VertexId>
where
G: Graph<Vertex = Vertex, Edge = Edge>,
{
let mut rng = StdRng::seed_from_u64(seed);
let vertex_count = size.vertex_count();
let mut vertex_ids = Vec::with_capacity(vertex_count);
for i in 0..vertex_count {
let vertex = Vertex::Person {
name: format!("Person-{}", i),
age: rng.random_range(18..80),
unique_id: Uuid::new_v4(),
username: format!("user_{}", i),
biography: format!("Social network user {}", i),
};
vertex_ids.push(graph.add_vertex(vertex));
}
for i in 0..vertex_count {
let num_connections = rng.random_range(5..15).min(vertex_count - 1);
let mut connections = HashSet::new();
while connections.len() < num_connections {
let target_idx = if rng.random_bool(0.8) {
let range = 50; let start = i.saturating_sub(range / 2);
let end = (i + range / 2).min(vertex_count - 1);
rng.random_range(start..=end)
} else {
rng.random_range(0..vertex_count)
};
if target_idx != i {
connections.insert(target_idx);
}
}
for target_idx in connections {
let target = vertex_ids[target_idx];
let year = rng.random_range(1980..2023);
graph.add_edge(vertex_ids[i], target, Edge::Knows { since: year });
}
}
vertex_ids
}
pub fn generate_project_graph<G>(graph: &mut G, size: GraphSize, seed: u64) -> Vec<G::VertexId>
where
G: Graph<Vertex = Vertex, Edge = Edge>,
{
let mut rng = StdRng::seed_from_u64(seed);
let project_count = size.vertex_count() / 3; let person_count = size.vertex_count() / 3; let rust_count = size.vertex_count() - project_count - person_count;
let mut project_ids = Vec::with_capacity(project_count);
for i in 0..project_count {
let vertex = Vertex::Project(Project {
name: format!("Project-{}", i),
});
project_ids.push(graph.add_vertex(vertex));
}
let mut person_ids = Vec::with_capacity(person_count);
for i in 0..person_count {
let vertex = Vertex::Person {
name: format!("Developer-{}", i),
age: rng.random_range(18..80),
unique_id: Uuid::new_v4(),
username: format!("dev_{}", i),
biography: format!("Developer working on project {}", i % project_count),
};
person_ids.push(graph.add_vertex(vertex));
}
let mut rust_ids = Vec::with_capacity(rust_count);
for _ in 0..rust_count {
rust_ids.push(graph.add_vertex(Vertex::Rust));
}
for (i, &project_id) in project_ids.iter().enumerate() {
let num_deps = rng.random_range(0..=5).min(project_count - 1);
let mut deps = HashSet::new();
while deps.len() < num_deps {
let target_idx = rng.random_range(0..project_count);
if target_idx != i {
deps.insert(target_idx);
}
}
for target_idx in deps {
let target = project_ids[target_idx];
graph.add_edge(project_id, target, Edge::Created);
}
}
for (i, &person_id) in person_ids.iter().enumerate() {
let num_projects = rng.random_range(1..=3).min(project_count);
let mut projects = HashSet::new();
projects.insert(i % project_count);
while projects.len() < num_projects {
projects.insert(rng.random_range(0..project_count));
}
for &project_idx in &projects {
let project_id = project_ids[project_idx];
graph.add_edge(person_id, project_id, Edge::Created);
}
}
for &project_id in &project_ids {
let num_langs = rng.random_range(1..=2).min(rust_count);
for _ in 0..num_langs {
let lang_idx = rng.random_range(0..rust_count);
let lang_id = rust_ids[lang_idx];
graph.add_edge(
project_id,
lang_id,
Edge::Language(Language {
name: "Rust".to_string(),
}),
);
}
}
let mut all_vertices = Vec::with_capacity(size.vertex_count());
all_vertices.extend(project_ids);
all_vertices.extend(person_ids);
all_vertices.extend(rust_ids);
all_vertices
}
fn random_biography(rng: &mut StdRng) -> String {
let adjectives = [
"creative",
"diligent",
"innovative",
"experienced",
"passionate",
];
let roles = ["developer", "engineer", "architect", "programmer", "coder"];
let interests = [
"graph databases",
"distributed systems",
"machine learning",
"web development",
"mobile apps",
];
format!(
"A {} {} interested in {}",
adjectives.choose(rng).unwrap(),
roles.choose(rng).unwrap(),
interests.choose(rng).unwrap()
)
}