Compile node graph description to GPU code

node-graph/gcore/src/lib.rs

@@ -14,7 +14,7 @@ pub mod ops;
 pub mod value;
 
 pub trait Node<'n> {
-    type Output: 'n; // TODO: replace with generic associated type
+    type Output; // TODO: replace with generic associated type
 
     fn eval(&'n self) -> Self::Output;
 }
@@ -27,12 +27,26 @@ impl<'n, N: Node<'n>> Node<'n> for &'n N {
     }
 }
 
+pub trait NodeInput {
+    type Nodes;
+
+    fn new(input: Self::Nodes) -> Self;
+}
+
+trait FQN {
+    fn fqn(&self) -> &'static str;
+}
+
+trait Input<I> {
+    unsafe fn input(&self, input: I);
+}
+
 #[cfg(feature = "async")]
 #[async_trait]
 pub trait AsyncNode<'n> {
-    type Output: 'n; // TODO: replace with generic associated type
+    type Output; // TODO: replace with generic associated type
 
-    async fn eval(&'n self) -> Self::Output;
+    async fn eval_async(&'n self) -> Self::Output;
 }
 
 #[cfg(feature = "async")]
@@ -40,7 +54,7 @@ pub trait AsyncNode<'n> {
 impl<'n, N: Node<'n> + Sync> AsyncNode<'n> for N {
     type Output = N::Output;
 
-    async fn eval(&'n self) -> Self::Output {
+    async fn eval_async(&'n self) -> Self::Output {
         Node::eval(self)
     }
 }

node-graph/gcore/src/ops.rs

@@ -1,9 +1,9 @@
 use core::{marker::PhantomData, ops::Add};
 
-use crate::Node;
+use crate::{Node, NodeInput};
 
 #[repr(C)]
-struct AddNode<'n, L: Add<R>, R, I1: Node<'n, Output = L>, I2: Node<'n, Output = R>>(
+pub struct AddNode<'n, L: Add<R>, R, I1: Node<'n, Output = L>, I2: Node<'n, Output = R>>(
     pub I1,
     pub I2,
     PhantomData<&'n (L, R)>,
@@ -16,6 +16,13 @@ impl<'n, L: Add<R>, R, I1: Node<'n, Output = L>, I2: Node<'n, Output = R>> Node<
         self.0.eval() + self.1.eval()
     }
 }
+impl<'n, L: Add<R>, R, I1: Node<'n, Output = L>, I2: Node<'n, Output = R>>
+    AddNode<'n, L, R, I1, I2>
+{
+    pub fn new(input: (I1, I2)) -> AddNode<'n, L, R, I1, I2> {
+        AddNode(input.0, input.1, PhantomData)
+    }
+}
 
 #[repr(C)]
 pub struct CloneNode<'n, N: Node<'n, Output = &'n O>, O: Clone + 'n>(pub N, PhantomData<&'n ()>);
@@ -25,6 +32,11 @@ impl<'n, N: Node<'n, Output = &'n O>, O: Clone> Node<'n> for CloneNode<'n, N, O>
         self.0.eval().clone()
     }
 }
+impl<'n, N: Node<'n, Output = &'n O>, O: Clone> CloneNode<'n, N, O> {
+    pub const fn new(node: N) -> CloneNode<'n, N, O> {
+        CloneNode(node, PhantomData)
+    }
+}
 
 #[repr(C)]
 pub struct FstNode<'n, N: Node<'n>>(pub N, PhantomData<&'n ()>);
@@ -56,13 +68,12 @@ impl<'n, N: Node<'n>> Node<'n> for DupNode<'n, N> {
         (self.0.eval(), self.0.eval()) //TODO: use Copy/Clone implementation
     }
 }
+impl<'n, N: Node<'n>> NodeInput for DupNode<'n, N> {
+    type Nodes = N;
 
-#[repr(C)]
-/// Return the unit value
-pub struct UnitNode;
-impl<'n> Node<'n> for UnitNode {
-    type Output = ();
-    fn eval(&'n self) -> Self::Output {}
+    fn new(input: Self::Nodes) -> Self {
+        Self(input, PhantomData)
+    }
 }
 
 #[repr(C)]
@@ -74,11 +85,18 @@ impl<'n, N: Node<'n>> Node<'n> for IdNode<'n, N> {
         self.0.eval()
     }
 }
+impl<'n, N: Node<'n>> NodeInput for IdNode<'n, N> {
+    type Nodes = N;
+
+    fn new(input: Self::Nodes) -> Self {
+        Self(input, PhantomData)
+    }
+}
 
 pub fn foo() {
-    let unit = UnitNode;
-    let value = IdNode(crate::value::ValueNode::new(2u32), PhantomData);
-    let value2 = crate::value::ValueNode::new(4u32);
+    let unit = crate::value::UnitNode;
+    let value = IdNode(crate::value::ValueNode(2u32), PhantomData);
+    let value2 = crate::value::ValueNode(4u32);
     let dup = DupNode(&value, PhantomData);
     fn int(_: (), state: &u32) -> &u32 {
         state
@@ -120,10 +138,16 @@ pub mod gpu {
         #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] y: &mut [(u32, u32)],
         #[spirv(push_constant)] push_consts: &PushConsts,
     ) {
+        fn node_graph(input: Input) -> Output {
+            let n0 = ValueNode::new(input);
+            let n1 = IdNode::new(n0);
+            let n2 = IdNode::new(n1);
+            return n2.eval();
+        }
         let gid = global_id.x as usize;
         // Only process up to n, which is the length of the buffers.
         if global_id.x < push_consts.n {
-            y[gid] = OPERATION.eval(a[gid]);
+            y[gid] = node_graph(a[gid]);
         }
     }
     #[allow(unused)]

node-graph/gcore/src/value.rs

@@ -1,4 +1,6 @@
 use core::marker::PhantomData;
+use core::mem::MaybeUninit;
+use core::sync::atomic::AtomicBool;
 
 use crate::Node;
 
@@ -18,8 +20,7 @@ impl<'n, T: 'n> Node<'n> for ValueNode<T> {
         &self.0
     }
 }
-
-impl<'n, T> ValueNode<T> {
+impl<T> ValueNode<T> {
     pub const fn new(value: T) -> ValueNode<T> {
         ValueNode(value)
     }
@@ -33,8 +34,29 @@ impl<'n, T: Default + 'n> Node<'n> for DefaultNode<T> {
         T::default()
     }
 }
-impl<T> DefaultNode<T> {
-    pub const fn new() -> DefaultNode<T> {
-        DefaultNode(PhantomData)
+
+#[repr(C)]
+/// Return the unit value
+pub struct UnitNode;
+impl<'n> Node<'n> for UnitNode {
+    type Output = ();
+    fn eval(&'n self) -> Self::Output {}
+}
+
+pub struct InputNode<T>(MaybeUninit<T>, AtomicBool);
+impl<'n, T: 'n> Node<'n> for InputNode<T> {
+    type Output = &'n T;
+    fn eval(&'n self) -> Self::Output {
+        if self.1.load(core::sync::atomic::Ordering::SeqCst) {
+            unsafe { self.0.assume_init_ref() }
+        } else {
+            panic!("tried to access an input before setting it")
+        }
+    }
+}
+
+impl<T> InputNode<T> {
+    pub const fn new() -> InputNode<T> {
+        InputNode(MaybeUninit::uninit(), AtomicBool::new(false))
     }
 }