Implement if then else

src/backend.rs

@@ -82,10 +82,6 @@ impl Registers {
     }
 }
 
-/// Label in code.
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub struct Label(DynamicLabel);
-
 /// Describes location of a argument.
 enum ArgLocation {
     /// Argument is passed via some register.
@@ -139,7 +135,7 @@ impl CodeGenSession {
             asm: &mut self.assembler,
             start: start_offset,
             regs: Registers::new(),
-            sp_depth: 0,
+            sp_depth: StackDepth(0),
         }
     }
 
@@ -167,9 +163,8 @@ pub struct Context<'a> {
     asm: &'a mut Assembler,
     start: AssemblyOffset,
     regs: Registers,
-    /// Offset from starting value of SP counted in words. Each push and pop 
+    /// Each push and pop on the value stack increments or decrements this value by 1 respectively.
-    /// on the value stack increments or decrements this value by 1 respectively.
+    sp_depth: StackDepth,
-    sp_depth: u32,
 }
 
 impl<'a> Context<'a> {
@@ -179,6 +174,10 @@ impl<'a> Context<'a> {
     }
 }
 
+/// Label in code.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct Label(DynamicLabel);
+
 /// Create a new undefined label.
 pub fn create_label(ctx: &mut Context) -> Label {
     Label(ctx.asm.new_dynamic_label())
@@ -192,10 +191,32 @@ pub fn define_label(ctx: &mut Context, label: Label) {
     ctx.asm.dynamic_label(label.0);
 }
 
+/// Offset from starting value of SP counted in words.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct StackDepth(u32);
+
+impl StackDepth {
+    pub fn increment(&mut self) {
+        self.0 += 1;
+    }
+
+    pub fn decrement(&mut self) {
+        self.0 -= 1;
+    }
+}
+
+pub fn current_stack_depth(ctx: &Context) -> StackDepth {
+    ctx.sp_depth
+}
+
+pub fn restore_stack_depth(ctx: &mut Context, stack_depth: StackDepth) {
+    ctx.sp_depth = stack_depth;
+}
+
 fn push_i32(ctx: &mut Context, gpr: GPR) {
     // For now, do an actual push (and pop below). In the future, we could
     // do on-the-fly register allocation here.
-    ctx.sp_depth += 1;
+    ctx.sp_depth.increment();
     dynasm!(ctx.asm
         ; push Rq(gpr)
     );
@@ -203,7 +224,7 @@ fn push_i32(ctx: &mut Context, gpr: GPR) {
 }
 
 fn pop_i32(ctx: &mut Context) -> GPR {
-    ctx.sp_depth -= 1;
+    ctx.sp_depth.decrement();
     let gpr = ctx.regs.take_scratch_gpr();
     dynasm!(ctx.asm
         ; pop Rq(gpr)
@@ -222,7 +243,7 @@ pub fn add_i32(ctx: &mut Context) {
 }
 
 fn sp_relative_offset(ctx: &mut Context, slot_idx: u32) -> i32 {
-    ((ctx.sp_depth as i32) + slot_idx as i32) * WORD_SIZE as i32
+    ((ctx.sp_depth.0 as i32) + slot_idx as i32) * WORD_SIZE as i32
 }
 
 pub fn get_local_i32(ctx: &mut Context, local_idx: u32) {
@@ -257,6 +278,24 @@ pub fn relop_eq_i32(ctx: &mut Context) {
     ctx.regs.release_scratch_gpr(right);
 }
 
+/// Pops i32 predicate and branches to the specified label
+/// if the predicate is equal to zero.
+pub fn pop_and_breq(ctx: &mut Context, label: Label) {
+    let predicate = pop_i32(ctx);
+    dynasm!(ctx.asm
+        ; test Rd(predicate), Rd(predicate)
+        ; je =>label.0
+    );
+    ctx.regs.release_scratch_gpr(predicate);
+}
+
+/// Branch unconditionally to the specified label.
+pub fn br(ctx: &mut Context, label: Label) {
+    dynasm!(ctx.asm
+        ; jmp =>label.0
+    );
+}
+
 pub fn prepare_return_value(ctx: &mut Context) {
     let ret_gpr = pop_i32(ctx);
     if ret_gpr != RAX {
@@ -303,11 +342,11 @@ pub fn prologue(ctx: &mut Context, stack_slots: u32) {
         ; mov rbp, rsp
         ; sub rsp, framesize
     );
-    ctx.sp_depth += aligned_stack_slots - stack_slots;
+    ctx.sp_depth.0 += aligned_stack_slots - stack_slots;
 }
 
 pub fn epilogue(ctx: &mut Context) {
-    assert_eq!(ctx.sp_depth, 0, "imbalanced pushes and pops detected");
+    assert_eq!(ctx.sp_depth, StackDepth(0), "imbalanced pushes and pops detected");
     dynasm!(ctx.asm
         ; mov rsp, rbp
         ; pop rbp

src/function_body.rs

@@ -1,6 +1,8 @@
 use backend::*;
 use error::Error;
-use wasmparser::{FunctionBody, Operator};
+use wasmparser::{FunctionBody, Operator, Type};
+
+// TODO: Use own declared `Type` enum.
 
 /// Type of a control frame.
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -10,6 +12,7 @@ enum ControlFrameKind {
     /// Can be used for an implicit function block.
     Block { end_label: Label },
     /// Loop frame (branching to the beginning of block).
+    #[allow(unused)]
     Loop { header: Label },
     /// True-subblock of if expression.
     IfTrue {
@@ -49,20 +52,34 @@ impl ControlFrameKind {
 struct ControlFrame {
     kind: ControlFrameKind,
     /// Boolean which signals whether value stack became polymorphic. Value stack starts in non-polymorphic state and
-	/// becomes polymorphic only after an instruction that never passes control further is executed,
+    /// becomes polymorphic only after an instruction that never passes control further is executed,
-	/// i.e. `unreachable`, `br` (but not `br_if`!), etc.
+    /// i.e. `unreachable`, `br` (but not `br_if`!), etc.
     stack_polymorphic: bool,
-    // TODO: type, stack height, etc
+    /// Relative stack depth at the beginning of the frame.
+    stack_depth: StackDepth,
+    ty: Type,
 }
 
 impl ControlFrame {
-    pub fn new(kind: ControlFrameKind) -> ControlFrame {
+    pub fn new(kind: ControlFrameKind, stack_depth: StackDepth, ty: Type) -> ControlFrame {
         ControlFrame {
             kind,
+            stack_depth,
+            ty,
             stack_polymorphic: false,
         }
     }
 
+    pub fn outgoing_stack_depth(&self) -> StackDepth {
+        let mut outgoing_stack_depth = self.stack_depth;
+        if self.ty != Type::EmptyBlockType {
+            // If there a return value then increment expected outgoing stack depth value
+            // to account for the result value.
+            outgoing_stack_depth.increment();
+        }
+        outgoing_stack_depth
+    }
+
     /// Marks this control frame as reached stack-polymorphic state.
     pub fn mark_stack_polymorphic(&mut self) {
         self.stack_polymorphic = true;
@@ -75,6 +92,7 @@ pub fn translate(session: &mut CodeGenSession, body: &FunctionBody) -> Result<()
     // Assume signature is (i32, i32) -> i32 for now.
     // TODO: Use a real signature
     const ARG_COUNT: u32 = 2;
+    let return_ty = Type::I32;
 
     let mut framesize = ARG_COUNT;
     for local in locals {
@@ -100,6 +118,8 @@ pub fn translate(session: &mut CodeGenSession, body: &FunctionBody) -> Result<()
         ControlFrameKind::Block {
             end_label: epilogue_label,
         },
+        current_stack_depth(&ctx),
+        return_ty,
     ));
 
     for op in operators {
@@ -114,24 +134,63 @@ pub fn translate(session: &mut CodeGenSession, body: &FunctionBody) -> Result<()
             Operator::If { ty } => {
                 let end_label = create_label(&mut ctx);
                 let if_not = create_label(&mut ctx);
-                control_frames.push(ControlFrame::new(
-                    ControlFrameKind::IfTrue {
-                        end_label,
-                        if_not,
-                    },
-                ));
 
-                // TODO: Generate code that pops a value and executes the if_true part if the value
+                pop_and_breq(&mut ctx, if_not);
-                // is not equal to zero.
+
+                control_frames.push(ControlFrame::new(
+                    ControlFrameKind::IfTrue { end_label, if_not },
+                    current_stack_depth(&ctx),
+                    ty,
+                ));
+            }
+            Operator::Else => {
+                match control_frames.pop() {
+                    Some(ControlFrame {
+                        kind: ControlFrameKind::IfTrue { if_not, end_label },
+                        ty,
+                        stack_depth,
+                        ..
+                    }) => {
+                        // Finalize if..else block by jumping to the `end_label`.
+                        br(&mut ctx, end_label);
+
+                        // Define `if_not` label here, so if the corresponding `if` block receives
+                        // 0 it will branch here.
+                        // After that reset stack depth to the value before entering `if` block.
+                        define_label(&mut ctx, if_not);
+                        restore_stack_depth(&mut ctx, stack_depth);
+
+                        // Carry over the `end_label`, so it will be resolved when the corresponding `end`
+                        // is encountered.
+                        //
+                        // Also note that we reset `stack_depth` to the value before entering `if` block.
+                        let mut frame = ControlFrame::new(
+                            ControlFrameKind::IfFalse { end_label },
+                            stack_depth,
+                            ty,
+                        );
+                        control_frames.push(frame);
+                    }
+                    Some(_) => panic!("else expects if block"),
+                    None => panic!("control stack is never empty"),
+                };
             }
             Operator::End => {
                 let control_frame = control_frames.pop().expect("control stack is never empty");
+
                 if !control_frame.kind.is_loop() {
                     // Branches to a control frame with block type directs control flow to the header of the loop
                     // and we don't need to resolve it here. Branching to other control frames always lead
                     // control flow to the corresponding `end`.
                     define_label(&mut ctx, control_frame.kind.br_destination());
                 }
+
+                if let ControlFrameKind::IfTrue { if_not, .. } = control_frame.kind {
+                    // this is `if .. end` construction. Define the `if_not` label here.
+                    define_label(&mut ctx, if_not);
+                }
+
+                restore_stack_depth(&mut ctx, control_frame.outgoing_stack_depth());
             }
             Operator::I32Eq => {
                 relop_eq_i32(&mut ctx);

src/tests.rs

@@ -53,4 +53,56 @@ fn relop_eq() {
     }
 }
 
+#[test]
+fn if_then_else() {
+    const CASES: &[(usize, usize, usize)] = &[
+        (0, 1, 1),
+        (0, 0, 0),
+        (1, 0, 0),
+        (1, 1, 1),
+        (1312, 1, 1),
+        (1312, 1312, 1312),
+    ];
+
+    let code = r#"
+(module
+  (func (param i32) (param i32) (result i32)
+    (if (result i32)
+      (i32.eq
+        (get_local 0)
+        (get_local 1)
+      )
+      (then (get_local 0))
+      (else (get_local 1))
+    )
+  )
+)
+    "#;
+
+    for (a, b, expected) in CASES {
+        assert_eq!(execute_wat(code, *a, *b), *expected, "{}, {}", a, b);
+    }
+}
+
+#[test]
+fn if_without_result() {
+    let code = r#"
+(module
+  (func (param i32) (param i32) (result i32)
+    (if
+      (i32.eq
+        (get_local 0)
+        (get_local 1)
+      )
+      (then (unreachable))
+    )
+
+    (get_local 0)
+  )
+)
+    "#;
+
+    assert_eq!(execute_wat(code, 2, 3), 2);
+}
+
 // TODO: Add a test that checks argument passing via the stack.