lang_2023/include/parse_tree.hpp

#pragma once

#include <string>

// for clangd
#include "tree_sitter/api.h"

extern "C" const TSLanguage* tree_sitter_LANG();

namespace parser {

class ParseTree {
public:
  class Node {
  public:
    Node() = default;
    Node(const TSNode &node, const std::string* source) : node_(node), source_(source) {}

    std::string GetType() {
        return ts_node_type(node_);
    }

    std::pair<size_t, size_t> GetStartPoint() {
        TSPoint point = ts_node_start_point(node_);
        return {point.row, point.column};
    }

    std::pair<size_t, size_t> GetEndPoint() {
        TSPoint point = ts_node_end_point(node_);
        return {point.row, point.column};
    }

    std::string GetAsSExpression() {
      return ts_node_string(node_);
    }

    std::string GetValue() { // from source
      size_t start = ts_node_start_byte(node_);
      size_t end = ts_node_end_byte(node_);
      return source_->substr(start, end - start); // TODO check
    }

    bool IsNull() {
      return ts_node_is_null(node_);
    }

    bool IsNamed() {
      return ts_node_is_named(node_);
    }

    bool IsMissing() {
      return ts_node_is_missing(node_);
    }

    bool IsExtra() { // comments, etc.
      return ts_node_is_extra(node_);
    }

    bool HasError() {
      return ts_node_has_error(node_);
    }

    Node NthChild(size_t n) {
      return Node(ts_node_child(node_, n), source_);
    }
    size_t ChildCount() {
      return ts_node_child_count(node_);
    }

    Node NthNamedChild(size_t n) {
      return Node(ts_node_named_child(node_, n), source_);
    }
    size_t NamedChildCount() {
      return ts_node_named_child_count(node_);
    }

    Node ChildByFieldName(const std::string& name) {
      return Node(ts_node_child_by_field_name(node_, name.c_str(), name.size()), source_);
    }

    // ?? use field id instaed of name ??
    // ?? node equality check needed ??
  private:
    TSNode node_;
    const std::string* source_ = nullptr;
  };
  class Cursor { // ?? needed ??
  public:
    Cursor(const Node& node);

    void ResetTo(const Node& node);

    Node GetCurrentNode();
    std::string GetCurrentNodeName();

    bool GoToParent();
    bool GoToNextSibling();
    bool GoToFirstChild();
  private:
    TSTreeCursor cursor_;
  };

  ParseTree(const std::string& source) : source_(source) {
    TSParser* parser = ts_parser_new();
    ts_parser_set_language(parser, tree_sitter_LANG());

    tree_ = ts_parser_parse_string(
      parser,
      NULL,
      source_.c_str(),
      source_.size());
  }

  Node GetRoot() const {
    return Node(ts_tree_root_node(tree_), &source_);
  }

private:
  TSTree* tree_;
  std::string source_; // for token value extraction
};

} // namespace parser