From 19eb37635f0db38867f2df9af67849be1515ee25 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <git@ozinenko.com>
Date: Mon, 23 Dec 2024 09:41:51 +0100
Subject: [PATCH] [mlir] add a simple pygments lexer

This enables syntax highlighting of MLIR using the Pygments package in Python,
which is in turn usable from LaTeX via the minted package.
---
 mlir/utils/pygments/README.md     | 45 +++++++++++++++++++++++++++++++
 mlir/utils/pygments/mlir_lexer.py | 38 ++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 mlir/utils/pygments/README.md
 create mode 100644 mlir/utils/pygments/mlir_lexer.py

diff --git a/mlir/utils/pygments/README.md b/mlir/utils/pygments/README.md
new file mode 100644
index 0000000000000..838faceb01b0f
--- /dev/null
+++ b/mlir/utils/pygments/README.md
@@ -0,0 +1,45 @@
+## Pygments Lexer for MLIR
+
+This file contains a simple Pygments lexer configuration for MLIR, derived from
+the version used in the original CGO paper. Pygments allows for advanced
+configurable syntax highlighting of any code. This lexer is known to be
+incomplete and support mostly core IR with a subset of built-in types.
+Additions and customizations are welcome.
+
+### Standalone Usage
+
+Install Pygments, e.g., by running `pip install Pygments` or a Python package
+manager of your choosing. Use the standalone `pygmentize` command by
+instructing it to load the custom lexer:
+
+```
+pygmentize -l /path/to/mlir_lexer.py:MlirLexer -x myfile.mlir
+```
+
+This will produce highlighted output in the terminal. Other output formats are
+available, see Pygments [documentation](https://pygments.org/docs/) for more
+information.
+
+### LaTeX Usage
+
+First, make sure your distribution includes the `minted` package and list in
+the preamble.
+
+```latex
+\usepackage{minted}
+```
+
+Place the `mlir_lexer.py` in a place where the `latex` binary can find it,
+typically in the working directory next to the main `.tex` file. Note that you
+will have to invoke `latex` with the `-shell-escape` flag. See the `minted` 
+package [documentation](https://ctan.org/pkg/minted?lang=en) for more
+information.
+
+Leverage the custom lexer facility of `minted` to use this lexer in your
+document as:
+
+```latex
+\begin{minted}{mlir_lexer.py:MlirLexer -x}
+   ... your code here ...
+\end{minted}
+```
diff --git a/mlir/utils/pygments/mlir_lexer.py b/mlir/utils/pygments/mlir_lexer.py
new file mode 100644
index 0000000000000..179a058e9110c
--- /dev/null
+++ b/mlir/utils/pygments/mlir_lexer.py
@@ -0,0 +1,38 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from pygments.lexer import RegexLexer
+from pygments.token import *
+
+
+class MlirLexer(RegexLexer):
+    name = "MLIR"
+    aliases = ["mlir"]
+    filenames = ["*.mlir"]
+
+    tokens = {
+        "root": [
+            (r"%[a-zA-Z0-9_]+", Name.Variable),
+            (r"@[a-zA-Z_][a-zA-Z0-9_]+", Name.Function),
+            (r"\^[a-zA-Z0-9_]+", Name.Label),
+            (r"#[a-zA-Z0-9_]+", Name.Constant),
+            (r"![a-zA-Z0-9_]+", Keyword.Type),
+            (r"[a-zA-Z_][a-zA-Z0-9_]*\.", Name.Entity),
+            (r"memref[^.]", Keyword.Type),
+            (r"index", Keyword.Type),
+            (r"i[0-9]+", Keyword.Type),
+            (r"f[0-9]+", Keyword.Type),
+            (r"[0-9]+", Number.Integer),
+            (r"[0-9]*\.[0-9]*", Number.Float),
+            (r'"[^"]*"', String.Double),
+            (r"affine_map", Keyword.Reserved),
+            # TODO: this should be within affine maps only
+            (r"\+-\*\/", Operator),
+            (r"floordiv", Operator.Word),
+            (r"ceildiv", Operator.Word),
+            (r"mod", Operator.Word),
+            (r"()\[\]<>,{}", Punctuation),
+            (r"\/\/.*\n", Comment.Single),
+        ]
+    }