fix: few regression fixes (#441)

core/src/cpu/air/branch.rs

@@ -47,6 +47,7 @@ impl CpuChip {
             // When we are branching, assert that local.pc <==> branch_columns.next_pc as Word.
             builder
                 .when_transition()
+                .when(next.is_real)
                 .when(local.branching)
                 .assert_eq(branch_cols.next_pc.reduce::<AB>(), next.pc);
 
@@ -61,9 +62,9 @@ impl CpuChip {
 
             // When we are not branching, assert that local.pc + 4 <==> next.pc.
             builder
-                .when(local.not_branching)
                 .when_transition()
                 .when(next.is_real)
+                .when(local.not_branching)
                 .assert_eq(local.pc + AB::Expr::from_canonical_u8(4), next.pc);
         }
 

core/src/cpu/air/mod.rs

@@ -200,6 +200,7 @@ impl CpuChip {
         // Verify that the word form of next.pc is correct for both jump instructions.
         builder
             .when_transition()
+            .when(next.is_real)
             .when(local.selectors.is_jal + local.selectors.is_jalr)
             .assert_eq(jump_columns.next_pc.reduce::<AB>(), next.pc);
 

core/src/stark/machine.rs

@@ -301,6 +301,7 @@ pub mod tests {
     use crate::runtime::tests::fibonacci_program;
     use crate::runtime::tests::simple_memory_program;
     use crate::runtime::tests::simple_program;
+    use crate::runtime::tests::ssz_withdrawals_program;
     use crate::runtime::Instruction;
     use crate::runtime::Opcode;
     use crate::runtime::Program;
@@ -458,4 +459,10 @@ pub mod tests {
         let program = simple_memory_program();
         run_test(program).unwrap();
     }
+
+    #[test]
+    fn test_ssz_withdrawal() {
+        let program = ssz_withdrawals_program();
+        run_test(program).unwrap();
+    }
 }