From fbf848295fc1b98186dd3e41be73f30f0c1e8809 Mon Sep 17 00:00:00 2001
From: Dustin Spicuzza <dustin@virtualroadside.com>
Date: Wed, 20 Dec 2023 02:31:36 -0500
Subject: [PATCH] Add wrapper for
 HALSIM_RegisterSimPeriodic{Before,After}Callback

---
 .../robotpy-hal/gen/simulation/MockHooks.yml  | 32 ++++++++++++++++--
 .../robotpy-hal/hal/simulation/main.cpp       |  8 +++++
 .../robotpy-hal/hal/simulation/sim_cb.h       | 33 +++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 subprojects/robotpy-hal/hal/simulation/sim_cb.h

diff --git a/subprojects/robotpy-hal/gen/simulation/MockHooks.yml b/subprojects/robotpy-hal/gen/simulation/MockHooks.yml
index 03960a91..77c95142 100644
--- a/subprojects/robotpy-hal/gen/simulation/MockHooks.yml
+++ b/subprojects/robotpy-hal/gen/simulation/MockHooks.yml
@@ -1,5 +1,9 @@
 ---
 
+extra_includes:
+- sim_cb.h
+- pybind11/functional.h
+
 strip_prefixes:
   - HALSIM_
 
@@ -23,11 +27,35 @@ functions:
     ignore: true
 
   HALSIM_RegisterSimPeriodicBeforeCallback:
-    ignore: true
+    param_override:
+      param:
+        ignore: true
+    cpp_code: |
+      [](std::function<void(void)> fn) -> std::unique_ptr<SimCB> {
+        auto cb = std::make_unique<SimCB>(fn, HALSIM_CancelSimPeriodicBeforeCallback);
+        auto uid = HALSIM_RegisterSimPeriodicBeforeCallback([](void *param) {
+          ((SimCB*)param)->m_fn();
+        }, cb.get());
+        cb->SetUID(uid);
+        return std::move(cb);
+      }
   HALSIM_CancelSimPeriodicBeforeCallback:
+    ignore: true
 
   HALSIM_RegisterSimPeriodicAfterCallback:
-    ignore: true
+    param_override:
+      param:
+        ignore: true
+    cpp_code: |
+      [](std::function<void(void)> fn) -> std::unique_ptr<SimCB> {
+        auto cb = std::make_unique<SimCB>(fn, HALSIM_CancelSimPeriodicAfterCallback);
+        auto uid = HALSIM_RegisterSimPeriodicAfterCallback([](void *param) {
+          ((SimCB*)param)->m_fn();
+        }, cb.get());
+        cb->SetUID(uid);
+        return std::move(cb);
+      }
   HALSIM_CancelSimPeriodicAfterCallback:
+    ignore: true
 
   HALSIM_CancelAllSimPeriodicCallbacks:
\ No newline at end of file
diff --git a/subprojects/robotpy-hal/hal/simulation/main.cpp b/subprojects/robotpy-hal/hal/simulation/main.cpp
index dda40a82..affebc05 100644
--- a/subprojects/robotpy-hal/hal/simulation/main.cpp
+++ b/subprojects/robotpy-hal/hal/simulation/main.cpp
@@ -1,9 +1,17 @@
 
 #include <rpygen_wrapper.hpp>
+#include <pybind11/functional.h>
+
+#include "sim_cb.h"
 
 void HALSIM_ResetGlobalHandles();
 
 RPYBUILD_PYBIND11_MODULE(m) {
+
+  py::class_<SimCB> cls_SimCB(m, "SimCB");
+  cls_SimCB.doc() = "Simulation callback handle";
+  cls_SimCB.def("cancel", &SimCB::Cancel, py::doc("Cancel the callback"));
+
   initWrapper(m);
 
   m.def(
diff --git a/subprojects/robotpy-hal/hal/simulation/sim_cb.h b/subprojects/robotpy-hal/hal/simulation/sim_cb.h
new file mode 100644
index 00000000..8b9d48d0
--- /dev/null
+++ b/subprojects/robotpy-hal/hal/simulation/sim_cb.h
@@ -0,0 +1,33 @@
+
+#pragma once
+
+class SimCB {
+public:
+
+    SimCB(std::function<void(void)> fn, std::function<void(int32_t)> cancel) : 
+        m_fn(fn),
+        m_cancel(cancel)
+    {}
+
+    void SetUID(int32_t uid) {
+        m_uid = uid;
+    }
+
+    ~SimCB() {
+        Cancel();
+    }
+
+    void Cancel() {
+        if (m_valid) {
+            m_cancel(m_uid);
+            m_valid = false;
+        }
+    }
+
+    std::function<void(void)> m_fn;
+
+private:
+    bool m_valid = true;
+    int32_t m_uid;
+    std::function<void(int32_t)> m_cancel;
+};
\ No newline at end of file