[compiler-v2] Fixing bug in stack balancing of file-format generator (a…

…ptos-labs#14310)

* [compiler-v2] Fixing bug in stack balancing of file-format generator

Fixes aptos-labs#13912

The `balance_stack_end_of_block` function in function_generator.rs was both dumb and incorrect. It first pushed the argument which is needed for a branch on the stack, then checked if the stack needs flush, and then flushed the value just pushed right away, leading to stupid code like

```

CopyLoc(temp)
StLoc(temp)
CopyLoc(temp)
BrFalse
```

This code is also wrong if `temp` is borrowed since `StLoc(temp)` is not allowed in this case, which was the bytecode verification failure hit by the bug.

The fix simpplifies the stack balance function to ensure `temp` is not flushed unnecessarily.

Existing tests have been migrated from reference-safety to file-format-generator, and a new test has been added, reflecting a regression first caused by the fix.

* Addressing reviewer feedback

* More reviewer feedback

third_party/move/move-compiler-v2/src/file_format_generator/function_generator.rs

@@ -341,22 +341,25 @@ impl<'a> FunctionGenerator<'a> {
     }
 
     /// Balance the stack such that it exactly contains the `result` temps and nothing else. This
-    /// is used for instructions like `return` or `abort` which terminate a block und must leave
-    /// the stack empty at end.
+    /// is used for instructions like `branch`, `return` or `abort` which terminate a block
+    /// and must leave the stack empty at end.
     fn balance_stack_end_of_block(
         &mut self,
         ctx: &BytecodeContext,
         result: impl AsRef<[TempIndex]>,
     ) {
         let result = result.as_ref();
-        // First ensure the arguments are on the stack.
-        self.abstract_push_args(ctx, result, None);
-        if self.stack.len() != result.len() {
-            // Unfortunately, there is more on the stack than needed.
-            // Need to flush and push again so the stack is empty after return.
+        // If the stack contains already exactly the result and none of the temps is used after,
+        // nothing to do.
+        let stack_ready = self.stack == result
+            && self
+                .stack
+                .iter()
+                .all(|temp| !ctx.is_alive_after(*temp, &[], false));
+        if !stack_ready {
+            // Flush the stack and push the result
             self.abstract_flush_stack_before(ctx, 0);
-            self.abstract_push_args(ctx, result.as_ref(), None);
-            assert_eq!(self.stack.len(), result.len())
+            self.abstract_push_args(ctx, result, None);
         }
     }
 

third_party/move/move-compiler-v2/tests/file-format-generator/bug_13952_stack_balance.exp

@@ -0,0 +1,38 @@
+
+============ disassembled file-format ==================
+// Move bytecode v7
+module cafe.Module0 {
+struct Struct0 has copy, drop {
+	x: bool
+}
+
+public function5(Arg0: bool, Arg1: bool) /* def_idx: 0 */ {
+L2:	loc0: &bool
+L3:	loc1: bool
+L4:	loc2: &bool
+B0:
+	0: ImmBorrowLoc[0](Arg0: bool)
+	1: StLoc[2](loc0: &bool)
+	2: CopyLoc[0](Arg0: bool)
+	3: BrFalse(7)
+B1:
+	4: LdTrue
+	5: StLoc[3](loc1: bool)
+	6: Branch(9)
+B2:
+	7: MoveLoc[1](Arg1: bool)
+	8: StLoc[3](loc1: bool)
+B3:
+	9: ImmBorrowLoc[3](loc1: bool)
+	10: StLoc[4](loc2: &bool)
+	11: MoveLoc[2](loc0: &bool)
+	12: MoveLoc[4](loc2: &bool)
+	13: Neq
+	14: MoveLoc[0](Arg0: bool)
+	15: Pack[0](Struct0)
+	16: Pop
+	17: Pop
+	18: Ret
+}
+}
+============ bytecode verification succeeded ========

third_party/move/move-compiler-v2/tests/reference-safety/bug_13952.move → third_party/move/move-compiler-v2/tests/file-format-generator/bug_13952_stack_balance.move

@@ -1,11 +1,10 @@
-// TODO(#13952): after fix, rename file to reflect issue (`bug_nnn.move` to `bug_nnn_<issue>.move`)
 module 0xCAFE::Module0 {
     struct Struct0 has drop, copy {
         x: bool,
     }
 
     public fun function5(var21: bool, var23: bool) {
-        let var67 =  (&(var21) != &((var21 || var23)));
+        let _var67 =  (&(var21) != &((var21 || var23)));
         Struct0 {
             x: var21
         };

third_party/move/move-compiler-v2/tests/file-format-generator/bug_13952_stack_balance.opt.exp

@@ -0,0 +1,38 @@
+
+============ disassembled file-format ==================
+// Move bytecode v7
+module cafe.Module0 {
+struct Struct0 has copy, drop {
+	x: bool
+}
+
+public function5(Arg0: bool, Arg1: bool) /* def_idx: 0 */ {
+L2:	loc0: &bool
+L3:	loc1: bool
+L4:	loc2: &bool
+B0:
+	0: ImmBorrowLoc[0](Arg0: bool)
+	1: StLoc[2](loc0: &bool)
+	2: CopyLoc[0](Arg0: bool)
+	3: BrFalse(7)
+B1:
+	4: LdTrue
+	5: StLoc[3](loc1: bool)
+	6: Branch(9)
+B2:
+	7: MoveLoc[1](Arg1: bool)
+	8: StLoc[3](loc1: bool)
+B3:
+	9: ImmBorrowLoc[3](loc1: bool)
+	10: StLoc[4](loc2: &bool)
+	11: MoveLoc[2](loc0: &bool)
+	12: MoveLoc[4](loc2: &bool)
+	13: Neq
+	14: MoveLoc[0](Arg0: bool)
+	15: Pack[0](Struct0)
+	16: Pop
+	17: Pop
+	18: Ret
+}
+}
+============ bytecode verification succeeded ========

third_party/move/move-compiler-v2/tests/file-format-generator/bug_13952_stack_balance_regression.exp

@@ -0,0 +1,87 @@
+
+============ disassembled file-format ==================
+// Move bytecode v7
+module 42.m {
+enum Data has drop {
+ V1{
+	x: u64
+ },
+ V2{
+	x: u64,
+	y: bool
+ },
+ V3{
+
+ }
+}
+
+test_v1(): bool /* def_idx: 0 */ {
+L0:	loc0: Data
+L1:	loc1: bool
+B0:
+	0: LdU64(43)
+	1: PackVariant[0](Data/V1)
+	2: StLoc[0](loc0: Data)
+	3: ImmBorrowLoc[0](loc0: Data)
+	4: TestVariant[0](Data/V1)
+	5: StLoc[1](loc1: bool)
+	6: CopyLoc[1](loc1: bool)
+	7: BrTrue(8)
+B1:
+	8: MoveLoc[1](loc1: bool)
+	9: Ret
+}
+test_v1v3(): bool /* def_idx: 1 */ {
+L0:	loc0: Data
+L1:	loc1: &Data
+L2:	loc2: bool
+L3:	loc3: Data
+L4:	loc4: &Data
+L5:	loc5: bool
+B0:
+	0: LdU64(43)
+	1: PackVariant[0](Data/V1)
+	2: StLoc[0](loc0: Data)
+	3: ImmBorrowLoc[0](loc0: Data)
+	4: StLoc[1](loc1: &Data)
+	5: CopyLoc[1](loc1: &Data)
+	6: TestVariant[0](Data/V1)
+	7: StLoc[2](loc2: bool)
+	8: CopyLoc[2](loc2: bool)
+	9: BrTrue(15)
+B1:
+	10: MoveLoc[1](loc1: &Data)
+	11: TestVariant[1](Data/V3)
+	12: StLoc[2](loc2: bool)
+	13: CopyLoc[2](loc2: bool)
+	14: BrTrue(15)
+B2:
+	15: PackVariant[1](Data/V3)
+	16: StLoc[3](loc3: Data)
+	17: MoveLoc[2](loc2: bool)
+	18: BrFalse(32)
+B3:
+	19: ImmBorrowLoc[3](loc3: Data)
+	20: StLoc[4](loc4: &Data)
+	21: CopyLoc[4](loc4: &Data)
+	22: TestVariant[0](Data/V1)
+	23: StLoc[5](loc5: bool)
+	24: CopyLoc[5](loc5: bool)
+	25: BrTrue(31)
+B4:
+	26: MoveLoc[4](loc4: &Data)
+	27: TestVariant[1](Data/V3)
+	28: StLoc[5](loc5: bool)
+	29: CopyLoc[5](loc5: bool)
+	30: BrTrue(31)
+B5:
+	31: Branch(34)
+B6:
+	32: LdFalse
+	33: StLoc[5](loc5: bool)
+B7:
+	34: MoveLoc[5](loc5: bool)
+	35: Ret
+}
+}
+============ bytecode verification succeeded ========

third_party/move/move-compiler-v2/tests/file-format-generator/bug_13952_stack_balance_regression.move

@@ -0,0 +1,21 @@
+// Test relates to #13952 since the fix made this case fail first
+module 0x42::m {
+
+  enum Data has drop {
+    V1{x: u64},
+    V2{x: u64, y: bool}
+    V3
+  }
+
+  fun test_v1(): bool {
+    let d = Data::V1{x: 43};
+    (d is V1)
+  }
+
+  fun test_v1v3(): bool {
+    let d = Data::V1{x: 43};
+    let t = (d is V1|V3);
+    let d = Data::V3{};
+    t && (d is V1|V3)
+  }
+}

third_party/move/move-compiler-v2/tests/file-format-generator/bug_13952_stack_balance_regression.opt.exp

@@ -0,0 +1,85 @@
+
+============ disassembled file-format ==================
+// Move bytecode v7
+module 42.m {
+enum Data has drop {
+ V1{
+	x: u64
+ },
+ V2{
+	x: u64,
+	y: bool
+ },
+ V3{
+
+ }
+}
+
+test_v1(): bool /* def_idx: 0 */ {
+L0:	loc0: Data
+L1:	loc1: bool
+B0:
+	0: LdU64(43)
+	1: PackVariant[0](Data/V1)
+	2: StLoc[0](loc0: Data)
+	3: ImmBorrowLoc[0](loc0: Data)
+	4: TestVariant[0](Data/V1)
+	5: StLoc[1](loc1: bool)
+	6: CopyLoc[1](loc1: bool)
+	7: BrTrue(8)
+B1:
+	8: MoveLoc[1](loc1: bool)
+	9: Ret
+}
+test_v1v3(): bool /* def_idx: 1 */ {
+L0:	loc0: Data
+L1:	loc1: &Data
+L2:	loc2: bool
+L3:	loc3: Data
+B0:
+	0: LdU64(43)
+	1: PackVariant[0](Data/V1)
+	2: StLoc[0](loc0: Data)
+	3: ImmBorrowLoc[0](loc0: Data)
+	4: StLoc[1](loc1: &Data)
+	5: CopyLoc[1](loc1: &Data)
+	6: TestVariant[0](Data/V1)
+	7: StLoc[2](loc2: bool)
+	8: CopyLoc[2](loc2: bool)
+	9: BrTrue(15)
+B1:
+	10: MoveLoc[1](loc1: &Data)
+	11: TestVariant[1](Data/V3)
+	12: StLoc[2](loc2: bool)
+	13: CopyLoc[2](loc2: bool)
+	14: BrTrue(15)
+B2:
+	15: PackVariant[1](Data/V3)
+	16: StLoc[3](loc3: Data)
+	17: MoveLoc[2](loc2: bool)
+	18: BrFalse(32)
+B3:
+	19: ImmBorrowLoc[3](loc3: Data)
+	20: StLoc[1](loc1: &Data)
+	21: CopyLoc[1](loc1: &Data)
+	22: TestVariant[0](Data/V1)
+	23: StLoc[2](loc2: bool)
+	24: CopyLoc[2](loc2: bool)
+	25: BrTrue(31)
+B4:
+	26: MoveLoc[1](loc1: &Data)
+	27: TestVariant[1](Data/V3)
+	28: StLoc[2](loc2: bool)
+	29: CopyLoc[2](loc2: bool)
+	30: BrTrue(31)
+B5:
+	31: Branch(34)
+B6:
+	32: LdFalse
+	33: StLoc[2](loc2: bool)
+B7:
+	34: MoveLoc[2](loc2: bool)
+	35: Ret
+}
+}
+============ bytecode verification succeeded ========