1 files changed, 85 insertions, 0 deletions
diff --git a/backend/include/rcc/CodeGen/PBQP/Solution.h b/backend/include/rcc/CodeGen/PBQP/Solution.h
new file mode 100644
index 0000000..512eb57
--- /dev/null
+++ b/backend/include/rcc/CodeGen/PBQP/Solution.h
@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef RCC_CODEGEN_PBQP_SOLUTION_H
+#define RCC_CODEGEN_PBQP_SOLUTION_H
+
+#include "rcc/CodeGen/PBQP/Graph.h"
+#include "rcc/Support/Error.h"
+
+#include <cstddef>
+#include <limits>
+#include <optional>
+#include <span>
+#include <vector>
+
+namespace rcc::pbqp {
+
+class Solution final {
+public:
+  explicit Solution(std::size_t NodeCount)
+      : Selections(NodeCount, std::numeric_limits<std::size_t>::max()) {}
+
+  [[nodiscard]] Expected<void> setSelection(NodeID ID, std::size_t Selection) {
+    if (!ID.isValid() || ID.value() >= Selections.size())
+      return makeUnexpected(ErrorCode::InvalidArgument, "PBQP solution received an unknown node");
+    Selections[ID.value()] = Selection;
+    return {};
+  }
+
+  [[nodiscard]] std::optional<std::size_t> selection(NodeID ID) const noexcept {
+    if (!ID.isValid() || ID.value() >= Selections.size() ||
+        Selections[ID.value()] == std::numeric_limits<std::size_t>::max())
+      return std::nullopt;
+    return Selections[ID.value()];
+  }
+
+  [[nodiscard]] std::span<const std::size_t> selections() const noexcept {
+    return Selections;
+  }
+  [[nodiscard]] Cost totalCost() const noexcept {
+    return TotalCost;
+  }
+  void setTotalCost(Cost Value) noexcept {
+    TotalCost = Value;
+  }
+
+  [[nodiscard]] Expected<void> verify(const Graph &Problem) const {
+    if (Selections.size() != Problem.nodeCount())
+      return makeUnexpected(ErrorCode::VerificationError, "PBQP solution has the wrong node count");
+    Cost Total = 0;
+    for (const Node &NodeValue : Problem.nodes()) {
+      const auto Selection = selection(NodeValue.ID);
+      if (!Selection)
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "PBQP solution omits a valid node selection");
+      const std::size_t SelectedOption = Selection.value();
+      if (SelectedOption >= NodeValue.Costs.size())
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "PBQP solution selects an invalid node option");
+      Total = addCosts(Total, NodeValue.Costs[SelectedOption]);
+    }
+    for (const Edge &EdgeValue : Problem.edges()) {
+      const auto FirstSelection = selection(EdgeValue.First);
+      const auto SecondSelection = selection(EdgeValue.Second);
+      if (!FirstSelection || !SecondSelection) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "PBQP solution omits an edge endpoint selection");
+      }
+      Total = addCosts(Total, EdgeValue.Costs.at(FirstSelection.value(), SecondSelection.value()));
+    }
+    if (isInfinite(Total))
+      return makeUnexpected(ErrorCode::RegisterAllocationFailed,
+                            "PBQP problem has no finite solution");
+    if (Total != TotalCost)
+      return makeUnexpected(ErrorCode::VerificationError, "PBQP solution cost is inconsistent");
+    return {};
+  }
+
+private:
+  std::vector<std::size_t> Selections;
+  Cost TotalCost = 0;
+};
+
+} // namespace rcc::pbqp
+
+#endif // RCC_CODEGEN_PBQP_SOLUTION_H