Branching

nanoqsh · Feb 16, 2024 · a7313ae · a7313ae
1 parent 0ff36e6
commit a7313ae
Show file tree

Hide file tree

Showing 12 changed files with 342 additions and 107 deletions.
diff --git a/dunge/tests/shader.rs b/dunge/tests/shader.rs
@@ -23,7 +23,7 @@ fn shader_calc() -> Result<(), Error> {
 
     let cx = helpers::block_on(dunge::context())?;
     let shader = cx.make_shader(compute);
-    assert_eq!(shader.debug_wgsl(), include_str!("shader_calc.wgsl"));
+    helpers::eq_lines(shader.debug_wgsl(), include_str!("shader_calc.wgsl"));
     Ok(())
 }
 
@@ -41,7 +41,25 @@ fn shader_if() -> Result<(), Error> {
 
     let cx = helpers::block_on(dunge::context())?;
     let shader = cx.make_shader(compute);
-    assert_eq!(shader.debug_wgsl(), include_str!("shader_if.wgsl"));
+    helpers::eq_lines(shader.debug_wgsl(), include_str!("shader_if.wgsl"));
+    Ok(())
+}
+
+#[test]
+fn shader_branch() -> Result<(), Error> {
+    use dunge::sl::{self, Out};
+
+    let compute = || Out {
+        place: sl::default(|| sl::splat_vec4(3.))
+            .when(true, || sl::splat_vec4(1.))
+            .when(false, || sl::splat_vec4(2.)),
+        color: sl::splat_vec4(1.),
+    };
+
+    let cx = helpers::block_on(dunge::context())?;
+    let shader = cx.make_shader(compute);
+    // helpers::eq_lines(shader.debug_wgsl(), include_str!("shader_branch.wgsl"));
+    _ = std::fs::write("tests/shader_branch.wgsl", shader.debug_wgsl());
     Ok(())
 }
 

diff --git a/dunge/tests/shader_branch.wgsl b/dunge/tests/shader_branch.wgsl
@@ -0,0 +1,25 @@
+struct VertexOutput {
+    @builtin(position) member: vec4<f32>,
+}
+
+@vertex 
+fn vs() -> VertexOutput {
+    var local: vec4<f32>;
+
+    if true {
+        local = vec4<f32>(1f, 1f, 1f, 1f);
+    } else {
+        if false {
+            local = vec4<f32>(2f, 2f, 2f, 2f);
+        } else {
+            local = vec4<f32>(3f, 3f, 3f, 3f);
+        }
+    }
+    let _e9: vec4<f32> = local;
+    return VertexOutput(_e9);
+}
+
+@fragment 
+fn fs(param: VertexOutput) -> @location(0) vec4<f32> {
+    return vec4<f32>(1f, 1f, 1f, 1f);
+}
diff --git a/dunge/tests/triangle_group.rs b/dunge/tests/triangle_group.rs
@@ -40,7 +40,7 @@ fn render() -> Result<(), Error> {
 
     let cx = helpers::block_on(dunge::context())?;
     let shader = cx.make_shader(triangle);
-    assert_eq!(shader.debug_wgsl(), include_str!("triangle_group.wgsl"));
+    helpers::eq_lines(shader.debug_wgsl(), include_str!("triangle_group.wgsl"));
 
     let map = {
         let texture = {

diff --git a/dunge/tests/triangle_index.rs b/dunge/tests/triangle_index.rs
@@ -32,7 +32,7 @@ fn render() -> Result<(), Error> {
 
     let cx = helpers::block_on(dunge::context())?;
     let shader = cx.make_shader(triangle);
-    assert_eq!(shader.debug_wgsl(), include_str!("triangle_index.wgsl"));
+    helpers::eq_lines(shader.debug_wgsl(), include_str!("triangle_index.wgsl"));
 
     let layer = cx.make_layer(&shader, Format::RgbAlpha);
     let view = {

diff --git a/dunge/tests/triangle_instance.rs b/dunge/tests/triangle_instance.rs
@@ -36,7 +36,7 @@ fn render() -> Result<(), Error> {
 
     let cx = helpers::block_on(dunge::context())?;
     let shader = cx.make_shader(triangle);
-    assert_eq!(shader.debug_wgsl(), include_str!("triangle_instance.wgsl"));
+    helpers::eq_lines(shader.debug_wgsl(), include_str!("triangle_instance.wgsl"));
 
     let layer = cx.make_layer(&shader, Format::RgbAlpha);
     let view = {

diff --git a/dunge/tests/triangle_vertex.png b/dunge/tests/triangle_vertex.png
diff --git a/dunge/tests/triangle_vertex.rs b/dunge/tests/triangle_vertex.rs
@@ -30,7 +30,7 @@ fn render() -> Result<(), Error> {
 
     let cx = helpers::block_on(dunge::context())?;
     let shader = cx.make_shader(triangle);
-    assert_eq!(shader.debug_wgsl(), include_str!("triangle_vertex.wgsl"));
+    helpers::eq_lines(shader.debug_wgsl(), include_str!("triangle_vertex.wgsl"));
 
     let layer = cx.make_layer(&shader, Format::RgbAlpha);
     let view = {

diff --git a/dunge_shader/src/branch.rs b/dunge_shader/src/branch.rs
@@ -0,0 +1,169 @@
+use {
+    crate::{
+        eval::{Branch, Eval, Expr, GetEntry},
+        ret::Ret,
+        types,
+    },
+    std::marker::PhantomData,
+};
+
+pub fn if_then_else<C, A, B, X, Y, E>(c: C, a: A, b: B) -> Ret<IfThenElse<C, A, B, E>, X::Out>
+where
+    C: Eval<E, Out = bool>,
+    A: FnOnce() -> X,
+    B: FnOnce() -> Y,
+    X: Eval<E>,
+    X::Out: types::Value,
+    Y: Eval<E, Out = X::Out>,
+{
+    Ret::new(IfThenElse {
+        c,
+        a,
+        b,
+        e: PhantomData,
+    })
+}
+
+pub struct IfThenElse<C, A, B, E> {
+    c: C,
+    a: A,
+    b: B,
+    e: PhantomData<E>,
+}
+
+impl<C, A, B, X, Y, E> Eval<E> for Ret<IfThenElse<C, A, B, E>, X::Out>
+where
+    C: Eval<E>,
+    A: FnOnce() -> X,
+    B: FnOnce() -> Y,
+    X: Eval<E>,
+    X::Out: types::Value,
+    Y: Eval<E>,
+    E: GetEntry,
+{
+    type Out = X::Out;
+
+    fn eval(self, en: &mut E) -> Expr {
+        let IfThenElse { c, a, b, .. } = self.get();
+        let c = c.eval(en);
+        let a = |en: &mut E| a().eval(en);
+        let b = |en: &mut E| Some(b().eval(en));
+        let valty = <X::Out as types::Value>::VALUE_TYPE;
+        let ty = en.get_entry().new_type(valty.ty());
+        let branch = Branch::new(en.get_entry(), ty);
+        branch.add(en, c, a, b);
+        branch.load(en.get_entry())
+    }
+}
+
+pub fn default<B, Y, E>(expr: B) -> Else<B>
+where
+    B: FnOnce() -> Y,
+    Y: Eval<E>,
+{
+    Else(expr)
+}
+
+pub struct Else<B>(B);
+
+impl<B> Else<B> {
+    pub fn when<C, A, X, Y, E>(self, cond: C, expr: A) -> Ret<When<C, A, B, E>, X::Out>
+    where
+        C: Eval<E, Out = bool>,
+        A: FnOnce() -> X,
+        B: FnOnce() -> Y,
+        X: Eval<E>,
+        X::Out: types::Value,
+        Y: Eval<E, Out = X::Out>,
+    {
+        Ret::new(When {
+            c: cond,
+            a: expr,
+            b: self.0,
+            e: PhantomData,
+        })
+    }
+}
+
+pub struct When<C, A, B, E> {
+    c: C,
+    a: A,
+    b: B,
+    e: PhantomData<E>,
+}
+
+impl<C, A, B, E, O> Ret<When<C, A, B, E>, O> {
+    #[allow(clippy::type_complexity)]
+    pub fn when<D, F, Z>(self, cond: D, expr: F) -> Ret<When<C, A, When<D, F, B, E>, E>, O>
+    where
+        D: Eval<E, Out = bool>,
+        F: FnOnce() -> Z,
+        Z: Eval<E, Out = O>,
+    {
+        let when = self.get();
+        Ret::new(When {
+            c: when.c,
+            a: when.a,
+            b: When {
+                c: cond,
+                a: expr,
+                b: when.b,
+                e: PhantomData,
+            },
+            e: PhantomData,
+        })
+    }
+}
+
+impl<C, A, B, X, E> Eval<E> for Ret<When<C, A, B, E>, X::Out>
+where
+    C: Eval<E>,
+    A: FnOnce() -> X,
+    B: EvalBranch<E>,
+    X: Eval<E>,
+    X::Out: types::Value,
+    E: GetEntry,
+{
+    type Out = X::Out;
+
+    fn eval(self, en: &mut E) -> Expr {
+        let when = self.get();
+        let valty = <X::Out as types::Value>::VALUE_TYPE;
+        let ty = en.get_entry().new_type(valty.ty());
+        let branch = Branch::new(en.get_entry(), ty);
+        when.eval_else(en, &branch);
+        branch.load(en.get_entry())
+    }
+}
+
+pub trait EvalBranch<E> {
+    fn eval_else(self, en: &mut E, branch: &Branch) -> Option<Expr>;
+}
+
+impl<F, R, E> EvalBranch<E> for F
+where
+    F: FnOnce() -> R,
+    R: Eval<E>,
+{
+    fn eval_else(self, en: &mut E, _: &Branch) -> Option<Expr> {
+        Some(self().eval(en))
+    }
+}
+
+impl<C, A, B, X, E> EvalBranch<E> for When<C, A, B, E>
+where
+    C: Eval<E>,
+    A: FnOnce() -> X,
+    X: Eval<E>,
+    B: EvalBranch<E>,
+    E: GetEntry,
+{
+    fn eval_else(self, en: &mut E, branch: &Branch) -> Option<Expr> {
+        let Self { c, a, b, .. } = self;
+        let c = c.eval(en);
+        let a = |en: &mut E| a().eval(en);
+        let b = |en: &mut E| b.eval_else(en, branch);
+        branch.add(en, c, a, b);
+        None
+    }
+}