/**
 * @file OpenQASM 3, and the OpenQASM 2 forms in common use, for tree-sitter.
 * @license Apache-2.0
 *
 * The grammar reads what is written: every statement of OpenQASM 3.1
 * (declarations of every type, gates and subroutines, modifiers, control
 * flow, timing, calibrations, pragmas and annotations) and OpenQASM 2's
 * `qreg`, `creg`, `opaque`, `measure q -> c` and `if(c==1)`. It does not say
 * what a program means; the engine's reader does.
 */

/// <reference types="tree-sitter-cli/dsl" />
// @ts-check

const PREC = {
  logical_or: 1,
  logical_and: 2,
  bit_or: 3,
  bit_xor: 4,
  bit_and: 5,
  equality: 6,
  comparison: 7,
  shift: 8,
  additive: 9,
  multiplicative: 10,
  unary: 11,
  power: 12,
  call: 13,
  index: 14,
};

const DECIMAL = /[0-9](_?[0-9])*/;
const FLOAT = choice(
  seq(DECIMAL, '.', optional(DECIMAL), optional(/[eE][+-]?[0-9]+/)),
  seq('.', DECIMAL, optional(/[eE][+-]?[0-9]+/)),
  seq(DECIMAL, /[eE][+-]?[0-9]+/),
);

/**
 * One or more of `rule`, separated by commas, with an optional trailing comma.
 * @param {RuleOrLiteral} rule
 */
function commaSep1(rule) {
  return seq(rule, repeat(seq(',', rule)), optional(','));
}

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

  extras: $ => [/\s/, $.line_comment, $.block_comment],

  word: $ => $.identifier,

  conflicts: $ => [
    // After a name and `[`: `g[100ns] q;` calls a gate for a duration,
    // `c[0] = …` assigns, and `c[0];` is an expression.
    [$.gate_call, $.indexed_identifier, $._expression],
    // `rx(θ)[100ns] q;` calls a gate for a duration; `f(x)[0];` indexes a call's value.
    [$.gate_call, $.call_expression],
  ],

  rules: {
    source_file: $ => seq(optional($.version), repeat($._statement_or_scope)),

    version: $ => seq('OPENQASM', field('number', choice($.integer_literal, $.float_literal)), ';'),

    _statement_or_scope: $ => choice($._statement, $.scope),

    scope: $ => seq('{', repeat($._statement_or_scope), '}'),

    _statement: $ => choice(
      $.pragma,
      $.annotation,
      $.include,
      $.calibration_grammar,
      $.quantum_declaration,
      $.old_style_declaration,
      $.classical_declaration,
      $.const_declaration,
      $.io_declaration,
      $.alias_declaration,
      $.gate_definition,
      $.opaque_declaration,
      $.subroutine_definition,
      $.extern_declaration,
      $.calibration,
      $.calibration_definition,
      $.gate_call,
      $.gphase_call,
      $.measure_statement,
      $.reset_statement,
      $.barrier_statement,
      $.delay_statement,
      $.box_statement,
      $.assignment_statement,
      $.expression_statement,
      $.if_statement,
      $.for_statement,
      $.while_statement,
      $.switch_statement,
      $.return_statement,
      $.break_statement,
      $.continue_statement,
      $.end_statement,
    ),

    pragma: _ => token(seq(choice('#pragma', 'pragma'), /[^\n]*/)),

    annotation: _ => token(seq('@', /[A-Za-z_][A-Za-z0-9_.]*/, /[^\n]*/)),

    include: $ => seq('include', field('path', $.string), ';'),

    calibration_grammar: $ => seq('defcalgrammar', field('name', $.string), ';'),

    // Declarations.

    quantum_declaration: $ => seq(field('type', $.qubit_type), field('name', $.identifier), ';'),

    qubit_type: $ => seq('qubit', optional($.designator)),

    old_style_declaration: $ => seq(choice('qreg', 'creg'), field('name', $.identifier), optional($.designator), ';'),

    classical_declaration: $ => seq(
      field('type', choice($.scalar_type, $.array_type)),
      field('name', $.identifier),
      optional(seq('=', field('value', $._declaration_value))),
      ';',
    ),

    const_declaration: $ => seq('const', field('type', $.scalar_type), field('name', $.identifier), '=', field('value', $._declaration_value), ';'),

    // `output` may name its value at once (`output int[32] n = 3;`), as the
    // engine's reader takes it; `input` takes its value when the program is read.
    io_declaration: $ => choice(
      seq('input', field('type', choice($.scalar_type, $.array_type)), field('name', $.identifier), ';'),
      seq('output', field('type', choice($.scalar_type, $.array_type)), field('name', $.identifier), optional(seq('=', field('value', $._expression))), ';'),
    ),

    alias_declaration: $ => seq('let', field('name', $.identifier), '=', field('value', $.alias_expression), ';'),

    _declaration_value: $ => choice($._expression, $.array_literal, $.measure_expression),

    designator: $ => seq('[', $._expression, ']'),

    scalar_type: $ => choice(
      seq(choice('bit', 'int', 'uint', 'float', 'angle'), optional($.designator)),
      'bool',
      'duration',
      'stretch',
      seq('complex', optional(seq('[', $.scalar_type, ']'))),
    ),

    array_type: $ => seq('array', '[', $.scalar_type, ',', commaSep1($._expression), ']'),

    array_reference_type: $ => seq(
      choice('readonly', 'mutable'),
      'array',
      '[',
      $.scalar_type,
      ',',
      choice(commaSep1($._expression), seq('#dim', '=', $._expression)),
      ']',
    ),

    // Gates, subroutines and calibrations.

    gate_definition: $ => seq(
      'gate',
      field('name', $.identifier),
      optional(field('parameters', $.gate_parameters)),
      field('qubits', $.identifier_list),
      field('body', $.scope),
    ),

    opaque_declaration: $ => seq('opaque', field('name', $.identifier), optional(field('parameters', $.gate_parameters)), field('qubits', $.identifier_list), ';'),

    gate_parameters: $ => seq('(', optional(commaSep1($.identifier)), ')'),

    identifier_list: $ => commaSep1($.identifier),

    subroutine_definition: $ => seq(
      'def',
      field('name', $.identifier),
      field('parameters', $.parameter_list),
      optional(field('return_type', $.return_signature)),
      field('body', $.scope),
    ),

    parameter_list: $ => seq('(', optional(commaSep1($.parameter)), ')'),

    parameter: $ => choice(
      seq(field('type', choice($.scalar_type, $.qubit_type, $.array_reference_type)), field('name', $.identifier)),
      seq(choice('creg', 'qreg'), field('name', $.identifier), optional($.designator)),
    ),

    return_signature: $ => seq('->', $.scalar_type),

    extern_declaration: $ => seq(
      'extern',
      field('name', $.identifier),
      '(',
      optional(commaSep1(choice($.scalar_type, $.array_reference_type, seq('creg', optional($.designator))))),
      ')',
      optional(field('return_type', $.return_signature)),
      ';',
    ),

    calibration: $ => seq('cal', field('body', $.calibration_block)),

    calibration_definition: $ => seq(
      'defcal',
      field('target', choice('measure', 'reset', 'delay', $.identifier)),
      optional(seq('(', optional(commaSep1(choice($.parameter, $._expression))), ')')),
      commaSep1(choice($.hardware_qubit, $.identifier)),
      optional(field('return_type', $.return_signature)),
      field('body', $.calibration_block),
    ),

    // A calibration's body is in its own grammar: kept as text, its braces matched.
    calibration_block: $ => seq('{', repeat(choice($.calibration_text, $.calibration_block)), '}'),

    calibration_text: _ => token(prec(-1, /[^{}\s][^{}]*/)),

    // Quantum statements.

    gate_call: $ => seq(
      repeat($.gate_modifier),
      field('name', $.identifier),
      optional(field('arguments', $.argument_list)),
      optional(field('duration', $.designator)),
      field('operands', $.operand_list),
      ';',
    ),

    gphase_call: $ => seq(
      repeat($.gate_modifier),
      'gphase',
      optional(field('arguments', $.argument_list)),
      optional(field('duration', $.designator)),
      optional(field('operands', $.operand_list)),
      ';',
    ),

    gate_modifier: $ => seq(
      choice(
        'inv',
        seq('pow', '(', $._expression, ')'),
        seq(choice('ctrl', 'negctrl'), optional(seq('(', $._expression, ')'))),
      ),
      '@',
    ),

    operand_list: $ => commaSep1($._operand),

    _operand: $ => choice($.indexed_identifier, $.hardware_qubit),

    indexed_identifier: $ => prec.left(seq($.identifier, repeat($.index))),

    index: $ => seq('[', choice($.set_expression, commaSep1(choice($._expression, $.range_expression))), ']'),

    measure_statement: $ => seq($.measure_expression, optional(seq('->', field('target', $.indexed_identifier))), ';'),

    measure_expression: $ => seq('measure', field('operand', $._operand)),

    reset_statement: $ => seq('reset', field('operand', $._operand), ';'),

    barrier_statement: $ => seq('barrier', optional($.operand_list), ';'),

    delay_statement: $ => seq('delay', field('duration', $.designator), optional($.operand_list), ';'),

    box_statement: $ => seq('box', optional(field('duration', $.designator)), field('body', $.scope)),

    // Classical statements.

    assignment_statement: $ => seq(
      field('left', $.indexed_identifier),
      field('operator', choice('=', '+=', '-=', '*=', '/=', '%=', '&=', '|=', '^=', '~=', '<<=', '>>=', '**=')),
      field('right', choice($._expression, $.measure_expression)),
      ';',
    ),

    expression_statement: $ => seq($._expression, ';'),

    if_statement: $ => prec.right(seq(
      'if',
      '(',
      field('condition', $._expression),
      ')',
      field('consequence', $._statement_or_scope),
      optional(seq('else', field('alternative', $._statement_or_scope))),
    )),

    for_statement: $ => seq(
      'for',
      optional(field('type', $.scalar_type)),
      field('variable', $.identifier),
      'in',
      field('iterable', choice($.set_expression, $.range, $.identifier)),
      field('body', $._statement_or_scope),
    ),

    range: $ => seq('[', $.range_expression, ']'),

    while_statement: $ => seq('while', '(', field('condition', $._expression), ')', field('body', $._statement_or_scope)),

    switch_statement: $ => seq('switch', '(', field('value', $._expression), ')', '{', repeat(choice($.case_clause, $.default_clause)), '}'),

    case_clause: $ => seq('case', commaSep1($._expression), field('body', $.scope)),

    default_clause: $ => seq('default', field('body', $.scope)),

    return_statement: $ => seq('return', optional(choice($._expression, $.measure_expression)), ';'),

    break_statement: _ => seq('break', ';'),

    continue_statement: _ => seq('continue', ';'),

    end_statement: _ => seq('end', ';'),

    // Expressions.

    _expression: $ => choice(
      $.identifier,
      $.hardware_qubit,
      $.integer_literal,
      $.float_literal,
      $.imaginary_literal,
      $.timing_literal,
      $.boolean_literal,
      $.string,
      $.parenthesized_expression,
      $.index_expression,
      $.call_expression,
      $.cast_expression,
      $.durationof_expression,
      $.unary_expression,
      $.binary_expression,
    ),

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

    index_expression: $ => prec(PREC.index, seq(field('value', $._expression), $.index)),

    call_expression: $ => seq(field('function', $.identifier), field('arguments', $.argument_list)),

    argument_list: $ => seq('(', optional(commaSep1($._expression)), ')'),

    cast_expression: $ => prec(PREC.call, seq(field('type', choice($.scalar_type, $.array_type)), '(', $._expression, ')')),

    durationof_expression: $ => seq('durationof', '(', $.scope, ')'),

    unary_expression: $ => prec(PREC.unary, seq(field('operator', choice('-', '+', '!', '~')), field('operand', $._expression))),

    binary_expression: $ => {
      /** @type {[string[], number, 'left' | 'right'][]} */
      const table = [
        [['**'], PREC.power, 'right'],
        [['*', '/', '%'], PREC.multiplicative, 'left'],
        [['+', '-'], PREC.additive, 'left'],
        [['<<', '>>'], PREC.shift, 'left'],
        [['<', '>', '<=', '>='], PREC.comparison, 'left'],
        [['==', '!='], PREC.equality, 'left'],
        [['&'], PREC.bit_and, 'left'],
        [['^'], PREC.bit_xor, 'left'],
        [['|'], PREC.bit_or, 'left'],
        [['&&'], PREC.logical_and, 'left'],
        [['||'], PREC.logical_or, 'left'],
      ];
      return choice(...table.map(([ops, p, assoc]) => {
        const op = ops.length === 1 ? ops[0] : choice(...ops);
        const rule = seq(field('left', $._expression), field('operator', op), field('right', $._expression));
        return assoc === 'right' ? prec.right(p, rule) : prec.left(p, rule);
      }));
    },

    alias_expression: $ => seq($._expression, repeat(seq('++', $._expression))),

    range_expression: $ => prec.right(seq(
      optional(field('start', $._expression)),
      ':',
      optional(field('step_or_end', $._expression)),
      optional(seq(':', optional(field('end', $._expression)))),
    )),

    set_expression: $ => seq('{', commaSep1($._expression), '}'),

    array_literal: $ => seq('{', optional(commaSep1(choice($._expression, $.array_literal))), '}'),

    // Literals and names.

    integer_literal: _ => token(choice(/0[bB][01](_?[01])*/, /0[oO][0-7](_?[0-7])*/, /0[xX][0-9a-fA-F](_?[0-9a-fA-F])*/, DECIMAL)),

    float_literal: _ => token(FLOAT),

    imaginary_literal: _ => token(seq(choice(FLOAT, DECIMAL), /[ \t]*/, 'im')),

    timing_literal: _ => token(seq(choice(FLOAT, DECIMAL), choice('dt', 'ns', 'us', 'µs', 'ms', 's'))),

    boolean_literal: _ => choice('true', 'false'),

    string: _ => token(choice(/"[^"\n]*"/, /'[^'\n]*'/)),

    hardware_qubit: _ => /\$[0-9]+/,

    identifier: _ => /[\p{L}_][\p{L}\p{N}_]*/u,

    line_comment: _ => token(seq('//', /[^\n]*/)),

    block_comment: _ => token(seq('/*', /[^*]*\*+([^/*][^*]*\*+)*/, '/')),
  },
});