[wgsl] Use function body behavior to detect missing return

We should not be using the behavior of the last statement in the
function, as it may be unreachable.

Fixed: 382053630
Change-Id: Iff47ea42f856de1f32ec064c0148cc54b656407b
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/217874
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Commit-Queue: James Price <jrprice@google.com>
diff --git a/src/tint/lang/wgsl/resolver/function_validation_test.cc b/src/tint/lang/wgsl/resolver/function_validation_test.cc
index 0abe53b..1a3c347 100644
--- a/src/tint/lang/wgsl/resolver/function_validation_test.cc
+++ b/src/tint/lang/wgsl/resolver/function_validation_test.cc
@@ -281,6 +281,19 @@
     EXPECT_EQ(r()->error(), R"(78:89 error: missing return at end of function)");
 }
 
+TEST_F(ResolverFunctionValidationTest, NonVoidFunctionEndWithDeadCodeAfterReturnStatement_Pass) {
+    // fn func() -> i32 { return 0; _ = 0; _ = 1; }
+
+    Func(Source{{12, 34}}, "func", tint::Empty, ty.i32(),
+         Block(Source{Source::Range{{45, 56}, {78, 90}}}, Vector{
+                                                              Return(0_a),
+                                                              Assign(Phony(), 0_a),
+                                                              Assign(Phony(), 1_a),
+                                                          }));
+
+    EXPECT_TRUE(r()->Resolve()) << r()->error();
+}
+
 TEST_F(ResolverFunctionValidationTest, FunctionTypeMustMatchReturnStatementType_Pass) {
     // fn func { return; }
 
diff --git a/src/tint/lang/wgsl/resolver/validator.cc b/src/tint/lang/wgsl/resolver/validator.cc
index 9e724b0..d5ce0a7 100644
--- a/src/tint/lang/wgsl/resolver/validator.cc
+++ b/src/tint/lang/wgsl/resolver/validator.cc
@@ -1238,10 +1238,7 @@
         }
 
         if (decl->body) {
-            sem::Behaviors behaviors{sem::Behavior::kNext};
-            if (auto* last = decl->body->Last()) {
-                behaviors = sem_.Get(last)->Behaviors();
-            }
+            auto behaviors = sem_.Get(decl->body)->Behaviors();
             if (behaviors.Contains(sem::Behavior::kNext)) {
                 auto end_source = decl->body->source.End();
                 end_source.range.begin.column--;