cyborg/ramen/src/lib.rs

325 lines
8.5 KiB
Rust

use slab::Slab;
use std::sync::Arc;
pub mod command;
pub mod node;
pub mod node_kind;
use command::*;
use node::{Edge, Node, SlotIndex};
use node_kind::*;
#[derive(Debug)]
pub enum GraphError {
InvalidReference,
MissingEdge,
Cyclic,
}
pub type GraphResult<T> = Result<T, GraphError>;
#[derive(Debug)]
pub struct Graph {
pub node_kinds: Arc<NodeKindStore>,
pub nodes: Slab<Node>,
}
impl Graph {
pub fn new(node_kinds: Arc<NodeKindStore>) -> Self {
Self {
node_kinds,
nodes: Slab::new(),
}
}
pub fn get_node(&self, index: usize) -> GraphResult<&Node> {
self.nodes.get(index).ok_or(GraphError::InvalidReference)
}
pub fn get_node_mut(&mut self, index: usize) -> GraphResult<&mut Node> {
self.nodes
.get_mut(index)
.ok_or(GraphError::InvalidReference)
}
// Return the edges that can be derived from the graph
pub fn edges(&self) -> Result<Vec<Edge>, GraphError> {
let topo_order = self.topological_order()?;
let mut edges = Vec::new();
for (node_idx, node) in topo_order.iter().enumerate() {
for (slot_idx, input) in self.nodes[*node].inputs.iter().enumerate() {
if let Some(input) = input {
let output = SlotIndex {
node: node_idx,
slot: slot_idx,
};
edges.push(Edge {
input: *input,
output,
});
}
}
}
Ok(edges)
}
pub fn topological_order(&self) -> Result<Vec<usize>, GraphError> {
// The number of nodes
let v = self.nodes.len();
// Forward-directed adjacency lists
let mut adjacency_lists: Vec<Vec<usize>> = vec![Vec::new(); v];
for (i, node) in self.nodes.iter() {
for j in node.inputs.iter().flatten() {
adjacency_lists[j.node].push(i);
}
}
// Create list of indegrees
let mut indegrees: Vec<usize> = vec![0; v];
for (i, node) in self.nodes.iter() {
indegrees[i] = node.inputs.len();
}
// Create a queue and initialize it with nodes with no indegrees
let mut queue: Vec<usize> = Vec::new();
for (i, count) in indegrees.iter().enumerate() {
if *count == 0 {
queue.push(i);
}
}
let mut sort_count = 0;
let mut topo_order: Vec<usize> = Vec::new();
while !queue.is_empty() {
let u = queue.pop().unwrap();
topo_order.push(u);
for i in &adjacency_lists[u] {
indegrees[*i] -= 1;
if indegrees[*i] == 0 {
queue.push(*i);
}
}
sort_count += 1;
}
// Check if there was a cycle
if sort_count != v {
Err(GraphError::Cyclic)
} else {
Ok(topo_order)
}
}
}
pub enum TestOpKind {
Add,
Sub,
Mul,
Div,
}
pub struct TestGraphBuilder {
pub graph: Graph,
pub literal_kind: usize,
pub add_kind: usize,
pub sub_kind: usize,
pub mul_kind: usize,
pub div_kind: usize,
}
impl TestGraphBuilder {
pub fn new() -> Self {
let mut kinds = NodeKindStore { kinds: vec![] };
let literal_kind = kinds.push(NodeKind {
title: "Literal".into(),
inputs: vec![],
outputs: vec![NodeKindOutput { label: "".into() }],
});
let mut make_arith_kind = |title: &str| -> usize {
kinds.push(NodeKind {
title: title.to_string(),
inputs: vec![
NodeKindInput {
label: "A".to_string(),
},
NodeKindInput {
label: "B".to_string(),
},
],
outputs: vec![NodeKindOutput {
label: "Result".to_string(),
}],
})
};
let add_kind = make_arith_kind("Add");
let sub_kind = make_arith_kind("Subtract");
let mul_kind = make_arith_kind("Multiply");
let div_kind = make_arith_kind("Divide");
let graph = Graph::new(Arc::new(kinds));
Self {
graph,
literal_kind,
add_kind,
sub_kind,
mul_kind,
div_kind,
}
}
pub fn add(&mut self, node: Node) -> usize {
self.graph.nodes.insert(node)
}
pub fn add_literal(&mut self, value: f32) -> usize {
self.graph.nodes.insert(self.make_literal(value))
}
pub fn make_literal(&self, _value: f32) -> Node {
Node {
kind: self.literal_kind,
pos: glam::Vec2::ZERO,
inputs: vec![],
}
}
pub fn add_op(&mut self, kind: TestOpKind, a: usize, b: usize) -> usize {
self.graph.nodes.insert(self.make_op(kind, a, b))
}
pub fn make_op(&self, kind: TestOpKind, a: usize, b: usize) -> Node {
let kind = match kind {
TestOpKind::Add => self.add_kind,
TestOpKind::Sub => self.sub_kind,
TestOpKind::Mul => self.mul_kind,
TestOpKind::Div => self.div_kind,
};
Node {
kind,
pos: glam::Vec2::ZERO,
inputs: vec![
Some(SlotIndex { node: a, slot: 0 }),
Some(SlotIndex { node: b, slot: 0 }),
],
}
}
pub fn three_node_add_cmds(&self) -> Vec<DynCommand> {
vec![
Box::new(CreateNode::from(self.make_literal(1.0))),
Box::new(CreateNode::from(self.make_literal(2.0))),
Box::new(CreateNode::from(Node {
kind: self.add_kind,
pos: glam::Vec2::ZERO,
inputs: vec![None; 2],
})),
Box::new(SetEdge {
edge: Edge {
input: SlotIndex { node: 0, slot: 0 },
output: SlotIndex { node: 2, slot: 0 },
},
}),
Box::new(SetEdge {
edge: Edge {
input: SlotIndex { node: 1, slot: 0 },
output: SlotIndex { node: 2, slot: 1 },
},
}),
]
}
}
#[cfg(test)]
mod tests {
use super::*;
impl CommandHistory {
pub fn push_multi(
&mut self,
graph: &mut Graph,
cmds: impl IntoIterator<Item = DynCommand>,
) -> GraphResult<()> {
for cmd in cmds {
println!("{:#?}", cmd);
self.push(graph, cmd)?;
println!("{:#?}", graph);
}
Ok(())
}
pub fn rewind(&mut self, graph: &mut Graph) -> GraphResult<()> {
while self.cursor > 0 {
self.cursor -= 1;
let undo = &self.commands[self.cursor].1;
println!("{:#?}", undo);
undo.apply(graph)?;
println!("{:#?}", graph);
}
Ok(())
}
}
#[test]
fn create_node() -> GraphResult<()> {
let mut builder = TestGraphBuilder::new();
let cmd = CreateNode {
kind: builder.literal_kind,
pos: glam::Vec2::ZERO,
inputs: vec![],
};
let undo = cmd.undo(&builder.graph)?;
cmd.apply(&mut builder.graph)?;
assert_eq!(builder.graph.nodes.len(), 1);
undo.apply(&mut builder.graph)?;
assert_eq!(builder.graph.nodes.len(), 0);
Ok(())
}
#[test]
fn three_node_add() -> GraphResult<()> {
let mut builder = TestGraphBuilder::new();
let cmds = builder.three_node_add_cmds();
for cmd in cmds.into_iter() {
println!("{:#?}", cmd);
cmd.apply(&mut builder.graph)?;
println!("{:#?}", builder.graph);
}
builder.graph.topological_order()?;
Ok(())
}
#[test]
fn three_node_add_rewind() -> GraphResult<()> {
let mut builder = TestGraphBuilder::new();
let mut history = CommandHistory::new();
let cmds = builder.three_node_add_cmds();
history.push_multi(&mut builder.graph, cmds)?;
builder.graph.topological_order()?;
history.rewind(&mut builder.graph)?;
assert_eq!(builder.graph.nodes.len(), 0);
Ok(())
}
}