ksys: Add Task utilities

1 files changed, 157 insertions, 0 deletions

diff --git a/src/KingSystem/Utils/Thread/TaskQueueBase.h b/src/KingSystem/Utils/Thread/TaskQueueBase.h
new file mode 100644
index 00000000..3d4b44cf
--- /dev/null
+++ b/src/KingSystem/Utils/Thread/TaskQueueBase.h
@@ -0,0 +1,157 @@
+#pragma once
+
+#include <basis/seadTypes.h>
+#include <container/seadBuffer.h>
+#include <container/seadOffsetList.h>
+#include <container/seadPtrArray.h>
+#include <prim/seadDelegate.h>
+#include <prim/seadRuntimeTypeInfo.h>
+#include <prim/seadTypedBitFlag.h>
+#include <time/seadTickSpan.h>
+#include "KingSystem/Utils/Thread/Event.h"
+#include "KingSystem/Utils/Types.h"
+
+namespace ksys::util {
+
+class Task;
+class TaskQueueLock;
+class TaskThread;
+
+struct TaskSelectionContext {
+    const auto& begin() const { return *it_begin; }
+    const auto& end() const { return *it_end; }
+    u8 lane_id;
+    const sead::OffsetList<Task>::iterator* it_begin;
+    const sead::OffsetList<Task>::iterator* it_end;
+};
+
+using TaskSelectionDelegate = sead::IDelegate1R<const TaskSelectionContext&, Task*>;
+
+class TaskQueueBase {
+    SEAD_RTTI_BASE(TaskQueueBase)
+public:
+    struct InitArg {
+        bool enable_locks;
+        /// Number of lanes.
+        u16 num_lanes;
+        /// Maximum number of threads that will be processing the queue.
+        u16 max_num_threads;
+        sead::Heap* heap;
+        TaskSelectionDelegate* task_selection_delegate;
+    };
+    KSYS_CHECK_SIZE_NX150(InitArg, 0x18);
+
+    struct PushArg {
+        u8 lane_id;
+        Task* task;
+    };
+    KSYS_CHECK_SIZE_NX150(PushArg, 0x10);
+
+    explicit TaskQueueBase(sead::Heap* heap);
+    virtual ~TaskQueueBase();
+
+    void clear();
+    bool init(const InitArg& arg);
+
+    bool addThread(TaskThread* thread);
+    void removeThread(TaskThread* thread);
+
+    s32 getNumActiveTasks() const;
+    s32 countTasksInLane(u16 id) const;
+    bool areNoThreadsBusy() const;
+    bool isAnyThreadBusy() const;
+    bool areAllThreadsPaused() const;
+
+    void waitForQueueToEmpty();
+    void waitForLaneToEmpty(u8 id);
+
+    void cancelTasks(u8 id);
+    bool isProcessingTask(u8 id) const;
+    void signalEmptyEventsIfNeeded();
+
+    void blockTasks(u8 id);
+    void blockTasksAndReloadThreads(u8 id);
+    void unblockTasks(u8 id);
+
+    void lock(TaskQueueLock* lock);
+
+    /// @returns the current thread if it is in the thread pool and nullptr otherwise.
+    TaskThread* getCurrentThread() const;
+
+    sead::OffsetList<Task>::iterator activeTasksBegin(TaskQueueLock* lock);
+    sead::OffsetList<Task>::robustIterator activeTasksRobustBegin(TaskQueueLock* lock);
+    sead::OffsetList<Task>::iterator activeTasksEnd() const;
+    sead::OffsetList<Task>::robustIterator activeTasksRobustEnd() const;
+
+    bool push(const PushArg& arg);
+    void removeTask(Task* task, bool b);
+    void fetchTask(Task** out_task);
+
+protected:
+    enum class Flag {
+        Lock = 0x1,
+        PreventThreadPoolChanges = 0x2,
+    };
+
+    struct Lane {
+        /// If true, tasks in this lane are not allowed to be fetched by any thread.
+        bool blocked = false;
+        /// First task in the lane. Tasks are also added to a linked list (mActiveTasks).
+        Task* head_task = nullptr;
+        Event* lane_empty_event = nullptr;
+    };
+    KSYS_CHECK_SIZE_NX150(Lane, 0x18);
+
+    class ScopedLock {
+    public:
+        explicit ScopedLock(const TaskQueueBase* queue) : mQueue(queue) { mQueue->lock(); }
+        ScopedLock(const ScopedLock&) = delete;
+        ~ScopedLock() { mQueue->unlock(); }
+        ScopedLock& operator=(const ScopedLock&) = delete;
+
+    private:
+        const TaskQueueBase* mQueue;
+    };
+
+    class ConditionalScopedLock {
+    public:
+        explicit ConditionalScopedLock(const TaskQueueBase* queue) : mQueue(queue) {
+            mQueue->lockIfNeeded();
+        }
+        ConditionalScopedLock(const ConditionalScopedLock&) = delete;
+        ~ConditionalScopedLock() { mQueue->unlockIfNeeded(); }
+        ConditionalScopedLock& operator=(const ConditionalScopedLock&) = delete;
+
+    private:
+        const TaskQueueBase* mQueue;
+    };
+
+    friend class TaskQueueLock;
+
+    virtual void lock() const {}
+    virtual void unlock() const {}
+
+    bool shouldLock() const { return mFlags.isOn(Flag::Lock); }
+
+    void lockIfNeeded() const {
+        if (shouldLock())
+            lock();
+    }
+
+    void unlockIfNeeded() const {
+        if (shouldLock())
+            unlock();
+    }
+
+    void notifyThreadsForNewTasks();
+
+    sead::TypedBitFlag<Flag, u8> mFlags;
+    sead::OffsetList<Task> mActiveTasks;
+    sead::Buffer<Lane> mLanes;
+    Event mQueueEmptyEvent;
+    sead::PtrArray<TaskThread> mThreads;
+    TaskSelectionDelegate* mTaskSelectionDelegate = nullptr;
+};
+KSYS_CHECK_SIZE_NX150(TaskQueueBase, 0x90);
+
+}  // namespace ksys::util