nac3core/codegen: refactor according to #23

nac3artiq/src/lib.rs

@@ -19,7 +19,7 @@ use rustpython_parser::{
 use parking_lot::{Mutex, RwLock};
 
 use nac3core::{
-    codegen::{CodeGenTask, WithCall, WorkerRegistry},
+    codegen::{CodeGenTask, DefaultCodeGenerator, WithCall, WorkerRegistry},
     symbol_resolver::SymbolResolver,
     toplevel::{composer::TopLevelComposer, DefinitionId, GenCall, TopLevelContext, TopLevelDef},
     typecheck::typedef::{FunSignature, FuncArg},
@@ -423,11 +423,13 @@ impl Nac3 {
                 .expect("couldn't write module to file");
         })));
         let thread_names: Vec<String> = (0..4).map(|i| format!("module{}", i)).collect();
-        let threads: Vec<_> = thread_names.iter().map(|s| s.as_str()).collect();
+        let threads: Vec<_> = thread_names
+            .iter()
+            .map(|s| Box::new(DefaultCodeGenerator::new(s.to_string())))
+            .collect();
 
         py.allow_threads(|| {
-            let (registry, handles) =
-                WorkerRegistry::create_workers(&threads, top_level.clone(), f);
+            let (registry, handles) = WorkerRegistry::create_workers(threads, top_level.clone(), f);
             registry.add_task(task);
             registry.wait_tasks_complete(handles);
         });

nac3core/src/codegen/generator.rs

@@ -0,0 +1,142 @@
+use crate::{
+    codegen::{expr::*, stmt::*, CodeGenContext},
+    toplevel::{DefinitionId, TopLevelDef},
+    typecheck::typedef::{FunSignature, Type},
+};
+use inkwell::values::{BasicValueEnum, PointerValue};
+use rustpython_parser::ast::{Expr, Stmt, StrRef};
+
+pub trait CodeGenerator {
+    /// Return the module name for the code generator.
+    fn get_name(&self) -> &str;
+
+    /// Generate function call and returns the function return value.
+    /// - obj: Optional object for method call.
+    /// - fun: Function signature and definition ID.
+    /// - params: Function parameters. Note that this does not include the object even if the
+    ///   function is a class method.
+    fn gen_call<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        obj: Option<(Type, BasicValueEnum<'ctx>)>,
+        fun: (&FunSignature, DefinitionId),
+        params: Vec<(Option<StrRef>, BasicValueEnum<'ctx>)>,
+    ) -> Option<BasicValueEnum<'ctx>> {
+        gen_call(self, ctx, obj, fun, params)
+    }
+
+    /// Generate object constructor and returns the constructed object.
+    /// - signature: Function signature of the contructor.
+    /// - def: Class definition for the constructor class.
+    /// - params: Function parameters.
+    fn gen_constructor<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        signature: &FunSignature,
+        def: &TopLevelDef,
+        params: Vec<(Option<StrRef>, BasicValueEnum<'ctx>)>,
+    ) -> BasicValueEnum<'ctx> {
+        gen_constructor(self, ctx, signature, def, params)
+    }
+
+    /// Generate a function instance.
+    /// - obj: Optional object for method call.
+    /// - fun: Function signature, definition ID and the substitution key.
+    /// - params: Function parameters. Note that this does not include the object even if the
+    ///   function is a class method.
+    /// Note that this function should check if the function is generated in another thread (due to
+    /// possible race condition), see the default implementation for an example.
+    fn gen_func_instance<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        obj: Option<(Type, BasicValueEnum<'ctx>)>,
+        fun: (&FunSignature, &mut TopLevelDef, String),
+    ) -> String {
+        gen_func_instance(ctx, obj, fun)
+    }
+
+    /// Generate the code for an expression.
+    fn gen_expr<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        expr: &Expr<Option<Type>>,
+    ) -> Option<BasicValueEnum<'ctx>> {
+        gen_expr(self, ctx, expr)
+    }
+
+    /// Allocate memory for a variable and return a pointer pointing to it.
+    /// The default implementation places the allocations at the start of the function.
+    fn gen_var_alloc<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        ty: Type,
+    ) -> PointerValue<'ctx> {
+        gen_var(ctx, ty)
+    }
+
+    /// Return a pointer pointing to the target of the expression.
+    fn gen_store_target<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        pattern: &Expr<Option<Type>>,
+    ) -> PointerValue<'ctx> {
+        gen_store_target(self, ctx, pattern)
+    }
+
+    /// Generate code for an assignment expression.
+    fn gen_assign<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        target: &Expr<Option<Type>>,
+        value: BasicValueEnum<'ctx>,
+    ) {
+        gen_assign(self, ctx, target, value)
+    }
+
+    /// Generate code for a while expression.
+    /// Return true if the while loop must early return
+    fn gen_while<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        stmt: &Stmt<Option<Type>>,
+    ) -> bool {
+        gen_while(self, ctx, stmt);
+        false
+    }
+
+    /// Generate code for an if expression.
+    /// Return true if the statement must early return
+    fn gen_if<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        stmt: &Stmt<Option<Type>>,
+    ) -> bool {
+        gen_if(self, ctx, stmt)
+    }
+
+    /// Generate code for a statement
+    /// Return true if the statement must early return
+    fn gen_stmt<'ctx, 'a>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, 'a>,
+        stmt: &Stmt<Option<Type>>,
+    ) -> bool {
+        gen_stmt(self, ctx, stmt)
+    }
+}
+
+pub struct DefaultCodeGenerator {
+    name: String,
+}
+
+impl DefaultCodeGenerator {
+    pub fn new(name: String) -> DefaultCodeGenerator {
+        DefaultCodeGenerator { name }
+    }
+}
+
+impl CodeGenerator for DefaultCodeGenerator {
+    fn get_name(&self) -> &str {
+        &self.name
+    }
+}

nac3core/src/codegen/mod.rs

@@ -27,11 +27,14 @@ use std::sync::{
 use std::thread;
 
 mod expr;
+mod generator;
 mod stmt;
 
 #[cfg(test)]
 mod test;
 
+pub use generator::{CodeGenerator, DefaultCodeGenerator};
+
 pub struct CodeGenContext<'ctx, 'a> {
     pub ctx: &'ctx Context,
     pub builder: Builder<'ctx>,
@@ -77,8 +80,8 @@ pub struct WorkerRegistry {
 }
 
 impl WorkerRegistry {
-    pub fn create_workers(
-        names: &[&str],
+    pub fn create_workers<G: CodeGenerator + Send + 'static>(
+        generators: Vec<Box<G>>,
         top_level_ctx: Arc<TopLevelContext>,
         f: Arc<WithCall>,
     ) -> (Arc<WorkerRegistry>, Vec<thread::JoinHandle<()>>) {
@@ -89,21 +92,20 @@ impl WorkerRegistry {
         let registry = Arc::new(WorkerRegistry {
             sender: Arc::new(sender),
             receiver: Arc::new(receiver),
-            thread_count: names.len(),
+            thread_count: generators.len(),
             panicked: AtomicBool::new(false),
             task_count,
             wait_condvar,
         });
 
         let mut handles = Vec::new();
-        for name in names.iter() {
+        for mut generator in generators.into_iter() {
             let top_level_ctx = top_level_ctx.clone();
             let registry = registry.clone();
             let registry2 = registry.clone();
-            let name = name.to_string();
             let f = f.clone();
             let handle = thread::spawn(move || {
-                registry.worker_thread(name, top_level_ctx, f);
+                registry.worker_thread(generator.as_mut(), top_level_ctx, f);
             });
             let handle = thread::spawn(move || {
                 if let Err(e) = handle.join() {
@@ -156,18 +158,19 @@ impl WorkerRegistry {
         self.sender.send(Some(task)).unwrap();
     }
 
-    fn worker_thread(
+    fn worker_thread<G: CodeGenerator>(
         &self,
-        module_name: String,
+        generator: &mut G,
         top_level_ctx: Arc<TopLevelContext>,
         f: Arc<WithCall>,
     ) {
         let context = Context::create();
         let mut builder = context.create_builder();
-        let mut module = context.create_module(&module_name);
+        let mut module = context.create_module(generator.get_name());
 
         while let Some(task) = self.receiver.recv().unwrap() {
-            let result = gen_func(&context, self, builder, module, task, top_level_ctx.clone());
+            let result =
+                gen_func(&context, generator, self, builder, module, task, top_level_ctx.clone());
             builder = result.0;
             module = result.1;
             *self.task_count.lock() -= 1;
@@ -243,8 +246,9 @@ fn get_llvm_type<'ctx>(
     })
 }
 
-pub fn gen_func<'ctx>(
+pub fn gen_func<'ctx, G: CodeGenerator + ?Sized>(
     context: &'ctx Context,
+    generator: &mut G,
     registry: &WorkerRegistry,
     builder: Builder<'ctx>,
     module: Module<'ctx>,
@@ -351,7 +355,7 @@ pub fn gen_func<'ctx>(
 
     let mut returned = false;
     for stmt in task.body.iter() {
-        returned = code_gen_context.gen_stmt(stmt);
+        returned = generator.gen_stmt(&mut code_gen_context, stmt);
         if returned {
             break;
         }

nac3core/src/codegen/stmt.rs

@@ -1,221 +1,250 @@
 use super::{
     expr::{assert_int_val, assert_pointer_val},
-    CodeGenContext,
+    CodeGenContext, CodeGenerator,
 };
 use crate::typecheck::typedef::Type;
 use inkwell::values::{BasicValue, BasicValueEnum, PointerValue};
 use rustpython_parser::ast::{Expr, ExprKind, Stmt, StmtKind};
 
-impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
-    fn gen_var(&mut self, ty: Type) -> PointerValue<'ctx> {
-        // put the alloca in init block
-        let current = self.builder.get_insert_block().unwrap();
-        // position before the last branching instruction...
-        self.builder.position_before(&self.init_bb.get_last_instruction().unwrap());
-        let ty = self.get_llvm_type(ty);
-        let ptr = self.builder.build_alloca(ty, "tmp");
-        self.builder.position_at_end(current);
-        ptr
-    }
+pub fn gen_var<'ctx, 'a>(ctx: &mut CodeGenContext<'ctx, 'a>, ty: Type) -> PointerValue<'ctx> {
+    // put the alloca in init block
+    let current = ctx.builder.get_insert_block().unwrap();
+    // position before the last branching instruction...
+    ctx.builder.position_before(&ctx.init_bb.get_last_instruction().unwrap());
+    let ty = ctx.get_llvm_type(ty);
+    let ptr = ctx.builder.build_alloca(ty, "tmp");
+    ctx.builder.position_at_end(current);
+    ptr
+}
 
-    fn parse_pattern(&mut self, pattern: &Expr<Option<Type>>) -> PointerValue<'ctx> {
-        // very similar to gen_expr, but we don't do an extra load at the end
-        // and we flatten nested tuples
-        match &pattern.node {
-            ExprKind::Name { id, .. } => {
-                self.var_assignment.get(id).cloned().unwrap_or_else(|| {
-                    let ptr = self.gen_var(pattern.custom.unwrap());
-                    self.var_assignment.insert(*id, ptr);
-                    ptr
-                })
-            }
-            ExprKind::Attribute { value, attr, .. } => {
-                let index = self.get_attr_index(value.custom.unwrap(), *attr);
-                let val = self.gen_expr(value).unwrap();
-                let ptr = if let BasicValueEnum::PointerValue(v) = val {
-                    v
-                } else {
-                    unreachable!();
-                };
-                unsafe {
-                    self.builder.build_in_bounds_gep(
-                        ptr,
-                        &[
-                            self.ctx.i32_type().const_zero(),
-                            self.ctx.i32_type().const_int(index as u64, false),
-                        ],
-                        "attr",
-                    )
-                }
-            }
-            ExprKind::Subscript { value, slice, .. } => {
-                let i32_type = self.ctx.i32_type();
-                let v = assert_pointer_val(self.gen_expr(value).unwrap());
-                let index = assert_int_val(self.gen_expr(slice).unwrap());
-                unsafe {
-                    let ptr_to_arr = self.builder.build_in_bounds_gep(
-                        v,
-                        &[i32_type.const_zero(), i32_type.const_int(1, false)],
-                        "ptr_to_arr",
-                    );
-                    let arr_ptr =
-                        assert_pointer_val(self.builder.build_load(ptr_to_arr, "loadptr"));
-                    self.builder.build_gep(arr_ptr, &[index], "loadarrgep")
-                }
-            }
-            _ => unreachable!(),
-        }
-    }
-
-    fn gen_assignment(&mut self, target: &Expr<Option<Type>>, value: BasicValueEnum<'ctx>) {
-        let i32_type = self.ctx.i32_type();
-        if let ExprKind::Tuple { elts, .. } = &target.node {
-            if let BasicValueEnum::PointerValue(ptr) = value {
-                for (i, elt) in elts.iter().enumerate() {
-                    unsafe {
-                        let t = self.builder.build_in_bounds_gep(
-                            ptr,
-                            &[i32_type.const_zero(), i32_type.const_int(i as u64, false)],
-                            "elem",
-                        );
-                        let v = self.builder.build_load(t, "tmpload");
-                        self.gen_assignment(elt, v);
-                    }
-                }
+pub fn gen_store_target<'ctx, 'a, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    pattern: &Expr<Option<Type>>,
+) -> PointerValue<'ctx> {
+    // very similar to gen_expr, but we don't do an extra load at the end
+    // and we flatten nested tuples
+    match &pattern.node {
+        ExprKind::Name { id, .. } => ctx.var_assignment.get(id).cloned().unwrap_or_else(|| {
+            let ptr = generator.gen_var_alloc(ctx, pattern.custom.unwrap());
+            ctx.var_assignment.insert(*id, ptr);
+            ptr
+        }),
+        ExprKind::Attribute { value, attr, .. } => {
+            let index = ctx.get_attr_index(value.custom.unwrap(), *attr);
+            let val = generator.gen_expr(ctx, value).unwrap();
+            let ptr = if let BasicValueEnum::PointerValue(v) = val {
+                v
             } else {
-                unreachable!()
+                unreachable!();
+            };
+            unsafe {
+                ctx.builder.build_in_bounds_gep(
+                    ptr,
+                    &[
+                        ctx.ctx.i32_type().const_zero(),
+                        ctx.ctx.i32_type().const_int(index as u64, false),
+                    ],
+                    "attr",
+                )
             }
-        } else {
-            let ptr = self.parse_pattern(target);
-            self.builder.build_store(ptr, value);
         }
-    }
-
-    // return true if it contains terminator
-    pub fn gen_stmt(&mut self, stmt: &Stmt<Option<Type>>) -> bool {
-        match &stmt.node {
-            StmtKind::Pass => {},
-            StmtKind::Expr { value } => {
-                self.gen_expr(value);
+        ExprKind::Subscript { value, slice, .. } => {
+            let i32_type = ctx.ctx.i32_type();
+            let v = assert_pointer_val(generator.gen_expr(ctx, value).unwrap());
+            let index = assert_int_val(generator.gen_expr(ctx, slice).unwrap());
+            unsafe {
+                let ptr_to_arr = ctx.builder.build_in_bounds_gep(
+                    v,
+                    &[i32_type.const_zero(), i32_type.const_int(1, false)],
+                    "ptr_to_arr",
+                );
+                let arr_ptr = assert_pointer_val(ctx.builder.build_load(ptr_to_arr, "loadptr"));
+                ctx.builder.build_gep(arr_ptr, &[index], "loadarrgep")
             }
-            StmtKind::Return { value } => {
-                let value = value.as_ref().map(|v| self.gen_expr(v).unwrap());
-                let value = value.as_ref().map(|v| v as &dyn BasicValue);
-                self.builder.build_return(value);
-                return true;
-            }
-            StmtKind::AnnAssign { target, value, .. } => {
-                if let Some(value) = value {
-                    let value = self.gen_expr(value).unwrap();
-                    self.gen_assignment(target, value);
-                }
-            }
-            StmtKind::Assign { targets, value, .. } => {
-                let value = self.gen_expr(value).unwrap();
-                for target in targets.iter() {
-                    self.gen_assignment(target, value);
-                }
-            }
-            StmtKind::Continue => {
-                self.builder.build_unconditional_branch(self.loop_bb.unwrap().0);
-                return true;
-            }
-            StmtKind::Break => {
-                self.builder.build_unconditional_branch(self.loop_bb.unwrap().1);
-                return true;
-            }
-            StmtKind::If { 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 mut cont_bb = None; // 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 = Some(self.ctx.append_basic_block(current, "cont"));
-                    cont_bb.unwrap()
-                } else {
-                    self.ctx.append_basic_block(current, "orelse")
-                };
-                self.builder.build_unconditional_branch(test_bb);
-                self.builder.position_at_end(test_bb);
-                let test = self.gen_expr(test).unwrap();
-                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);
-                let mut exited = false;
-                for stmt in body.iter() {
-                    exited = self.gen_stmt(stmt);
-                    if exited {
-                        break;
-                    }
-                }
-                if !exited {
-                    if cont_bb.is_none() {
-                        cont_bb = Some(self.ctx.append_basic_block(current, "cont"));
-                    }
-                    self.builder.build_unconditional_branch(cont_bb.unwrap());
-                }
-                if !orelse.is_empty() {
-                    exited = false;
-                    self.builder.position_at_end(orelse_bb);
-                    for stmt in orelse.iter() {
-                        exited = self.gen_stmt(stmt);
-                        if exited {
-                            break;
-                        }
-                    }
-                    if !exited {
-                        if cont_bb.is_none() {
-                            cont_bb = Some(self.ctx.append_basic_block(current, "cont"));
-                        }
-                        self.builder.build_unconditional_branch(cont_bb.unwrap());
-                    }
-                }
-                if let Some(cont_bb) = cont_bb {
-                    self.builder.position_at_end(cont_bb);
-                }
-            }
-            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).unwrap();
-                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!("{:?}", stmt),
-        };
-        false
+        }
+        _ => unreachable!(),
     }
 }
+
+pub fn gen_assign<'ctx, 'a, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    target: &Expr<Option<Type>>,
+    value: BasicValueEnum<'ctx>,
+) {
+    let i32_type = ctx.ctx.i32_type();
+    if let ExprKind::Tuple { elts, .. } = &target.node {
+        if let BasicValueEnum::PointerValue(ptr) = value {
+            for (i, elt) in elts.iter().enumerate() {
+                unsafe {
+                    let t = ctx.builder.build_in_bounds_gep(
+                        ptr,
+                        &[i32_type.const_zero(), i32_type.const_int(i as u64, false)],
+                        "elem",
+                    );
+                    let v = ctx.builder.build_load(t, "tmpload");
+                    generator.gen_assign(ctx, elt, v);
+                }
+            }
+        } else {
+            unreachable!()
+        }
+    } else {
+        let ptr = generator.gen_store_target(ctx, target);
+        ctx.builder.build_store(ptr, value);
+    }
+}
+
+pub fn gen_while<'ctx, 'a, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    stmt: &Stmt<Option<Type>>,
+) {
+    if let StmtKind::While { test, body, orelse } = &stmt.node {
+        let current = ctx.builder.get_insert_block().unwrap().get_parent().unwrap();
+        let test_bb = ctx.ctx.append_basic_block(current, "test");
+        let body_bb = ctx.ctx.append_basic_block(current, "body");
+        let cont_bb = ctx.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 { ctx.ctx.append_basic_block(current, "orelse") };
+        // store loop bb information and restore it later
+        let loop_bb = ctx.loop_bb.replace((test_bb, cont_bb));
+        ctx.builder.build_unconditional_branch(test_bb);
+        ctx.builder.position_at_end(test_bb);
+        let test = generator.gen_expr(ctx, test).unwrap();
+        if let BasicValueEnum::IntValue(test) = test {
+            ctx.builder.build_conditional_branch(test, body_bb, orelse_bb);
+        } else {
+            unreachable!()
+        };
+        ctx.builder.position_at_end(body_bb);
+        for stmt in body.iter() {
+            generator.gen_stmt(ctx, stmt);
+        }
+        ctx.builder.build_unconditional_branch(test_bb);
+        if !orelse.is_empty() {
+            ctx.builder.position_at_end(orelse_bb);
+            for stmt in orelse.iter() {
+                generator.gen_stmt(ctx, stmt);
+            }
+            ctx.builder.build_unconditional_branch(cont_bb);
+        }
+        ctx.builder.position_at_end(cont_bb);
+        ctx.loop_bb = loop_bb;
+    } else {
+        unreachable!()
+    }
+}
+
+pub fn gen_if<'ctx, 'a, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    stmt: &Stmt<Option<Type>>,
+) -> bool {
+    if let StmtKind::If { test, body, orelse } = &stmt.node {
+        let current = ctx.builder.get_insert_block().unwrap().get_parent().unwrap();
+        let test_bb = ctx.ctx.append_basic_block(current, "test");
+        let body_bb = ctx.ctx.append_basic_block(current, "body");
+        let mut cont_bb = None;
+        // if there is no orelse, we just go to cont_bb
+        let orelse_bb = if orelse.is_empty() {
+            cont_bb = Some(ctx.ctx.append_basic_block(current, "cont"));
+            cont_bb.unwrap()
+        } else {
+            ctx.ctx.append_basic_block(current, "orelse")
+        };
+        ctx.builder.build_unconditional_branch(test_bb);
+        ctx.builder.position_at_end(test_bb);
+        let test = generator.gen_expr(ctx, test).unwrap();
+        if let BasicValueEnum::IntValue(test) = test {
+            ctx.builder.build_conditional_branch(test, body_bb, orelse_bb);
+        } else {
+            unreachable!()
+        };
+        ctx.builder.position_at_end(body_bb);
+        let mut exited = false;
+        for stmt in body.iter() {
+            exited = generator.gen_stmt(ctx, stmt);
+            if exited {
+                break;
+            }
+        }
+        if !exited {
+            if cont_bb.is_none() {
+                cont_bb = Some(ctx.ctx.append_basic_block(current, "cont"));
+            }
+            ctx.builder.build_unconditional_branch(cont_bb.unwrap());
+        }
+        let then_exited = exited;
+        let else_exited = if !orelse.is_empty() {
+            exited = false;
+            ctx.builder.position_at_end(orelse_bb);
+            for stmt in orelse.iter() {
+                exited = generator.gen_stmt(ctx, stmt);
+                if exited {
+                    break;
+                }
+            }
+            if !exited {
+                if cont_bb.is_none() {
+                    cont_bb = Some(ctx.ctx.append_basic_block(current, "cont"));
+                }
+                ctx.builder.build_unconditional_branch(cont_bb.unwrap());
+            }
+            exited
+        } else {
+            false
+        };
+        if let Some(cont_bb) = cont_bb {
+            ctx.builder.position_at_end(cont_bb);
+        }
+        then_exited && else_exited
+    } else {
+        unreachable!()
+    }
+}
+
+pub fn gen_stmt<'ctx, 'a, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    stmt: &Stmt<Option<Type>>,
+) -> bool {
+    match &stmt.node {
+        StmtKind::Pass => {}
+        StmtKind::Expr { value } => {
+            generator.gen_expr(ctx, value);
+        }
+        StmtKind::Return { value } => {
+            let value = value.as_ref().map(|v| generator.gen_expr(ctx, v).unwrap());
+            let value = value.as_ref().map(|v| v as &dyn BasicValue);
+            ctx.builder.build_return(value);
+            return true;
+        }
+        StmtKind::AnnAssign { target, value, .. } => {
+            if let Some(value) = value {
+                let value = generator.gen_expr(ctx, value).unwrap();
+                generator.gen_assign(ctx, target, value);
+            }
+        }
+        StmtKind::Assign { targets, value, .. } => {
+            let value = generator.gen_expr(ctx, value).unwrap();
+            for target in targets.iter() {
+                generator.gen_assign(ctx, target, value);
+            }
+        }
+        StmtKind::Continue => {
+            ctx.builder.build_unconditional_branch(ctx.loop_bb.unwrap().0);
+            return true;
+        }
+        StmtKind::Break => {
+            ctx.builder.build_unconditional_branch(ctx.loop_bb.unwrap().1);
+            return true;
+        }
+        StmtKind::If { .. } => return generator.gen_if(ctx, stmt),
+        StmtKind::While { .. } => return generator.gen_while(ctx, stmt),
+        _ => unimplemented!()
+    };
+    false
+}

nac3core/src/codegen/test.rs

@@ -1,5 +1,5 @@
 use crate::{
-    codegen::{CodeGenTask, WithCall, WorkerRegistry, CodeGenContext},
+    codegen::{CodeGenTask, WithCall, WorkerRegistry, CodeGenContext, DefaultCodeGenerator},
     location::Location,
     symbol_resolver::SymbolResolver,
     toplevel::{
@@ -72,7 +72,7 @@ fn test_primitives() {
         class_names: Default::default(),
     }) as Arc<dyn SymbolResolver + Send + Sync>;
 
-    let threads = ["test"];
+    let threads = vec![DefaultCodeGenerator::new("test".into()).into()];
     let signature = FunSignature {
         args: vec![
             FuncArg { name: "a".into(), ty: primitives.int32, default_value: None },
@@ -186,7 +186,7 @@ fn test_primitives() {
         .trim();
         assert_eq!(expected, module.print_to_string().to_str().unwrap().trim());
     })));
-    let (registry, handles) = WorkerRegistry::create_workers(&threads, top_level, f);
+    let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, f);
     registry.add_task(task);
     registry.wait_tasks_complete(handles);
 }
@@ -245,7 +245,7 @@ fn test_simple_call() {
         unreachable!()
     }
 
-    let threads = ["test"];
+    let threads = vec![DefaultCodeGenerator::new("test".into()).into()];
     let mut function_data = FunctionData {
         resolver: resolver.clone(),
         bound_variables: Vec::new(),
@@ -351,7 +351,7 @@ fn test_simple_call() {
         .trim();
         assert_eq!(expected, module.print_to_string().to_str().unwrap().trim());
     })));
-    let (registry, handles) = WorkerRegistry::create_workers(&threads, top_level, f);
+    let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, f);
     registry.add_task(task);
     registry.wait_tasks_complete(handles);
 }

nac3standalone/src/main.rs

@@ -1,12 +1,16 @@
-use std::env;
-use std::fs;
-use inkwell::{OptimizationLevel, passes::{PassManager, PassManagerBuilder}, targets::*};
+use inkwell::{
+    passes::{PassManager, PassManagerBuilder},
+    targets::*,
+    OptimizationLevel,
+};
 use nac3core::typecheck::type_inferencer::PrimitiveStore;
 use rustpython_parser::parser;
+use std::env;
+use std::fs;
 use std::{collections::HashMap, path::Path, sync::Arc, time::SystemTime};
 
 use nac3core::{
-    codegen::{CodeGenTask, WithCall, WorkerRegistry},
+    codegen::{CodeGenTask, DefaultCodeGenerator, WithCall, WorkerRegistry},
     symbol_resolver::SymbolResolver,
     toplevel::{composer::TopLevelComposer, TopLevelDef},
     typecheck::typedef::FunSignature,
@@ -17,7 +21,10 @@ use basic_symbol_resolver::*;
 
 fn main() {
     let demo_name = env::args().nth(1).unwrap();
-    let threads: u32 = env::args().nth(2).map(|s| str::parse(&s).unwrap()).unwrap_or(1);
+    let threads: u32 = env::args()
+        .nth(2)
+        .map(|s| str::parse(&s).unwrap())
+        .unwrap_or(1);
 
     let start = SystemTime::now();
 
@@ -29,7 +36,7 @@ fn main() {
             println!("Cannot open input file: {}", err);
             return;
         }
-   };
+    };
 
     let primitive: PrimitiveStore = TopLevelComposer::make_primitives().0;
     let (mut composer, builtins_def, builtins_ty) = TopLevelComposer::new(vec![]);
@@ -38,23 +45,27 @@ fn main() {
         id_to_type: builtins_ty.into(),
         id_to_def: builtins_def.into(),
         class_names: Default::default(),
-    }.into();
-    let resolver = Arc::new(
-        Resolver(internal_resolver.clone())
-    ) as Arc<dyn SymbolResolver + Send + Sync>;
+    }
+    .into();
+    let resolver =
+        Arc::new(Resolver(internal_resolver.clone())) as Arc<dyn SymbolResolver + Send + Sync>;
     let setup_time = SystemTime::now();
-    println!("setup time: {}ms", setup_time.duration_since(start).unwrap().as_millis());
+    println!(
+        "setup time: {}ms",
+        setup_time.duration_since(start).unwrap().as_millis()
+    );
 
     let parser_result = parser::parse_program(&program).unwrap();
     let parse_time = SystemTime::now();
-    println!("parse time: {}ms", parse_time.duration_since(setup_time).unwrap().as_millis());
+    println!(
+        "parse time: {}ms",
+        parse_time.duration_since(setup_time).unwrap().as_millis()
+    );
 
     for stmt in parser_result.into_iter() {
-        let (name, def_id, ty) = composer.register_top_level(
-            stmt,
-            Some(resolver.clone()),
-            "__main__".into(),
-        ).unwrap();
+        let (name, def_id, ty) = composer
+            .register_top_level(stmt, Some(resolver.clone()), "__main__".into())
+            .unwrap();
 
         internal_resolver.add_id_def(name, def_id);
         if let Some(ty) = ty {
@@ -64,7 +75,13 @@ fn main() {
 
     composer.start_analysis(true).unwrap();
     let analysis_time = SystemTime::now();
-    println!("analysis time: {}ms", analysis_time.duration_since(parse_time).unwrap().as_millis());
+    println!(
+        "analysis time: {}ms",
+        analysis_time
+            .duration_since(parse_time)
+            .unwrap()
+            .as_millis()
+    );
 
     let top_level = Arc::new(composer.make_top_level_context());
 
@@ -119,19 +136,32 @@ fn main() {
             )
             .expect("couldn't create target machine");
         target_machine
-            .write_to_file(module, FileType::Object, Path::new(&format!("{}.o", module.get_name().to_str().unwrap())))
+            .write_to_file(
+                module,
+                FileType::Object,
+                Path::new(&format!("{}.o", module.get_name().to_str().unwrap())),
+            )
             .expect("couldn't write module to file");
 
         // println!("IR:\n{}", module.print_to_string().to_str().unwrap());
-
     })));
-    let threads: Vec<String> = (0..threads).map(|i| format!("module{}", i)).collect();
-    let threads: Vec<_> = threads.iter().map(|s| s.as_str()).collect();
-    let (registry, handles) = WorkerRegistry::create_workers(&threads, top_level, f);
+    let threads = (0..threads)
+        .map(|i| Box::new(DefaultCodeGenerator::new(format!("module{}", i))))
+        .collect();
+    let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, f);
     registry.add_task(task);
     registry.wait_tasks_complete(handles);
 
     let final_time = SystemTime::now();
-    println!("codegen time (including LLVM): {}ms", final_time.duration_since(analysis_time).unwrap().as_millis());
-    println!("total time: {}ms", final_time.duration_since(start).unwrap().as_millis());
+    println!(
+        "codegen time (including LLVM): {}ms",
+        final_time
+            .duration_since(analysis_time)
+            .unwrap()
+            .as_millis()
+    );
+    println!(
+        "total time: {}ms",
+        final_time.duration_since(start).unwrap().as_millis()
+    );
 }