Update mem2reg test

examples/mem2reg/mem2reg.input.ir

@@ -15,14 +15,17 @@ init:
     %l0:i32:x
 
 block b0:
-  %b0:i0:unit = store 42:i32 %l0:*i32
+  %b0:i0:i32 = load %l0:*i32
-  %b0:i1:i32 = load %l0:*i32
 
-  %b0:i2:unit = store 37:i32 %l0:*i32
+  %b0:i1:unit = store 42:i32 %l0:*i32
-  %b0:i3:i32 = load %l0:*i32
+  %b0:i2:i32 = load %l0:*i32
 
-  %b0:i4:unit = call @sink(%b0:i1:i32)
+  %b0:i3:unit = store 37:i32 %l0:*i32
-  %b0:i5:unit = call @sink(%b0:i3:i32)
+  %b0:i4:i32 = load %l0:*i32
+
+  %b0:i5:unit = call @sink(%b0:i0:i32)
+  %b0:i6:unit = call @sink(%b0:i2:i32)
+  %b0:i7:unit = call @sink(%b0:i4:i32)
 
   ret 0:i32
 }
@@ -34,17 +37,117 @@ init:
     %l0:i32:x
 
 block b0:
-  %b0:i0:unit = store 42:i32 %l0:*i32
+  %b0:i0:i32 = load %l0:*i32
-  %b0:i1:i32 = load %l0:*i32
 
-  %b0:i2:unit = store 37:i32 %l0:*i32
+  %b0:i1:unit = store 42:i32 %l0:*i32
-  %b0:i3:i32 = load %l0:*i32
+  %b0:i2:i32 = load %l0:*i32
+
+  %b0:i3:unit = store 37:i32 %l0:*i32
+  %b0:i4:i32 = load %l0:*i32
 
   j b1()
 
 block b1:
-  %b1:i0:unit = call @sink(%b0:i1:i32)
+  %b1:i0:i32 = load %l0:*i32
-  %b1:i1:unit = call @sink(%b0:i3:i32)
+
+  %b1:i1:unit = call @sink(%b0:i0:i32)
+  %b1:i2:unit = call @sink(%b0:i2:i32)
+  %b1:i3:unit = call @sink(%b0:i4:i32)
+  %b1:i4:unit = call @sink(%b1:i0:i32)
+
+  ret 0:i32
+}
+
+fun i32 @single_location {
+init:
+  bid: b666
+  allocations:
+    %l0:i32:x
+
+block b666:
+  j b0()
+
+block b0:
+  %b0:i0:i32 = load %l0:*i32
+  %b0:i1:unit = call @sink(%b0:i0:i32)
+  %b0:i2:unit = store 0:i32 %l0:*i32
+  switch undef:i32 default b1() [
+    2:i32 b2()
+    3:i32 b3()
+  ]
+
+block b1:
+  %b1:i0:i32 = load %l0:*i32
+  %b1:i1:unit = call @sink(%b1:i0:i32)
+  %b1:i2:unit = store 100:i32 %l0:*i32
+  j b2()
+
+block b2:
+  %b2:i0:i32 = load %l0:*i32
+  %b2:i1:unit = call @sink(%b2:i0:i32)
+  %b2:i2:unit = store 200:i32 %l0:*i32
+  j b3()
+
+block b3:
+  %b3:i0:unit = store 300:i32 %l0:*i32
+  j b0()
+}
+
+fun i32 @multi_location {
+init:
+  bid: b666
+  allocations:
+    %l0:i32:x
+    %l1:i16:y
+    %l2:i64:z
+
+block b666:
+  j b0()
+
+block b0:
+  %b0:i0:i32 = load %l0:*i32
+  %b0:i1:unit = call @sink(%b0:i0:i32)
+  %b0:i2:unit = store 0:i32 %l0:*i32
+
+  %b0:i3:i16 = load %l1:*i16
+  %b0:i4:unit = call @sink(%b0:i3:i16)
+  %b0:i5:unit = store 1:i16 %l1:*i16
+
+  %b0:i6:i64 = load %l2:*i64
+  %b0:i7:unit = call @sink(%b0:i6:i64)
+  %b0:i8:unit = store 2:i64 %l2:*i64
+
+  switch undef:i32 default b1() [
+    2:i32 b2()
+    3:i32 b3()
+  ]
+
+block b1:
+  %b1:i0:i32 = load %l0:*i32
+  %b1:i1:unit = call @sink(%b1:i0:i32)
+  %b1:i2:unit = store 100:i32 %l0:*i32
+
+  %b1:i3:i64 = load %l2:*i64
+  %b1:i4:unit = call @sink(%b1:i3:i64)
+  %b1:i5:unit = store 102:i64 %l2:*i64
+
+  j b2()
+
+block b2:
+  %b2:i0:i32 = load %l0:*i32
+  %b2:i1:unit = call @sink(%b2:i0:i32)
+  %b2:i2:unit = store 200:i32 %l0:*i32
+
+  %b2:i3:i16 = load %l1:*i16
+  %b2:i4:unit = call @sink(%b2:i3:i16)
+  %b2:i5:unit = store 201:i16 %l1:*i16
+
+  j b3()
+
+block b3:
+  %b3:i0:unit = store 300:i32 %l0:*i32
+  %b3:i1:unit = store 301:i16 %l1:*i16
+  %b3:i2:unit = store 302:i64 %l2:*i64
 
   j b0()
 }

examples/mem2reg/mem2reg.output.ir

@@ -15,14 +15,17 @@ init:
     %l0:i32:x
 
 block b0:
-  %b0:i0:unit = nop
+  %b0:i0:i32 = load %l0:*i32
-  %b0:i1:i32 = load %l0:*i32
 
-  %b0:i2:unit = nop
+  %b0:i1:unit = nop
-  %b0:i3:i32 = load %l0:*i32
+  %b0:i2:i32 = load %l0:*i32
 
-  %b0:i4:unit = call @sink(42:i32)
+  %b0:i3:unit = nop
-  %b0:i5:unit = call @sink(37:i32)
+  %b0:i4:i32 = load %l0:*i32
+
+  %b0:i5:unit = call @sink(undef:i32)
+  %b0:i6:unit = call @sink(42:i32)
+  %b0:i7:unit = call @sink(37:i32)
 
   ret 0:i32
 }
@@ -34,17 +37,134 @@ init:
     %l0:i32:x
 
 block b0:
-  %b0:i0:unit = nop
+  %b0:i0:i32 = load %l0:*i32
-  %b0:i1:i32 = load %l0:*i32
 
-  %b0:i2:unit = nop
+  %b0:i1:unit = nop
-  %b0:i3:i32 = load %l0:*i32
+  %b0:i2:i32 = load %l0:*i32
 
-  j b1()
+  %b0:i3:unit = nop
+  %b0:i4:i32 = load %l0:*i32
+
+  j b1(37:i32)
 
 block b1:
-  %b1:i0:unit = call @sink(42:i32)
+  %b1:p0:i32:x
-  %b1:i1:unit = call @sink(37:i32)
 
-  j b0()
+  %b1:i0:i32 = load %l0:*i32
+
+  %b1:i1:unit = call @sink(undef:i32)
+  %b1:i2:unit = call @sink(42:i32)
+  %b1:i3:unit = call @sink(37:i32)
+  %b1:i4:unit = call @sink(%b1:p0:i32)
+
+  ret 0:i32
+}
+
+fun i32 @single_location {
+init:
+  bid: b666
+  allocations:
+    %l0:i32:x
+
+block b666:
+  j b0(undef:i32)
+
+block b0:
+  %b0:p0:i32:x
+
+  %b0:i0:i32 = load %l0:*i32
+  %b0:i1:unit = call @sink(%b0:p0:i32)
+  %b0:i2:unit = nop
+  switch undef:i32 default b1(0:i32) [
+    2:i32 b2(0:i32)
+    3:i32 b3()
+  ]
+
+block b1:
+  %b1:p0:i32:x
+  %b1:i0:i32 = load %l0:*i32
+  %b1:i1:unit = call @sink(%b1:p0:i32)
+  %b1:i2:unit = nop
+  j b2(100:i32)
+
+block b2:
+  %b2:p0:i32:x
+  %b2:i0:i32 = load %l0:*i32
+  %b2:i1:unit = call @sink(%b2:p0:i32)
+  %b2:i2:unit = nop
+  j b3()
+
+block b3:
+  %b3:i0:unit = nop
+  j b0(300:i32)
+}
+
+fun i32 @multi_location {
+init:
+  bid: b666
+  allocations:
+    %l0:i32:x
+    %l1:i16:y
+    %l2:i64:z
+
+block b666:
+  j b0(undef:i32, undef:i16, undef:i64)
+
+block b0:
+  %b0:p0:i32:x
+  %b0:p1:i16:y
+  %b0:p2:i64:z
+
+  %b0:i0:i32 = load %l0:*i32
+  %b0:i1:unit = call @sink(%b0:p0:i32)
+  %b0:i2:unit = nop
+
+  %b0:i3:i16 = load %l1:*i16
+  %b0:i4:unit = call @sink(%b0:p1:i16)
+  %b0:i5:unit = nop
+
+  %b0:i6:i64 = load %l2:*i64
+  %b0:i7:unit = call @sink(%b0:p2:i64)
+  %b0:i8:unit = nop
+
+  switch undef:i32 default b1(0:i32, 1:i16, 2:i64) [
+    2:i32 b2(0:i32, 1:i16)
+    3:i32 b3()
+  ]
+
+block b1:
+  %b1:p0:i32:x
+  %b1:p1:i16:y
+  %b1:p2:i64:z
+
+  %b1:i0:i32 = load %l0:*i32
+  %b1:i1:unit = call @sink(%b1:p0:i32)
+  %b1:i2:unit = nop
+
+  %b1:i3:i64 = load %l2:*i64
+  %b1:i4:unit = call @sink(%b1:p2:i64)
+  %b1:i5:unit = nop
+
+  j b2(100:i32, %b1:p1:i16)
+
+block b2:
+  %b2:p0:i32:x
+  %b2:p1:i16:y
+
+  %b2:i0:i32 = load %l0:*i32
+  %b2:i1:unit = call @sink(%b2:p0:i32)
+  %b2:i2:unit = nop
+
+  %b2:i3:i16 = load %l1:*i16
+  %b2:i4:unit = call @sink(%b2:p1:i16)
+  %b2:i5:unit = nop
+
+  j b3()
+
+block b3:
+  %b3:i0:unit = nop
+  %b3:i1:unit = nop
+  %b3:i2:unit = nop
+
+  j b0(300:i32, 301:i16, 302:i64)
 }

src/ir/mod.rs

@@ -321,7 +321,7 @@ impl fmt::Display for JumpArg {
             self.bid,
             self.args
                 .iter()
-                .map(|a| a.to_string())
+                .map(|a| format!("{}:{}", a, a.dtype()))
                 .collect::<Vec<_>>()
                 .join(", ")
         )

src/ir/write_ir.rs

@@ -259,12 +259,12 @@ impl WriteString for BlockExit {
                 default,
                 cases,
             } => format!(
-                "switch {} default: {} [\n{}\n  ]",
+                "switch {} default {} [\n{}\n  ]",
                 value.write_string(),
                 default,
                 cases
                     .iter()
-                    .map(|(v, b)| format!("    {}:{}, {}", v, v.dtype(), b))
+                    .map(|(v, b)| format!("    {}:{} {}", v, v.dtype(), b))
                     .collect::<Vec<_>>()
                     .join("\n")
             ),