From 126d7b6825e04c0342067b1ce7dc84d7cc4a14a0 Mon Sep 17 00:00:00 2001
From: Pat Hickey <pat@moreproductive.org>
Date: Wed, 26 Feb 2020 12:37:54 -0800
Subject: [PATCH] argument marshaling for unions

---
 crates/generate/src/funcs.rs | 43 ++++++++++++++++++------------------
 1 file changed, 21 insertions(+), 22 deletions(-)

diff --git a/crates/generate/src/funcs.rs b/crates/generate/src/funcs.rs
index 8814df19ab..50344f876e 100644
--- a/crates/generate/src/funcs.rs
+++ b/crates/generate/src/funcs.rs
@@ -153,6 +153,25 @@ fn marshal_arg(
         }
     };
 
+    let read_conversion = {
+        let pointee_type = names.type_ref(tref, anon_lifetime());
+        let arg_name = names.func_ptr_binding(&param.name);
+        let name = names.func_param(&param.name);
+        quote! {
+            let #name = match memory.ptr::<#pointee_type>(#arg_name as u32) {
+                Ok(p) => match p.read() {
+                    Ok(r) => r,
+                    Err(e) => {
+                        #error_handling
+                    }
+                },
+                Err(e) => {
+                    #error_handling
+                }
+            };
+        }
+    };
+
     match &*tref.type_() {
         witx::Type::Enum(_e) => try_into_conversion,
         witx::Type::Flags(_f) => try_into_conversion,
@@ -257,24 +276,7 @@ fn marshal_arg(
                 };
             }
         }
-        witx::Type::Struct(s) if !struct_is_copy(&s) => {
-            let pointee_type = names.type_ref(tref, anon_lifetime());
-            let arg_name = names.func_ptr_binding(&param.name);
-            let name = names.func_param(&param.name);
-            quote! {
-                let #name = match memory.ptr::<#pointee_type>(#arg_name as u32) {
-                    Ok(p) => match p.read() {
-                        Ok(r) => r,
-                        Err(e) => {
-                            #error_handling
-                        }
-                    },
-                    Err(e) => {
-                        #error_handling
-                    }
-                };
-            }
-        }
+        witx::Type::Struct(s) if !struct_is_copy(&s) => read_conversion,
         witx::Type::Array(arr) => {
             let pointee_type = names.type_ref(arr, anon_lifetime());
             let ptr_name = names.func_ptr_binding(&param.name);
@@ -305,10 +307,7 @@ fn marshal_arg(
                 };
             }
         }
-        witx::Type::Union(_u) => {
-            let name = names.func_param(&param.name);
-            quote!(let #name = unimplemented!("union argument marshaling");)
-        }
+        witx::Type::Union(_u) => read_conversion,
         _ => unimplemented!("argument type marshalling"),
     }
 }