Improper parsing of memref syntax

[3ulalia]

Hello,

I'm trying to do some operations on the MLIR generated for a simple C function:

void test(int *x, int *y) {
  *x = *y + 2;
}

Polygeist yields the following MLIR for this:

❯ ./bin/cgeist test/first.c -function=test -S
module { // (module attributes elided for space, and because mlir-opt chokes on them for some reason)
  func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2_i32 = arith.constant 2 : i32
    %0 = affine.load %arg1[0] : memref<?xi32>
    %1 = arith.addi %0, %c2_i32 : i32
    affine.store %1, %arg0[0] : memref<?xi32>
    return
  }
}

Which can then be lowered by mlir-opt:

❯ mlir-opt --lower-affine first.mlir
module {
  func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2_i32 = arith.constant 2 : i32
    %c0 = arith.constant 0 : index
    %0 = memref.load %arg1[%c0] : memref<?xi32>
    %1 = arith.addi %0, %c2_i32 : i32
    %c0_0 = arith.constant 0 : index
    memref.store %1, %arg0[%c0_0] : memref<?xi32>
    return
  }
}

However, attempting to parse this with pymlir makes it choke on the braces after the memref.load argument:

❯ python3
Python 3.12.9 (main, Feb  4 2025, 14:38:38) [GCC 14.2.1 20241116] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import mlir
>>> mlir.parse_string("""module {
...   func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
...     %c2_i32 = arith.constant 2 : i32
...     %c0 = arith.constant 0 : index
...     %0 = memref.load %arg1[%c0] : memref<?xi32>
...     %1 = arith.addi %0, %c2_i32 : i32
...     %c0_0 = arith.constant 0 : index
...     memref.store %1, %arg0[%c0_0] : memref<?xi32>
...     return
...   }
... }""")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/nix/store/s0d7723ifv9s3i57rgysj3b4sj1c2lc6-python3.12-pymlir-0.5/lib/python3.12/site-packages/mlir/parser.py", line 133, in parse_string
    return parser.parse(code)
           ^^^^^^^^^^^^^^^^^^
  File "/nix/store/s0d7723ifv9s3i57rgysj3b4sj1c2lc6-python3.12-pymlir-0.5/lib/python3.12/site-packages/mlir/parser.py", line 93, in parse
    tree = self.parser.parse(code)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/lark.py", line 655, in parse
    return self.parser.parse(text, start=start, on_error=on_error)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parser_frontends.py", line 104, in parse
    return self.parser.parse(stream, chosen_start, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/earley.py", line 280, in parse
    to_scan = self._parse(lexer, columns, to_scan, start_symbol)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/xearley.py", line 152, in _parse
    to_scan = scan(i, to_scan)
              ^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/xearley.py", line 125, in scan
    raise UnexpectedCharacters(stream, i, text_line, text_column, {item.expect.name for item in to_scan},
lark.exceptions.UnexpectedCharacters: No terminal matches '[' in the current parser context, at line 5 col 27

    %0 = memref.load %arg1[%c0] : memref<?xi32>
                          ^
Expected one of:
	* __ANON_4
	* __ANON_8
	* __ANON_0
	* COMMA
	* HASH
	* COLON
	* __ANON_7

Is this me being silly or is it an actual parsing error? I haven't spent enough time with pymlir to be able to tell.

Thanks!