From 8308a1965cca0c762a34c86f6fa5d541ffccd36a Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Sat, 23 Sep 2023 09:40:32 -0400
Subject: [PATCH] fixed escape! in last stmt in block

---
 luisa_compute/examples/mpm.rs  |  1 -
 luisa_compute_track/src/lib.rs | 12 ++++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/luisa_compute/examples/mpm.rs b/luisa_compute/examples/mpm.rs
index b731b8e..0500bda 100644
--- a/luisa_compute/examples/mpm.rs
+++ b/luisa_compute/examples/mpm.rs
@@ -130,7 +130,6 @@ fn main() {
                     grid_m.var().atomic_fetch_add(idx, weight * P_MASS);
                 });
             });
-            let _ = (); // WHAT?
         }),
     );
 
diff --git a/luisa_compute_track/src/lib.rs b/luisa_compute_track/src/lib.rs
index 2619897..4f130a3 100644
--- a/luisa_compute_track/src/lib.rs
+++ b/luisa_compute_track/src/lib.rs
@@ -15,6 +15,7 @@ use pretty_assertions::assert_eq;
 struct TraceVisitor {
     trait_path: TokenStream,
     flow_path: TokenStream,
+    is_last_stmt: bool,
 }
 
 impl VisitMut for TraceVisitor {
@@ -22,16 +23,22 @@ impl VisitMut for TraceVisitor {
         let len = node.stmts.len();
         if len > 0 {
             for stmt in node.stmts[0..len - 1].iter_mut() {
+                let old = self.is_last_stmt;
+                self.is_last_stmt = false;
                 self.visit_stmt_mut(stmt);
+                self.is_last_stmt = old;
             }
-            visit_stmt_mut(self, node.stmts.last_mut().unwrap());
+            let old = self.is_last_stmt;
+            self.is_last_stmt = true;
+            self.visit_stmt_mut(node.stmts.last_mut().unwrap());
+            self.is_last_stmt = old;
         }
     }
     fn visit_stmt_mut(&mut self, node: &mut Stmt) {
         let span = node.span();
         match node {
             Stmt::Expr(_, semi) => {
-                if semi.is_none() {
+                if !self.is_last_stmt && semi.is_none() {
                     *semi = Some(Token![;](span));
                 }
             }
@@ -337,6 +344,7 @@ pub fn tracked(
 
 fn track_impl(mut ast: Expr) -> TokenStream {
     (TraceVisitor {
+        is_last_stmt: false,
         flow_path: quote!(::luisa_compute::lang::control_flow),
         trait_path: quote!(::luisa_compute::lang::ops),
     })