@felixtensor/tree-sitter-mlir 0.1.4 → 0.1.6

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):
@@ -65,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),
@@ -284,10 +309,17 @@ 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,
@@ -361,6 +393,8 @@ 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,
@@ -449,50 +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), ")"),
@@ -562,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
@@ -608,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),
 
     // =========================================================================
@@ -666,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.4",
+  "version": "0.1.6",
   "description": "MLIR grammar for tree-sitter",
   "type": "module",
   "repository": {
@@ -36,8 +36,8 @@
     "node-gyp-build": "^4.8.4"
   },
   "devDependencies": {
+    "node-gyp": "^12.4.0",
     "prebuildify": "^6.0.1",
-    "tree-sitter": "^0.25.0",
     "tree-sitter-cli": "^0.26.9"
   },
   "peerDependencies": {

package/queries/highlights.scm

@@ -11,11 +11,17 @@
 (func_operation name: _ @function.builtin)
 (module_operation name: _ @function.builtin)
 (func_operation ["private" "public" "attributes"] @keyword)
+(function_specifier) @keyword
 (module_operation "attributes" @keyword)
 
 ;; Dialect operations (e.g., arith.addi)
 (custom_op_name) @function.builtin
-(generic_operation (string_literal) @function.builtin)
+(custom_operation ["array" "sparse" "tensor" "vector"] @keyword)
+(custom_operation ["+" "-" "*" "/" "&" "|" "~"] @operator)
+(custom_operation "?" @punctuation.special)
+(custom_operation "loc" @keyword)
+(custom_operation "module(" @keyword)
+(custom_operation ">" @punctuation.bracket)
 
 ;; Symbols (@name)
 (symbol_ref_id) @string.special.symbol
@@ -42,20 +48,48 @@
 (dim_list "x" @punctuation.delimiter)
 (dimension_separator) @punctuation.delimiter
 (vector_dim_list "x" @punctuation.delimiter)
+(dim_list ["?" "*"] @punctuation.special)
+(dialect_dim_list ["?" "*"] @punctuation.special)
 
 [(attribute_alias) (attribute_alias_def) (dialect_attribute) (builtin_attribute) (dictionary_attribute)] @attribute
 
 ;; Specific attribute content
+(properties ["<{" "}>"] @punctuation.bracket)
+(affine_map "affine_map" @keyword)
+(affine_set "affine_set" @keyword)
 (affine_map ["max" "min" "symbol"] @keyword)
 (affine_set ["max" "min" "symbol"] @keyword)
+(affine_map
+  ["dense" "sparse" "compressed" "singleton" "loose_compressed" "n_out_of_m"]
+  @keyword)
+(affine_set
+  ["dense" "sparse" "compressed" "singleton" "loose_compressed" "n_out_of_m"]
+  @keyword)
+(affine_map ["+" "-" "*" "==" ">=" "<="] @operator)
+(affine_set ["+" "-" "*" "==" ">=" "<="] @operator)
+(strided_layout "strided" @keyword)
 (strided_layout "offset" @keyword)
-["ceildiv" "floordiv" "mod"] @keyword.operator
+(strided_layout ["?" "*"] @punctuation.special)
+(distinct_attribute "distinct" @keyword)
+(dense_resource_literal "dense_resource" @keyword)
+["ceildiv" "floordiv" "mod"] @operator
+(pretty_dialect_item_body
+  ["array" "dense" "opaque" "sparse" "tensor" "vector"] @keyword)
+(pretty_dialect_item_body ["?" "*"] @punctuation.special)
 
 ;; ── Literals ────────────────────────────────────────────────────────────────
 [(integer_literal) (float_literal) (complex_literal)] @number
 (bool_literal) @boolean
-[(tensor_literal) (dense_resource_literal) (array_literal) (unit_literal) (uninitialized_literal)] @constant.builtin
+[(tensor_literal) (array_literal) (unit_literal) (uninitialized_literal)] @constant.builtin
+(tensor_literal ["dense" "sparse"] @keyword)
+(array_literal "array" @keyword)
 (string_literal) @string
+(generic_operation (string_literal) @function.builtin)
+
+;; Escape sequences inside strings (\n, \t, \", \\, \HH) overlay on @string;
+;; malformed escapes are flagged distinctly rather than silently colored.
+(escape_sequence) @string.escape
+(invalid_escape) @error
 
 ;; ── SSA Variables (%name) ───────────────────────────────────────────────────
 ;; General uses and results (catch-all, overridden by more specific rules below)
@@ -67,17 +101,31 @@
 (block_arg_list (value_use) @variable.parameter)
 
 ;; ── Control Flow ────────────────────────────────────────────────────────────
-(caret_id) @label
+(caret_id) @tag
 (trailing_location "loc" @keyword)
+(callsite_location ["callsite" "at"] @keyword)
+(fused_location "fused" @keyword)
+(location "to" @keyword)
+(unknown_location) @constant.builtin
 (variadic) @punctuation.special
 
+;; ── External Resource Blocks ───────────────────────────────────────────────
+(external_resources ["{-#" "#-}"] @punctuation.bracket)
+
 ;; ── Punctuation ─────────────────────────────────────────────────────────────
 ["(" ")" "{" "}" "[" "]" "<" ">"] @punctuation.bracket
 ["," ":"] @punctuation.delimiter
 ["=" "->" "::"] @operator
 
-;; Catch-all for bare keywords in Op bodies (ins, outs, etc.)
-(bare_id) @keyword
+;; Catch-all for bare keywords in custom operation bodies (ins, outs, etc.).
+;; Keep this scoped to direct custom_operation children so attribute keys,
+;; affine dimensions, and pretty dialect payload identifiers do not all become
+;; keywords.
+(custom_operation (bare_id) @keyword)
+
+;; Dense resource handles override the bare_id catch-all above.
+(dense_resource_literal (bare_id) @constant.builtin)
 
 ;; Dictionary attribute keys override the bare_id catch-all above
 (attribute_entry (bare_id) @attribute)
+(attribute_entry (string_literal) @attribute)