From cf50278dad0cbeb95fe744844ee78e1d71038b1f Mon Sep 17 00:00:00 2001
From: mwojcik <mw@m-labs.hk>
Date: Fri, 1 Nov 2024 17:14:23 +0800
Subject: [PATCH] await sk message before receiving it

---
 nac3artiq/src/subkernels.rs | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/nac3artiq/src/subkernels.rs b/nac3artiq/src/subkernels.rs
index f2c777b4..25f2fe6b 100644
--- a/nac3artiq/src/subkernels.rs
+++ b/nac3artiq/src/subkernels.rs
@@ -188,7 +188,7 @@ impl Subkernels {
     
         let arg_length = args.len() + usize::from(obj.is_some());
     
-        let stackptr = call_stacksave(ctx, Some("rpc.stack"));
+        let stackptr = call_stacksave(ctx, Some("subkernel.stack"));
         let args_ptr = ctx
             .builder
             .build_array_alloca(
@@ -237,7 +237,7 @@ impl Subkernels {
                 ctx.builder.build_gep(
                     args_ptr,
                     &[int32.const_int(i as u64, false)],
-                    &format!("rpc.arg{i}"),
+                    &format!("subkernel.arg{i}"),
                 )
             }
             .unwrap();
@@ -260,7 +260,7 @@ impl Subkernels {
             )
         });
         ctx.builder
-            .build_call_or_invoke(subkernel_send, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "rpc.send")
+            .build_call_or_invoke(subkernel_send, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "subkernel.send")
             .unwrap();
         // reclaim stack space used by arguments
         call_stackrestore(ctx, stackptr);
@@ -279,6 +279,31 @@ impl Subkernels {
         args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
         generator: &mut dyn CodeGenerator,
     ) -> Result<Option<BasicValueEnum<'ctx>>, String>  {
+
+        // await for the message first
+        let int8 = ctx.ctx.i8_type();
+        let int64 = ctx.ctx.i64_type();
+        let size_type = generator.get_size_type(ctx.ctx);
+        let ptr_type = int8.ptr_type(AddressSpace::default());
+        let tag_ptr_type = ctx.ctx.struct_type(&[ptr_type.into(), size_type.into()], false);
+    
+        let subkernel_await_message = ctx.module.get_function("subkernel_await_message").unwrap_or_else(|| {
+            ctx.module.add_function(
+                "subkernel_await_message",
+                ctx.ctx.void_type().fn_type(
+                    &[
+                        ctx.ctx.void_type().fn_type(&[int8.into(), int64.into(), tag_ptr_type.into(), int8.into(), int8.into()], false),
+                    ],
+                    false,
+                ),
+                None,
+            )
+        });
+
+        ctx.builder
+            .build_call_or_invoke(subkernel_await_message, &[subkernel_id.into(), destination.into(), run_type.const_int(1, false)], "subkernel.run")
+            .unwrap();
+        
         // -- receive value:
         // T result = {
         //   void *ret_ptr = alloca(sizeof(T));