From 83dfd6e6185926c57d53e9e87b56abba3595ef60 Mon Sep 17 00:00:00 2001
From: Guojie Luo <gluo@pku.edu.cn>
Date: Tue, 22 Mar 2022 10:27:05 +0800
Subject: [PATCH] Add a preliminary codegen of case statement

---
 src/cst/case_statement.rs | 105 ++++++++++++++++++++++++++++++++++----
 1 file changed, 96 insertions(+), 9 deletions(-)

diff --git a/src/cst/case_statement.rs b/src/cst/case_statement.rs
index 9cc5b21..e681d7a 100644
--- a/src/cst/case_statement.rs
+++ b/src/cst/case_statement.rs
@@ -1,4 +1,4 @@
-use crate::cst::{Expression, ModuleContext, Statement, Tag, Tools};
+use crate::cst::{Expression, ModuleContext, Statement, Tag, Tools, UnitContext};
 use json::JsonValue;
 use llhd::ir::Block;
 
@@ -11,21 +11,108 @@ impl CaseStatement {
     pub fn codegen(json: &JsonValue, context: &mut ModuleContext) -> Option<Block> {
         let (json_expression, json_case_items, json_default_statement) = Self::analyze(json);
 
-        Expression::codegen(json_expression, context);
+        if context.unit_ctx.data.is_none() {
+            Expression::codegen(json_expression, context);
 
-        for (json_case_expressions, json_statement) in json_case_items {
-            for json_case_expression in json_case_expressions {
-                Expression::codegen(json_case_expression, context);
+            for (json_case_expressions, json_statement) in json_case_items {
+                for json_case_expression in json_case_expressions {
+                    Expression::codegen(json_case_expression, context);
+                }
+                Self::visit_case_body(json_statement, context);
             }
-            Statement::codegen(json_statement, context);
+
+            if json_default_statement.is_some() {
+                Self::visit_case_body(json_default_statement.unwrap(), context);
+            }
+
+            return None;
         }
 
+        let expr = {
+            let e = Expression::codegen(json_expression, context);
+            assert!(e.is_some(), "FIXME: CST node '{}'", json_expression);
+            e.unwrap()
+        };
+
+        let mut bb_next = {
+            let unit_ctx = &mut context.unit_ctx;
+            let mut builder = UnitContext::builder(&mut unit_ctx.data);
+            builder.block()
+        };
+
+        let case_tails: Vec<_> = json_case_items
+            .iter()
+            .map(|(json_case_expressions, json_statement)| {
+                let case_expressions: Vec<_> = json_case_expressions
+                    .iter()
+                    .map(|x| Expression::codegen(x, context).unwrap())
+                    .collect();
+
+                let bb_curr = {
+                    let unit_ctx = &mut context.unit_ctx;
+                    let mut builder = UnitContext::builder(&mut unit_ctx.data);
+
+                    let bb_curr = bb_next;
+                    bb_next = builder.block();
+
+                    builder.append_to(unit_ctx.bb_head.pop().unwrap());
+                    unit_ctx.bb_head.push(bb_next); // basic block for the next case
+
+                    let case_expr = case_expressions
+                        .into_iter()
+                        .reduce(|x, y| builder.ins().or(x, y))
+                        .unwrap();
+                    let cond = builder.ins().eq(expr, case_expr);
+                    builder.ins().br_cond(cond, bb_next, bb_curr);
+
+                    bb_curr
+                };
+
+                context.unit_ctx.bb_head.push(bb_curr);
+                let stmt = Self::visit_case_body(json_statement, context);
+                context.unit_ctx.bb_head.pop();
+
+                stmt.unwrap()
+            })
+            .collect();
+
         if json_default_statement.is_some() {
-            Statement::codegen(json_default_statement.unwrap(), context);
+            let stmt = Self::visit_case_body(json_default_statement.unwrap(), context);
+            let unit_ctx = &mut context.unit_ctx;
+            unit_ctx.bb_head.pop();
+            unit_ctx.bb_head.push(stmt.unwrap()); // basic block for the future statement
+            bb_next = stmt.unwrap();
         }
 
-        warn!("TODO: finish codegen()");
-        None
+        let unit_ctx = &mut context.unit_ctx;
+        if unit_ctx.data.is_some() {
+            let mut builder = UnitContext::builder(&mut unit_ctx.data);
+
+            // merge dangling basic blocks at the end of each case
+            for case_tail in case_tails {
+                builder.replace_block_use(case_tail, bb_next);
+                builder.delete_block(case_tail);
+            }
+        }
+
+        Some(bb_next)
+    }
+
+    fn visit_case_body(json: &JsonValue, context: &mut ModuleContext) -> Option<Block> {
+        let stmt = Statement::codegen(json, context);
+
+        if context.unit_ctx.data.is_some() {
+            let unit_ctx = &mut context.unit_ctx;
+            let mut builder = UnitContext::builder(&mut unit_ctx.data);
+
+            let bb_dangled = builder.block();
+            builder.append_to(stmt.unwrap());
+            builder.ins().br(bb_dangled);
+
+            Some(bb_dangled)
+        } else {
+            None
+        }
     }
 
     fn analyze(