From fefc3c733a807e1d01c941d7ba63a76078c568aa Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Fri, 22 Sep 2023 15:24:43 -0400
Subject: [PATCH] fixed printer

---
 luisa_compute/examples/printer.rs   | 12 ++++++------
 luisa_compute/examples/ray_query.rs |  2 +-
 luisa_compute/src/lang/types.rs     |  7 +++++++
 luisa_compute/src/printer.rs        |  3 ++-
 luisa_compute_derive/src/lib.rs     |  5 ++++-
 5 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/luisa_compute/examples/printer.rs b/luisa_compute/examples/printer.rs
index 752c206..af173ae 100644
--- a/luisa_compute/examples/printer.rs
+++ b/luisa_compute/examples/printer.rs
@@ -21,14 +21,14 @@ fn main() {
         "cpu"
     });
     let printer = Printer::new(&device, 65536);
-    let kernel = device.create_kernel::<fn()>(&|| {
+    let kernel = device.create_kernel::<fn()>(track!(&|| {
         let id = dispatch_id().xy();
-        if_!(id.x().cmpeq(id.y()), {
+        if id.x == id.y {
             lc_info!(printer, "id = {:?}", id);
-        }, else {
-            lc_info!(printer, "not equal!, id = [{} {}]", id.x(), id.y());
-        });
-    });
+        } else {
+            lc_info!(printer, "not equal!, id = [{} {}]", id.x, id.y);
+        }
+    }));
     device.default_stream().with_scope(|s| {
         s.reset_printer(&printer);
         s.submit([kernel.dispatch_async([4, 4, 1])]);
diff --git a/luisa_compute/examples/ray_query.rs b/luisa_compute/examples/ray_query.rs
index 0fe78a1..1b6d63f 100644
--- a/luisa_compute/examples/ray_query.rs
+++ b/luisa_compute/examples/ray_query.rs
@@ -163,7 +163,7 @@ fn main() {
                     let t = Var::<f32>::zeroed();
 
                     for _ in 0..100 {
-                        let dist = (o + d * t.load() - (sphere.center + translate.expr())).length()
+                        let dist = (o + d * t - (sphere.center + translate.expr())).length()
                             - sphere.radius;
                         if dist < 0.001 {
                             if (px == Uint2::expr(400, 400)).all() {
diff --git a/luisa_compute/src/lang/types.rs b/luisa_compute/src/lang/types.rs
index da75a46..ddcf62e 100644
--- a/luisa_compute/src/lang/types.rs
+++ b/luisa_compute/src/lang/types.rs
@@ -154,6 +154,10 @@ pub struct Expr<T: Value> {
     proxy: *mut ExprProxyData<T>,
 }
 
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub struct TypeTag<T: Value>(PhantomData<T>);
+
 /// A variable within a [`Kernel`](crate::runtime::Kernel) or
 /// [`Callable`](crate::runtime::Callable). Created using [`Expr::var`]
 /// and [`Value::var`].
@@ -294,6 +298,9 @@ impl<T: Value> Expr<T> {
         let node = __current_scope(|s| s.bitcast(self.node(), ty));
         Expr::<S>::from_node(node)
     }
+    pub fn _type_tag(&self)->TypeTag<T> {
+        TypeTag(PhantomData)
+    }
 }
 
 impl<T: Value> Var<T> {
diff --git a/luisa_compute/src/printer.rs b/luisa_compute/src/printer.rs
index f0520f6..cde50f6 100644
--- a/luisa_compute/src/printer.rs
+++ b/luisa_compute/src/printer.rs
@@ -8,6 +8,7 @@ pub use log as _log;
 
 use crate::internal_prelude::*;
 
+use crate::lang::types::TypeTag;
 use crate::lang::{pack_to, packed_size};
 
 pub use crate::{lc_debug, lc_error, lc_info, lc_warn};
@@ -102,7 +103,7 @@ macro_rules! lc_error {
         $crate::lc_log!($printer, log::Level::Error, $fmt, $($arg)*);
     };
 }
-pub fn _unpack_from_expr<V: Value>(data: *const u32, _: Expr<V>) -> V {
+pub fn _unpack_from_expr<V: Value>(data: *const u32, _: TypeTag<V>) -> V {
     unsafe { std::ptr::read_unaligned(data as *const V) }
 }
 impl Printer {
diff --git a/luisa_compute_derive/src/lib.rs b/luisa_compute_derive/src/lib.rs
index 6ca8ee3..2b220fc 100644
--- a/luisa_compute_derive/src/lib.rs
+++ b/luisa_compute_derive/src/lib.rs
@@ -68,7 +68,7 @@ pub fn _log(item: TokenStream) -> TokenStream {
     quote! {
         {
             #(
-                let #arg_idents = #args;
+                let #arg_idents = #args._type_tag();
             )*;
             let mut __log_priv_i = 0;
             let log_fn = Box::new(move |args: &[*const u32]| -> () {
@@ -81,6 +81,9 @@ pub fn _log(item: TokenStream) -> TokenStream {
                     }
                 ), *);
             });
+            #(
+                let #arg_idents = #args;
+            )*;
             let mut printer_args = luisa_compute::printer::PrinterArgs::new();
             #(
                 printer_args.append(#arg_idents);