1#include "tree_sitter/parser.h"
  2#include <stdlib.h>
  3#include <string.h>
  4#include <wctype.h>
  5
  6enum TokenType {
  7  DOLLAR_QUOTED_STRING_START_TAG,
  8  DOLLAR_QUOTED_STRING_END_TAG,
  9  DOLLAR_QUOTED_STRING
 10};
 11
 12#define MALLOC_STRING_SIZE 1024
 13
 14typedef struct LexerState {
 15  char* start_tag;
 16} LexerState;
 17
 18void *tree_sitter_sql_external_scanner_create() {
 19  LexerState *state = malloc(sizeof(LexerState));
 20  state->start_tag = NULL;
 21  return state;
 22}
 23
 24void *tree_sitter_sql_external_scanner_destroy(void *payload) {
 25  LexerState *state = (LexerState*)payload;
 26  if (state->start_tag != NULL) {
 27    free(state->start_tag);
 28    state->start_tag = NULL;
 29  }
 30  free(payload);
 31  return NULL;
 32}
 33
 34char* add_char(char* text, size_t* text_size, char c, int index) {
 35  if (text == NULL) {
 36    text = malloc(sizeof(char) * MALLOC_STRING_SIZE);
 37    *text_size = MALLOC_STRING_SIZE;
 38  }
 39
 40  // will break when indexes advances more than MALLOC_STRING_SIZE
 41  if (index + 1 >= *text_size) {
 42    *text_size += MALLOC_STRING_SIZE;
 43    char* tmp = malloc(*text_size * sizeof(char));
 44    strncpy(tmp, text, *text_size);
 45    free(text);
 46    text = tmp;
 47  }
 48
 49  text[index] = c;
 50  text[index + 1] = '\0';
 51  return text;
 52}
 53
 54char* scan_dollar_string_tag(TSLexer *lexer) {
 55  char* tag = NULL;
 56  int index = 0;
 57  size_t* text_size = malloc(sizeof(size_t));
 58  *text_size = 0;
 59  if (lexer->lookahead == '$') {
 60    tag = add_char(tag, text_size, '$', index);
 61    lexer->advance(lexer, false);
 62  } else {
 63    free(text_size);
 64    return NULL;
 65  }
 66
 67  while (lexer->lookahead != '$' && !iswspace(lexer->lookahead) && !lexer->eof(lexer)) {
 68    tag = add_char(tag, text_size, lexer->lookahead, ++index);
 69    lexer->advance(lexer, false);
 70  }
 71
 72  if (lexer->lookahead == '$') {
 73    tag = add_char(tag, text_size, lexer->lookahead, ++index);
 74    lexer->advance(lexer, false);
 75    free(text_size);
 76    return tag;
 77  } else {
 78    free(tag);
 79    free(text_size);
 80    return NULL;
 81  }
 82}
 83
 84bool tree_sitter_sql_external_scanner_scan(void *payload, TSLexer *lexer, const bool *valid_symbols) {
 85  LexerState *state = (LexerState*)payload;
 86  if (valid_symbols[DOLLAR_QUOTED_STRING_START_TAG] && state->start_tag == NULL) {
 87    while (iswspace(lexer->lookahead)) lexer->advance(lexer, true);
 88
 89    char* start_tag = scan_dollar_string_tag(lexer);
 90    if (start_tag == NULL) {
 91      return false;
 92    }
 93    if (state->start_tag != NULL) {
 94      free(state->start_tag);
 95      state->start_tag = NULL;
 96    }
 97    state->start_tag = start_tag;
 98    lexer->result_symbol = DOLLAR_QUOTED_STRING_START_TAG;
 99    return true;
100  }
101
102  if (valid_symbols[DOLLAR_QUOTED_STRING_END_TAG] && state->start_tag != NULL) {
103    while (iswspace(lexer->lookahead)) lexer->advance(lexer, true);
104
105    char* end_tag = scan_dollar_string_tag(lexer);
106    if (end_tag != NULL && strcmp(end_tag, state->start_tag) == 0) {
107      free(state->start_tag);
108      state->start_tag = NULL;
109      lexer->result_symbol = DOLLAR_QUOTED_STRING_END_TAG;
110      free(end_tag);
111      return true;
112    }
113    if (end_tag != NULL) {
114      free(end_tag);
115    }
116    return false;
117  }
118
119  if (valid_symbols[DOLLAR_QUOTED_STRING]) {
120    lexer->mark_end(lexer);
121    while (iswspace(lexer->lookahead)) lexer->advance(lexer, true);
122
123    char* start_tag = scan_dollar_string_tag(lexer);
124    if (start_tag == NULL) {
125      return false;
126    }
127
128    if (state->start_tag != NULL && strcmp(state->start_tag, start_tag) == 0) {
129      return false;
130    }
131
132    char* end_tag = NULL;
133    while (true) {
134      if (lexer->eof(lexer)) {
135        free(start_tag);
136        free(end_tag);
137        return false;
138      }
139
140      end_tag = scan_dollar_string_tag(lexer);
141      if (end_tag == NULL) {
142        lexer->advance(lexer, false);
143        continue;
144      }
145
146      if (strcmp(end_tag, start_tag) == 0) {
147        free(start_tag);
148        free(end_tag);
149        lexer->mark_end(lexer);
150        lexer->result_symbol = DOLLAR_QUOTED_STRING;
151        return true;
152      }
153
154      free(end_tag);
155      end_tag = NULL;
156    }
157  }
158
159  return false;
160}
161
162unsigned tree_sitter_sql_external_scanner_serialize(void *payload, char *buffer) {
163  LexerState *state = (LexerState *)payload;
164  if (state == NULL || state->start_tag == NULL) {
165    return 0;
166  }
167  // + 1 for the '\0'
168  int tag_length = strlen(state->start_tag) + 1;
169  if (tag_length >= TREE_SITTER_SERIALIZATION_BUFFER_SIZE) {
170    return 0;
171  }
172
173  memcpy(buffer, state->start_tag, tag_length);
174  if (state->start_tag != NULL) {
175    free(state->start_tag);
176    state->start_tag = NULL;
177  }
178  return tag_length;
179}
180
181void tree_sitter_sql_external_scanner_deserialize(void *payload, const char *buffer, unsigned length) {
182  LexerState *state = (LexerState *)payload;
183  state->start_tag = NULL;
184  // A length of 1 can't exists.
185  if (length > 1) {
186    state->start_tag = malloc(length);
187    memcpy(state->start_tag, buffer, length);
188  }
189}