Fix bug in sampling a kernel with mz-reset-mz pattern (overwriting sample_result register)

runtime/common/MeasureCounts.cpp

@@ -186,7 +186,13 @@ sample_result::sample_result(double preComputedExp,
 }
 
 void sample_result::append(ExecutionResult &result) {
-  sampleResults.insert({result.registerName, result});
+  // If given a result corresponding to the same register name,
+  // replace the existing one if in the map.
+  auto iter = sampleResults.find(result.registerName);
+  if (iter != sampleResults.end())
+    sampleResults[result.registerName] = result;
+  else
+    sampleResults.insert({result.registerName, result});
   if (!totalShots)
     for (auto &[bits, count] : result.counts)
       totalShots += count;

unittests/CMakeLists.txt

@@ -33,6 +33,7 @@ set(CUDAQ_RUNTIME_TEST_SOURCES
   integration/observe_result_tester.cpp
   integration/noise_tester.cpp
   integration/get_state_tester.cpp
+  integration/measure_reset_tester.cpp
   qir/NVQIRTester.cpp
   qis/QubitQISTester.cpp
   integration/kernels_tester.cpp

unittests/integration/measure_reset_tester.cpp

@@ -0,0 +1,45 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include "CUDAQTestUtils.h"
+
+#include <cudaq.h>
+#include <iostream>
+
+TEST(MeasureResetTester, checkBug980) {
+  auto foo = []() __qpu__ {
+    cudaq::qubit a;
+    cudaq::mz(a);
+    cudaq::reset(a); // properly reset the qubit!
+    cudaq::h(a);
+    cudaq::mz(a);
+  };
+
+  auto bar = []() __qpu__ {
+    cudaq::qubit a;
+    cudaq::x(a);
+    [[maybe_unused]] auto a0 = cudaq::mz(a);
+    cudaq::reset(a); // properly reset the qubit!
+    cudaq::h(a);
+    [[maybe_unused]] auto a1 = cudaq::mz(a);
+  };
+
+  std::cout << "Foo:\n";
+  auto result = cudaq::sample(foo);
+  result.dump();
+  EXPECT_EQ(2, result.size());
+  EXPECT_TRUE(result.count("0") > 0);
+  EXPECT_TRUE(result.count("1") > 0);
+
+  std::cout << "Bar:\n";
+  result = cudaq::sample(bar);
+  result.dump();
+  EXPECT_EQ(2, result.size());
+  EXPECT_TRUE(result.count("0") > 0);
+  EXPECT_TRUE(result.count("1") > 0);
+}