use nac3core::{ codegen::{expr::gen_call, stmt::gen_with, CodeGenContext, CodeGenerator}, toplevel::DefinitionId, typecheck::typedef::{FunSignature, Type}, }; use rustpython_parser::ast::{Expr, ExprKind, Located, Stmt, StmtKind, StrRef}; use inkwell::values::BasicValueEnum; use crate::timeline::TimeFns; pub struct ArtiqCodeGenerator<'a> { name: String, name_counter: u32, start: Option>>, end: Option>>, timeline: &'a (dyn TimeFns + Sync), } impl<'a> ArtiqCodeGenerator<'a> { pub fn new(name: String, timeline: &'a (dyn TimeFns + Sync)) -> ArtiqCodeGenerator<'a> { ArtiqCodeGenerator { name, name_counter: 0, start: None, end: None, timeline, } } } impl<'b> CodeGenerator for ArtiqCodeGenerator<'b> { fn get_name(&self) -> &str { &self.name } fn gen_call<'ctx, 'a>( &mut self, ctx: &mut CodeGenContext<'ctx, 'a>, obj: Option<(Type, BasicValueEnum<'ctx>)>, fun: (&FunSignature, DefinitionId), params: Vec<(Option, BasicValueEnum<'ctx>)>, ) -> Option> { let result = gen_call(self, ctx, obj, fun, params); if let Some(end) = self.end.clone() { let old_end = self.gen_expr(ctx, &end).unwrap(); let now = self.timeline.emit_now_mu(ctx); let smax = ctx.module.get_function("llvm.smax.i64").unwrap_or_else(|| { let i64 = ctx.ctx.i64_type(); ctx.module.add_function( "llvm.smax.i64", i64.fn_type(&[i64.into(), i64.into()], false), None, ) }); let max = ctx .builder .build_call(smax, &[old_end, now], "smax") .try_as_basic_value() .left() .unwrap(); let end_store = self.gen_store_target(ctx, &end); ctx.builder.build_store(end_store, max); } if let Some(start) = self.start.clone() { let start_val = self.gen_expr(ctx, &start).unwrap(); self.timeline.emit_at_mu(ctx, start_val); } result } fn gen_with<'ctx, 'a>( &mut self, ctx: &mut CodeGenContext<'ctx, 'a>, stmt: &Stmt>, ) -> bool { if let StmtKind::With { items, body, .. } = &stmt.node { if items.len() == 1 && items[0].optional_vars.is_none() { let item = &items[0]; // Behavior of parallel and sequential: // Each function call (indirectly, can be inside a sequential block) within a parallel // block will update the end variable to the maximum now_mu in the block. // Each function call directly inside a parallel block will reset the timeline after // execution. A parallel block within a sequential block (or not within any block) will // set the timeline to the max now_mu within the block (and the outer max now_mu will also // be updated). // // Implementation: We track the start and end separately. // - If there is a start variable, it indicates that we are directly inside a // parallel block and we have to reset the timeline after every function call. // - If there is a end variable, it indicates that we are (indirectly) inside a // parallel block, and we should update the max end value. if let ExprKind::Name { id, ctx: name_ctx } = &item.context_expr.node { if id == &"parallel".into() { let old_start = self.start.take(); let old_end = self.end.take(); let now = if let Some(old_start) = &old_start { self.gen_expr(ctx, old_start).unwrap() } else { self.timeline.emit_now_mu(ctx) }; // Emulate variable allocation, as we need to use the CodeGenContext // HashMap to store our variable due to lifetime limitation // Note: we should be able to store variables directly if generic // associative type is used by limiting the lifetime of CodeGenerator to // the LLVM Context. // The name is guaranteed to be unique as users cannot use this as variable // name. self.start = old_start.clone().or_else(|| { let start = format!("with-{}-start", self.name_counter).into(); let start_expr = Located { // location does not matter at this point location: stmt.location, node: ExprKind::Name { id: start, ctx: name_ctx.clone(), }, custom: Some(ctx.primitives.int64), }; let start = self.gen_store_target(ctx, &start_expr); ctx.builder.build_store(start, now); Some(start_expr) }); let end = format!("with-{}-end", self.name_counter).into(); let end_expr = Located { // location does not matter at this point location: stmt.location, node: ExprKind::Name { id: end, ctx: name_ctx.clone(), }, custom: Some(ctx.primitives.int64), }; let end = self.gen_store_target(ctx, &end_expr); ctx.builder.build_store(end, now); self.end = Some(end_expr); self.name_counter += 1; let mut exited = false; for stmt in body.iter() { if self.gen_stmt(ctx, stmt) { exited = true; break; } } // set duration let end_expr = self.end.take().unwrap(); let end_val = self.gen_expr(ctx, &end_expr).unwrap(); // inside an sequential block if old_start.is_none() { self.timeline.emit_at_mu(ctx, end_val); } // inside a parallel block, should update the outer max now_mu if let Some(old_end) = &old_end { let outer_end_val = self.gen_expr(ctx, old_end).unwrap(); let smax = ctx.module.get_function("llvm.smax.i64").unwrap_or_else(|| { let i64 = ctx.ctx.i64_type(); ctx.module.add_function( "llvm.smax.i64", i64.fn_type(&[i64.into(), i64.into()], false), None, ) }); let max = ctx .builder .build_call(smax, &[end_val, outer_end_val], "smax") .try_as_basic_value() .left() .unwrap(); let outer_end = self.gen_store_target(ctx, old_end); ctx.builder.build_store(outer_end, max); } self.start = old_start; self.end = old_end; return exited; } else if id == &"sequential".into() { let start = self.start.take(); for stmt in body.iter() { if self.gen_stmt(ctx, stmt) { self.start = start; return true; } } self.start = start; return false; } } } // not parallel/sequential gen_with(self, ctx, stmt) } else { unreachable!() } } }