#include <catch.hpp>

#include <pls/pls.h>

using namespace pls;

class once_sub_task : public task {
  std::atomic<int> *counter_;
  int children_;

 protected:
  void execute_internal() override {
    (*counter_)++;
    for (int i = 0; i < children_; i++) {
      spawn_child<once_sub_task>(counter_, children_ - 1);
    }
  }

 public:
  explicit once_sub_task(std::atomic<int> *counter, int children) :
      task{},
      counter_{counter},
      children_{children} {}
};

class force_steal_sub_task : public task {
  std::atomic<int> *parent_counter_;
  std::atomic<int> *overall_counter_;

 protected:
  void execute_internal() override {
    (*overall_counter_)--;
    if (overall_counter_->load() > 0) {
      std::atomic<int> counter{1};
      spawn_child<force_steal_sub_task>(&counter, overall_counter_);
      while (counter.load() > 0); // Spin...
    }

    (*parent_counter_)--;
  }

 public:
  explicit force_steal_sub_task(std::atomic<int> *parent_counter, std::atomic<int> *overall_counter) :
      task{},
      parent_counter_{parent_counter},
      overall_counter_{overall_counter} {}
};

TEST_CASE("tbb task are scheduled correctly", "[internal/scheduling/fork_join_task.h]") {
  malloc_scheduler_memory my_scheduler_memory{8, 2 << 12};

  SECTION("tasks are executed exactly once") {
    scheduler my_scheduler{&my_scheduler_memory, 2};
    int start_counter = 4;
    int total_tasks = 1 + 4 + 4 * 3 + 4 * 3 * 2 + 4 * 3 * 2 * 1;
    std::atomic<int> counter{0};

    my_scheduler.perform_work([&]() {
      scheduler::spawn_child<once_sub_task>(&counter, start_counter);
    });

    REQUIRE(counter.load() == total_tasks);
  }

  SECTION("tasks can be stolen") {
    scheduler my_scheduler{&my_scheduler_memory, 8};
    my_scheduler.perform_work([&]() {
      std::atomic<int> dummy_parent{1}, overall_counter{8};
      scheduler::spawn_child<force_steal_sub_task>(&dummy_parent, &overall_counter);

      // Required, as child operates on our stack's memory!!!
      scheduler::wait_for_all();
    });
  }
}