refactor(moonbit): derive list borrows from lowering

Author: peter-jerry-ye

crates/moonbit/src/lib.rs

@@ -49,64 +49,6 @@ pub(crate) fn is_list_canonical(resolve: &Resolve, element: &Type) -> bool {
     }
 }
 
-fn collect_direct_canonical_lists(
-    resolve: &Resolve,
-    ty: &Type,
-    expression: String,
-    wasm_param: usize,
-    expressions: &mut HashSet<String>,
-    wasm_params: &mut BTreeMap<usize, Type>,
-) {
-    let Type::Id(id) = ty else {
-        return;
-    };
-    match &resolve.types[*id].kind {
-        TypeDefKind::Type(ty) => collect_direct_canonical_lists(
-            resolve,
-            ty,
-            expression,
-            wasm_param,
-            expressions,
-            wasm_params,
-        ),
-        TypeDefKind::List(element) => {
-            if is_list_canonical(resolve, element) {
-                expressions.insert(expression);
-                wasm_params.insert(wasm_param, *element);
-            }
-        }
-        TypeDefKind::Record(record) => {
-            let mut field_wasm_param = wasm_param;
-            for field in &record.fields {
-                collect_direct_canonical_lists(
-                    resolve,
-                    &field.ty,
-                    format!("({expression}).{}", field.name.to_moonbit_ident()),
-                    field_wasm_param,
-                    expressions,
-                    wasm_params,
-                );
-                field_wasm_param += abi::flat_types(resolve, &field.ty, None).unwrap().len();
-            }
-        }
-        TypeDefKind::Tuple(tuple) => {
-            let mut field_wasm_param = wasm_param;
-            for (index, field) in tuple.types.iter().enumerate() {
-                collect_direct_canonical_lists(
-                    resolve,
-                    field,
-                    format!("({expression}).{index}"),
-                    field_wasm_param,
-                    expressions,
-                    wasm_params,
-                );
-                field_wasm_param += abi::flat_types(resolve, field, None).unwrap().len();
-            }
-        }
-        _ => {}
-    }
-}
-
 // Assumptions:
 // - Data: u8 -> Byte, s8 | s32 -> Int, s16 -> Int16, u16 -> UInt16, u32 -> UInt, s64 -> Int64, u64 -> UInt64, f32 -> Float, f64 -> Double, address -> Int
 // - Encoding: UTF16
@@ -685,25 +627,16 @@ impl InterfaceGenerator<'_> {
         let endpoint_plan = self.import_async_function_plan(self.interface, func);
         let wasm_sig = self.resolve.wasm_signature(variant, func);
         let mbt_sig = self.world_gen.pkg_resolver.mbt_sig(self.name, func, false);
-        let mut direct_canonical_lists = HashSet::new();
-        let mut direct_canonical_wasm_params = BTreeMap::new();
-        if !async_plan.is_async() && !wasm_sig.indirect_params {
-            let mut wasm_param = 0;
-            for Param { name, ty, .. } in &func.params {
-                collect_direct_canonical_lists(
-                    self.resolve,
-                    ty,
-                    name.to_moonbit_ident(),
-                    wasm_param,
-                    &mut direct_canonical_lists,
-                    &mut direct_canonical_wasm_params,
-                );
-                wasm_param += abi::flat_types(self.resolve, ty, None).unwrap().len();
-            }
-        }
-        let (src, needs_cleanup_list, endpoint_state) = if async_plan.is_async() {
+        let (src, needs_cleanup_list, endpoint_state, direct_canonical_wasm_params) = if async_plan
+            .is_async()
+        {
             let body = self.generate_async_import_body(&endpoint_plan, func, &mbt_sig, &wasm_sig);
-            (body.src, body.needs_cleanup_list, body.state)
+            (
+                body.src,
+                body.needs_cleanup_list,
+                body.state,
+                BTreeMap::new(),
+            )
         } else {
             let mut bindgen = FunctionBindgen::new(
                 self,
@@ -712,13 +645,7 @@ impl InterfaceGenerator<'_> {
                     .map(|Param { name, .. }| name.to_moonbit_ident())
                     .collect(),
             )
-            .with_direct_canonical_list_params(
-                direct_canonical_lists,
-                direct_canonical_wasm_params
-                    .keys()
-                    .map(|param| param + 1)
-                    .collect(),
-            )
+            .with_direct_list_borrows(!wasm_sig.indirect_params)
             .with_async_state(endpoint_plan.state());
             if endpoint_plan.has_endpoints() {
                 bindgen = bindgen.with_sync_import_commit(
@@ -734,7 +661,12 @@ impl InterfaceGenerator<'_> {
                 &mut bindgen,
                 false,
             );
-            (bindgen.src, bindgen.needs_cleanup_list, bindgen.async_state)
+            (
+                bindgen.src,
+                bindgen.needs_cleanup_list,
+                bindgen.async_state,
+                bindgen.direct_canonical_wasm_params,
+            )
         };
 
         let cleanup_list = if needs_cleanup_list {
@@ -763,8 +695,6 @@ impl InterfaceGenerator<'_> {
                 if let Some(element) = direct_canonical_wasm_params.get(&i) {
                     let element = self.world_gen.pkg_resolver.type_name(self.name, element);
                     format!("p{i} : FixedArray[{element}]")
-                } else if i > 0 && direct_canonical_wasm_params.contains_key(&(i - 1)) {
-                    format!("p{i}? : Int = p{}.length()", i - 1)
                 } else {
                     format!("p{i} : {}", wasm_type(*param))
                 }
@@ -1537,8 +1467,9 @@ struct FunctionBindgen<'a, 'b> {
     sync_endpoint_drop: bool,
     commit_endpoints: bool,
     sync_import_argument_types: Option<Vec<Type>>,
-    direct_canonical_list_params: HashSet<String>,
-    direct_canonical_list_length_params: HashSet<usize>,
+    direct_list_borrows: bool,
+    direct_list_borrow_candidates: HashMap<String, (String, Type)>,
+    direct_canonical_wasm_params: BTreeMap<usize, Type>,
     async_state: AsyncFunctionState,
 }
 
@@ -1569,19 +1500,15 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             sync_endpoint_drop: false,
             commit_endpoints: false,
             sync_import_argument_types: None,
-            direct_canonical_list_params: HashSet::new(),
-            direct_canonical_list_length_params: HashSet::new(),
+            direct_list_borrows: false,
+            direct_list_borrow_candidates: HashMap::new(),
+            direct_canonical_wasm_params: BTreeMap::new(),
             async_state: AsyncFunctionState::default(),
         }
     }
 
-    fn with_direct_canonical_list_params(
-        mut self,
-        params: HashSet<String>,
-        length_params: HashSet<usize>,
-    ) -> Self {
-        self.direct_canonical_list_params = params;
-        self.direct_canonical_list_length_params = length_params;
+    fn with_direct_list_borrows(mut self, enabled: bool) -> Self {
+        self.direct_list_borrows = enabled;
         self
     }
 
@@ -2135,22 +2062,29 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             )),
 
             Instruction::ListCanonLower { element, realloc } => {
-                let element: &Type = element;
-                let element = match element {
+                let original_element: &Type = element;
+                let element = match original_element {
                     Type::Id(id) => match &resolve.types[dealias(resolve, *id)].kind {
                         TypeDefKind::Type(element) => element,
                         _ => unreachable!("unsupported list element type"),
                     },
-                    _ => element,
+                    _ => original_element,
                 };
+                let op = &operands[0];
+                // A canonical list can be passed as a typed borrow only when its
+                // pointer and length flow directly to the Wasm call. Lists lowered
+                // inside blocks are merged into variants or written into enclosing
+                // list storage, so those still need their integer pointer.
+                if self.direct_list_borrows && realloc.is_none() && self.block_storage.is_empty() {
+                    let length = format!("{op}.length()");
+                    self.direct_list_borrow_candidates
+                        .insert(op.clone(), (length.clone(), *original_element));
+                    results.push(op.clone());
+                    results.push(length);
+                    return;
+                }
                 match element {
                     Type::U8 => {
-                        let op = &operands[0];
-                        if realloc.is_none() && self.direct_canonical_list_params.contains(op) {
-                            results.push(op.clone());
-                            results.push(format!("{op}.length()"));
-                            return;
-                        }
                         let ptr = self.locals.tmp("ptr");
                         self.use_ffi(ffi::BYTES2PTR);
                         uwriteln!(
@@ -2174,12 +2108,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                     | Type::S64
                     | Type::F32
                     | Type::F64 => {
-                        let op = &operands[0];
-                        if realloc.is_none() && self.direct_canonical_list_params.contains(op) {
-                            results.push(op.clone());
-                            results.push(format!("{op}.length()"));
-                            return;
-                        }
                         let ptr = self.locals.tmp("ptr");
                         let (owned_ffi, ty) = match element {
                             Type::Bool => (ffi::BOOL_ARRAY2PTR, "bool"),
@@ -2441,6 +2369,13 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 };
 
                 let func_name = name.to_upper_camel_case();
+                for (index, operand) in operands.iter().enumerate() {
+                    if let Some((length, element)) = self.direct_list_borrow_candidates.get(operand)
+                    {
+                        assert_eq!(operands.get(index + 1), Some(length));
+                        self.direct_canonical_wasm_params.insert(index, *element);
+                    }
+                }
                 let call_operands = if self.sync_import_argument_types.is_some() {
                     operands
                         .iter()
@@ -2453,15 +2388,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 } else {
                     operands.clone()
                 };
-                let arguments = call_operands
-                    .iter()
-                    .enumerate()
-                    .filter_map(|(i, operand)| {
-                        (!self.direct_canonical_list_length_params.contains(&i))
-                            .then_some(operand.as_str())
-                    })
-                    .collect::<Vec<_>>()
-                    .join(", ");
+                let arguments = call_operands.join(", ");
                 // TODO: handle this to support async functions
                 uwriteln!(self.src, "{assignment} wasmImport{func_name}({arguments});");
                 self.commit_sync_import_arguments(sig, &call_operands);
@@ -3745,9 +3672,11 @@ mod tests {
 
         assert!(
             top.contains(
-                "wasmImportSend(bytes, unsigned_shorts, signed_shorts, words, (envelope).words, booleans"
+                "wasmImportSend(bytes, bytes.length(), unsigned_shorts, unsigned_shorts.length(), \
+                 signed_shorts, signed_shorts.length(), words, words.length(), \
+                 (envelope).words, (envelope).words.length(), booleans, booleans.length())"
             ),
-            "direct canonical lists must let the FFI supply their lengths: {top}"
+            "direct canonical lists must retain their canonical pointer/length operands: {top}"
         );
         assert!(
             !top.contains("mbt_ffi_borrowed_array2ptr(bytes)")
@@ -3756,27 +3685,18 @@ mod tests {
                 && !top.contains("mbt_ffi_borrowed_array2ptr(words)"),
             "direct list parameters must not round-trip through Int: {top}"
         );
-        assert!(
-            !top.contains("bytes.length()")
-                && !top.contains("unsigned_shorts.length()")
-                && !top.contains("signed_shorts.length()")
-                && !top.contains("words.length()")
-                && !top.contains("(envelope).words.length()")
-                && !top.contains("booleans.length()"),
-            "direct canonical list lengths must be supplied by FFI defaults: {top}"
-        );
         assert!(!top.contains("mbt_ffi_borrowed_array2ptr"), "{top}");
         assert!(
             ffi.contains(
                 "#unsafe_skip_stub_check\n#borrow(p0, p2, p4, p6, p8, p10)\nfn wasmImportSend(\
-                 p0 : FixedArray[Byte], p1? : Int = p0.length(), \
-                 p2 : FixedArray[UInt16], p3? : Int = p2.length(), \
-                 p4 : FixedArray[Int16], p5? : Int = p4.length(), \
-                 p6 : FixedArray[UInt], p7? : Int = p6.length(), \
-                 p8 : FixedArray[UInt], p9? : Int = p8.length(), \
-                 p10 : FixedArray[Bool], p11? : Int = p10.length())"
+                 p0 : FixedArray[Byte], p1 : Int, \
+                 p2 : FixedArray[UInt16], p3 : Int, \
+                 p4 : FixedArray[Int16], p5 : Int, \
+                 p6 : FixedArray[UInt], p7 : Int, \
+                 p8 : FixedArray[UInt], p9 : Int, \
+                 p10 : FixedArray[Bool], p11 : Int)"
             ),
-            "the imported FFI must derive borrowed array lengths by default: {ffi}"
+            "the imported FFI must borrow arrays while retaining explicit lengths: {ffi}"
         );
         assert!(
             !ffi.contains("mbt_ffi_borrowed_array2ptr"),
@@ -3824,6 +3744,74 @@ mod tests {
         );
     }
 
+    #[test]
+    fn imported_conditional_canonical_lists_retain_pointer_lowering() {
+        let files = generate(
+            r#"
+            package a:b;
+
+            interface api {
+                variant optional-bytes {
+                    none,
+                    some(list<u8>),
+                }
+                send: func(bytes: optional-bytes);
+            }
+
+            world client { import api; }
+            "#,
+            "client",
+        );
+        let top = file(&files, "interface/a/b/api/top.mbt");
+        let ffi = file(&files, "interface/a/b/api/ffi.mbt");
+
+        assert!(
+            top.contains("mbt_ffi_bytes2ptr"),
+            "conditional list pointers must be lowered before merging variant cases: {top}"
+        );
+        assert!(
+            !ffi.contains("#borrow"),
+            "a conditionally occupied pointer parameter cannot be a typed borrow: {ffi}"
+        );
+    }
+
+    #[test]
+    fn imported_indirect_canonical_lists_retain_pointer_lowering() {
+        let files = generate(
+            r#"
+            package a:b;
+
+            interface api {
+                send: func(
+                    a: list<u8>,
+                    b: list<u8>,
+                    c: list<u8>,
+                    d: list<u8>,
+                    e: list<u8>,
+                    f: list<u8>,
+                    g: list<u8>,
+                    h: list<u8>,
+                    i: list<u8>,
+                );
+            }
+
+            world client { import api; }
+            "#,
+            "client",
+        );
+        let top = file(&files, "interface/a/b/api/top.mbt");
+        let ffi = file(&files, "interface/a/b/api/ffi.mbt");
+
+        assert!(
+            top.contains("mbt_ffi_bytes2ptr(a)"),
+            "lists stored in an indirect argument area still need integer pointers: {top}"
+        );
+        assert!(
+            ffi.contains("fn wasmImportSend(p0 : Int)") && !ffi.contains("#borrow"),
+            "an indirect import receives only its argument-area pointer: {ffi}"
+        );
+    }
+
     #[test]
     fn async_export_background_group_name_is_deconflicted() {
         let files = generate(