154 lines
4.1 KiB
Rust
154 lines
4.1 KiB
Rust
use slab::Slab;
|
|
|
|
pub mod command;
|
|
pub mod node;
|
|
|
|
use node::{AtomOp, Edge, Node, NodeKind, SlotIndex};
|
|
|
|
#[derive(Debug)]
|
|
pub enum GraphError {
|
|
InvalidReference,
|
|
MissingEdge,
|
|
Cyclic,
|
|
}
|
|
|
|
pub type GraphResult<T> = Result<T, GraphError>;
|
|
|
|
#[derive(Debug)]
|
|
pub struct Graph {
|
|
pub nodes: Slab<Node>,
|
|
}
|
|
|
|
impl Graph {
|
|
pub fn new() -> Self {
|
|
Self { 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)
|
|
}
|
|
|
|
pub fn compile(&self) -> GraphResult<()> {
|
|
// TODO: Only compile the nodes that are actually used (set an output node)
|
|
|
|
// Iterate through the nodes in topological order
|
|
for key in self.topological_order()? {
|
|
let node = &self.nodes[key];
|
|
match node.kind {
|
|
NodeKind::Lit(x) => {
|
|
println!("let n{}: f32 = {:?};", key, x);
|
|
}
|
|
NodeKind::Op(op) => {
|
|
let symbol = match op {
|
|
AtomOp::Add => '+',
|
|
AtomOp::Sub => '-',
|
|
AtomOp::Mul => '*',
|
|
AtomOp::Div => '/',
|
|
};
|
|
|
|
let inputs = node.unwrap_inputs()?;
|
|
let lhs = inputs[0].node;
|
|
let rhs = inputs[1].node;
|
|
println!("let n{}: f32 = n{} {} n{};", key, lhs, symbol, rhs);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<Vec<Node>> for Graph {
|
|
fn from(vec: Vec<Node>) -> Self {
|
|
let mut nodes = Slab::with_capacity(vec.len());
|
|
|
|
for (idx, node) in vec.into_iter().enumerate() {
|
|
let new_idx = nodes.insert(node);
|
|
assert_eq!(new_idx, idx);
|
|
}
|
|
|
|
Self { nodes }
|
|
}
|
|
}
|