xref: /llvm-project/mlir/utils/pygments/mlir_lexer.py (revision c7d237085bf9102ecf0c9105d8cc7fd94b752a3a)
1# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2# See https://llvm.org/LICENSE.txt for license information.
3# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5from pygments.lexer import RegexLexer
6from pygments.token import *
7
8
9class MlirLexer(RegexLexer):
10    name = "MLIR"
11    aliases = ["mlir"]
12    filenames = ["*.mlir"]
13
14    tokens = {
15        "root": [
16            (r"%[a-zA-Z0-9_]+", Name.Variable),
17            (r"@[a-zA-Z_][a-zA-Z0-9_]+", Name.Function),
18            (r"\^[a-zA-Z0-9_]+", Name.Label),
19            (r"#[a-zA-Z0-9_]+", Name.Constant),
20            (r"![a-zA-Z0-9_]+", Keyword.Type),
21            (r"[a-zA-Z_][a-zA-Z0-9_]*\.", Name.Entity),
22            (r"memref[^.]", Keyword.Type),
23            (r"index", Keyword.Type),
24            (r"i[0-9]+", Keyword.Type),
25            (r"f[0-9]+", Keyword.Type),
26            (r"[0-9]+", Number.Integer),
27            (r"[0-9]*\.[0-9]*", Number.Float),
28            (r'"[^"]*"', String.Double),
29            (r"affine_map", Keyword.Reserved),
30            # TODO: this should be within affine maps only
31            (r"\+-\*\/", Operator),
32            (r"floordiv", Operator.Word),
33            (r"ceildiv", Operator.Word),
34            (r"mod", Operator.Word),
35            (r"()\[\]<>,{}", Punctuation),
36            (r"\/\/.*\n", Comment.Single),
37        ]
38    }
39