lang/include/nodes/expression_nodes.hpp

#pragma once

#include "basic_nodes.hpp"
#include "type_nodes.hpp"

#include <optional>
#include <variant>
#include <vector>

namespace nodes {

class TypedNode : public Node {
public:
  TypedNode(Node node) : Node(node) {}

  void set_expression_type(TypeProxy expression_type) {
    expression_type_ = expression_type;
  }

  std::optional<TypeProxy> get_expression_type() const {
    return expression_type_;
  }

protected:
  std::optional<TypeProxy> expression_type_;
};

class Expression;
class ExpressionStorage;

class ExpressionProxy {
  friend ExpressionStorage;

public:
  ExpressionProxy() = delete;

  Expression *get();

  const Expression *get() const;

private:
  ExpressionProxy(ExpressionStorage &expression_storage, size_t id)
      : expression_storage_(&expression_storage), id_(id) {}

private:
  ExpressionStorage *expression_storage_;
  size_t id_;
};

namespace utils {

inline std::optional<nodes::Expression *>
proxy_to_expr_optional(std::optional<nodes::ExpressionProxy> &proxy) {
  if (proxy.has_value()) {
    return proxy.value().get();
  }
  return std::nullopt;
}

inline std::optional<const nodes::Expression *>
proxy_to_expr_optional(const std::optional<nodes::ExpressionProxy> &proxy) {
  if (proxy.has_value()) {
    return proxy.value().get();
  }
  return std::nullopt;
}

} // namespace utils

// --- flow control

class Match : public TypedNode {
public:
  class Case : public TypedNode {
  public:
    enum CaseType {
      PATTERN_VALUE,
      VALUE_PATTERN,
    };

    Case(Node node, CaseType case_type, ExpressionProxy value,
         std::optional<ExpressionProxy> condition = std::nullopt,
         std::optional<ExpressionProxy> expression = std::nullopt)
        : TypedNode(node), case_type_(case_type), value_(value),
          condition_(condition), expression_(expression) {}

    CaseType case_type() const { return case_type_; }

    Expression *get_value() { return value_.get(); }

    const Expression *get_value() const { return value_.get(); }

    std::optional<Expression *> get_condition() {
      return utils::proxy_to_expr_optional(condition_);
    }

    std::optional<const Expression *> get_condition() const {
      return utils::proxy_to_expr_optional(condition_);
    }

    std::optional<Expression *> get_expression() {
      return utils::proxy_to_expr_optional(expression_);
    }

    std::optional<const Expression *> get_expression() const {
      return utils::proxy_to_expr_optional(expression_);
    }

  private:
    CaseType case_type_;
    ExpressionProxy value_;
    std::optional<ExpressionProxy> condition_;
    std::optional<ExpressionProxy> expression_;
  };

  Match(Node node, ExpressionProxy value, std::vector<Case> &&cases)
      : TypedNode(node), value_(value), cases_(std::move(cases)) {}

  Match(Node node, ExpressionProxy value, const std::vector<Case> &cases)
      : TypedNode(node), value_(value), cases_(cases) {}

  Expression *get_value() { return value_.get(); }

  const Expression *get_value() const { return value_.get(); }

  size_t cases_size() const { return cases_.size(); }

  Case *get_case(size_t id) { return &cases_.at(id); }

  const Case *get_case(size_t id) const { return &cases_.at(id); }

private:
  ExpressionProxy value_;
  std::vector<Case> cases_;
};

class Condition : public Node {
public:
  Condition(Node node,
            std::vector<std::pair<ExpressionProxy, ExpressionProxy>> &&cases,
            std::optional<ExpressionProxy> else_case = std::nullopt)
      : Node(node), cases_(std::move(cases)), else_case_(else_case) {}

  Condition(
      Node node,
      const std::vector<std::pair<ExpressionProxy, ExpressionProxy>> &cases,
      std::optional<ExpressionProxy> else_case = std::nullopt)
      : Node(node), cases_(cases), else_case_(else_case) {}

  size_t cases_size() const { return cases_.size(); }

  std::pair<Expression *, Expression *> get_case(size_t id) {
    return {cases_.at(id).first.get(), cases_[id].second.get()};
  }

  std::pair<const Expression *, const Expression *> get_case(size_t id) const {
    return {cases_.at(id).first.get(), cases_[id].second.get()};
  }

  std::optional<Expression *> get_else_case() {
    return utils::proxy_to_expr_optional(else_case_);
  }

  std::optional<const Expression *> get_else_case() const {
    return utils::proxy_to_expr_optional(else_case_);
  }

private:
  std::vector<std::pair<ExpressionProxy, ExpressionProxy>> cases_;
  std::optional<ExpressionProxy> else_case_;
};

class Loop : public TypedNode {
public:
  enum LoopType {
    LOOP,
    WHILE,
    FOR,
  };

  // LOOP
  Loop(Node node, ExpressionProxy expression)
      : TypedNode(node), loop_type_(LOOP), expression_(expression) {}

  // WHILE
  Loop(Node node, ExpressionProxy condition, ExpressionProxy expression)
      : TypedNode(node), loop_type_(WHILE), expression_(expression),
        condition_(condition) {}

  // FOR
  Loop(Node node, ExpressionProxy variable, ExpressionProxy interval,
       ExpressionProxy expression)
      : TypedNode(node), loop_type_(FOR), expression_(expression),
        variable_(variable), interval_(interval) {}

  LoopType get_type() const { return loop_type_; }

  Expression *get_expression() { return expression_.get(); }

  const Expression *get_expression() const { return expression_.get(); }

  std::optional<Expression *> get_condition() {
    return utils::proxy_to_expr_optional(condition_);
  }

  std::optional<const Expression *> get_condition() const {
    return utils::proxy_to_expr_optional(condition_);
  }

  std::optional<Expression *> get_variable() {
    return utils::proxy_to_expr_optional(variable_);
  }

  std::optional<const Expression *> get_variable() const {
    return utils::proxy_to_expr_optional(variable_);
  }

  std::optional<Expression *> get_interval() {
    return utils::proxy_to_expr_optional(interval_);
  }

  std::optional<const Expression *> get_interval() const {
    return utils::proxy_to_expr_optional(interval_);
  }

private:
  LoopType loop_type_;

  ExpressionProxy expression_;

  std::optional<ExpressionProxy> condition_;
  std::optional<ExpressionProxy> variable_;
  std::optional<ExpressionProxy> interval_;
};

// --- containers

class Container : public TypedNode {
public:
  enum ContainerType {
    BLOCK,
    ARRAY,
  };

  Container(Node node, ContainerType type,
            std::vector<ExpressionProxy> &&expressions)
      : TypedNode(node), container_type_(type),
        expressions_(std::move(expressions)) {}

  Container(Node node, ContainerType type,
            const std::vector<ExpressionProxy> &expressions)
      : TypedNode(node), container_type_(type), expressions_(expressions) {}

  ContainerType get_type() const { return container_type_; }

  size_t expressions_size() const { return expressions_.size(); }

  Expression *get_expression(size_t id) { return expressions_.at(id).get(); }

  const Expression *get_expression(size_t id) const {
    return expressions_.at(id).get();
  }

private:
  ContainerType container_type_;

  std::vector<ExpressionProxy> expressions_;
};

// --- modifiers

class Return : public Node {
public:
  enum ReturnType {
    RETURN,
    BRING,
  };

  Return(Node node, ReturnType type, ExpressionProxy expression)
      : Node(node), return_type_(type), expression_(expression) {}

  ReturnType get_type() const { return return_type_; }

  Expression *get_expression() { return expression_.get(); }

  const Expression *get_expression() const { return expression_.get(); }

private:
  ReturnType return_type_;

  ExpressionProxy expression_;
};

class NameDefinition : public TypedNode {
public:
  enum Modifier {
    LET, // %
    VAR, // $
  };

  NameDefinition(Node node, Modifier modifier, Identifier &&name)
      : TypedNode(node), modifier_(modifier), name_(std::move(name)) {}

  NameDefinition(Node node, Modifier modifier, const Identifier &name)
      : TypedNode(node), modifier_(modifier), name_(name) {}

  Modifier get_modifier() const { return modifier_; }

  Identifier *get_name() { return &name_; }

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

private:
  Modifier modifier_;

  Identifier name_;
};

class Access : public TypedNode {
public:
  enum AccessType {
    ARRAY,
    TUPLE, // only number literal index allowed
  };

  Access(Node node, AccessType type, ExpressionProxy value,
         ExpressionProxy index)
      : TypedNode(node), access_type_(type), value_(value), index_(index) {}

  AccessType get_type() const { return access_type_; }

  Expression *get_value() { return value_.get(); }

  const Expression *get_value() const { return value_.get(); }

  Expression *get_index() { return index_.get(); }

  const Expression *get_index() const { return index_.get(); }

private:
  AccessType access_type_;

  ExpressionProxy value_;
  ExpressionProxy index_;
};

class LoopControl : public TypedNode {
public:
  enum LoopControlType {
    BREAK,
    CONTINUE,
  };

  LoopControl(Node node, LoopControlType type)
      : TypedNode(node), loop_control_type_(type) {}

  LoopControlType get_type() const { return loop_control_type_; }

private:
  LoopControlType loop_control_type_;
};

class ModifierExpression : public TypedNode {
public:
  ModifierExpression(Node node, Modifier modifier, ExpressionProxy expression)
      : TypedNode(node), modifier_(modifier), expression_(expression) {}

  Modifier get_modifier() const { return modifier_; }

  Expression *get_expression() { return expression_.get(); }

  const Expression *get_expression() const { return expression_.get(); }

private:
  Modifier modifier_;

  ExpressionProxy expression_;
};

// --- other

class NameExpression : public TypedNode {
public:
  NameExpression(Node node, Identifier &&name)
      : TypedNode(node), name_(std::move(name)) {}

  NameExpression(Node node, const Identifier &name)
      : TypedNode(node), name_(name) {}

  NameExpression(
      Node node, Identifier &&name,
      std::vector<std::pair<std::optional<std::string>, ExpressionProxy>>
          &&arguments,
      std::optional<TypeProxy> &&prefix, bool is_point_call = false,
      bool is_operator_call = false)
      : TypedNode(node), name_(std::move(name)),
        arguments_(std::move(arguments)), prefix_(std::move(prefix)),
        is_point_call_(is_point_call), is_operator_call_(is_operator_call) {}

  Identifier *get_name() { return &name_; }

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

  std::optional<Type *> get_prefix() {
    if (prefix_.has_value()) {
      return prefix_.value().get();
    }
    return std::nullopt;
  }

  std::optional<const Type *> get_prefix() const {
    if (prefix_.has_value()) {
      return prefix_.value().get();
    }
    return std::nullopt;
  }

  size_t arguments_size() const { return arguments_.size(); }

  Expression *get_argument_value(size_t id) {
    return arguments_.at(id).second.get();
  }

  const Expression *get_argument_value(size_t id) const {
    return arguments_.at(id).second.get();
  }

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

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

  bool is_point_call() const { return is_point_call_; }

  bool is_operator_call() const { return is_operator_call_; }

private:
  Identifier name_;
  // universal function call syntax
  std::vector<std::pair<std::optional<std::string>, ExpressionProxy>>
      arguments_;
  std::optional<TypeProxy> prefix_;
  // for static methods
  bool is_point_call_ = false;    // x.f ... or f x ...
  bool is_operator_call_ = false; // ... operator ...
};

class Constructor : public TypedNode {
public:
  Constructor(
      Node node, TypeProxy type,
      std::vector<std::pair<std::optional<std::string>, ExpressionProxy>>
          &&arguments)
      : TypedNode(node), constructor_type_(type),
        arguments_(std::move(arguments)) {}

  Constructor(
      Node node, TypeProxy type,
      const std::vector<std::pair<std::optional<std::string>, ExpressionProxy>>
          &arguments)
      : TypedNode(node), constructor_type_(type), arguments_(arguments) {}

  Type *get_type() { return constructor_type_.get(); }

  const Type *get_type() const { return constructor_type_.get(); }

  size_t arguments_size() const { return arguments_.size(); }

  Expression *get_argument_value(size_t id) {
    return arguments_.at(id).second.get();
  }

  const Expression *get_argument_value(size_t id) const {
    return arguments_.at(id).second.get();
  }

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

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

private:
  TypeProxy constructor_type_;
  std::vector<std::pair<std::optional<std::string>, ExpressionProxy>>
      arguments_;
};

class Lambda : public TypedNode {
public:
  Lambda(Node node, std::vector<Identifier> &&arguments,
         ExpressionProxy expression)
      : TypedNode(node), arguments_(std::move(arguments)),
        expression_(expression) {}

  Lambda(Node node, const std::vector<Identifier> &arguments,
         ExpressionProxy expression)
      : TypedNode(node), arguments_(arguments), expression_(expression) {}

  size_t arguments_size() const { return arguments_.size(); }

  Identifier *get_argument(size_t id) { return &arguments_.at(id); }

  const Identifier *get_argument(size_t id) const { return &arguments_.at(id); }

  Expression *get_expression() { return expression_.get(); }

  const Expression *get_expression() const { return expression_.get(); }

private:
  std::vector<Identifier> arguments_;
  ExpressionProxy expression_;
};

class Expression {
public:
  template <typename T>
  Expression(T &&expression, bool is_scoped)
      : expression_(std::forward<T>(expression)), is_scoped_(is_scoped) {}

  template <typename T> std::optional<T *> get() {
    if (std::holds_alternative<T>(expression_)) {
      return &std::get<T>(expression_);
    }
    return std::nullopt;
  }

  template <typename T> std::optional<const T *> get() const {
    if (std::holds_alternative<T>(expression_)) {
      return &std::get<T>(expression_);
    }
    return std::nullopt;
  }

  auto get_any() { return &expression_; }

  auto get_any() const { return &expression_; }

  bool is_scoped() const { return is_scoped_; }

private:
  std::variant<
      // --- flow control
      Match, Condition, Loop,

      // --- operators
      // CommaExpression is OperatorExpression for operator ','
      // OperatorExpression is NameExpression with two arguments

      // --- containers
      // Block value is brought value ("bring x")
      // Array value is array of expression's values
      Container,

      // --- modifiers
      Return, NameDefinition,
      // TupleAccess - Access with number Literal index
      Access, LoopControl,
      // Reference or SuffixExpression is ModifierExpression
      ModifierExpression,

      // --- other
      NameExpression, Constructor, Lambda,

      // --- literal
      Literal,

      // --- extra
      Extra,

      // --- empty lines
      EmptyLines>
      expression_;

  bool is_scoped_ = false;
};

class ExpressionStorage {
  friend ExpressionProxy;

public:
  ExpressionProxy add_expression(const Expression &expression) {
    storage_.push_back(expression);
    return ExpressionProxy(*this, storage_.size() - 1);
  }

  ExpressionProxy add_expression(Expression &&expression) {
    storage_.push_back(std::move(expression));
    return ExpressionProxy(*this, storage_.size() - 1);
  }

private:
  Expression *get_expression(size_t id) { return &storage_.at(id); }

  const Expression *get_expression(size_t id) const { return &storage_.at(id); }

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

} // namespace nodes