ast.hpp

#pragma once
#include <memory>
#include <ostream>
#include <string>
#include <utility>
#include <vector>

// ---------- helpers ----------
inline void pad(std::ostream &os, int n) {
  for (int i = 0; i < n; i++)
    os << ' ';
}

struct SourceLoc {
  int line = 0;
  int col = 0;
};

enum class TypeKind { Unknown, Int, Bool, Str, Array, Error };

inline const char *typeName(TypeKind t) {
  switch (t) {
  case TypeKind::Unknown:
    return "unknown";
  case TypeKind::Int:
    return "int";
  case TypeKind::Bool:
    return "bool";
  case TypeKind::Str:
    return "str";
  case TypeKind::Array:
    return "array";
  case TypeKind::Error:
    return "error";
  }
  return "error";
}

// ---------- Expr ----------
struct Expr {
  virtual ~Expr() = default;
  virtual void dump(std::ostream &os, int indent = 0) const = 0;
};

struct NumberExpr : Expr {
  long long value;
  explicit NumberExpr(long long v) : value(v) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Number(" << value << ")\n";
  }
};

struct StringExpr : Expr {
  std::string value;
  explicit StringExpr(std::string v) : value(std::move(v)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "String(\"" << value << "\")\n";
  }
};

struct BoolExpr : Expr {
  bool value;
  explicit BoolExpr(bool v) : value(v) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Bool(" << (value ? "true" : "false") << ")\n";
  }
};

struct IdentExpr : Expr {
  std::string name;
  SourceLoc loc;
  explicit IdentExpr(std::string n, SourceLoc l) : name(std::move(n)), loc(l) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Ident(" << name << ") @" << loc.line << ":" << loc.col << "\n";
  }
};

struct UnaryExpr : Expr {
  std::string op;
  std::unique_ptr<Expr> rhs;
  UnaryExpr(std::string o, std::unique_ptr<Expr> r)
      : op(std::move(o)), rhs(std::move(r)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Unary(" << op << ")\n";
    rhs->dump(os, indent + 2);
  }
};

struct BinaryExpr : Expr {
  std::string op;
  std::unique_ptr<Expr> lhs, rhs;
  BinaryExpr(std::string o, std::unique_ptr<Expr> l, std::unique_ptr<Expr> r)
      : op(std::move(o)), lhs(std::move(l)), rhs(std::move(r)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Binary(" << op << ")\n";
    lhs->dump(os, indent + 2);
    rhs->dump(os, indent + 2);
  }
};

struct CallExpr : Expr {
  std::string callee;
  SourceLoc loc;
  std::vector<std::unique_ptr<Expr>> args;

  CallExpr(std::string c, SourceLoc l, std::vector<std::unique_ptr<Expr>> a)
      : callee(std::move(c)), loc(l), args(std::move(a)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Call(" << callee << ") @" << loc.line << ":" << loc.col << "\n";
    for (auto &arg : args)
      arg->dump(os, indent + 2);
  }
};

struct IndexExpr : Expr {
  std::string name;
  SourceLoc loc;
  std::unique_ptr<Expr> index;

  IndexExpr(std::string n, SourceLoc l, std::unique_ptr<Expr> i)
      : name(std::move(n)), loc(l), index(std::move(i)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Index(" << name << ") @" << loc.line << ":" << loc.col << "\n";
    index->dump(os, indent + 2);
  }
};

struct AssignExpr : Expr {
  std::string name;
  SourceLoc loc;
  std::unique_ptr<Expr> value;

  AssignExpr(std::string n, SourceLoc l, std::unique_ptr<Expr> v)
      : name(std::move(n)), loc(l), value(std::move(v)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Assign(" << name << ") @" << loc.line << ":" << loc.col << "\n";
    value->dump(os, indent + 2);
  }
};

struct ArrayAssignExpr : Expr {
  std::string name;
  SourceLoc loc;
  std::unique_ptr<Expr> index;
  std::unique_ptr<Expr> value;

  ArrayAssignExpr(std::string n, SourceLoc l, std::unique_ptr<Expr> i,
                  std::unique_ptr<Expr> v)
      : name(std::move(n)), loc(l), index(std::move(i)), value(std::move(v)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "ArrayAssign(" << name << ") @" << loc.line << ":" << loc.col << "\n";

    pad(os, indent + 2);
    os << "Index:\n";
    index->dump(os, indent + 4);

    pad(os, indent + 2);
    os << "Value:\n";
    value->dump(os, indent + 4);
  }
};

// ---------- Stmt ----------
struct BlockStmt;

struct Stmt {
  virtual ~Stmt() = default;
  virtual void dump(std::ostream &os, int indent = 0) const = 0;
};

struct FuncStmt : Stmt {
  std::string name;
  SourceLoc loc;
  std::vector<std::string> params;
  std::unique_ptr<Stmt> body;

  FuncStmt(std::string n, SourceLoc l, std::vector<std::string> p,
           std::unique_ptr<Stmt> b)
      : name(std::move(n)), loc(l), params(std::move(p)), body(std::move(b)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Func(" << name << ") @" << loc.line << ":" << loc.col << "\n";
    pad(os, indent + 2);
    os << "Params:";
    for (const auto &p : params)
      os << " " << p;
    os << "\n";
    body->dump(os, indent + 2);
  }
};

struct ReturnStmt : Stmt {
  SourceLoc loc;
  std::unique_ptr<Expr> value;

  ReturnStmt(SourceLoc l, std::unique_ptr<Expr> v)
      : loc(l), value(std::move(v)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Return @" << loc.line << ":" << loc.col << "\n";
    if (value)
      value->dump(os, indent + 2);
  }
};

struct VarStmt : Stmt {
  std::string name;
  SourceLoc loc;
  std::unique_ptr<Expr> init;
  std::size_t arraySize = 0;

  VarStmt(std::string n, SourceLoc l, std::unique_ptr<Expr> i,
          std::size_t size = 0)
      : name(std::move(n)), loc(l), init(std::move(i)), arraySize(size) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);

    if (arraySize > 0) {
      os << "ArrayVar(" << name << "[" << arraySize << "]) @" << loc.line << ":"
         << loc.col << "\n";
    } else {
      os << "Var(" << name << ") @" << loc.line << ":" << loc.col << "\n";
    }

    if (init) {
      init->dump(os, indent + 2);
    } else if (arraySize == 0) {
      pad(os, indent + 2);
      os << "(no init)\n";
    }
  }
};
struct PrintStmt : Stmt {
  std::unique_ptr<Expr> expr;
  explicit PrintStmt(std::unique_ptr<Expr> e) : expr(std::move(e)) {}
  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Print\n";
    expr->dump(os, indent + 2);
  }
};

struct ExprStmt : Stmt {
  std::unique_ptr<Expr> expr;
  explicit ExprStmt(std::unique_ptr<Expr> e) : expr(std::move(e)) {}
  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "ExprStmt\n";
    expr->dump(os, indent + 2);
  }
};

struct BlockStmt : Stmt {
  std::vector<std::unique_ptr<Stmt>> stmts;
  explicit BlockStmt(std::vector<std::unique_ptr<Stmt>> s)
      : stmts(std::move(s)) {}
  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "Block\n";
    for (auto &st : stmts)
      st->dump(os, indent + 2);
  }
};

struct IfStmt : Stmt {
  std::unique_ptr<Expr> cond;
  std::unique_ptr<Stmt> thenBranch;
  std::unique_ptr<Stmt> elseBranch;
  IfStmt(std::unique_ptr<Expr> c, std::unique_ptr<Stmt> t,
         std::unique_ptr<Stmt> e)
      : cond(std::move(c)), thenBranch(std::move(t)), elseBranch(std::move(e)) {
  }

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "If\n";
    pad(os, indent + 2);
    os << "Cond:\n";
    cond->dump(os, indent + 4);
    pad(os, indent + 2);
    os << "Then:\n";
    thenBranch->dump(os, indent + 4);
    if (elseBranch) {
      pad(os, indent + 2);
      os << "Else:\n";
      elseBranch->dump(os, indent + 4);
    }
  }
};

struct WhileStmt : Stmt {
  std::unique_ptr<Expr> cond;
  std::unique_ptr<Stmt> body;
  WhileStmt(std::unique_ptr<Expr> c, std::unique_ptr<Stmt> b)
      : cond(std::move(c)), body(std::move(b)) {}

  void dump(std::ostream &os, int indent) const override {
    pad(os, indent);
    os << "While\n";
    pad(os, indent + 2);
    os << "Cond:\n";
    cond->dump(os, indent + 4);
    pad(os, indent + 2);
    os << "Body:\n";
    body->dump(os, indent + 4);
  }
};