From 7a90ff5791e5801bc76f680da4fea8082b3b7792 Mon Sep 17 00:00:00 2001
From: pca006132 <john.lck40@gmail.com>
Date: Mon, 9 Aug 2021 16:10:17 +0800
Subject: [PATCH] while loop constructs

---
 nac3core/src/codegen/stmt.rs | 42 ++++++++++++++++++++++++++++++++++++
 nac3core/src/top_level.rs    |  8 ++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/nac3core/src/codegen/stmt.rs b/nac3core/src/codegen/stmt.rs
index 1fa650605..897cf1b7b 100644
--- a/nac3core/src/codegen/stmt.rs
+++ b/nac3core/src/codegen/stmt.rs
@@ -82,6 +82,48 @@ impl<'ctx> CodeGenContext<'ctx> {
                     self.gen_assignment(target, value);
                 }
             }
+            StmtKind::Continue => {
+                self.builder.build_unconditional_branch(self.loop_bb.unwrap().0);
+            }
+            StmtKind::Break => {
+                self.builder.build_unconditional_branch(self.loop_bb.unwrap().1);
+            }
+            StmtKind::While { test, body, orelse } => {
+                let current = self.builder.get_insert_block().unwrap().get_parent().unwrap();
+                let test_bb = self.ctx.append_basic_block(current, "test");
+                let body_bb = self.ctx.append_basic_block(current, "body");
+                let cont_bb = self.ctx.append_basic_block(current, "cont");
+                // if there is no orelse, we just go to cont_bb
+                let orelse_bb = if orelse.is_empty() {
+                    cont_bb
+                } else {
+                    self.ctx.append_basic_block(current, "orelse")
+                };
+                // store loop bb information and restore it later
+                let loop_bb = self.loop_bb.replace((test_bb, cont_bb));
+                self.builder.build_unconditional_branch(test_bb);
+                self.builder.position_at_end(test_bb);
+                let test = self.gen_expr(test);
+                if let BasicValueEnum::IntValue(test) = test {
+                    self.builder.build_conditional_branch(test, body_bb, orelse_bb);
+                } else {
+                    unreachable!()
+                };
+                self.builder.position_at_end(body_bb);
+                for stmt in body.iter() {
+                    self.gen_stmt(stmt);
+                }
+                self.builder.build_unconditional_branch(test_bb);
+                if !orelse.is_empty() {
+                    self.builder.position_at_end(orelse_bb);
+                    for stmt in orelse.iter() {
+                        self.gen_stmt(stmt);
+                    }
+                    self.builder.build_unconditional_branch(cont_bb);
+                }
+                self.builder.position_at_end(cont_bb);
+                self.loop_bb = loop_bb;
+            }
             _ => unimplemented!(),
         }
     }
diff --git a/nac3core/src/top_level.rs b/nac3core/src/top_level.rs
index e205af288..cd4af4d84 100644
--- a/nac3core/src/top_level.rs
+++ b/nac3core/src/top_level.rs
@@ -3,7 +3,10 @@ use std::{collections::HashMap, sync::Arc};
 use super::typecheck::type_inferencer::PrimitiveStore;
 use super::typecheck::typedef::{SharedUnifier, Type, Unifier};
 use crate::symbol_resolver::SymbolResolver;
-use inkwell::{builder::Builder, context::Context, module::Module, types::BasicTypeEnum, values::PointerValue};
+use inkwell::{
+    basic_block::BasicBlock, builder::Builder, context::Context, module::Module,
+    types::BasicTypeEnum, values::PointerValue,
+};
 use parking_lot::RwLock;
 use rustpython_parser::ast::Stmt;
 
@@ -65,4 +68,7 @@ pub struct CodeGenContext<'ctx> {
     pub var_assignment: HashMap<String, PointerValue<'ctx>>,
     pub type_cache: HashMap<Type, BasicTypeEnum<'ctx>>,
     pub primitives: PrimitiveStore,
+    // where continue and break should go to respectively
+    // the first one is the test_bb, and the second one is bb after the loop
+    pub loop_bb: Option<(BasicBlock<'ctx>, BasicBlock<'ctx>)>,
 }