lang/include/type_nodes.hpp

#pragma once

#include "basic_nodes.hpp"

#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <variant>
#include <vector>

namespace nodes {

class Type;
class TypeStorage;

class TypeProxy {
  friend TypeStorage;

public:
  Type *get();

  const Type *get() const;

private:
  TypeProxy(TypeStorage &type_storage, size_t id)
      : type_storage_(&type_storage), id_(id) {}

private:
  TypeStorage *type_storage_;
  size_t id_;
};

// can't have both optional and result modifiers ??
class Type : public Node {
public:
  Type(Node node, Identifier &&identifier, bool is_on_heap = false,
       Modifier modifier = Modifier::NONE)
      : Node(node), name_(std::move(identifier)), is_on_heap_(is_on_heap),
        modifier_(modifier) {}

  Type(Node node, const Identifier &identifier, bool is_on_heap = false,
       Modifier modifier = Modifier::NONE)
      : Node(node), name_(identifier), is_on_heap_(is_on_heap),
        modifier_(modifier) {}

  Type(Node node, Identifier &&identifier, std::vector<TypeProxy> &&parameters,
       bool is_on_heap = false, Modifier modifier = Modifier::NONE)
      : Node(node), name_(std::move(identifier)),
        parameters_(std::move(parameters)), is_on_heap_(is_on_heap),
        modifier_(modifier) {}

  Type(Node node, const Identifier &identifier,
       std::vector<TypeProxy> &&parameters, bool is_on_heap = false,
       Modifier modifier = Modifier::NONE)
      : Node(node), name_(identifier), parameters_(std::move(parameters)),
        is_on_heap_(is_on_heap), modifier_(modifier) {}

  Identifier *get_name() { return &name_; }

  const Identifier *get_name() const { return &name_; }

  size_t get_parametrs_size() const { return parameters_.size(); }

  Type *get_parameter(size_t id) { return parameters_.at(id).get(); }

  const Type *get_parameter(size_t id) const {
    return parameters_.at(id).get();
  }

  bool is_on_heap() const { return is_on_heap_; }

  Modifier get_modifier() const { return modifier_; }

private:
  Identifier name_;
  std::vector<TypeProxy> parameters_;
  // or use allocator ??
  bool is_on_heap_ = false;
  Modifier modifier_ = Modifier::NONE; // optional, result or none
};

class TypeStorage {
  friend TypeProxy;

public:
  TypeProxy add_type(const Type &type) {
    storage_.push_back(type);
    return TypeProxy(*this, storage_.size() - 1);
  }

  TypeProxy add_type(Type &&type) {
    storage_.push_back(std::move(type));
    return TypeProxy(*this, storage_.size() - 1);
  }

private:
  Type *get_type(size_t id) { return &storage_.at(id); }

  const Type *get_type(size_t id) const { return &storage_.at(id); }

private:
  std::vector<Type> storage_;
};

class TupleType : public Node {
public:
  TupleType(Node node,
            const std::vector<std::pair<std::optional<std::string>, TypeProxy>>
                &fields)
      : Node(node) {
    annotations_.reserve(fields.size());
    fields_.reserve(fields.size());
    for (auto &field : fields) {
      if (field.first.has_value()) {
        annotation_fields_[field.first.value()] = fields_.size();
      }
      annotations_.push_back(field.first);
      fields_.push_back(field.second);
    }
  }

  size_t size() const { return fields_.size(); }

  std::optional<std::string *> get_annotation(size_t id) {
    if (annotations_.at(id).has_value()) {
      return &annotations_[id].value();
    }
    return std::nullopt;
  }

  std::optional<const std::string *> get_annotation(size_t id) const {
    if (annotations_.at(id).has_value()) {
      return &annotations_[id].value();
    }
    return std::nullopt;
  }

  Type *get(size_t id) { return fields_.at(id).get(); }

  const Type *get(size_t id) const { return fields_.at(id).get(); }

  Type *get(const std::string &annotation) {
    return fields_.at(annotation_fields_.at(annotation)).get();
  }

  const Type *get(const std::string &annotation) const {
    return fields_.at(annotation_fields_.at(annotation)).get();
  }

  std::vector<std::string> get_all_annotations() {
    std::vector<std::string> annotations;

    annotations.reserve(annotation_fields_.size());
    for (auto &annotation_with_field : annotation_fields_) {
      annotations.push_back(annotation_with_field.first);
    }

    return annotations;
  }

private:
  std::unordered_map<std::string, size_t> annotation_fields_;
  std::vector<TypeProxy> fields_;
  std::vector<std::optional<std::string>> annotations_;
};

class VariantType : public Node {
public:
  VariantType(Node node, std::vector<TupleType> &&constructors_)
      : Node(node), constructors_(std::move(constructors_)) {}

  VariantType(Node node, const std::vector<TupleType> &constructors_)
      : Node(node), constructors_(constructors_) {}

  size_t size() const { return constructors_.size(); }

  TupleType *get(size_t id) { return &constructors_.at(id); }

  const TupleType *get(size_t id) const { return &constructors_.at(id); }

private:
  // named constructors ??
  std::vector<TupleType> constructors_;
};

} // namespace nodes