dialects: (varith) fix switch parsing (#3351)

Apparently putting generic syntax in a roundtrip test never tests the parsing.
xdslproject · Oct 29, 2024 · c556a6a · c556a6a
1 parent f9123d0
commit c556a6a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/tests/filecheck/dialects/varith/varith_ops.mlir b/tests/filecheck/dialects/varith/varith_ops.mlir
@@ -30,11 +30,16 @@
 // CHECK:  %x6 = varith.mul %ta, %tb, %tc, %td : tensor<10xf32>
 // CHECK-GENERIC: %x6 = "varith.mul"(%ta, %tb, %tc, %td) : (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
 
-%x7 = "varith.switch"(%ia, %fa, %fb, %fc, %fd) <{"case_values" = dense<[0, 1, 2]> : vector<3xi32>}> : (i32, f32, f32, f32, f32) -> f32
+%x7 = varith.switch %ia : i32 -> f32, [
+  default: %fa,
+  0: %fb,
+  1: %fc,
+  2: %fd
+]
 // CHECK:      %x7 = varith.switch %ia : i32 -> f32, [
 // CHECK-NEXT:   default: %fa,
-// CHECK-NEXT:   0: %fb
-// CHECK-NEXT:   1: %fc
+// CHECK-NEXT:   0: %fb,
+// CHECK-NEXT:   1: %fc,
 // CHECK-NEXT:   2: %fd
 // CHECK-NEXT: ]
 

diff --git a/xdsl/dialects/varith.py b/xdsl/dialects/varith.py
@@ -134,13 +134,14 @@ def parse(cls, parser: Parser) -> Self:
         parser.parse_punctuation(",")
         parser.parse_punctuation("[")
         parser.parse_keyword("default")
+        parser.parse_punctuation(":")
         default_arg = parser.resolve_operand(
             parser.parse_unresolved_operand(), return_type
         )
 
         values: list[int] = []
         args: list[SSAValue] = []
-        while parser.parse_punctuation(","):
+        while parser.parse_optional_punctuation(","):
             values.append(parser.parse_integer())
             parser.parse_punctuation(":")
             args.append(