From ed3a8b01e672daf8b0d75932365499747ab4097e Mon Sep 17 00:00:00 2001
From: Titan Yuan <tryuan99@gmail.com>
Date: Mon, 24 Feb 2025 12:05:34 -0800
Subject: [PATCH] Add the assignment plugin

---
 BUILD.bazel                     |  6 ++++++
 assignment/BUILD.bazel          |  6 ++++++
 assignment/assignment_plugin.cc | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+)
 create mode 100644 assignment/assignment_plugin.cc

diff --git a/BUILD.bazel b/BUILD.bazel
index 4640842..0816433 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -2,6 +2,11 @@ load("@rules_pkg//pkg:tar.bzl", "pkg_tar")
 
 package(default_visibility = ["//visibility:public"])
 
+cc_shared_library(
+    name = "assignment",
+    deps = ["//assignment:assignment_plugin"],
+)
+
 cc_shared_library(
     name = "example",
     deps = ["//example:example_plugin"],
@@ -10,6 +15,7 @@ cc_shared_library(
 pkg_tar(
     name = "plugins",
     srcs = [
+        ":assignment",
         ":example",
     ],
     extension = ".tar.gz",
diff --git a/assignment/BUILD.bazel b/assignment/BUILD.bazel
index 051b207..9f118a5 100644
--- a/assignment/BUILD.bazel
+++ b/assignment/BUILD.bazel
@@ -29,3 +29,9 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+cc_library(
+    name = "assignment_plugin",
+    srcs = ["assignment_plugin.cc"],
+    deps = [":covered_assignment"],
+)
diff --git a/assignment/assignment_plugin.cc b/assignment/assignment_plugin.cc
new file mode 100644
index 0000000..dca1f5a
--- /dev/null
+++ b/assignment/assignment_plugin.cc
@@ -0,0 +1,33 @@
+// The assignment plugin includes various assignment algorithms.
+
+#include <vector>
+
+#include "assignment/covered_assignment.h"
+
+extern "C" {
+// Assign the agents to the tasks using a covered assignment.
+int Assignment_CoveredAssignment_Assign(const int num_agents,
+                                        const int num_tasks, const float* costs,
+                                        int* assigned_agents,
+                                        int* assigned_tasks) {
+  // Create the cost matrix.
+  std::vector<std::vector<double>> cost_matrix(num_agents,
+                                               std::vector<double>(num_tasks));
+  for (int i = 0; i < num_agents; ++i) {
+    for (int j = 0; j < num_tasks; ++j) {
+      cost_matrix[i][j] = *(costs + num_tasks * i + j);
+    }
+  }
+
+  // Perform the assignment.
+  assignment::CoveredAssignment assignment(num_agents, num_tasks, cost_matrix);
+  const auto assignments = assignment.Assign();
+
+  // Record the assignments.
+  for (int i = 0; i < assignments.size(); ++i) {
+    assigned_agents[i] = assignments[i].first;
+    assigned_tasks[i] = assignments[i].second;
+  }
+  return assignments.size();
+}
+}