lang_2023/include/global_info.hpp

#pragma once

#include <string>
#include <functional>
#include <unordered_map>
#include <unordered_set>

// for clangd
#include "definitions.hpp"
#include "interpreter_tree.hpp"
#include "typeclass_graph.hpp"
#include "utils.hpp"

namespace info {

// TODO: add classes / structs and functions module interface
class GlobalInfo {
  friend class NamespaceVisitor;
public:
  GlobalInfo() {
    namespaces_.emplace_back(); // global namespace
  }

  struct PartitionInfo {
    std::vector<std::string> path;
    std::string name;
    interpreter::tokens::PartitionStatement* node = nullptr;
  };

  class NamespaceVisitor {
    friend GlobalInfo;
    NamespaceVisitor() = delete;
  public:

    struct Path {
      std::vector<std::optional<utils::IdType>> path_types;
      definition::Namespace* result;
    };

    void AddImport(definition::Import&& import_info, const std::optional<std::string>& name = std::nullopt);

    void AddEnterNamespace(const std::string& name,
                           utils::ClassInternalsModifier modifier,
                           std::optional<interpreter::tokens::Namespace*> node,
                           const interpreter::tokens::BaseNode& base_node);

    void EnterNamespace(const std::string& name,
                        utils::ClassInternalsModifier modifier);

    void ExitNamespace();

    void ToGlobalNamespace();

    utils::IdType AddFunctionDeclaration(const std::string& name,
                                         definition::FunctionDeclaration&& function_declaration_info,
                                         const interpreter::tokens::BaseNode& base_node);

    utils::IdType AddFunctionDefinition(const std::string& name,
                                        definition::FunctionDefinition&& function_definition_info,
                                        const interpreter::tokens::BaseNode& base_node);

    utils::IdType AddType(const std::string& type,
                          definition::Type&& type_info,
                          const interpreter::tokens::BaseNode& base_node);

    utils::IdType AddAbstractType(const std::string& abstract_type,
                                  definition::AbstractType&& abstract_type_info,
                                  const interpreter::tokens::BaseNode& base_node);

    utils::IdType AddTypeclass(const std::string& typeclass,
                               definition::Typeclass&& typeclass_info,
                               const interpreter::tokens::BaseNode& base_node);

    utils::IdType AddConstructor(const std::string& constructor,
                                 definition::Constructor&& constructor_info,
                                 const interpreter::tokens::BaseNode& base_node);

    utils::IdType AddPartition(const std::vector<std::string>& path, // including name
                               interpreter::tokens::PartitionStatement* node);

    std::optional<utils::IdType> FindNamespaceId(const std::optional<std::vector<std::string>>& path);

    std::optional<utils::IdType> FindFunctionId(const std::optional<std::vector<std::string>>& path,
                                                const std::string& name);

    std::optional<utils::IdType> FindMethodId(const std::optional<std::vector<std::string>>& path,
                                              const std::string& type,
                                              const std::string& name,
                                              utils::IsConstModifier modifier);

    std::optional<utils::IdType> FindTypeId(const std::optional<std::vector<std::string>>& path,
                                            const std::string& type);

    std::optional<utils::IdType> FindLocalTypeId(const std::string& type);

    std::optional<utils::IdType> FindAbstractTypeId(const std::string& abstract_type);

    std::optional<utils::IdType> FindTypeclassId(const std::string& typeclass);

    std::optional<utils::IdType> FindConstructorId(const std::optional<std::vector<std::string>>& path,
                                                 const std::string& constructor);

    std::optional<definition::Namespace*> FindNamespace(const std::optional<std::vector<std::string>>& path) {
      std::optional<utils::IdType> id = FindNamespaceId(path);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetNamespaceInfo(id.value());
    }

    std::optional<definition::Function*> FindFunction(const std::optional<std::vector<std::string>>& path,
                                                      const std::string& name) {
      std::optional<utils::IdType> id = FindFunctionId(path, name);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetFunctionInfo(id.value());
    }

    std::optional<definition::Function*> FindMethod(const std::optional<std::vector<std::string>>& path,
                                                    const std::string& type,
                                                    const std::string& name,
                                                    utils::IsConstModifier modifier) {
      std::optional<utils::IdType> id = FindFunctionId(path, name);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetFunctionInfo(id.value());
    }

    std::optional<definition::Type*> FindAnyType(const std::optional<std::vector<std::string>>& path,
                                              const std::string& type) {
      std::optional<utils::IdType> id = FindTypeId(path, type);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetAnyTypeInfo(id.value());
    }

    template<typename T>
    std::optional<T*> FindType(const std::optional<std::vector<std::string>>& path,
                                              const std::string& type) {
      std::optional<utils::IdType> id = FindTypeId(path, type);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return global_info_.GetTypeInfo<T>(id.value());
    }

    std::optional<definition::Type*> FindAnyLocalType(const std::string& type) {
      std::optional<utils::IdType> id = FindLocalTypeId(type);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetAnyTypeInfo(id.value());
    }

    template<typename T>
    std::optional<T*> FindLocalType(const std::string& type) {
      std::optional<utils::IdType> id = FindLocalTypeId(type);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return global_info_.GetTypeInfo<T>(id.value());
    }

    std::optional<definition::AbstractType*> FindAbstractType(const std::string& abstract_type) {
      std::optional<utils::IdType> id = FindAbstractTypeId(abstract_type);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetAbstractTypeInfo(id.value());
    }

    std::optional<definition::Typeclass*> FindTypeclass(const std::string& typeclass) {
      std::optional<utils::IdType> id = FindTypeclassId(typeclass);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetTypeclassInfo(id.value());
    }

    std::optional<definition::Constructor*> FindConstructor(const std::optional<std::vector<std::string>>& path,
                                                            const std::string& constructor) {
      std::optional<utils::IdType> id = FindConstructorId(path, constructor);
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetConstructorInfo(id.value());
    }

    NamespaceVisitor CreateVisitor() {
      return global_info_.CreateVisitor();
    }

    GlobalInfo* GetGlobalInfo() {
      return &global_info_;
    }

    TypeclassGraph* GetTypeclassGraph() {
      return global_info_.GetTypeclassGraph();
    }

    const std::vector<std::string>& GetCurrentPath() {
      return current_path_;
    }

    utils::IdType GetCurrentNamespaceId() {
      return namespace_stack_.back();
    }

    definition::Namespace* GetCurrentNamespace() {
      return &global_info_.namespaces_[GetCurrentNamespaceId()];
    }

    std::optional<utils::IdType> GetCurrentNamespaceTypeId() {
      if (!GetCurrentNamespace()->any_node.has_value()) {
        return std::nullopt;
      }
      return GetCurrentNamespace()->any_node.value()->link_type_id_;
    }

    std::optional<definition::Type*> GetCurrentNamespaceAnyType() {
      std::optional<utils::IdType> id = GetCurrentNamespaceTypeId();
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetAnyTypeInfo(id.value());
    }

    template<typename T>
    std::optional<T*> GetCurrentNamespaceType() {
      std::optional<utils::IdType> id = GetCurrentNamespaceTypeId();
      if (!id.has_value()) {
        return std::nullopt;
      }
      return global_info_.GetTypeInfo<T>(id.value());
    }

    std::optional<utils::IdType> GetCurrentNamespaceTypeclassId() {
      if (!GetCurrentNamespace()->any_node.has_value()) {
        return std::nullopt;
      }
      return GetCurrentNamespace()->any_node.value()->link_typeclass_id_;
    }

    std::optional<definition::Typeclass*> GetCurrentNamespaceTypeclass() {
      std::optional<utils::IdType> id = GetCurrentNamespaceTypeclassId();
      if (!id.has_value()) {
        return std::nullopt;
      }
      return &global_info_.GetTypeclassInfo(id.value());
    }

    bool IsInGlobalNamespace() {
      return namespace_stack_.size() == 1;
    }
  private:
    NamespaceVisitor(GlobalInfo& global_info) : global_info_(global_info),
                                                namespace_stack_ {global_info.GlobalNamespaceId} {}

    template<typename T>
    std::optional<T> FindSomething(
      const std::optional<std::vector<std::string>>& path,
      std::function<std::optional<T>(utils::IdType)> search_func);

    std::optional<utils::IdType> FindNamespaceIn(
      utils::IdType current_namespace,
      const std::vector<std::string>& path);

    std::unordered_map<std::string, utils::IdType>*
      ChooseNamespaces(utils::ClassInternalsModifier modifier,
                       utils::IdType namespace_id);
  private:
    GlobalInfo& global_info_;

    std::vector<utils::IdType> namespace_stack_;
    std::vector<std::string> current_path_;
  };

  NamespaceVisitor CreateVisitor() {
    return NamespaceVisitor(*this);
  }

  definition::Function& GetFunctionInfo(utils::IdType id) {
    return functions_.at(id);
  }

  definition::Type& GetAnyTypeInfo(utils::IdType id) {
    return types_.at(id);
  }

  template<typename T>
  std::optional<T*> GetTypeInfo(utils::IdType id) {
    if (!std::holds_alternative<T>(types_.at(id).type)) {
      return std::nullopt;
    }
    return &std::get<T>(types_[id].type);
  }

  definition::AbstractType& GetAbstractTypeInfo(utils::IdType id) {
    return abstract_types_.at(id);
  }

  definition::Typeclass& GetTypeclassInfo(utils::IdType id) {
    return typeclasses_.at(id);
  }

  definition::Constructor& GetConstructorInfo(utils::IdType id) {
    return constructors_.at(id);
  }

  definition::Namespace& GetNamespaceInfo(utils::IdType id) {
    return namespaces_.at(id);
  }

  PartitionInfo& GetPartitionInfo(utils::IdType id) {
    return partitions_.at(id);
  }

  std::optional<utils::IdType> FindPartition(const std::vector<std::string>& path) {
    auto trie_ans = partitions_trie_.Find(path);
    return trie_ans.has_value() ? std::optional<utils::IdType>(*trie_ans.value()) : std::nullopt;
  }

  std::vector<utils::IdType> FindPartitionsByPrefix(const std::vector<std::string>& path) { // optimize ??
    auto trie_ans = partitions_trie_.FindByPrefix(path);

    std::vector<utils::IdType> ans(trie_ans.size());

    for (size_t i = 0; i < ans.size(); ++i) {
      ans[i] = *trie_ans[i];
    }

    return ans;
  }

  TypeclassGraph* GetTypeclassGraph() {
    return &typeclass_graph_;
  }

  // cache ??
  std::unordered_set<utils::IdType> GetAnnotatedTypeTypeclassesSet(interpreter::tokens::AnnotatedType* node) {
    std::unordered_set<utils::IdType> typeclasses;

    for (auto& typeclass : node->typeclasses) {
      utils::IdType graph_id = typeclasses_[typeclass->typeclass_id_].graph_id_;

      std::vector<utils::IdType> dependencies = typeclass_graph_.GetTypeclassDependencies(graph_id);
      for (auto& dependency : dependencies) {
        typeclasses.insert(dependency);
      }

      typeclasses.insert(graph_id);
    }
    return typeclasses;
  }

  std::vector<utils::IdType> GetAnnotatedTypeTypeclassesVector(interpreter::tokens::AnnotatedType* node) {
    std::unordered_set<utils::IdType> typeclasses_set = GetAnnotatedTypeTypeclassesSet(node);

    std::vector<utils::IdType> typeclasses_vector;
    typeclasses_vector.reserve(typeclasses_vector.size());
    for (auto& typeclass : typeclasses_set) {
      typeclasses_vector.push_back(typeclass);
    }
    return typeclasses_vector;
  }

  std::optional<utils::IdType> AddTypeclassToGraph(utils::IdType typeclass);

private:
  void CollectFunctionInfo(
    utils::IdType current_namespace,
    utils::ClassInternalsModifier modifier,
    std::vector<std::pair<std::string, std::pair<utils::ClassInternalsModifier, interpreter::tokens::FunctionDeclaration*>>>& function_declarations,
    std::vector<std::pair<std::string, interpreter::tokens::FunctionDefinitionStatement*>>& function_definitions);

private:
  const utils::IdType GlobalNamespaceId = 0;

  std::vector<definition::Function> functions_;
  std::vector<definition::Type> types_;
  std::vector<definition::AbstractType> abstract_types_;
  std::vector<definition::Typeclass> typeclasses_;
  std::vector<definition::Constructor> constructors_;

  std::vector<definition::Namespace> namespaces_;

  std::unordered_map<std::string, utils::IdType> name_to_typeclass_;
  std::unordered_map<std::string, utils::IdType> name_to_abstract_type_;

  std::vector<PartitionInfo> partitions_;
  utils::Trie<std::string, utils::IdType> partitions_trie_;

  std::vector<definition::Import> imports_;
  std::unordered_map<std::string, definition::Import> usages_;

  TypeclassGraph typeclass_graph_;
};

} // namespace info