From ba4e842d31e5d9eea3c8d5bd11b052f25ab388c3 Mon Sep 17 00:00:00 2001
From: David Mak <csdavidmak@ust.hk>
Date: Mon, 23 Oct 2023 15:58:13 +0800
Subject: [PATCH] artiq: Fix timeline tracking for parallel blocks

This commit fixes timeline tracking by updating the maximum mu and
restoring the start mu on every top-level statement within the parallel
block as opposed to on each function call.

Fixes #298.
---
 nac3artiq/src/codegen.rs | 85 +++++++++++++++++++++++-----------------
 1 file changed, 50 insertions(+), 35 deletions(-)

diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs
index ccf11eb..0d915b9 100644
--- a/nac3artiq/src/codegen.rs
+++ b/nac3artiq/src/codegen.rs
@@ -1,7 +1,6 @@
 use nac3core::{
     codegen::{
-        expr::gen_call,
-        stmt::gen_with,
+        stmt::{gen_block, gen_with},
         CodeGenContext, CodeGenerator,
     },
     symbol_resolver::ValueEnum,
@@ -67,39 +66,52 @@ impl<'b> CodeGenerator for ArtiqCodeGenerator<'b> {
         }
     }
 
-    fn gen_call<'ctx, 'a>(
-        &mut self,
-        ctx: &mut CodeGenContext<'ctx, 'a>,
-        obj: Option<(Type, ValueEnum<'ctx>)>,
-        fun: (&FunSignature, DefinitionId),
-        params: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
-    ) -> Result<Option<BasicValueEnum<'ctx>>, String> {
-        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().to_basic_value_enum(ctx, self, end.custom.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.into(), now.into()], "smax")
-                .try_as_basic_value()
-                .left()
-                .unwrap();
-            let end_store = self.gen_store_target(ctx, &end, Some("end_store.addr"))?;
-            ctx.builder.build_store(end_store, max);
-        }
+    fn gen_block<'ctx, 'a, 'c, I: Iterator<Item=&'c Stmt<Option<Type>>>>(&mut self, ctx: &mut CodeGenContext<'ctx, 'a>, stmts: I) -> Result<(), String> where Self: Sized {
         if let Some(start) = self.start.clone() {
-            let start_val = self.gen_expr(ctx, &start)?.unwrap().to_basic_value_enum(ctx, self, start.custom.unwrap())?;
-            self.timeline.emit_at_mu(ctx, start_val);
+            // This block is a direct child of a `with parallel` block.
+            //
+            // We generate each statement as-is, update the maximum now_mu, then reset now_mu back
+            // to the original value.
+            for stmt in stmts {
+                self.gen_stmt(ctx, stmt)?;
+
+                if ctx.is_terminated() {
+                    break;
+                }
+
+                // gen_stmt should not change that this block is the body of a `with parallel` block
+                assert_eq!(self.start.is_some(), true);
+                // The `end` value should always be defined within a `with parallel` context
+                assert_eq!(self.end.is_some(), true);
+
+                let Some(end) = self.end.clone() else { unreachable!() };
+                let old_end = self.gen_expr(ctx, &end)?.unwrap().to_basic_value_enum(ctx, self, end.custom.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.into(), now.into()], "smax")
+                    .try_as_basic_value()
+                    .left()
+                    .unwrap();
+                let end_store = self.gen_store_target(ctx, &end, Some("end_store.addr"))?;
+                ctx.builder.build_store(end_store, max);
+
+                let start_val = self.gen_expr(ctx, &start)?.unwrap().to_basic_value_enum(ctx, self, start.custom.unwrap())?;
+                self.timeline.emit_at_mu(ctx, start_val);
+            }
+
+            Ok(())
+        } else {
+            gen_block(self, ctx, stmts)
         }
-        Ok(result)
     }
 
     fn gen_with<'ctx, 'a>(
@@ -196,12 +208,14 @@ impl<'b> CodeGenerator for ArtiqCodeGenerator<'b> {
                             .unwrap()
                             .to_basic_value_enum(ctx, self, end_expr.custom.unwrap())?;
 
-                        // inside a sequential block
+                        // Our parallel block is a direct descendant of a non-parallel block - Mark
+                        // the end mu of this entire parallel 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
+                        // Our parallel block is an indirect descendant of a parallel block -
+                        // Propagate the end mu of this parallel block to our parent.
                         if let Some(old_end) = &old_end {
                             let outer_end_val = self
                                 .gen_expr(ctx, old_end)?
@@ -226,6 +240,7 @@ impl<'b> CodeGenerator for ArtiqCodeGenerator<'b> {
                             ctx.builder.build_store(outer_end, max);
                         }
 
+                        // Restore the start/end mu of our parent with block (if any)
                         self.start = old_start;
                         self.end = old_end;