From 2b35890bbbc1df3ed81691ad04f253bc7085de82 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Fri, 19 Feb 2016 15:41:30 -0800
Subject: [PATCH] Add ts_node_symbols() function

---
 include/tree_sitter/runtime.h | 16 ++++++++++++----
 spec/runtime/node_spec.cc     | 29 +++++++++++++++++++++++++++++
 src/runtime/language.c        |  2 +-
 src/runtime/node.c            | 22 ++++++++++++++++++++++
 4 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/include/tree_sitter/runtime.h b/include/tree_sitter/runtime.h
index 18f55ed4..fb24d6f4 100644
--- a/include/tree_sitter/runtime.h
+++ b/include/tree_sitter/runtime.h
@@ -8,6 +8,10 @@ extern "C" {
 #include <stdlib.h>
 #include <stdbool.h>
 
+typedef unsigned short TSSymbol;
+typedef struct TSLanguage TSLanguage;
+typedef struct TSDocument TSDocument;
+
 typedef enum {
   TSInputEncodingUTF8,
   TSInputEncodingUTF16,
@@ -46,9 +50,11 @@ typedef struct {
   size_t offset[3];
 } TSNode;
 
-typedef unsigned short TSSymbol;
-typedef struct TSLanguage TSLanguage;
-typedef struct TSDocument TSDocument;
+typedef struct {
+  TSSymbol value;
+  bool done;
+  void *data;
+} TSSymbolIterator;
 
 size_t ts_node_start_char(TSNode);
 size_t ts_node_start_byte(TSNode);
@@ -57,6 +63,8 @@ size_t ts_node_end_char(TSNode);
 size_t ts_node_end_byte(TSNode);
 TSPoint ts_node_end_point(TSNode);
 TSSymbol ts_node_symbol(TSNode);
+TSSymbolIterator ts_node_symbols(TSNode);
+void ts_symbol_iterator_next(TSSymbolIterator *);
 const char *ts_node_name(TSNode, const TSDocument *);
 char *ts_node_string(TSNode, const TSDocument *);
 bool ts_node_eq(TSNode, TSNode);
@@ -90,7 +98,7 @@ TSNode ts_document_root_node(const TSDocument *);
 size_t ts_document_parse_count(const TSDocument *);
 
 size_t ts_language_symbol_count(const TSLanguage *);
-const char *ts_language_symbol_name(TSLanguage *, TSSymbol);
+const char *ts_language_symbol_name(const TSLanguage *, TSSymbol);
 
 #define ts_builtin_sym_error 0
 #define ts_builtin_sym_end 1
diff --git a/spec/runtime/node_spec.cc b/spec/runtime/node_spec.cc
index cfd1bc1a..007ba4a2 100644
--- a/spec/runtime/node_spec.cc
+++ b/spec/runtime/node_spec.cc
@@ -133,6 +133,35 @@ describe("Node", []() {
     });
   });
 
+  describe("symbols()", [&]() {
+    it("returns an iterator that yields each of the node's symbols", [&]() {
+      const TSLanguage *language = ts_document_language(document);
+
+      TSSymbolIterator iterator = ts_node_symbols(array_node);
+      AssertThat(iterator.done, Equals(false));
+      AssertThat(ts_language_symbol_name(language, iterator.value), Equals("array"));
+
+      ts_symbol_iterator_next(&iterator);
+      AssertThat(iterator.done, Equals(false));
+      AssertThat(ts_language_symbol_name(language, iterator.value), Equals("_value"));
+
+      ts_symbol_iterator_next(&iterator);
+      AssertThat(iterator.done, Equals(true));
+
+      TSNode false_node = ts_node_descendant_for_range(array_node, false_index, false_index + 1);
+      iterator = ts_node_symbols(false_node);
+      AssertThat(iterator.done, Equals(false));
+      AssertThat(ts_language_symbol_name(language, iterator.value), Equals("false"));
+
+      ts_symbol_iterator_next(&iterator);
+      AssertThat(iterator.done, Equals(false));
+      AssertThat(ts_language_symbol_name(language, iterator.value), Equals("_value"));
+
+      ts_symbol_iterator_next(&iterator);
+      AssertThat(iterator.done, Equals(true));
+    });
+  });
+
   describe("child_count(), child(i)", [&]() {
     it("returns the child node at the given index, including anonymous nodes", [&]() {
       AssertThat(ts_node_child_count(array_node), Equals<size_t>(7));
diff --git a/src/runtime/language.c b/src/runtime/language.c
index 0ad4c08a..f56d0893 100644
--- a/src/runtime/language.c
+++ b/src/runtime/language.c
@@ -22,6 +22,6 @@ size_t ts_language_symbol_count(const TSLanguage *language) {
   return language->symbol_count;
 }
 
-const char *ts_language_symbol_name(TSLanguage *language, TSSymbol symbol) {
+const char *ts_language_symbol_name(const TSLanguage *language, TSSymbol symbol) {
   return language->symbol_names[symbol];
 }
diff --git a/src/runtime/node.c b/src/runtime/node.c
index 41072499..a18c7745 100644
--- a/src/runtime/node.c
+++ b/src/runtime/node.c
@@ -204,6 +204,28 @@ TSSymbol ts_node_symbol(TSNode self) {
   return ts_node__tree(self)->symbol;
 }
 
+TSSymbolIterator ts_node_symbols(TSNode self) {
+  const TSTree *tree = ts_node__tree(self);
+  return (TSSymbolIterator){
+    .value = tree->symbol,
+    .done = false,
+    .data = (void *)tree,
+  };
+}
+
+void ts_symbol_iterator_next(TSSymbolIterator *self) {
+  const TSTree *tree = (const TSTree *)self->data;
+  const TSTree *parent = tree->context.parent;
+  if (!self->done && parent) {
+    if (parent->child_count == 1 && !parent->visible) {
+      self->value = parent->symbol;
+      self->data = (void *)parent;
+      return;
+    }
+  }
+  self->done = true;
+}
+
 const char *ts_node_name(TSNode self, const TSDocument *document) {
   return document->parser.language->symbol_names[ts_node__tree(self)->symbol];
 }