goap.hpp

#pragma once
 
#include <unordered_set>
#include <functional>
#include <coroutine>
#include <optional>
#include <memory>
#include <vector>
#include <queue>
#include <stack>
 
#include <type_traits>
#include <concepts>
 
namespace aitoolkit::goap {
  template <typename T>
  concept blackboard_trait = requires(const T &a, const T &b) {
    { a == b } -> std::convertible_to<bool>;
    { std::hash<T>{}(a) } -> std::convertible_to<size_t>;
  };
 
  template <blackboard_trait T>
  class action {
    public:
      virtual ~action() = default;
 
      virtual float cost(const T& blackboard) const = 0;
 
      virtual bool check_preconditions(const T& blackboard) const = 0;
 
      virtual void apply_effects(T& blackboard, bool dry_run) const = 0;
  };
 
  template <blackboard_trait T>
  using action_ptr = std::unique_ptr<action<T>>;
 
  template <typename A, typename T>
  concept action_trait = std::derived_from<A, action<T>>;
 
  template <blackboard_trait T, action_trait<T>... Actions>
  std::vector<action_ptr<T>> action_list(Actions... actions) {
    auto action_list = std::vector<action_ptr<T>>{};
    action_list.reserve(sizeof...(Actions));
    (action_list.push_back(std::make_unique<Actions>(std::move(actions))), ...);
    return action_list;
  }
 
  template <blackboard_trait T>
  class plan {
    public:
      plan() = default;
 
      size_t size() const {
        return m_plan.size();
      }
 
      operator bool() const {
        return !m_plan.empty();
      }
 
      void run_next(T& blackboard) {
        if (!m_plan.empty()) {
          auto action_idx = m_plan.top();
          m_plan.pop();
 
          auto& action = m_actions[action_idx];
          action->apply_effects(blackboard, false);
        }
      }
 
    private:
      std::stack<size_t> m_plan;
      std::vector<action_ptr<T>> m_actions;
 
      friend plan<T> planner<T>(
        std::vector<action_ptr<T>> actions,
        T initital_blackboard,
        T goal_blackboard,
        size_t max_iterations
      );
  };
 
  template <blackboard_trait T>
  plan<T> planner(
    std::vector<action_ptr<T>> actions,
    T initital_blackboard,
    T goal_blackboard,
    size_t max_iterations = 0
  ) {
    struct node_type {
      T blackboard;
      float cost;
 
      std::optional<size_t> action_taken_idx;
      std::shared_ptr<node_type> parent;
    };
 
    using node_ptr = std::shared_ptr<node_type>;
 
    struct node_compare {
      bool operator()(const node_ptr& a, const node_ptr& b) const {
        return a->cost > b->cost;
      }
    };
 
    std::priority_queue<node_ptr, std::vector<node_ptr>, node_compare> open_set;
    std::unordered_set<T> closed_set;
 
    open_set.push(std::make_shared<node_type>(node_type{
      .blackboard = initital_blackboard,
      .cost = 0.0f,
      .action_taken_idx = std::nullopt,
      .parent = nullptr
    }));
 
    for (
      size_t iteration = 0;
      !open_set.empty() && (max_iterations == 0 || iteration < max_iterations);
      ++iteration
    ) {
      auto current_node = open_set.top();
      open_set.pop();
 
      if (current_node->blackboard == goal_blackboard) {
        auto p = plan<T>();
        p.m_actions = std::move(actions);
 
        while (current_node->parent != nullptr) {
          auto action_idx = current_node->action_taken_idx.value();
          p.m_plan.push(action_idx);
          current_node = current_node->parent;
        }
 
        return p;
      }
 
      if (!closed_set.contains(current_node->blackboard)) {
        closed_set.insert(current_node->blackboard);
 
        for (size_t action_idx = 0; action_idx < actions.size(); action_idx++) {
          auto& action = actions[action_idx];
 
          if (action->check_preconditions(current_node->blackboard)) {
            auto next_blackboard = current_node->blackboard;
            action->apply_effects(next_blackboard, true);
            auto next_cost = current_node->cost + action->cost(current_node->blackboard);
 
            if (!closed_set.contains(next_blackboard)) {
              open_set.push(std::make_shared<node_type>(node_type{
                .blackboard = next_blackboard,
                .cost = next_cost,
                .action_taken_idx = action_idx,
                .parent = current_node
              }));
            }
          }
        }
      }
    }
 
    return plan<T>();
  }
}