1 files changed, 104 insertions, 0 deletions
diff --git a/src/snarl_ext.rs b/src/snarl_ext.rs
new file mode 100644
index 0000000..7d0464d
--- /dev/null
+++ b/src/snarl_ext.rs
@@ -0,0 +1,104 @@
+use egui_snarl::{InPinId, NodeId, OutPinId, Snarl};
+use std::fmt;
+use topological_sort::TopologicalSort;
+
+use crate::node::{AnyNode, ConcreteNode};
+use crate::types::{Type, TypeSignature};
+
+pub trait WithSignatures {
+    fn in_pin_type(&self, pin: InPinId) -> Option<Type>;
+    fn out_pin_type(&self, pin: OutPinId) -> Type;
+    fn get_input_types(&self, node: NodeId) -> Vec<Option<Type>>;
+    fn get_node_signature(&self, node: NodeId) -> TypeSignature;
+}
+
+pub trait Compilable {
+    fn compile_node(&self, node: NodeId, f: &mut dyn fmt::Write) -> fmt::Result;
+    fn compile(&self, f: &mut dyn fmt::Write) -> fmt::Result;
+}
+
+impl WithSignatures for Snarl<AnyNode> {
+    fn in_pin_type(&self, pin: InPinId) -> Option<Type> {
+        let sig = &self.get_node_signature(pin.node);
+        sig.inputs.get(pin.input).copied()
+    }
+
+    fn out_pin_type(&self, pin: OutPinId) -> Type {
+        let sig = &self.get_node_signature(pin.node);
+        sig.outputs[pin.output]
+    }
+
+    fn get_input_types(&self, node: NodeId) -> Vec<Option<Type>> {
+        (0..self[node].max_inputs())
+            .map(
+                move |input| match &*self.in_pin(InPinId { node, input }).remotes {
+                    [] => None,
+                    [out_pin] => Some(self.out_pin_type(*out_pin)),
+                    _ => unreachable!("cannot connect to multiple inputs"),
+                },
+            )
+            .collect()
+    }
+
+    fn get_node_signature(&self, node: NodeId) -> TypeSignature {
+        self[node].signature(&self.get_input_types(node))
+    }
+}
+
+fn compile_output(pin: OutPinId) -> String {
+    let node_id = pin.node.0;
+    let out_id = pin.output;
+    format!("n{node_id}_o{out_id}")
+}
+
+impl Compilable for Snarl<AnyNode> {
+    fn compile(&self, f: &mut dyn fmt::Write) -> fmt::Result {
+        let mut order = TopologicalSort::<NodeId>::new();
+        for (out, inp) in self.wires() {
+            order.add_dependency(out.node, inp.node);
+        }
+
+        write!(
+            f,
+            "
+#version 450
+in vec2 in_uv;
+out vec4 out_color;
+
+layout(std140, binding = 0) uniform Uniforms {{
+    vec2 in_resolution;
+    float in_time;
+    float _pad3;
+}} uniforms;
+
+void main() {{
+"
+        )?;
+        while let Some(id) = order.pop() {
+            self.compile_node(id, f)?;
+        }
+        writeln!(f, "}}")?;
+
+        Ok(())
+    }
+
+    fn compile_node(&self, node: NodeId, f: &mut dyn fmt::Write) -> fmt::Result {
+        let signature = self.get_node_signature(node);
+
+        let inputs: Result<Vec<String>, fmt::Error> = (0..signature.inputs.len())
+            .map(
+                |input| match &*self.in_pin(InPinId { node, input }).remotes {
+                    [] => Err(fmt::Error),
+                    [pin] => Ok(compile_output(*pin)),
+                    _ => unreachable!("cannot connect to multiple inputs"),
+                },
+            )
+            .collect();
+
+        let outputs: Vec<String> = (0..signature.outputs.len())
+            .map(|output| compile_output(OutPinId { node, output }))
+            .collect();
+
+        self[node].compile(signature, inputs?, outputs, f)
+    }
+}