Handle input chunks that end within multi-byte characters

src/runtime/lexer.c

@@ -8,11 +8,15 @@
 #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); \
+    self->logger.log(self->logger.payload, TSLogTypeLex, self->debug_buffer);         \
   }
 
 #define LOG_CHARACTER(message, character) \
-  LOG(character < 255 ? message " character:'%c'" : message " character:%d", character)
+  LOG(                                    \
+    32 <= character && character < 127 ?  \
+    message " character:'%c'" :           \
+    message " character:%d", character    \
+  )
 
 static const char empty_chunk[3] = { 0, 0 };
 
@@ -27,6 +31,12 @@ static void ts_lexer__get_chunk(Lexer *self) {
   if (!self->chunk_size) self->chunk = empty_chunk;
 }
 
+typedef utf8proc_ssize_t (*DecodeFunction)(
+  const utf8proc_uint8_t *,
+  utf8proc_ssize_t,
+  utf8proc_int32_t *
+);
+
 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;
@@ -38,15 +48,22 @@ static void ts_lexer__get_lookahead(Lexer *self) {
     return;
   }
 
-  if (self->input.encoding == TSInputEncodingUTF8) {
-    int64_t lookahead_size = utf8proc_iterate(chunk, size, &self->data.lookahead);
-    if (lookahead_size < 0) {
-      self->lookahead_size = 1;
-    } else {
-      self->lookahead_size = lookahead_size;
-    }
-  } else {
-    self->lookahead_size = utf16_iterate(chunk, size, &self->data.lookahead);
+  DecodeFunction decode =
+    self->input.encoding == TSInputEncodingUTF8 ? utf8proc_iterate : utf16_iterate;
+
+  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) {
+    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) {
+    self->lookahead_size = 1;
   }
 }
 

src/runtime/utf16.c

@@ -1,6 +1,10 @@
 #include "runtime/utf16.h"
 
-int utf16_iterate(const uint8_t *string, size_t length, int32_t *code_point) {
+utf8proc_ssize_t utf16_iterate(
+  const utf8proc_uint8_t *string,
+  utf8proc_ssize_t length,
+  utf8proc_int32_t *code_point
+) {
   if (length < 2) {
     *code_point = -1;
     return 0;

src/runtime/utf16.h

@@ -7,11 +7,12 @@ extern "C" {
 
 #include <stdint.h>
 #include <stdlib.h>
+#include "utf8proc.h"
 
 // Analogous to utf8proc's utf8proc_iterate function. Reads one code point from
-// the given string and stores it in the location pointed to by `code_point`.
+// the given UTF16 string and stores it in the location pointed to by `code_point`.
 // Returns the number of bytes in `string` that were read.
-int utf16_iterate(const uint8_t *string, size_t length, int32_t *code_point);
+utf8proc_ssize_t utf16_iterate(const utf8proc_uint8_t *, utf8proc_ssize_t, utf8proc_int32_t *);
 
 #ifdef __cplusplus
 }