Repository: spcl/pymlir

Issue:
Conditional branches in generic form
The following MLIR:
```
module  {
    func @mlir_entry(%a: i32) -> i32 {
        %1 = constant 1 : i32

        %isOne = cmpi "sle", %a, %1 : i32
        cond_br %isOne, ^one, ^else

        ^one: 
            return %1 : i32

        ^else: 
            return %1 : i32
    }
}
```
results in the following generic form:

```
// out.mlir
"module"() ( {
  "func"() ( {
  ^bb0(%arg0: i32):  // no predecessors
    %c1_i32 = "std.constant"() {value = 1 : i32} : () -> i32
    %0 = "std.cmpi"(%arg0, %c1_i32) {predicate = 3 : i64} : (i32, i32) -> i1
    "std.cond_br"(%0)[^bb1, ^bb2] {operand_segment_sizes = dense<[1, 0, 0]> : vector<3xi32>} : (i1) -> ()
  ^bb1:  // pred: ^bb0
    "std.return"(%c1_i32) : (i32) -> ()
  ^bb2:  // pred: ^bb0
    "std.return"(%c1_i32) : (i32) -> ()
  }) {sym_name = "mlir_entry", type = (i32) -> i32} : () -> ()
}) : () -> ()
``` 

Executing:
```
import mlir
ast1 = mlir.parse_path('out.mlir')
```

throws this exception:

```
lark.exceptions.UnexpectedCharacters: No terminal matches '[' in the current parser context, at line 6 col 22

    "std.cond_br"(%0)[^bb1, ^bb2] {operand_segment_sizes = de
                     ^
Expected one of: 
        * LBRACE
        * COLON
```


Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl) Report
> Merging [#16](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl) (92457c1) into [master](https://codecov.io/gh/spcl/pymlir/commit/5ce958e917c3249aff0076fd98e428d9f26e1dfd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl) (5ce958e) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/spcl/pymlir/pull/16/graphs/tree.svg?width=650&height=150&src=pr&token=hbsNVG5iGp&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl)](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl)

```diff
@@           Coverage Diff           @@
##           master      #16   +/-   ##
=======================================
  Coverage   94.43%   94.44%           
=======================================
  Files          15       15           
  Lines        1779     1782    +3     
=======================================
+ Hits         1680     1683    +3     
  Misses         99       99           
```


| [Impacted Files](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl) | Coverage Δ | |
|---|---|---|
| [mlir/astnodes.py](https://codecov.io/gh/spcl/pymlir/pull/16/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl#diff-bWxpci9hc3Rub2Rlcy5weQ==) | `93.37% <100.00%> (+0.03%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl). Last update [5ce958e...92457c1](https://codecov.io/gh/spcl/pymlir/pull/16?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spcl).


Patch:
diff --git a/AUTHORS b/AUTHORS
index 5047613..6583143 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -5,3 +5,4 @@ ETH Zurich
 Tal Ben-Nun
 Kaushik Kulkarni
 Mehdi Amini
+Berke Ates
\ No newline at end of file
diff --git a/mlir/astnodes.py b/mlir/astnodes.py
index e15c71f..db8b57c 100644
--- a/mlir/astnodes.py
+++ b/mlir/astnodes.py
@@ -508,6 +508,7 @@ class Op(Node):
 class GenericOperation(Op):
     name: str
     args: Optional[List[SsaId]]
+    successors: Optional[List[BlockId]]
     attributes: Optional[AttributeDict]
     type: List[Type]
 
@@ -519,6 +520,8 @@ class GenericOperation(Op):
             result += ', '.join(dump_or_value(arg, indent) for arg in self.args)
 
         result += ')'
+        if self.successors:
+            result += '[' + dump_or_value(self.successors, indent) + ']'
         if self.attributes:
             result += ' ' + dump_or_value(self.attributes, indent)
         if isinstance(self.type, list):
diff --git a/mlir/lark/mlir.lark b/mlir/lark/mlir.lark
index 8636fc9..d54dffe 100644
--- a/mlir/lark/mlir.lark
+++ b/mlir/lark/mlir.lark
@@ -167,7 +167,7 @@ location : string_literal ":" decimal_literal ":" decimal_literal
 trailing_location : ("loc" "(" location ")")
 
 // Undefined operations in all dialects
-generic_operation : string_literal "(" optional_ssa_use_list ")" optional_attr_dict trailing_type
+generic_operation : string_literal "(" optional_ssa_use_list ")" optional_successor_list optional_attr_dict trailing_type
 custom_operation  : bare_id "." bare_id optional_ssa_use_list trailing_type
 
 // Final operation definition
@@ -184,6 +184,7 @@ ssa_id_and_type_list : ssa_id_and_type ("," ssa_id_and_type)*
 operation_list: operation+
 
 block_label     : block_id optional_block_arg_list ":"
+successor_list   : "[" block_id? ("," block_id)* "]"
 
 block           : optional_block_label operation_list
 region : "{" block* "}"
@@ -211,6 +212,7 @@ region : "{" block* "}"
 ?optional_memory_space                  : ("," memory_space)?                  -> optional
 ?optional_block_label                   : block_label?                         -> optional
 ?optional_symbol_use_list               : symbol_use_list?                     -> optional
+?optional_successor_list                : successor_list?                      -> optional
 // ----------------------------------------------------------------------
 // Modules and functions
 


Test patch:
diff --git a/tests/test_syntax.py b/tests/test_syntax.py
index a785b23..37208de 100644
--- a/tests/test_syntax.py
+++ b/tests/test_syntax.py
@@ -215,6 +215,24 @@ def test_generic_dialect_std(parser: Optional[Parser] = None):
     module = parser.parse(code)
     print(module.pretty())
 
+def test_generic_dialect_std_cond_br(parser: Optional[Parser] = None):
+    code = '''
+"module"() ( {
+"func"() ( {
+^bb0(%arg0: i32):  // no predecessors
+    %c1_i32 = "std.constant"() {value = 1 : i32} : () -> i32
+    %0 = "std.cmpi"(%arg0, %c1_i32) {predicate = 3 : i64} : (i32, i32) -> i1
+    "std.cond_br"(%0)[^bb1, ^bb2] {operand_segment_sizes = dense<[1, 0, 0]> : vector<3xi32>} : (i1) -> ()
+^bb1:  // pred: ^bb0
+    "std.return"(%c1_i32) : (i32) -> ()
+^bb2:  // pred: ^bb0
+    "std.return"(%c1_i32) : (i32) -> ()
+}) {sym_name = "mlir_entry", type = (i32) -> i32} : () -> ()
+}) : () -> ()
+    '''
+    parser = parser or Parser()
+    module = parser.parse(code)
+    print(module.pretty())
 
 def test_generic_dialect_llvm(parser: Optional[Parser] = None):
     code = '''
@@ -255,5 +273,6 @@ if __name__ == '__main__':
     test_affine(p)
     test_definitions(p)
     test_generic_dialect_std(p)
+    test_generic_dialect_std_cond_br(p)
     test_generic_dialect_llvm(p)
     test_integer_sign(p)