Address the comments in CL 169943

Bug: dawn:2184
Change-Id: Ifa984b70bc980239e9634d83b13b664f710868f4
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/171060
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Peng Huang <penghuang@chromium.org>
diff --git a/src/dawn/native/ShaderModule.cpp b/src/dawn/native/ShaderModule.cpp
index a50ea41..193c992 100644
--- a/src/dawn/native/ShaderModule.cpp
+++ b/src/dawn/native/ShaderModule.cpp
@@ -1292,14 +1292,19 @@
 }
 
 ShaderModuleBase::ScopedUseTintProgram ShaderModuleBase::UseTintProgram() {
-    mTintData.Use([&](auto tintData) {
-        // When the ShaderModuleBase is not referenced externally, and not used for initializing any
-        // pipeline, the mTintProgram will be released. However the ShaderModuleBase itself may
-        // still alive due to being referenced by some pipelines. In this case, when
-        // DeviceBase::APICreateShaderModule() with the same shader source code, Dawn will look up
-        // from the cache and return the same ShaderModuleBase. In this case, we have to recreate
-        // mTintProgram, when the mTintProgram is required for initializing new pipelines.
-        if (tintData->tintProgram == nullptr) {
+    return mTintData.Use(
+        [&](auto tintData) {
+            if (tintData->tintProgram) {
+                return ScopedUseTintProgram(this);
+            }
+
+            // When the ShaderModuleBase is not referenced externally, and not used for initializing
+            // any pipeline, the mTintProgram will be released. However the ShaderModuleBase itself
+            // may still alive due to being referenced by some pipelines. In this case, when
+            // DeviceBase::APICreateShaderModule() with the same shader source code, Dawn will look
+            // up from the cache and return the same ShaderModuleBase. In this case, we have to
+            // recreate mTintProgram, when the mTintProgram is required for initializing new
+            // pipelines.
             ShaderModuleDescriptor descriptor;
             ShaderModuleWGSLDescriptor wgslDescriptor;
             ShaderModuleSPIRVDescriptor sprivDescriptor;
@@ -1326,10 +1331,9 @@
 
             tintData->tintProgram = std::move(parseResult.tintProgram);
             tintData->tintProgramRecreateCount++;
-        }
-    });
 
-    return ScopedUseTintProgram(this);
+            return ScopedUseTintProgram(this);
+        });
 }
 
 Ref<TintProgram> ShaderModuleBase::GetTintProgram() const {
diff --git a/src/dawn/tests/unittests/RefCountedTests.cpp b/src/dawn/tests/unittests/RefCountedTests.cpp
index 6ae82e8..4222243 100644
--- a/src/dawn/tests/unittests/RefCountedTests.cpp
+++ b/src/dawn/tests/unittests/RefCountedTests.cpp
@@ -61,11 +61,21 @@
 };
 
 TEST(RefCount, Increment) {
-    RefCount refCount(/*initCount=*/0, /*payload=*/0);
+    RefCount refCount0(/*initCount=*/0, /*payload=*/0);
     // Previous count is 0
-    EXPECT_TRUE(refCount.Increment());
+    EXPECT_TRUE(refCount0.Increment());
     // Previous count is 1
-    EXPECT_FALSE(refCount.Increment());
+    EXPECT_FALSE(refCount0.Increment());
+
+    EXPECT_FALSE(refCount0.Decrement());
+    EXPECT_TRUE(refCount0.Decrement());
+
+    RefCount refCount1(/*initCount=*/1, /*payload=*/0);
+    // Previous count is 1
+    EXPECT_FALSE(refCount1.Increment());
+
+    EXPECT_FALSE(refCount1.Decrement());
+    EXPECT_TRUE(refCount1.Decrement());
 }
 
 // Test that RCs start with one ref, and removing it destroys the object.