From 04c3e792ddc95d160e5627e3b4c3181fa4bf34cc Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Tue, 12 Dec 2023 23:40:45 -0500
Subject: [PATCH] minor

---
 luisa_compute/src/lang/control_flow.rs | 24 ++++++++++++++++--------
 luisa_compute/src/lang/debug.rs        | 13 +++++++++----
 luisa_compute/src/lang/ops/impls.rs    | 14 ++++++++++++++
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/luisa_compute/src/lang/control_flow.rs b/luisa_compute/src/lang/control_flow.rs
index 2c43a61..25f0b14 100644
--- a/luisa_compute/src/lang/control_flow.rs
+++ b/luisa_compute/src/lang/control_flow.rs
@@ -3,6 +3,8 @@ use std::ffi::CString;
 use crate::internal_prelude::*;
 use ir::SwitchCase;
 
+use super::debug::__unreachable_typed;
+
 /**
  * If you want rustfmt to format your code, use if_!(cond, { .. }, { .. })
  * or if_!(cond, { .. }, else, {...}) instead of if_!(cond, { .. }, else
@@ -426,14 +428,20 @@ impl<R: Aggregate> SwitchBuilder<R> {
                 s.push(IrBuilder::new(pools));
             });
             for i in 0..phi_count {
-                let msg = CString::new("unreachable code in switch statement!").unwrap();
-                let default_node = __current_scope(|b| {
-                    b.call(
-                        Func::Unreachable(CBoxedSlice::from(msg)),
-                        &[],
-                        case_phis[0][i].type_().clone(),
-                    )
-                });
+                // let msg = CString::new("unreachable code in switch statement!").unwrap();
+                // let default_node = __current_scope(|b| {
+                //     b.call(
+                //         Func::Unreachable(CBoxedSlice::from(msg)),
+                //         &[],
+                //         case_phis[0][i].type_().clone(),
+                //     )
+                // });
+                let default_node = __unreachable_typed(
+                    case_phis[0][i].type_().clone(),
+                    file!(),
+                    line!(),
+                    column!(),
+                );
                 default_nodes.push(default_node);
             }
             __pop_scope()
diff --git a/luisa_compute/src/lang/debug.rs b/luisa_compute/src/lang/debug.rs
index 038023c..106b5b9 100644
--- a/luisa_compute/src/lang/debug.rs
+++ b/luisa_compute/src/lang/debug.rs
@@ -71,8 +71,13 @@ pub fn __env_need_backtrace() -> bool {
         Err(_) => false,
     }
 }
-
+#[doc(hidden)]
 pub fn __unreachable(file: &str, line: u32, col: u32) {
+    __unreachable_typed(Type::void(), file, line, col);
+}
+
+#[doc(hidden)]
+pub fn __unreachable_typed(ty: CArc<Type>, file: &str, line: u32, col: u32) -> NodeRef {
     let path = std::path::Path::new(file);
     let pretty_filename: String;
     if path.exists() {
@@ -102,9 +107,9 @@ pub fn __unreachable(file: &str, line: u32, col: u32) {
                 CString::new(msg).unwrap().into_bytes_with_nul(),
             )),
             &[],
-            Type::void(),
-        );
-    });
+            ty,
+        )
+    })
 }
 
 pub fn __assert(cond: impl Into<Expr<bool>>, msg: &str, file: &str, line: u32, col: u32) {
diff --git a/luisa_compute/src/lang/ops/impls.rs b/luisa_compute/src/lang/ops/impls.rs
index c91c88a..36f828e 100644
--- a/luisa_compute/src/lang/ops/impls.rs
+++ b/luisa_compute/src/lang/ops/impls.rs
@@ -417,6 +417,14 @@ impl<R: Aggregate> SelectMaybeExpr<R> for Expr<bool> {
     }
 }
 
+impl<R: Aggregate> SelectMaybeExpr<R> for Var<bool> {
+    fn if_then_else(self, on: impl Fn() -> R, off: impl Fn() -> R) -> R {
+        crate::lang::control_flow::if_then_else(**self, on, off)
+    }
+    fn select(self, on: R, off: R) -> R {
+        crate::lang::control_flow::select(**self, on, off)
+    }
+}
 impl ActivateMaybeExpr for bool {
     fn activate(self, then: impl Fn()) {
         if self {
@@ -430,6 +438,12 @@ impl ActivateMaybeExpr for Expr<bool> {
     }
 }
 
+impl ActivateMaybeExpr for Var<bool> {
+    fn activate(self, then: impl Fn()) {
+        crate::lang::control_flow::if_then_else(self.load(), then, || {})
+    }
+}
+
 impl LoopMaybeExpr for bool {
     fn while_loop(mut cond: impl FnMut() -> Self, mut body: impl FnMut()) {
         while cond() {