lang/include/tree_sitter_wrapper.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(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);
    }

    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_);
    }

    Node PreviousSibling() {
      return Node(ts_node_prev_sibling(node_), source_);
    }

    Node PreviousNamedSibling() {
      return Node(ts_node_prev_named_sibling(node_), source_);
    }

    Node NextSibling() { return Node(ts_node_next_sibling(node_), source_); }

    Node NextNamedSibling() {
      return Node(ts_node_next_named_sibling(node_), source_);
    }

  private:
    TSNode node_;
    const std::string *source_;
  };

  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, nullptr, source_.c_str(),
                                   source_.size());

    ts_parser_delete(parser);
  }

  ParseTree(const ParseTree &parse_tree)
      : tree_(ts_tree_copy(parse_tree.tree_)), source_(parse_tree.source_) {}

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

  ~ParseTree() { ts_tree_delete(tree_); }

  bool IsProperlyParsed() { // TODO: find place
    return !GetRoot().HasError();
  }

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

} // namespace parser