4 files changed, 29 insertions, 16 deletions

diff --git a/arch/x86_64/arch/devices/legacy_pit.cpp b/arch/x86_64/arch/devices/legacy_pit.cpp
index b2a30ef9..b24f87ea 100644
--- a/arch/x86_64/arch/devices/legacy_pit.cpp
+++ b/arch/x86_64/arch/devices/legacy_pit.cpp
@@ -63,7 +63,7 @@ namespace arch::devices
 
     auto divisor = static_cast<std::uint16_t>(base_frequency / data->frequency_in_hz);
 
-    kapi::interrupts::register_handler(data->irq_number, *this);
+    kapi::interrupts::register_handler(data->irq_number, *this, data);
 
     command_port::write<std::uint8_t>(square_wave_mode);
     io::wait();
@@ -89,17 +89,20 @@ namespace arch::devices
     return pit_names;
   }
 
-  auto pit_driver::handle_interrupt(std::uint32_t irq_number) -> kapi::interrupts::status
+  auto pit_driver::handle_interrupt(std::uint32_t irq_number, kstd::weak_ptr<void> context) -> kapi::interrupts::status
   {
-    // auto data = static_pointer_cast<driver_data>(device.driver_data());
+    auto data = static_pointer_cast<driver_data>(context.lock());
+    if (!data)
+    {
+      return kapi::interrupts::status::unhandled;
+    }
 
-    if (irq_number != 0)
+    if (irq_number != data->irq_number)
     {
       return kapi::interrupts::status::unhandled;
     }
 
-    // TODO: upgrade kapi interrupts to gain access to the device.
-    // ++m_ticks;
+    ++data->ticks;
 
     return kapi::interrupts::status::handled;
   }
diff --git a/arch/x86_64/arch/devices/legacy_pit.hpp b/arch/x86_64/arch/devices/legacy_pit.hpp
index f0a18ecb..31c02207 100644
--- a/arch/x86_64/arch/devices/legacy_pit.hpp
+++ b/arch/x86_64/arch/devices/legacy_pit.hpp
@@ -6,6 +6,7 @@
 #include <kapi/devices.hpp>
 #include <kapi/interrupts.hpp>
 
+#include <kstd/memory.hpp>
 #include <kstd/result.hpp>
 
 #include <cstdint>
@@ -42,7 +43,7 @@ namespace arch::devices
 
     [[nodiscard]] auto supported_names() const -> std::span<std::string_view const> override;
 
-    auto handle_interrupt(std::uint32_t irq_number) -> kapi::interrupts::status override;
+    auto handle_interrupt(std::uint32_t irq_number, kstd::weak_ptr<void> context) -> kapi::interrupts::status override;
 
   protected:
     auto query_interface(kapi::devices::interface interface) -> void * override;
diff --git a/kapi/kapi/interrupts.hpp b/kapi/kapi/interrupts.hpp
index 4ba06848..41710c2b 100644
--- a/kapi/kapi/interrupts.hpp
+++ b/kapi/kapi/interrupts.hpp
@@ -1,6 +1,8 @@
 #ifndef TEACHOS_KAPI_INTERRUPTS_HPP
 #define TEACHOS_KAPI_INTERRUPTS_HPP
 
+#include <kstd/memory.hpp>
+
 #include <cstdint>
 
 namespace kapi::interrupts
@@ -29,8 +31,9 @@ namespace kapi::interrupts
     //! task quickly and must take care when acquiring globally shared locks.
     //!
     //! @param irq_number The identifier of the interrupt request that triggered the handler.
+    //! @param context Additional context data associated during handler registration.
     //! @return status::handled if the handler successfully handled the interrupt, status::unhandled otherwise.
-    virtual auto handle_interrupt(std::uint32_t irq_number) -> status = 0;
+    virtual auto handle_interrupt(std::uint32_t irq_number, kstd::weak_ptr<void> context) -> status = 0;
   };
 
   //! @}
@@ -42,7 +45,8 @@ namespace kapi::interrupts
   //!
   //! @param irq_number The IRQ number to register the handler for.
   //! @param handler The interrupt handler to register.
-  auto register_handler(std::uint32_t irq_number, handler & handler) -> void;
+  //! @param context Additional context data to be passed back to the handler when the interrupt occurs.
+  auto register_handler(std::uint32_t irq_number, handler & handler, kstd::weak_ptr<void> context) -> void;
 
   //! Unregister a previously registered interrupt handler.
   //!
diff --git a/kernel/kapi/interrupts.cpp b/kernel/kapi/interrupts.cpp
index 18a1b29d..c5925058 100644
--- a/kernel/kapi/interrupts.cpp
+++ b/kernel/kapi/interrupts.cpp
@@ -1,37 +1,42 @@
 #include <kapi/interrupts.hpp>
 
 #include <kstd/flat_map.hpp>
+#include <kstd/memory.hpp>
 #include <kstd/print.hpp>
 #include <kstd/vector.hpp>
 
 #include <algorithm>
 #include <cstdint>
+#include <utility>
 
 namespace kapi::interrupts
 {
 
   namespace
   {
-    auto constinit handlers = kstd::flat_map<std::uint32_t, kstd::vector<handler *>>{};
+    auto constinit handlers = kstd::flat_map<std::uint32_t, kstd::vector<std::pair<handler *, kstd::weak_ptr<void>>>>{};
   }  // namespace
 
-  auto register_handler(std::uint32_t irq_number, handler & handler) -> void
+  auto register_handler(std::uint32_t irq_number, handler & handler, kstd::weak_ptr<void> context) -> void
   {
     if (handlers.contains(irq_number))
     {
       auto & handler_list = handlers.at(irq_number);
-      handler_list.push_back(&handler);
+      handler_list.push_back(std::pair{&handler, context});
     }
     else
     {
-      handlers.emplace(irq_number, kstd::vector{&handler});
+      handlers.emplace(irq_number, kstd::vector{
+                                     std::pair{&handler, context}
+      });
     }
   }
 
   auto unregister_handler(std::uint32_t irq_number, handler & handler) -> void
   {
     auto & handler_list = handlers.at(irq_number);
-    auto [first, last] = std::ranges::remove(handler_list, &handler);
+    auto [first, last] =
+        std::ranges::remove_if(handler_list, [&](auto const & entry) { return entry.first == &handler; });
     handler_list.erase(first, last);
   }
 
@@ -51,9 +56,9 @@ namespace kapi::interrupts
       return status::unhandled;
     }
 
-    for (auto handler : handler_list)
+    for (auto const & [handler, context] : handler_list)
     {
-      if (handler && handler->handle_interrupt(irq_number) == status::handled)
+      if (handler && handler->handle_interrupt(irq_number, context) == status::handled)
       {
         return status::handled;
       }