WIP - [artiq] Consume each dim one-by-one

nac3artiq/src/codegen.rs

@@ -600,16 +600,46 @@ fn format_rpc_ret<'ctx>(
                 unreachable!();
             };
             llvm_ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false));
-            llvm_ndarray.create_dim_sizes(ctx, llvm_usize, llvm_usize.const_int(ndims, false));
+            llvm_ndarray.create_dim_sizes(ctx, llvm_usize, llvm_ndarray.load_ndims(ctx));
 
-            let dims_buf_sz =
-                ctx.builder.build_int_mul(llvm_ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap();
+            let buffer = ctx.builder.build_array_alloca(llvm_i8, llvm_usize_sizeof, "rpc.buffer").unwrap();
+            let buffer = ArraySliceValue::from_ptr_val(buffer, llvm_usize_sizeof, Some("rpc.buffer"));
 
-            let buffer_size =
-                ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap();
+            gen_for_callback_incrementing(
+                generator,
+                ctx,
+                None,
+                llvm_usize.const_zero(),
+                (llvm_ndarray.load_ndims(ctx), false),
+                |generator, ctx, _, i| {
+                    ctx
+                        .build_call_or_invoke(rpc_recv, &[buffer.base_ptr(ctx, generator).into()], "rpc.size.next")
+                        .map(BasicValueEnum::into_int_value)
+                        .unwrap();
 
-            let buffer = ctx.builder.build_array_alloca(llvm_i8, buffer_size, "rpc.buffer").unwrap();
-            let buffer = ArraySliceValue::from_ptr_val(buffer, buffer_size, Some("rpc.buffer"));
+                    let offset = ctx.builder.build_int_add(
+                        llvm_pdata_sizeof,
+                        ctx.builder.build_int_mul(i, llvm_usize_sizeof, "").unwrap(),
+                        "",
+                    ).unwrap();
+                    let dims_ptr = unsafe { llvm_ndarray.dim_sizes().ptr_offset_unchecked(ctx, generator, &offset, None) };
+                    let pbuffer_dims_ptr =
+                        unsafe { buffer.ptr_offset_unchecked(ctx, generator, &offset, None) };
+                    call_memcpy_generic(
+                        ctx,
+                        dims_ptr,
+                        pbuffer_dims_ptr,
+                        llvm_usize_sizeof,
+                        llvm_i1.const_zero(),
+                    ); 
+
+                    Ok(())
+                },
+                llvm_usize.const_int(1, false),
+            ).unwrap();
+
+            let buffer = ctx.builder.build_array_alloca(llvm_i8, llvm_pdata_sizeof, "rpc.buffer").unwrap();
+            let buffer = ArraySliceValue::from_ptr_val(buffer, llvm_pdata_sizeof, Some("rpc.buffer"));
 
             // recv [*data, dim_sz[..]]
             let alloc_size = ctx
@@ -627,16 +657,6 @@ fn format_rpc_ret<'ctx>(
                 llvm_i1.const_zero(),
             );
 
-            let pbuffer_dims_begin =
-                unsafe { buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None) };
-            call_memcpy_generic(
-                ctx,
-                llvm_ndarray.dim_sizes().base_ptr(ctx, generator),
-                pbuffer_dims_begin,
-                dims_buf_sz,
-                llvm_i1.const_zero(),
-            );
-
             let is_done = ctx
                 .builder
                 .build_int_compare(IntPredicate::EQ, llvm_i32.const_zero(), alloc_size, "rpc.done")
@@ -644,12 +664,13 @@ fn format_rpc_ret<'ctx>(
 
             let alloc_ptr = ctx.builder.build_array_alloca(llvm_pi8, alloc_size, "rpc.alloc").unwrap();
             let alloc_ptr = ctx.builder.build_bitcast(alloc_ptr, llvm_pi8, "rpc.alloc.ptr").map(BasicValueEnum::into_pointer_value).unwrap();
+            let prehead_metadata_bb = ctx.builder.get_insert_block().unwrap();
 
             ctx.builder.build_conditional_branch(is_done, tail_bb, head_bb).unwrap();
 
             ctx.builder.position_at_end(head_bb);
             let phi = ctx.builder.build_phi(llvm_pi8, "rpc.ptr").unwrap();
-            phi.add_incoming(&[(&alloc_ptr, prehead_bb)]);
+            phi.add_incoming(&[(&alloc_ptr, prehead_metadata_bb)]);
             let alloc_size = ctx
                 .build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next")
                 .unwrap()