use std::collections::HashMap;
use panproto_schema::Edge;
use serde::{Deserialize, Serialize};
use crate::error::RestrictError;
use crate::metadata::Node;
use crate::value::Value;
use crate::wtype::CompiledMigration;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct GInstance {
pub nodes: HashMap<u32, Node>,
pub edges: Vec<(u32, u32, Edge)>,
pub values: HashMap<u32, Value>,
}
impl GInstance {
#[must_use]
pub fn new() -> Self {
Self {
nodes: HashMap::new(),
edges: Vec::new(),
values: HashMap::new(),
}
}
#[must_use]
pub fn with_node(mut self, node: Node) -> Self {
self.nodes.insert(node.id, node);
self
}
#[must_use]
pub fn with_edge(mut self, src: u32, tgt: u32, edge: Edge) -> Self {
self.edges.push((src, tgt, edge));
self
}
#[must_use]
pub fn with_value(mut self, node_id: u32, value: Value) -> Self {
self.values.insert(node_id, value);
self
}
#[must_use]
pub fn node_count(&self) -> usize {
self.nodes.len()
}
#[must_use]
pub fn edge_count(&self) -> usize {
self.edges.len()
}
}
impl Default for GInstance {
fn default() -> Self {
Self::new()
}
}
pub fn graph_restrict(
instance: &GInstance,
migration: &CompiledMigration,
) -> Result<GInstance, RestrictError> {
let mut new_nodes = HashMap::new();
let mut new_edges = Vec::new();
let mut new_values = HashMap::new();
for (id, node) in &instance.nodes {
if let Some(new_anchor) = migration.vertex_remap.get(&node.anchor) {
let mut new_node = node.clone();
new_anchor.clone_into(&mut new_node.anchor);
new_nodes.insert(*id, new_node);
if let Some(value) = instance.values.get(id) {
new_values.insert(*id, value.clone());
}
}
}
for (src, tgt, edge) in &instance.edges {
if new_nodes.contains_key(src) && new_nodes.contains_key(tgt) {
if let Some(new_edge) = migration.edge_remap.get(edge) {
new_edges.push((*src, *tgt, new_edge.clone()));
} else if migration.surviving_edges.contains(edge) {
new_edges.push((*src, *tgt, edge.clone()));
}
}
}
Ok(GInstance {
nodes: new_nodes,
edges: new_edges,
values: new_values,
})
}
pub fn graph_extend(
instance: &GInstance,
migration: &CompiledMigration,
) -> Result<GInstance, RestrictError> {
let mut new_nodes = HashMap::new();
let mut new_edges = Vec::new();
let mut new_values = HashMap::new();
for (&id, node) in &instance.nodes {
let mut new_node = node.clone();
if let Some(remapped) = migration.vertex_remap.get(&node.anchor) {
new_node.anchor.clone_from(remapped);
} else if !migration.surviving_verts.contains(&node.anchor) {
continue;
}
new_nodes.insert(id, new_node);
if let Some(value) = instance.values.get(&id) {
new_values.insert(id, value.clone());
}
}
for (src, tgt, edge) in &instance.edges {
if !new_nodes.contains_key(src) || !new_nodes.contains_key(tgt) {
continue;
}
if let Some(new_edge) = migration.edge_remap.get(edge) {
new_edges.push((*src, *tgt, new_edge.clone()));
} else if migration.surviving_edges.contains(edge) {
new_edges.push((*src, *tgt, edge.clone()));
}
}
Ok(GInstance {
nodes: new_nodes,
edges: new_edges,
values: new_values,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_graph_instance() {
let g = GInstance::new();
assert_eq!(g.node_count(), 0);
assert_eq!(g.edge_count(), 0);
}
#[test]
fn build_graph_instance() {
let g = GInstance::new()
.with_node(Node::new(0, "person"))
.with_node(Node::new(1, "person"))
.with_edge(
0,
1,
Edge {
src: "person".into(),
tgt: "person".into(),
kind: "knows".into(),
name: None,
},
)
.with_value(0, Value::Str("Alice".into()))
.with_value(1, Value::Str("Bob".into()));
assert_eq!(g.node_count(), 2);
assert_eq!(g.edge_count(), 1);
assert_eq!(g.values.len(), 2);
}
#[test]
#[allow(clippy::expect_used)]
fn graph_restrict_drops_unmapped_nodes() {
let g = GInstance::new()
.with_node(Node::new(0, "a"))
.with_node(Node::new(1, "b"))
.with_node(Node::new(2, "c"));
let migration = CompiledMigration {
surviving_verts: std::iter::once("a_new".into()).collect(),
surviving_edges: std::collections::HashSet::new(),
vertex_remap: std::iter::once(("a".into(), "a_new".into())).collect(),
edge_remap: HashMap::new(),
resolver: HashMap::new(),
hyper_resolver: HashMap::new(),
field_transforms: HashMap::new(),
};
let restricted = graph_restrict(&g, &migration).expect("graph_restrict should succeed");
assert_eq!(restricted.node_count(), 1);
assert_eq!(restricted.nodes[&0].anchor, "a_new");
}
}