grammar.js

const list = (item, sep = ',') => seq(item, repeat(seq(sep, item)));

const prec_op = {
  or: 1,
  and: 2,
  is: 3,
  comp: 4,
  bor: 5,
  bnot: 6,
  band: 7,
  bshift: 8,
  concat: 9,
  plus: 10,
  mult: 11,
  unary: 12,
  power: 13,

  as: 100,

  index: 1000,

  expr_for_statement: 10000,
};

const prec_type_op = {
  union: 100,
  index: 1000,
};

module.exports = grammar({
  name: 'teal',

  extras: $ => [
    $.comment,
    /[\s\n]/,
  ],

  externals: $ => [
    $.comment,

    $._long_string_start,
    $._long_string_char,
    $._long_string_end,

    $._short_string_start,
    $._short_string_char,
    $._short_string_end,
  ],

  inline: $ => [
    $.scope,
    $._partypelist,
    $._parnamelist,
    $._type_list_maybe_vararg,
  ],

  conflicts: $ => [
    // conflict lies in table entries with type annotations vs array entry of a method call
    // ex: { foo: bar = nil } vs { foo:bar() }
    [$._var, $.table_entry],
    [$.return_type],
  ],

  word: $ => $.identifier,

  rules: {
    program: $ => seq(
      alias(optional(token.immediate(/#![^\n\r]*/)), $.shebang_comment),
      repeat($._statement),
    ),

    _statement: $ => prec(1, seq(choice(
      $.var_declaration,
      $.var_assignment,
      $.type_declaration,
      $.record_declaration,
      $.interface_declaration,
      $.enum_declaration,
      $.return_statement,
      $.break,
      $._for_statement,
      $.do_statement,
      $.while_statement,
      $.repeat_statement,
      $.function_call,
      $.function_statement,
      $.macroexp_statement,
      $.if_statement,
      seq('::', alias($.identifier, $.label), '::'),
      $.goto,
    ), optional(';'))),

    return_statement: $ => prec.right(1, seq(
      'return',
      optional(choice(
        list($._expression),
        seq('(', list($._expression), ')'),
      )),
    )),

    break: $ => 'break',

    if_statement: $ => seq(
      'if', field('condition', $._expression), 'then',
      repeat($._statement),
      repeat($.elseif_block),
      optional($.else_block),
      'end',
    ),

    elseif_block: $ => seq('elseif', field('condition', $._expression), 'then', repeat($._statement)),
    else_block: $ => seq('else', repeat($._statement)),

    numeric_for_statement: $ => seq(
      'for', field('variable', $.identifier), '=', field('initializer', $._expression),
      ',', field('target', $._expression),
      optional(seq(',', field('step', $._expression))), field('body', alias($.do_statement, $.for_body)),
    ),

    generic_for_statement: $ => seq(
      'for', field('variable', list($.identifier)), 'in', field('iterator', list($._expression)), field('body', alias($.do_statement, $.for_body)),
    ),

    _for_statement: $ => choice(
      $.numeric_for_statement,
      $.generic_for_statement,
    ),

    while_statement: $ => seq(
      'while', field('condition', $._expression), alias($.do_statement, $.while_body),
    ),

    repeat_statement: $ => seq(
      'repeat', repeat($._statement),
      'until', field('condition', $._expression),
    ),

    do_statement: $ => seq(
      'do',
      repeat($._statement),
      'end',
    ),

    _expression: $ => choice(
      $.identifier,
      $._var,
      $.number,
      $.string,
      $.boolean,
      $.nil,
      $.table_constructor,
      $.anon_function,
      $.function_call,
      $._prefix_expression,
      $.bin_op,
      $.unary_op,
      $.varargs,
      $.type_cast,
      $.parenthesized_expression,
    ),

    parenthesized_expression: $ => seq('(', $._expression, ')'),

    unary_op: $ => prec.left(prec_op.unary, seq(
      field('op', alias(choice('not', '#', '-', '~'), $.op)),
      field('right', $._expression),
    )),

    bin_op: $ => choice(
      ...[
        ['or', prec_op.or],
        ['and', prec_op.and],
        ['<', prec_op.comp],
        ['<=', prec_op.comp],
        ['==', prec_op.comp],
        ['~=', prec_op.comp],
        ['>=', prec_op.comp],
        ['>', prec_op.comp],
        ['|', prec_op.bor],
        ['~', prec_op.bnot],
        ['&', prec_op.band],
        ['<<', prec_op.bshift],
        ['>>', prec_op.bshift],
        ['+', prec_op.plus],
        ['-', prec_op.plus],
        ['*', prec_op.mult],
        ['/', prec_op.mult],
        ['//', prec_op.mult],
        ['%', prec_op.mult],
      ].map(([operator, precedence]) => prec.left(precedence, seq(
        field('left', $._expression),
        field('op', alias(operator, $.op)),
        field('right', $._expression),
      ))),
      ...[
        ['..', prec_op.concat],
        ['^', prec_op.power],
      ].map(([operator, precedence]) => prec.right(precedence, seq(
        field('left', $._expression),
        field('op', alias(operator, $.op)),
        field('right', $._expression),
      ))),
      prec.right(prec_op.is, seq(
        field('left', $._expression),
        field('op', alias('is', $.op)),
        field('right', $._type),
      )),
    ),

    type_cast: $ => prec.right(prec_op.as, seq(
      $._expression,
      'as',
      choice($._type, $.type_tuple),
    )),

    type_tuple: $ => prec(10, seq('(', list($._type), ')')),
    type_union: $ => prec.left(prec_type_op.union, seq($._type, '|', $._type)),

    var_declarator: $ => seq(
      field('name', $.identifier),
      optional(seq('<', field('attribute', alias($.identifier, $.attribute)), '>')),
    ),

    var_declarators: $ => list(alias($.var_declarator, $.var)),

    expressions: $ => list($._expression),

    var_declaration: $ => seq(
      $.scope,
      field('declarators', $.var_declarators),
      optional(field('type_annotation', $.type_annotation)),
      optional(seq('=', field('initializers', $.expressions))),
    ),

    _type_def: $ => seq(
      'type',
      field('name', $.identifier),
      '=',
      field('value', choice(
        $._type, $._newtype,
      )),
    ),

    type_declaration: $ => choice(
      seq(
        field('scope', $.scope),
        $._type_def,
      ),
      $.record_declaration,
      $.interface_declaration,
      $.enum_declaration,
    ),

    assignment_variables: $ => list(alias($._var, $.var)),

    var_assignment: $ => seq(
      field('variables', $.assignment_variables), '=', field('expressions', $.expressions),
    ),

    _prefix_expression: $ => prec(10, choice(
      $._var,
      $.function_call,
      $.parenthesized_expression,
    )),
    method_index: $ => seq($._prefix_expression, ':', field('key', $.identifier)),
    arguments: $ => choice(
      seq('(', optional(list($._expression)), ')'),
      $.string,
      $.table_constructor,
    ),
    function_call: $ => prec(10, seq(
      field('called_object', choice(
        $._prefix_expression,
        $.method_index,
      )),
      field('arguments', $.arguments),
    )),

    table_entry: $ => choice(
      seq('[', field('expr_key', $._expression), ']', '=', field('value', $._expression)),
      prec(1, seq(
        field('key', $.identifier),
        field('type', optional(seq(':', $._type))),
        '=',
        field('value', $._expression),
      )),
      prec(1, field('value', $._expression)),
    ),

    table_constructor: $ => seq(
      '{',
      optional(list($.table_entry, choice(',', ';'))),
      optional(choice(',', ';')),
      '}',
    ),

    function_name: $ => seq(
      field('base', $.identifier),
      choice(
        seq(
          repeat(seq('.', field('entry', $.identifier))),
          seq(':', field('method', $.identifier)),
        ),
        repeat1(seq('.', field('entry', $.identifier))),
      ),
    ),

    scope: $ => field('scope', choice('local', 'global')),

    function_statement: $ => choice(
      seq(
        $.scope,
        'function',
        field('name', $.identifier),
        field('signature', $.function_signature),
        field('body', $.function_body),
      ),
      seq(
        'function',
        field('name', $.function_name),
        field('signature', $.function_signature),
        field('body', $.function_body),
      ),
    ),

    macroexp_statement: $ => seq(
      'local',
      $._macroexp_def,
    ),

    _macroexp_def: $ => seq(
      'macroexp',
      field('name', $.identifier),
      field('signature', alias($.function_signature, $.macroexp_signature)),
      field('body', $.macroexp_body),
    ),

    macroexp_body: $ => seq(
      $.return_statement,
      'end',
    ),

    variadic_type: $ => prec(prec_type_op.index - 1, seq($._type, '...')),
    _type_list_maybe_vararg: $ => seq(list($._type), optional(seq(',', $.variadic_type))),
    return_type: $ => choice(
      $.variadic_type,
      prec.dynamic(1, $._type_list_maybe_vararg),
      prec.dynamic(10, seq('(', $._type_list_maybe_vararg, ')')),
    ),

    _partypelist: $ => list(alias($._partype, $.arg)),
    _partype: $ => seq(optional(seq(field('name', $.identifier), ':')), field('type', $._type)),
    _parnamelist: $ => list(alias($._parname, $.arg)),
    _parname: $ => seq(
      field('name', $.identifier),
      optional(field('optional_marker', $.optional_marker)),
      optional(seq(':', field('type', $._type))),
    ),
    typeargs: $ => seq('<', list($.identifier), '>'),

    anon_function: $ => seq(
      'function',
      field('signature', $.function_signature),
      field('body', $.function_body),
    ),

    optional_marker: $ => '?',

    _annotated_var_arg: $ => seq('...', ':', field('type', $._type)),
    signature_arguments: $ => seq('(',
      optional(choice(
        $._parnamelist,
        seq(
          $._parnamelist,
          ',',
          alias($._annotated_var_arg, $.varargs),
        ),
        alias($._annotated_var_arg, $.varargs),
      )),
      ')',
    ),

    function_signature: $ => seq(
      field('typeargs', optional($.typeargs)),
      field('arguments', alias($.signature_arguments, $.arguments)),
      optional(seq(':', field('return_type', $.return_type))),
    ),

    function_body: $ => seq(
      repeat($._statement),
      'end',
    ),

    record_field: $ => choice(
      seq(
        field('key', $.identifier),
        ':', field('type', $._type),
      ),
      seq(
        '[', field('key', $.string), ']',
        ':', field('type', $._type),
      ),
      // TODO: there has to be a way around doing this, but I can't figure it out
      ...['type', 'record', 'interface', 'enum', 'userdata', 'metamethod'].map((reserved_id) => seq(
        field('key', alias(reserved_id, $.identifier)),
        ':', field('type', $._type),
      )),
    ),

    _record_entry: $ => choice(
      alias($.record_field, $.field),
      alias($._type_def, $.typedef),
      alias($._interface_def, $.interface_declaration),
      alias($._record_def, $.record_declaration),
      alias($._enum_def, $.enum_declaration),
      alias($._interface_def, $.interface_declaration),
      alias($._macroexp_def, $.macroexp_declaration),
    ),

    metamethod_annotation: $ => seq(
      'metamethod', field('name', $.identifier), ':', field('type', $._type),
    ),

    record_body: $ => seq(
      optional(seq(
        field('is_types', seq('is', list($._type))),
        optional(field('where', seq('where', $._expression))),
      )),
      optional(seq('{', alias($._type, $.record_array_type), '}')),
      repeat(choice(
        $._record_entry,
        alias('userdata', $.userdata),
        field('metamethod', alias($.metamethod_annotation, $.metamethod)),
      )),
      'end',
    ),

    _record_def: $ => seq(
      'record',
      field('name', $.identifier),
      field('typeargs', optional($.typeargs)),
      field('record_body', $.record_body),
    ),

    _interface_def: $ => seq(
      'interface',
      field('name', $.identifier),
      field('typeargs', optional($.typeargs)),
      field('interface_body', alias($.record_body, $.interface_body)),
    ),

    interface_declaration: $ => seq(
      $.scope,
      $._interface_def,
    ),

    record_declaration: $ => seq(
      $.scope,
      $._record_def,
    ),

    enum_body: $ => seq(
      repeat($.string),
      'end',
    ),

    _enum_def: $ => seq(
      'enum',
      field('name', $.identifier),
      field('enum_body', $.enum_body),
    ),

    enum_declaration: $ => seq(
      $.scope,
      $._enum_def,
    ),

    anon_interface: $ => seq(
      'interface',
      optional($.typeargs),
      field('interface_body', alias($.record_body, $.interface_body)),
    ),

    anon_record: $ => seq(
      'record',
      optional($.typeargs),
      field('record_body', $.record_body),
    ),

    _anon_enum: $ => seq(
      'enum',
      $.enum_body,
    ),

    _newtype: $ => choice(
      $._anon_enum,
      $.anon_record,
      $.anon_interface,
    ),

    type_annotation: $ => seq(
      ':',
      list($._type),
    ),

    _type: $ => prec(2, choice(
      $.simple_type,
      $.type_index,
      $.table_type,
      $.function_type,
      $.type_union,
      seq('(', $._type, ')'),
    )),

    typearg_params: $ => prec.right(1000, seq(
      '<', list($._type), '>',
    )),

    type_index: $ => prec.right(prec_type_op.index, seq(
      choice($.identifier, $.type_index), '.', $.identifier,
      optional(alias($.typearg_params, $.typeargs)),
    )),

    simple_type: $ => prec.right(1000, seq(
      field('name', $.identifier),
      optional(alias($.typearg_params, $.typeargs)),
    )),

    table_type: $ => seq(
      '{',
      choice(
        field('value_type', $._type),
        seq(
          field('key_type', $._type),
          ':',
          field('value_type', $._type),
        ),
        seq(
          field('tuple_type', $._type),
          repeat1(seq(',', field('tuple_type', $._type))),
        ),
      ),
      '}',
    ),

    function_type_args: $ => seq(
      '(',
      optional(choice(
        seq(
          $._partypelist,
          optional(seq(',', alias($._annotated_var_arg, $.varargs))),
        ),
        alias($._annotated_var_arg, $.varargs),
      )),
      ')',
    ),

    function_type: $ => prec.right(1, seq(
      'function',
      field('typeargs', optional($.typeargs)),
      optional(seq(
        field('arguments', alias($.function_type_args, $.arguments)),
        optional(seq(
          ':',
          field('return_type', $.return_type),
        )),
      )),
    )),

    goto: $ => seq('goto', $.identifier),

    index: $ => prec(1, seq(
      $._prefix_expression,
      choice(
        field('key', seq('.', $.identifier)),
        field('expr_key', seq('[', $._expression, ']')),
      ),
    )),

    _var: $ => prec(1, choice(
      $.identifier,
      $.index,
    )),

    varargs: $ => '...',
    identifier: $ => /[a-zA-Z_][a-zA-Z_0-9]*/,
    number: $ => choice(
      /\d+(\.\d+)?(e\d+)?/i,
      /\.\d+(e\d+)?/i,
      /0x[0-9a-fA-F]+(\.[0-9a-fA-F]+)?(p\d+)?/,
    ),
    boolean: $ => choice('true', 'false'),

    _short_string_content: $ => alias(repeat1(choice(
      $.format_specifier,
      $.escape_sequence,
      token.immediate(prec(1, '%')),
      prec(0, $._short_string_char),
    )), $.string_content),

    _long_string_content: $ => alias(repeat1(choice(
      $.format_specifier,
      $._long_string_char,
      token.immediate(prec(1, '%')),
    )), $.string_content),

    string: $ => prec(2, choice(
      seq(
        field('start', alias($._short_string_start, 'short_string_start')),
        field('content', optional($._short_string_content)),
        field('end', alias($._short_string_end, 'short_string_end')),
      ),

      seq(
        field('start', alias($._long_string_start, 'long_string_start')),
        field('content', optional($._long_string_content)),
        field('end', alias($._long_string_end, 'long_string_end')),
      ),
    )),

    format_specifier: $ => token.immediate(prec(3, seq(
      '%',
      choice(
        '%',
        seq(
          optional(choice(
            '+', '-',
          )),
          optional(' '),
          optional('#'),
          optional(/[0-9]+/),
          optional('.'),
          optional(/[0-9]+/),
          /[AaEefGgcdiouXxpqs]/,
        ),
      ),
    ))),

    escape_sequence: $ => token.immediate(prec(3, seq(
      '\\',
      choice(
        /[abfnrtvz"'\\]/,
        seq('x', /[0-9a-fA-F]{2}/),
        seq('d', /[0-7]{3}/),
        seq('u{', /[0-9a-fA-F]{1,8}/, '}'),
      ),
    ))),

    nil: $ => 'nil',
  },
});