Fix undefined behavior in cudaq::qec::to_parity_matrix (#87)

Using `const_cast` to remove the const-ness of a parameter results in
undefined behavior if the incoming parameter was truly a `const` value.
Given that the `to_parity_matrix` function was called by `const` class
member functions with member variables as parameters, this was indeed
unsafe behavior. While I have not seen any direct badness result from
this, I think it's best to pre-emptively clean this up.

---------

Signed-off-by: Ben Howe <[email protected]>

Author: bmhowe23

libs/qec/lib/codes/repetition.cpp

@@ -44,6 +44,10 @@ repetition::repetition(const heterogeneous_map &options) : code() {
   Lz = Lz * cudaq::spin::i(get_num_data_qubits() - 1);
 
   m_pauli_observables.push_back(Lz);
+
+  // Sort now to avoid repeated sorts later.
+  sortStabilizerOps(m_stabilizers);
+  sortStabilizerOps(m_pauli_observables);
 }
 
 /// @brief Register the repetition code type

libs/qec/lib/codes/steane.cpp

@@ -30,6 +30,10 @@ steane::steane(const heterogeneous_map &options) : code() {
   m_stabilizers = fromPauliWords(
       {"XXXXIII", "IXXIXXI", "IIXXIXX", "ZZZZIII", "IZZIZZI", "IIZZIZZ"});
   m_pauli_observables = fromPauliWords({"IIIIXXX", "IIIIZZZ"});
+
+  // Sort now to avoid repeated sorts later.
+  sortStabilizerOps(m_stabilizers);
+  sortStabilizerOps(m_pauli_observables);
 }
 
 /// @brief Register the Steane code type

libs/qec/lib/codes/surface_code.cpp

@@ -358,6 +358,10 @@ surface_code::surface_code(const heterogeneous_map &options) : code() {
   m_stabilizers = grid.get_spin_op_stabilizers();
   m_pauli_observables = grid.get_spin_op_observables();
 
+  // Sort now to avoid repeated sorts later.
+  sortStabilizerOps(m_stabilizers);
+  sortStabilizerOps(m_pauli_observables);
+
   operation_encodings.insert(std::make_pair(operation::x, x));
   operation_encodings.insert(std::make_pair(operation::z, z));
   operation_encodings.insert(std::make_pair(operation::cx, cx));

libs/qec/lib/stabilizer_utils.cpp

@@ -1,5 +1,5 @@
-/****************************************************************-*- C++ -*-****
- * Copyright (c) 2024 NVIDIA Corporation & Affiliates.                         *
+/*******************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
@@ -8,29 +8,36 @@
 #include "cudaq/qec/stabilizer_utils.h"
 
 namespace cudaq::qec {
+// Sort the stabilizers by the occurrence of 'Z' first, then by 'X'.
+// An earlier occurrence is considered "less".
+// Example: a = "ZZX", b = "XXZ" -> a < b
+static bool spinOpComparator(const cudaq::spin_op &a, const cudaq::spin_op &b) {
+  auto astr = a.to_string(false);
+  auto bstr = b.to_string(false);
+  auto zIdxA = astr.find_first_of("Z");
+  auto zIdxB = bstr.find_first_of("Z");
+  if (zIdxA == std::string::npos) {
+    if (zIdxB != std::string::npos)
+      return false;
+
+    // No Z in either, must both contain a X
+    auto xIdxA = astr.find_first_of("X");
+    return xIdxA < bstr.find_first_of("X");
+  }
+
+  // First contains a Z
+  if (zIdxB == std::string::npos)
+    return true;
+
+  return zIdxA < zIdxB;
+}
+
+static bool isStabilizerSorted(const std::vector<cudaq::spin_op> &ops) {
+  return std::is_sorted(ops.begin(), ops.end(), spinOpComparator);
+}
+
 void sortStabilizerOps(std::vector<cudaq::spin_op> &ops) {
-  // Sort the stabilizers, z first, then x
-  std::sort(ops.begin(), ops.end(),
-            [](const cudaq::spin_op &a, const cudaq::spin_op &b) {
-              auto astr = a.to_string(false);
-              auto bstr = b.to_string(false);
-              auto zIdxA = astr.find_first_of("Z");
-              auto zIdxB = bstr.find_first_of("Z");
-              if (zIdxA == std::string::npos) {
-                if (zIdxB != std::string::npos)
-                  return false;
-
-                // No Z in either, must both contain a X
-                auto xIdxA = astr.find_first_of("X");
-                return xIdxA < bstr.find_first_of("X");
-              }
-
-              // First contains a Z
-              if (zIdxB == std::string::npos)
-                return true;
-
-              return zIdxA < zIdxB;
-            });
+  std::sort(ops.begin(), ops.end(), spinOpComparator);
 }
 
 // Need to push into the form
@@ -42,14 +49,25 @@ to_parity_matrix(const std::vector<cudaq::spin_op> &stabilizers,
   if (stabilizers.empty())
     return cudaqx::tensor<uint8_t>();
 
-  sortStabilizerOps(const_cast<std::vector<cudaq::spin_op> &>(stabilizers));
+  // Stabilizers must be sorted prior to use, but they are "const" into this
+  // function. First check to see if they are already sorted, and if so, proceed
+  // without making any copies. Otherwise make a copy, sort the copy, and make
+  // the downstream code use the sorted copy.
+  const std::vector<cudaq::spin_op> *p_stabilizers = &stabilizers;
+  std::vector<cudaq::spin_op> stabilizers_copy; // only use if necessary
+  if (!isStabilizerSorted(stabilizers)) {
+    // Make a copy and sort them
+    stabilizers_copy = stabilizers;
+    sortStabilizerOps(stabilizers_copy);
+    p_stabilizers = &stabilizers_copy;
+  }
 
   if (type == stabilizer_type::XZ) {
-    auto numQubits = stabilizers[0].num_qubits();
-    cudaqx::tensor<uint8_t> t({stabilizers.size(), 2 * numQubits});
+    auto numQubits = (*p_stabilizers)[0].num_qubits();
+    cudaqx::tensor<uint8_t> t({p_stabilizers->size(), 2 * numQubits});
     // Start by counting the number of Z spin_ops
     std::size_t numZRows = 0;
-    for (auto &op : stabilizers)
+    for (auto &op : *p_stabilizers)
       if (op.to_string(false).find("Z") != std::string::npos)
         numZRows++;
       else
@@ -58,16 +76,16 @@ to_parity_matrix(const std::vector<cudaq::spin_op> &stabilizers,
     // Need to shift Z bits left
     for (std::size_t row = 0; row < numZRows; row++) {
       for (std::size_t i = numQubits; i < 2 * numQubits; i++) {
-        if (stabilizers[row].get_raw_data().first[0][i])
+        if ((*p_stabilizers)[row].get_raw_data().first[0][i])
           t.at({row, i - numQubits}) = 1;
       }
     }
 
-    auto numXRows = stabilizers.size() - numZRows;
+    auto numXRows = p_stabilizers->size() - numZRows;
 
     for (std::size_t row = 0; row < numXRows; row++) {
       for (std::size_t i = 0; i < numQubits; i++) {
-        if (stabilizers[numZRows + row].get_raw_data().first[0][i])
+        if ((*p_stabilizers)[numZRows + row].get_raw_data().first[0][i])
           t.at({numZRows + row, i + numQubits}) = 1;
       }
     }
@@ -76,10 +94,10 @@ to_parity_matrix(const std::vector<cudaq::spin_op> &stabilizers,
   }
 
   if (type == stabilizer_type::Z) {
-    auto numQubits = stabilizers[0].num_qubits();
+    auto numQubits = (*p_stabilizers)[0].num_qubits();
     // Start by counting the number of Z spin_ops
     std::size_t numZRows = 0;
-    for (auto &op : stabilizers)
+    for (auto &op : *p_stabilizers)
       if (op.to_string(false).find("Z") != std::string::npos)
         numZRows++;
       else
@@ -91,32 +109,32 @@ to_parity_matrix(const std::vector<cudaq::spin_op> &stabilizers,
     cudaqx::tensor<uint8_t> ret({numZRows, numQubits});
     for (std::size_t row = 0; row < numZRows; row++) {
       for (std::size_t i = numQubits; i < 2 * numQubits; i++) {
-        if (stabilizers[row].get_raw_data().first[0][i])
+        if ((*p_stabilizers)[row].get_raw_data().first[0][i])
           ret.at({row, i - numQubits}) = 1;
       }
     }
 
     return ret;
   }
 
-  auto numQubits = stabilizers[0].num_qubits();
+  auto numQubits = (*p_stabilizers)[0].num_qubits();
   // Start by counting the number of Z spin_ops
   std::size_t numZRows = 0;
-  for (auto &op : stabilizers)
+  for (auto &op : *p_stabilizers)
     if (op.to_string(false).find("Z") != std::string::npos)
       numZRows++;
     else
       break;
 
-  auto numXRows = stabilizers.size() - numZRows;
+  auto numXRows = p_stabilizers->size() - numZRows;
 
   if (numXRows == 0)
     return cudaqx::tensor<uint8_t>();
 
   cudaqx::tensor<uint8_t> ret({numXRows, numQubits});
   for (std::size_t row = 0; row < numXRows; row++) {
     for (std::size_t i = 0; i < numQubits; i++) {
-      if (stabilizers[numZRows + row].get_raw_data().first[0][i])
+      if ((*p_stabilizers)[numZRows + row].get_raw_data().first[0][i])
         ret.at({row, i}) = 1;
     }
   }
@@ -130,6 +148,7 @@ cudaqx::tensor<uint8_t> to_parity_matrix(const std::vector<std::string> &words,
   std::vector<cudaq::spin_op> ops;
   for (auto &os : words)
     ops.emplace_back(cudaq::spin_op::from_word(os));
+  sortStabilizerOps(ops);
   return to_parity_matrix(ops, type);
 }
 } // namespace cudaq::qec