@felixtensor/tree-sitter-mlir 0.1.3 → 0.1.5

package/grammar.js

@@ -5,20 +5,45 @@ export default grammar({
   name: "mlir",
   extras: ($) => [/[\s\x00]/, $.comment],
   conflicts: ($) => [
+    // Core MLIR overlaps: shaped dimensions, aliases, pretty dialect payloads,
+    // value/type lists, and affine syntax share prefixes by design.
     [$._static_dim_list, $._static_dim_list],
-    [$.dictionary_attribute, $.region],
-    [$.custom_op_name, $.attribute_entry],
     [$.type_alias, $.dialect_namespace],
     [$.dialect_namespace, $.attribute_alias],
     [$.pretty_dialect_item],
-    [$.array_literal, $._custom_body_element_base],
-    [$._custom_body_element_base, $.tensor_type],
-    [$._custom_body_element_base, $._custom_body_arrow],
-    [$._generic_custom_operation_with_location_attr_dict, $.custom_op_name],
-    [$._custom_body_dict_key, $.attribute_entry],
     [$._value_use_list, $._value_use_and_type],
     [$._type_list_no_parens, $._type_or_func_type],
     [$._type_list_parens, $._multi_dim_affine_expr_parens],
+
+    // Custom operation fallback overlaps: loose body syntax must preserve
+    // dialect keywords, dictionary-looking payloads, and loc-sensitive forms
+    // without enumerating every upstream dialect operation.
+    [$.custom_op_name, $.attribute_entry],
+    [$.array_literal, $._custom_body_array_keyword],
+    [$._custom_body_tensor_keyword, $.tensor_type],
+    [$._generic_custom_operation_with_location_attr_dict, $.custom_op_name],
+    [$._custom_body_dict_key, $.attribute_entry],
+  ],
+  inline: ($) => [
+    $._tier1_custom_operation,
+    $._tier2_custom_operation,
+    $._custom_body_reference_element,
+    $._custom_body_type_element,
+    $._custom_body_attribute_or_braced_element,
+    $._custom_body_dialect_marker,
+    $._custom_body_group,
+    $._custom_body_atom,
+    $._custom_body_literal_element,
+    $._custom_body_reserved_keyword,
+    $._custom_body_affine_keyword,
+    $._custom_body_brace_payload,
+    $._custom_body_punctuation,
+    $._custom_body_separator_punctuation,
+    $._custom_body_operator_punctuation,
+    $._pretty_dialect_structural_content,
+    $._pretty_dialect_value_content,
+    $._pretty_dialect_keyword_content,
+    $._pretty_dialect_punctuation_content,
   ],
 
   // Token-level precedence constants (higher wins the token race):
@@ -47,7 +72,17 @@ export default grammar({
         $.external_resources,
       ),
     external_resources: ($) =>
-      seq("{-#", repeat($._pretty_dialect_item_contents), "#-}"),
+      seq("{-#", repeat($._external_resource_content), "#-}"),
+    _external_resource_content: ($) =>
+      choice(
+        $.string_literal,
+        // File metadata is YAML-like, not a dialect pretty body. Keep it loose
+        // so resource keys such as dense_resource_test_2xi32 are not split by
+        // dimension-list/type rules.
+        token(prec(-1, /[^#"\/]+/)),
+        "#",
+        "/",
+      ),
 
     // =========================================================================
     // Common syntax (lang-ref)
@@ -55,7 +90,7 @@ export default grammar({
     //  decimal-literal ::= digit+
     //  hexadecimal-literal ::= `0x` hex_digit+
     //  float-literal ::= [-+]?[0-9]+[.][0-9]*([eE][-+]?[0-9]+)?
-    //  string-literal  ::= `"` [^"\n\f\v\r]* `"`
+    //  string-literal  ::= `"` (char | escape-sequence | invalid-escape)* `"`
     // =========================================================================
     _digit: ($) => /[0-9]/,
     integer_literal: ($) => choice($._decimal_literal, $._hexadecimal_literal),
@@ -274,10 +309,18 @@ export default grammar({
     // Tier 2: _generic_custom_operation — all other dialect.op_name patterns
     // =========================================================================
     custom_operation: ($) =>
+      choice($._tier1_custom_operation, $._tier2_custom_operation),
+
+    _tier1_custom_operation: ($) =>
       choice(
         prec(2, $.func_operation),
         prec(2, $.module_operation),
         prec(2, $._affine_for_operation),
+      ),
+
+    _tier2_custom_operation: ($) =>
+      choice(
+        $._pdl_interp_record_match_operation,
         $._generic_custom_operation_with_location_attr_dict,
         $._generic_custom_operation,
       ),
@@ -350,6 +393,27 @@ export default grammar({
         ),
       ),
 
+    // These stay as specialized operation forms because `loc(...)` can be
+    // either custom body syntax or the operation-level trailing location.
+    _pdl_interp_record_match_operation: ($) =>
+      prec.dynamic(
+        -1,
+        prec.right(
+          seq(
+            field(
+              "name",
+              alias(
+                token(prec(20, "pdl_interp.record_match")),
+                $.custom_op_name,
+              ),
+            ),
+            repeat($._custom_body_element),
+            $._custom_body_location_list,
+            repeat($._custom_body_element),
+          ),
+        ),
+      ),
+
     // Tier 2: Generic custom operation — dialect.op_name + structural body
     // Negative dynamic precedence makes the parser prefer ending the body
     // and starting a new operation (with _op_result_list) over extending
@@ -419,48 +483,116 @@ export default grammar({
     _custom_body_element: ($) =>
       choice(
         $._custom_body_element_base,
+        // Kept out of _custom_body_element_base because nested delimiter bodies
+        // should still treat `>` as an angle-group boundary, not a loose marker.
         $._custom_body_successor_marker, // >^bb1 (WasmSSA if continuation)
       ),
 
     _custom_body_element_base: ($) =>
+      choice(
+        $._custom_body_reference_element,
+        $._custom_body_type_element,
+        $._custom_body_attribute_or_braced_element,
+        $._custom_body_dialect_marker,
+        $._custom_body_group,
+        $._custom_body_atom,
+        $._custom_body_punctuation,
+      ),
+
+    _custom_body_reference_element: ($) =>
       choice(
         $.value_use, // %foo, %0
         $.symbol_ref_id, // @sym, @"string"
         $.successor, // ^bb0, ^bb0(%arg : type)
-        $._custom_body_complex_label, // complex: %value (IRDL operand label)
-        prec(2, $.type), // !type, i32, memref<...>, etc.
+      ),
+
+    _custom_body_type_element: ($) =>
+      prec(2, $.type), // !type, i32, memref<...>, etc.
+
+    // Attribute includes dictionary_attribute, so keep it adjacent to the
+    // custom-body `{...}` payloads while preserving the public wrapper.
+    _custom_body_attribute_or_braced_element: ($) =>
+      choice(
         $.attribute, // #attr, {dict}, affine_map<...>
+        $._custom_body_brace_payload,
+      ),
+
+    _custom_body_brace_payload: ($) =>
+      choice(
         $._custom_body_tuple_group, // {(%v), (%w)}
         $.region, // { ... } (regions with operations)
         $._custom_body_value_group, // {%v : type, ...}
         $._custom_body_ssa_dict, // {"attr" = %value, ...} / options with SSA values
+      ),
+
+    _custom_body_dialect_marker: ($) =>
+      choice(
+        $._custom_body_arrow, // <- (OpenMP loop transform mapped-from marker)
+        $._custom_body_complex_label, // complex: %value (IRDL operand label)
+        $._custom_body_module_symbol_arg, // module(@sym) kernel attr
         $._custom_body_sparse_operand, // sparse(%idx : type)
+      ),
+
+    _custom_body_group: ($) =>
+      choice(
         $._custom_body_paren, // ( ... )
         $._custom_body_bracket, // [ ... ]
         $._custom_body_angle_group, // < ... >
+      ),
+
+    _custom_body_atom: ($) =>
+      choice(
+        $._custom_body_literal_element,
+        $._custom_body_reserved_keyword,
+        $.bare_id, // keywords: to, from, step, ins, outs, etc.
+      ),
+
+    _custom_body_literal_element: ($) =>
+      choice(
+        $.variadic, // custom assembly ellipsis marker
         $._literal, // 42, 3.14, "string", true, dense<...>
-        "array", // property names may collide with array<...>
+      ),
+
+    _custom_body_reserved_keyword: ($) =>
+      choice(
+        $._custom_body_array_keyword, // property names may collide with array<...>
         "vector", // OpenACC keyword may collide with vector<...>
-        "tensor", // AMDGPU/NVGPU keyword may collide with tensor<...>
-        "ceildiv",
-        "floordiv",
-        "mod", // inline affine keywords
-        $.bare_id, // keywords: to, from, step, ins, outs, etc.
-        $._custom_body_arrow, // <- (mapped-from, e.g. omp.fuse <- (...))
+        $._custom_body_tensor_keyword, // AMDGPU/NVGPU keyword may collide with tensor<...>
+        $._custom_body_affine_keyword, // inline affine keywords
+      ),
+
+    _custom_body_array_keyword: ($) => "array",
+    _custom_body_tensor_keyword: ($) => "tensor",
+    _custom_body_affine_keyword: ($) =>
+      choice("ceildiv", "floordiv", "mod"),
+
+    _custom_body_punctuation: ($) =>
+      choice(
+        $._custom_body_separator_punctuation,
+        $._custom_body_operator_punctuation,
+      ),
+
+    _custom_body_separator_punctuation: ($) =>
+      choice(
         ",",
         "=",
         ":",
         "->",
+      ),
+
+    _custom_body_operator_punctuation: ($) =>
+      choice(
         "*",
         "?",
         $.dimension_separator,
         "+",
-        "-",
+        $._custom_body_minus_punctuation,
         "/",
         "&",
         "|",
         "~",
       ),
+    _custom_body_minus_punctuation: ($) => "-",
 
     _custom_body_paren: ($) =>
       seq("(", repeat($._nested_custom_body_element), ")"),
@@ -513,8 +645,12 @@ export default grammar({
     _custom_body_ssa_value_array: ($) =>
       seq("[", $.value_use, repeat(seq(",", $.value_use)), "]"),
     _custom_body_successor_marker: ($) => seq(">", $.successor),
+    _custom_body_module_symbol_arg: ($) =>
+      seq(token(prec(20, "module(")), $.symbol_ref_id, ")"),
     _custom_body_sparse_operand: ($) =>
       seq($._sparse_keyword, $._custom_body_paren),
+    _custom_body_location_list: ($) =>
+      seq(token("loc"), "(", "[", optional($._value_use_list), "]", ")"),
     _custom_body_angle_group: ($) =>
       seq("<", repeat($._nested_custom_body_element), ">"),
     // "Mapped-from" arrow used by OpenMP loop-transform ops, e.g.
@@ -526,7 +662,7 @@ export default grammar({
     // mis-lex negative payloads like `#smt.bv<-1>` as `<-` `1`. Keeping the
     // tokens separate leaves the lexer unchanged; GLR distinguishes the arrow
     // from `_custom_body_angle_group` (which requires a closing '>').
-    _custom_body_arrow: ($) => seq("<", "-"),
+    _custom_body_arrow: ($) => prec(1, seq("<", "-")),
     // Only nested groups accept `trailing_location` as a body element.
     // At the top level it is omitted on purpose so the operation rule
     // captures a trailing `loc(...)` as the operation's location instead of
@@ -572,7 +708,8 @@ export default grammar({
     //   region      ::= `{` entry-block? block* `}`
     //   entry-block ::= operation+
     // =========================================================================
-    region: ($) => seq("{", optional($.entry_block), repeat($.block), "}"),
+    region: ($) =>
+      prec(1, seq("{", optional($.entry_block), repeat($.block), "}")),
     entry_block: ($) => repeat1($.operation),
 
     // =========================================================================
@@ -630,37 +767,48 @@ export default grammar({
     _pretty_dialect_item_contents: ($) =>
       prec.left(
         choice(
-          $.pretty_dialect_item_body,
-          $._pretty_dialect_bang_body_token,
-          $._pretty_dialect_body_attribute,
-          $.dialect_dim_list,
-          $.type,
-          prec(2, $.attribute),
-          $._literal,
-          $._dense_keyword,
-          $._sparse_keyword,
-          "array",
-          "vector",
-          "tensor",
-          "opaque",
-          $.bare_id,
-          ",",
-          ":",
-          "=",
-          "->",
-          "(",
-          ")",
-          "[",
-          "]",
-          "{",
-          "}",
-          "*",
-          "?",
-          "@",
-          "#",
+          $._pretty_dialect_structural_content,
+          $._pretty_dialect_value_content,
+          $._pretty_dialect_keyword_content,
+          $._pretty_dialect_punctuation_content,
           token(prec(-1, /[^<>]/)),
         ),
       ),
+    _pretty_dialect_structural_content: ($) =>
+      choice(
+        $.pretty_dialect_item_body,
+        $._pretty_dialect_bang_body_token,
+        $._pretty_dialect_body_attribute,
+      ),
+    _pretty_dialect_value_content: ($) =>
+      choice($.dialect_dim_list, $.type, prec(2, $.attribute), $._literal),
+    _pretty_dialect_keyword_content: ($) =>
+      choice(
+        $._dense_keyword,
+        $._sparse_keyword,
+        "array",
+        "vector",
+        "tensor",
+        "opaque",
+        $.bare_id,
+      ),
+    _pretty_dialect_punctuation_content: ($) =>
+      choice(
+        ",",
+        ":",
+        "=",
+        "->",
+        "(",
+        ")",
+        "[",
+        "]",
+        "{",
+        "}",
+        "*",
+        "?",
+        "@",
+        "#",
+      ),
     _pretty_dialect_bang_body_token: ($) =>
       token(prec(1, seq("!", /[^a-zA-Z_<>]/))),
     _pretty_dialect_body_attribute: ($) =>

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@felixtensor/tree-sitter-mlir",
-  "version": "0.1.3",
+  "version": "0.1.5",
   "description": "MLIR grammar for tree-sitter",
   "type": "module",
   "repository": {
@@ -55,9 +55,10 @@
     "compile": "tree-sitter generate",
     "update-tests": "tree-sitter test --update",
     "test": "tree-sitter test",
+    "test:bench": "node --test test/*.test.mjs",
     "test:bindings": "node --test bindings/node/*_test.js",
-    "test:examples": "tree-sitter parse --quiet --stat examples/**/*.mlir",
-    "test:all": "npm run test && npm run test:examples",
+    "test:examples": "tree-sitter parse --quiet --stat \"examples/**/*.mlir\"",
+    "test:all": "npm run test && npm run test:bench && npm run test:examples",
     "bench": "node bench.mjs",
     "prebuildify": "prebuildify --napi --strip"
   },