Merge pull request #11113 from bfredl/tree-sitter-query

tree-sitter step 2: query API and  highlighting prototype

26 files changed, 4653 insertions, 229 deletions

diff --git a/src/tree_sitter/api.h b/src/tree_sitter/api.h
index d39d0521ee..40187e3db0 100644
--- a/src/tree_sitter/api.h
+++ b/src/tree_sitter/api.h
@@ -14,7 +14,19 @@ extern "C" {
 /* Section - ABI Versioning */
 /****************************/
 
+/**
+ * The latest ABI version that is supported by the current version of the
+ * library. When Languages are generated by the Tree-sitter CLI, they are
+ * assigned an ABI version number that corresponds to the current CLI version.
+ * The Tree-sitter library is generally backwards-compatible with languages
+ * generated using older CLI versions, but is not forwards-compatible.
+ */
 #define TREE_SITTER_LANGUAGE_VERSION 11
+
+/**
+ * The earliest ABI version that is supported by the current version of the
+ * library.
+ */
 #define TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION 9
 
 /*******************/
@@ -26,6 +38,8 @@ typedef uint16_t TSFieldId;
 typedef struct TSLanguage TSLanguage;
 typedef struct TSParser TSParser;
 typedef struct TSTree TSTree;
+typedef struct TSQuery TSQuery;
+typedef struct TSQueryCursor TSQueryCursor;
 
 typedef enum {
   TSInputEncodingUTF8,
@@ -87,6 +101,37 @@ typedef struct {
   uint32_t context[2];
 } TSTreeCursor;
 
+typedef struct {
+  TSNode node;
+  uint32_t index;
+} TSQueryCapture;
+
+typedef struct {
+  uint32_t id;
+  uint16_t pattern_index;
+  uint16_t capture_count;
+  const TSQueryCapture *captures;
+} TSQueryMatch;
+
+typedef enum {
+  TSQueryPredicateStepTypeDone,
+  TSQueryPredicateStepTypeCapture,
+  TSQueryPredicateStepTypeString,
+} TSQueryPredicateStepType;
+
+typedef struct {
+  TSQueryPredicateStepType type;
+  uint32_t value_id;
+} TSQueryPredicateStep;
+
+typedef enum {
+  TSQueryErrorNone = 0,
+  TSQueryErrorSyntax,
+  TSQueryErrorNodeType,
+  TSQueryErrorField,
+  TSQueryErrorCapture,
+} TSQueryError;
+
 /********************/
 /* Section - Parser */
 /********************/
@@ -119,7 +164,7 @@ bool ts_parser_set_language(TSParser *self, const TSLanguage *language);
 const TSLanguage *ts_parser_language(const TSParser *self);
 
 /**
- * Set the spans of text that the parser should include when parsing.
+ * Set the ranges of text that the parser should include when parsing.
  *
  * By default, the parser will always include entire documents. This function
  * allows you to parse only a *portion* of a document but still return a syntax
@@ -226,14 +271,16 @@ TSTree *ts_parser_parse_string_encoding(
  * by default, it will resume where it left off on the next call to
  * `ts_parser_parse` or other parsing functions. If you don't want to resume,
  * and instead intend to use this parser to parse some other document, you must
- * call this `ts_parser_reset` first.
+ * call `ts_parser_reset` first.
  */
 void ts_parser_reset(TSParser *self);
 
 /**
  * Set the maximum duration in microseconds that parsing should be allowed to
- * take before halting. If parsing takes longer than this, it will halt early,
- * returning NULL. See `ts_parser_parse` for more information.
+ * take before halting.
+ *
+ * If parsing takes longer than this, it will halt early, returning NULL.
+ * See `ts_parser_parse` for more information.
  */
 void ts_parser_set_timeout_micros(TSParser *self, uint64_t timeout);
 
@@ -243,10 +290,11 @@ void ts_parser_set_timeout_micros(TSParser *self, uint64_t timeout);
 uint64_t ts_parser_timeout_micros(const TSParser *self);
 
 /**
- * Set the parser's current cancellation flag pointer. If a non-null pointer is
- * assigned, then the parser will periodically read from this pointer during
- * parsing. If it reads a non-zero value, it will halt early, returning NULL.
- * See `ts_parser_parse` for more information.
+ * Set the parser's current cancellation flag pointer.
+ *
+ * If a non-null pointer is assigned, then the parser will periodically read
+ * from this pointer during parsing. If it reads a non-zero value, it will
+ * halt early, returning NULL. See `ts_parser_parse` for more information.
  */
 void ts_parser_set_cancellation_flag(TSParser *self, const size_t *flag);
 
@@ -322,22 +370,22 @@ const TSLanguage *ts_tree_language(const TSTree *);
 void ts_tree_edit(TSTree *self, const TSInputEdit *edit);
 
 /**
- * Compare a new syntax tree to a previous syntax tree representing the same
+ * Compare an old edited syntax tree to a new syntax tree representing the same
  * document, returning an array of ranges whose syntactic structure has changed.
  *
  * For this to work correctly, the old syntax tree must have been edited such
  * that its ranges match up to the new tree. Generally, you'll want to call
- * this function right after calling one of the `ts_parser_parse` functions,
- * passing in the new tree that was returned from `ts_parser_parse` and the old
- * tree that was passed as a parameter.
+ * this function right after calling one of the `ts_parser_parse` functions.
+ * You need to pass the old tree that was passed to parse, as well as the new
+ * tree that was returned from that function.
  *
  * The returned array is allocated using `malloc` and the caller is responsible
  * for freeing it using `free`. The length of the array will be written to the
  * given `length` pointer.
  */
 TSRange *ts_tree_get_changed_ranges(
-  const TSTree *self,
   const TSTree *old_tree,
+  const TSTree *new_tree,
   uint32_t *length
 );
 
@@ -409,8 +457,8 @@ bool ts_node_is_named(TSNode);
 bool ts_node_is_missing(TSNode);
 
 /**
- * Check if the node is *missing*. Missing nodes are inserted by the parser in
- * order to recover from certain kinds of syntax errors.
+ * Check if the node is *extra*. Extra nodes represent things like comments,
+ * which are not required the grammar, but can appear anywhere.
  */
 bool ts_node_is_extra(TSNode);
 
@@ -542,7 +590,7 @@ TSTreeCursor ts_tree_cursor_new(TSNode);
 void ts_tree_cursor_delete(TSTreeCursor *);
 
 /**
- * Re-initialize a tree cursor to start at a different ndoe.
+ * Re-initialize a tree cursor to start at a different node.
  */
 void ts_tree_cursor_reset(TSTreeCursor *, TSNode);
 
@@ -584,7 +632,7 @@ bool ts_tree_cursor_goto_parent(TSTreeCursor *);
 bool ts_tree_cursor_goto_next_sibling(TSTreeCursor *);
 
 /**
- * Move the cursor to the first schild of its current node.
+ * Move the cursor to the first child of its current node.
  *
  * This returns `true` if the cursor successfully moved, and returns `false`
  * if there were no children.
@@ -592,7 +640,7 @@ bool ts_tree_cursor_goto_next_sibling(TSTreeCursor *);
 bool ts_tree_cursor_goto_first_child(TSTreeCursor *);
 
 /**
- * Move the cursor to the first schild of its current node that extends beyond
+ * Move the cursor to the first child of its current node that extends beyond
  * the given byte offset.
  *
  * This returns the index of the child node if one was found, and returns -1
@@ -602,6 +650,156 @@ int64_t ts_tree_cursor_goto_first_child_for_byte(TSTreeCursor *, uint32_t);
 
 TSTreeCursor ts_tree_cursor_copy(const TSTreeCursor *);
 
+/*******************/
+/* Section - Query */
+/*******************/
+
+/**
+ * Create a new query from a string containing one or more S-expression
+ * patterns. The query is associated with a particular language, and can
+ * only be run on syntax nodes parsed with that language.
+ *
+ * If all of the given patterns are valid, this returns a `TSQuery`.
+ * If a pattern is invalid, this returns `NULL`, and provides two pieces
+ * of information about the problem:
+ * 1. The byte offset of the error is written to the `error_offset` parameter.
+ * 2. The type of error is written to the `error_type` parameter.
+ */
+TSQuery *ts_query_new(
+  const TSLanguage *language,
+  const char *source,
+  uint32_t source_len,
+  uint32_t *error_offset,
+  TSQueryError *error_type
+);
+
+/**
+ * Delete a query, freeing all of the memory that it used.
+ */
+void ts_query_delete(TSQuery *);
+
+/**
+ * Get the number of patterns, captures, or string literals in the query.
+ */
+uint32_t ts_query_pattern_count(const TSQuery *);
+uint32_t ts_query_capture_count(const TSQuery *);
+uint32_t ts_query_string_count(const TSQuery *);
+
+/**
+ * Get the byte offset where the given pattern starts in the query's source.
+ *
+ * This can be useful when combining queries by concatenating their source
+ * code strings.
+ */
+uint32_t ts_query_start_byte_for_pattern(const TSQuery *, uint32_t);
+
+/**
+ * Get all of the predicates for the given pattern in the query.
+ *
+ * The predicates are represented as a single array of steps. There are three
+ * types of steps in this array, which correspond to the three legal values for
+ * the `type` field:
+ * - `TSQueryPredicateStepTypeCapture` - Steps with this type represent names
+ *    of captures. Their `value_id` can be used with the
+ *   `ts_query_capture_name_for_id` function to obtain the name of the capture.
+ * - `TSQueryPredicateStepTypeString` - Steps with this type represent literal
+ *    strings. Their `value_id` can be used with the
+ *    `ts_query_string_value_for_id` function to obtain their string value.
+ * - `TSQueryPredicateStepTypeDone` - Steps with this type are *sentinels*
+ *    that represent the end of an individual predicate. If a pattern has two
+ *    predicates, then there will be two steps with this `type` in the array.
+ */
+const TSQueryPredicateStep *ts_query_predicates_for_pattern(
+  const TSQuery *self,
+  uint32_t pattern_index,
+  uint32_t *length
+);
+
+/**
+ * Get the name and length of one of the query's captures, or one of the
+ * query's string literals. Each capture and string is associated with a
+ * numeric id based on the order that it appeared in the query's source.
+ */
+const char *ts_query_capture_name_for_id(
+  const TSQuery *,
+  uint32_t id,
+  uint32_t *length
+);
+const char *ts_query_string_value_for_id(
+  const TSQuery *,
+  uint32_t id,
+  uint32_t *length
+);
+
+/**
+ * Disable a certain capture within a query. This prevents the capture
+ * from being returned in matches, and also avoids any resource usage
+ * associated with recording the capture.
+ */
+void ts_query_disable_capture(TSQuery *, const char *, uint32_t);
+
+/**
+ * Create a new cursor for executing a given query.
+ *
+ * The cursor stores the state that is needed to iteratively search
+ * for matches. To use the query cursor, first call `ts_query_cursor_exec`
+ * to start running a given query on a given syntax node. Then, there are
+ * two options for consuming the results of the query:
+ * 1. Repeatedly call `ts_query_cursor_next_match` to iterate over all of the
+ *    the *matches* in the order that they were found. Each match contains the
+ *    index of the pattern that matched, and an array of captures. Because
+ *    multiple patterns can match the same set of nodes, one match may contain
+ *    captures that appear *before* some of the captures from a previous match.
+ * 2. Repeatedly call `ts_query_cursor_next_capture` to iterate over all of the
+ *    individual *captures* in the order that they appear. This is useful if
+ *    don't care about which pattern matched, and just want a single ordered
+ *    sequence of captures.
+ *
+ * If you don't care about consuming all of the results, you can stop calling
+ * `ts_query_cursor_next_match` or `ts_query_cursor_next_capture` at any point.
+ *  You can then start executing another query on another node by calling
+ *  `ts_query_cursor_exec` again.
+ */
+TSQueryCursor *ts_query_cursor_new(void);
+
+/**
+ * Delete a query cursor, freeing all of the memory that it used.
+ */
+void ts_query_cursor_delete(TSQueryCursor *);
+
+/**
+ * Start running a given query on a given node.
+ */
+void ts_query_cursor_exec(TSQueryCursor *, const TSQuery *, TSNode);
+
+/**
+ * Set the range of bytes or (row, column) positions in which the query
+ * will be executed.
+ */
+void ts_query_cursor_set_byte_range(TSQueryCursor *, uint32_t, uint32_t);
+void ts_query_cursor_set_point_range(TSQueryCursor *, TSPoint, TSPoint);
+
+/**
+ * Advance to the next match of the currently running query.
+ *
+ * If there is a match, write it to `*match` and return `true`.
+ * Otherwise, return `false`.
+ */
+bool ts_query_cursor_next_match(TSQueryCursor *, TSQueryMatch *match);
+void ts_query_cursor_remove_match(TSQueryCursor *, uint32_t id);
+
+/**
+ * Advance to the next capture of the currently running query.
+ *
+ * If there is a capture, write its match to `*match` and its index within
+ * the matche's capture list to `*capture_index`. Otherwise, return `false`.
+ */
+bool ts_query_cursor_next_capture(
+  TSQueryCursor *,
+  TSQueryMatch *match,
+  uint32_t *capture_index
+);
+
 /**********************/
 /* Section - Language */
 /**********************/
@@ -619,7 +817,12 @@ const char *ts_language_symbol_name(const TSLanguage *, TSSymbol);
 /**
  * Get the numerical id for the given node type string.
  */
-TSSymbol ts_language_symbol_for_name(const TSLanguage *, const char *);
+TSSymbol ts_language_symbol_for_name(
+  const TSLanguage *self,
+  const char *string,
+  uint32_t length,
+  bool is_named
+);
 
 /**
  * Get the number of distinct field names in the language.
diff --git a/src/tree_sitter/bits.h b/src/tree_sitter/bits.h
new file mode 100644
index 0000000000..3bec455dd1
--- /dev/null
+++ b/src/tree_sitter/bits.h
@@ -0,0 +1,29 @@
+#ifndef TREE_SITTER_BITS_H_
+#define TREE_SITTER_BITS_H_
+
+#include <stdint.h>
+
+static inline uint32_t bitmask_for_index(uint16_t id) {
+  return (1u << (31 - id));
+}
+
+#ifdef _WIN32
+
+#include <intrin.h>
+
+static inline uint32_t count_leading_zeros(uint32_t x) {
+  if (x == 0) return 32;
+  uint32_t result;
+  _BitScanReverse(&result, x);
+  return 31 - result;
+}
+
+#else
+
+static inline uint32_t count_leading_zeros(uint32_t x) {
+  if (x == 0) return 32;
+  return __builtin_clz(x);
+}
+
+#endif
+#endif  // TREE_SITTER_BITS_H_
diff --git a/src/tree_sitter/language.c b/src/tree_sitter/language.c
index 1bfb1a8d03..e240ef2a53 100644
--- a/src/tree_sitter/language.c
+++ b/src/tree_sitter/language.c
@@ -3,8 +3,28 @@
 #include "./error_costs.h"
 #include <string.h>
 
-void ts_language_table_entry(const TSLanguage *self, TSStateId state,
-                             TSSymbol symbol, TableEntry *result) {
+uint32_t ts_language_symbol_count(const TSLanguage *self) {
+  return self->symbol_count + self->alias_count;
+}
+
+uint32_t ts_language_version(const TSLanguage *self) {
+  return self->version;
+}
+
+uint32_t ts_language_field_count(const TSLanguage *self) {
+  if (self->version >= TREE_SITTER_LANGUAGE_VERSION_WITH_FIELDS) {
+    return self->field_count;
+  } else {
+    return 0;
+  }
+}
+
+void ts_language_table_entry(
+  const TSLanguage *self,
+  TSStateId state,
+  TSSymbol symbol,
+  TableEntry *result
+) {
   if (symbol == ts_builtin_sym_error || symbol == ts_builtin_sym_error_repeat) {
     result->action_count = 0;
     result->is_reusable = false;
@@ -19,48 +39,72 @@ void ts_language_table_entry(const TSLanguage *self, TSStateId state,
   }
 }
 
-uint32_t ts_language_symbol_count(const TSLanguage *language) {
-  return language->symbol_count + language->alias_count;
-}
-
-uint32_t ts_language_version(const TSLanguage *language) {
-  return language->version;
-}
-
-TSSymbolMetadata ts_language_symbol_metadata(const TSLanguage *language, TSSymbol symbol) {
+TSSymbolMetadata ts_language_symbol_metadata(
+  const TSLanguage *self,
+  TSSymbol symbol
+) {
   if (symbol == ts_builtin_sym_error)  {
     return (TSSymbolMetadata){.visible = true, .named = true};
   } else if (symbol == ts_builtin_sym_error_repeat) {
     return (TSSymbolMetadata){.visible = false, .named = false};
   } else {
-    return language->symbol_metadata[symbol];
+    return self->symbol_metadata[symbol];
+  }
+}
+
+TSSymbol ts_language_public_symbol(
+  const TSLanguage *self,
+  TSSymbol symbol
+) {
+  if (symbol == ts_builtin_sym_error) return symbol;
+  if (self->version >= TREE_SITTER_LANGUAGE_VERSION_WITH_SYMBOL_DEDUPING) {
+    return self->public_symbol_map[symbol];
+  } else {
+    return symbol;
   }
 }
 
-const char *ts_language_symbol_name(const TSLanguage *language, TSSymbol symbol) {
+const char *ts_language_symbol_name(
+  const TSLanguage *self,
+  TSSymbol symbol
+) {
   if (symbol == ts_builtin_sym_error) {
     return "ERROR";
   } else if (symbol == ts_builtin_sym_error_repeat) {
     return "_ERROR";
   } else {
-    return language->symbol_names[symbol];
+    return self->symbol_names[symbol];
   }
 }
 
-TSSymbol ts_language_symbol_for_name(const TSLanguage *self, const char *name) {
-  if (!strcmp(name, "ERROR")) return ts_builtin_sym_error;
-
+TSSymbol ts_language_symbol_for_name(
+  const TSLanguage *self,
+  const char *string,
+  uint32_t length,
+  bool is_named
+) {
+  if (!strncmp(string, "ERROR", length)) return ts_builtin_sym_error;
   uint32_t count = ts_language_symbol_count(self);
   for (TSSymbol i = 0; i < count; i++) {
-    if (!strcmp(self->symbol_names[i], name)) {
-      return i;
+    TSSymbolMetadata metadata = ts_language_symbol_metadata(self, i);
+    if (!metadata.visible || metadata.named != is_named) continue;
+    const char *symbol_name = self->symbol_names[i];
+    if (!strncmp(symbol_name, string, length) && !symbol_name[length]) {
+      if (self->version >= TREE_SITTER_LANGUAGE_VERSION_WITH_SYMBOL_DEDUPING) {
+        return self->public_symbol_map[i];
+      } else {
+        return i;
+      }
     }
   }
   return 0;
 }
 
-TSSymbolType ts_language_symbol_type(const TSLanguage *language, TSSymbol symbol) {
-  TSSymbolMetadata metadata = ts_language_symbol_metadata(language, symbol);
+TSSymbolType ts_language_symbol_type(
+  const TSLanguage *self,
+  TSSymbol symbol
+) {
+  TSSymbolMetadata metadata = ts_language_symbol_metadata(self, symbol);
   if (metadata.named) {
     return TSSymbolTypeRegular;
   } else if (metadata.visible) {
@@ -70,15 +114,10 @@ TSSymbolType ts_language_symbol_type(const TSLanguage *language, TSSymbol symbol
   }
 }
 
-uint32_t ts_language_field_count(const TSLanguage *self) {
-  if (self->version >= TREE_SITTER_LANGUAGE_VERSION_WITH_FIELDS) {
-    return self->field_count;
-  } else {
-    return 0;
-  }
-}
-
-const char *ts_language_field_name_for_id(const TSLanguage *self, TSFieldId id) {
+const char *ts_language_field_name_for_id(
+  const TSLanguage *self,
+  TSFieldId id
+) {
   uint32_t count = ts_language_field_count(self);
   if (count) {
     return self->field_names[id];
@@ -96,7 +135,8 @@ TSFieldId ts_language_field_id_for_name(
   for (TSSymbol i = 1; i < count + 1; i++) {
     switch (strncmp(name, self->field_names[i], name_length)) {
       case 0:
-        return i;
+        if (self->field_names[i][name_length] == 0) return i;
+        break;
       case -1:
         return 0;
       default:
diff --git a/src/tree_sitter/language.h b/src/tree_sitter/language.h
index 0741486a1b..d7e17c3d70 100644
--- a/src/tree_sitter/language.h
+++ b/src/tree_sitter/language.h
@@ -10,6 +10,7 @@ extern "C" {
 
 #define ts_builtin_sym_error_repeat (ts_builtin_sym_error - 1)
 #define TREE_SITTER_LANGUAGE_VERSION_WITH_FIELDS 10
+#define TREE_SITTER_LANGUAGE_VERSION_WITH_SYMBOL_DEDUPING 11
 #define TREE_SITTER_LANGUAGE_VERSION_WITH_SMALL_STATES 11
 
 typedef struct {
@@ -22,6 +23,8 @@ void ts_language_table_entry(const TSLanguage *, TSStateId, TSSymbol, TableEntry
 
 TSSymbolMetadata ts_language_symbol_metadata(const TSLanguage *, TSSymbol);
 
+TSSymbol ts_language_public_symbol(const TSLanguage *, TSSymbol);
+
 static inline bool ts_language_is_symbol_external(const TSLanguage *self, TSSymbol symbol) {
   return 0 < symbol && symbol < self->external_token_count + 1;
 }
diff --git a/src/tree_sitter/lexer.c b/src/tree_sitter/lexer.c
index fdc127466f..e2ca851973 100644
--- a/src/tree_sitter/lexer.c
+++ b/src/tree_sitter/lexer.c
@@ -2,26 +2,58 @@
 #include "./lexer.h"
 #include "./subtree.h"
 #include "./length.h"
-#include "./utf16.h"
-#include "utf8proc.h"
-
-#define LOG(...)                                                                      \
-  if (self->logger.log) {                                                             \
-    snprintf(self->debug_buffer, TREE_SITTER_SERIALIZATION_BUFFER_SIZE, __VA_ARGS__); \
-    self->logger.log(self->logger.payload, TSLogTypeLex, self->debug_buffer);         \
+#include "./unicode.h"
+
+#define LOG(message, character)              \
+  if (self->logger.log) {                    \
+    snprintf(                                \
+      self->debug_buffer,                    \
+      TREE_SITTER_SERIALIZATION_BUFFER_SIZE, \
+      32 <= character && character < 127 ?   \
+        message " character:'%c'" :          \
+        message " character:%d",             \
+      character                              \
+    );                                       \
+    self->logger.log(                        \
+      self->logger.payload,                  \
+      TSLogTypeLex,                          \
+      self->debug_buffer                     \
+    );                                       \
   }
 
-#define LOG_CHARACTER(message, character) \
-  LOG(                                    \
-    32 <= character && character < 127 ?  \
-    message " character:'%c'" :           \
-    message " character:%d", character    \
-  )
+static const int32_t BYTE_ORDER_MARK = 0xFEFF;
 
-static const char empty_chunk[3] = { 0, 0 };
+static const TSRange DEFAULT_RANGE = {
+  .start_point = {
+    .row = 0,
+    .column = 0,
+  },
+  .end_point = {
+    .row = UINT32_MAX,
+    .column = UINT32_MAX,
+  },
+  .start_byte = 0,
+  .end_byte = UINT32_MAX
+};
 
-static const int32_t BYTE_ORDER_MARK = 0xFEFF;
+// Check if the lexer has reached EOF. This state is stored
+// by setting the lexer's `current_included_range_index` such that
+// it has consumed all of its available ranges.
+static bool ts_lexer__eof(const TSLexer *_self) {
+  Lexer *self = (Lexer *)_self;
+  return self->current_included_range_index == self->included_range_count;
+}
+
+// Clear the currently stored chunk of source code, because the lexer's
+// position has changed.
+static void ts_lexer__clear_chunk(Lexer *self) {
+  self->chunk = NULL;
+  self->chunk_size = 0;
+  self->chunk_start = 0;
+}
 
+// Call the lexer's input callback to obtain a new chunk of source code
+// for the current position.
 static void ts_lexer__get_chunk(Lexer *self) {
   self->chunk_start = self->current_position.bytes;
   self->chunk = self->input.read(
@@ -30,15 +62,15 @@ static void ts_lexer__get_chunk(Lexer *self) {
     self->current_position.extent,
     &self->chunk_size
   );
-  if (!self->chunk_size) self->chunk = empty_chunk;
+  if (!self->chunk_size) {
+    self->current_included_range_index = self->included_range_count;
+    self->chunk = NULL;
+  }
 }
 
-typedef utf8proc_ssize_t (*DecodeFunction)(
-  const utf8proc_uint8_t *,
-  utf8proc_ssize_t,
-  utf8proc_int32_t *
-);
-
+// Decode the next unicode character in the current chunk of source code.
+// This assumes that the lexer has already retrieved a chunk of source
+// code that spans the current position.
 static void ts_lexer__get_lookahead(Lexer *self) {
   uint32_t position_in_chunk = self->current_position.bytes - self->chunk_start;
   const uint8_t *chunk = (const uint8_t *)self->chunk + position_in_chunk;
@@ -50,29 +82,37 @@ static void ts_lexer__get_lookahead(Lexer *self) {
     return;
   }
 
-  DecodeFunction decode =
-    self->input.encoding == TSInputEncodingUTF8 ? utf8proc_iterate : utf16_iterate;
+  UnicodeDecodeFunction decode = self->input.encoding == TSInputEncodingUTF8
+    ? ts_decode_utf8
+    : ts_decode_utf16;
 
   self->lookahead_size = decode(chunk, size, &self->data.lookahead);
 
   // If this chunk ended in the middle of a multi-byte character,
   // try again with a fresh chunk.
-  if (self->data.lookahead == -1 && size < 4) {
+  if (self->data.lookahead == TS_DECODE_ERROR && size < 4) {
     ts_lexer__get_chunk(self);
     chunk = (const uint8_t *)self->chunk;
     size = self->chunk_size;
     self->lookahead_size = decode(chunk, size, &self->data.lookahead);
   }
 
-  if (self->data.lookahead == -1) {
+  if (self->data.lookahead == TS_DECODE_ERROR) {
     self->lookahead_size = 1;
   }
 }
 
-static void ts_lexer__advance(TSLexer *payload, bool skip) {
-  Lexer *self = (Lexer *)payload;
-  if (self->chunk == empty_chunk)
-    return;
+// Advance to the next character in the source code, retrieving a new
+// chunk of source code if needed.
+static void ts_lexer__advance(TSLexer *_self, bool skip) {
+  Lexer *self = (Lexer *)_self;
+  if (!self->chunk) return;
+
+  if (skip) {
+    LOG("skip", self->data.lookahead);
+  } else {
+    LOG("consume", self->data.lookahead);
+  }
 
   if (self->lookahead_size) {
     self->current_position.bytes += self->lookahead_size;
@@ -84,53 +124,65 @@ static void ts_lexer__advance(TSLexer *payload, bool skip) {
     }
   }
 
-  TSRange *current_range = &self->included_ranges[self->current_included_range_index];
-  if (self->current_position.bytes == current_range->end_byte) {
-    self->current_included_range_index++;
-    if (self->current_included_range_index == self->included_range_count) {
-      self->data.lookahead = '\0';
-      self->lookahead_size = 1;
-      return;
-    } else {
-      current_range++;
-      self->current_position = (Length) {
-        current_range->start_byte,
-        current_range->start_point,
-      };
+  const TSRange *current_range = NULL;
+  if (self->current_included_range_index < self->included_range_count) {
+    current_range = &self->included_ranges[self->current_included_range_index];
+    if (self->current_position.bytes == current_range->end_byte) {
+      self->current_included_range_index++;
+      if (self->current_included_range_index < self->included_range_count) {
+        current_range++;
+        self->current_position = (Length) {
+          current_range->start_byte,
+          current_range->start_point,
+        };
+      } else {
+        current_range = NULL;
+      }
     }
   }
 
-  if (skip) {
-    LOG_CHARACTER("skip", self->data.lookahead);
-    self->token_start_position = self->current_position;
-  } else {
-    LOG_CHARACTER("consume", self->data.lookahead);
-  }
+  if (skip) self->token_start_position = self->current_position;
 
-  if (self->current_position.bytes >= self->chunk_start + self->chunk_size) {
-    ts_lexer__get_chunk(self);
+  if (current_range) {
+    if (self->current_position.bytes >= self->chunk_start + self->chunk_size) {
+      ts_lexer__get_chunk(self);
+    }
+    ts_lexer__get_lookahead(self);
+  } else {
+    ts_lexer__clear_chunk(self);
+    self->data.lookahead = '\0';
+    self->lookahead_size = 1;
   }
-
-  ts_lexer__get_lookahead(self);
 }
 
-static void ts_lexer__mark_end(TSLexer *payload) {
-  Lexer *self = (Lexer *)payload;
-  TSRange *current_included_range = &self->included_ranges[self->current_included_range_index];
-  if (self->current_included_range_index > 0 &&
-      self->current_position.bytes == current_included_range->start_byte) {
-    TSRange *previous_included_range = current_included_range - 1;
-    self->token_end_position = (Length) {
-      previous_included_range->end_byte,
-      previous_included_range->end_point,
-    };
-  } else {
-    self->token_end_position = self->current_position;
+// Mark that a token match has completed. This can be called multiple
+// times if a longer match is found later.
+static void ts_lexer__mark_end(TSLexer *_self) {
+  Lexer *self = (Lexer *)_self;
+  if (!ts_lexer__eof(&self->data)) {
+    // If the lexer is right at the beginning of included range,
+    // then the token should be considered to end at the *end* of the
+    // previous included range, rather than here.
+    TSRange *current_included_range = &self->included_ranges[
+      self->current_included_range_index
+    ];
+    if (
+      self->current_included_range_index > 0 &&
+      self->current_position.bytes == current_included_range->start_byte
+    ) {
+      TSRange *previous_included_range = current_included_range - 1;
+      self->t