Fix NifTaggedEnum derived Encoder impl for named-field variants (#547)

Use map_from_term_arrays in Nif{Map,Struct,Exception,TaggedEnum} encoders

rustler/src/types/map.rs

@@ -27,7 +27,7 @@ impl<'a> Term<'a> {
     /// ```elixir
     /// keys = ["foo", "bar"]
     /// values = [1, 2]
-    /// List.zip(keys, values) |> Map.new()
+    /// Enum.zip(keys, values) |> Map.new()
     /// ```
     pub fn map_from_arrays(
         env: Env<'a>,
@@ -47,6 +47,29 @@ impl<'a> Term<'a> {
         }
     }
 
+    /// Construct a new map from two vectors of terms.
+    ///
+    /// It is identical to map_from_arrays, but requires the keys and values to
+    /// be encoded already - this is useful for constructing maps whose values
+    /// or keys are different Rust types, with the same performance as map_from_arrays.
+    pub fn map_from_term_arrays(
+        env: Env<'a>,
+        keys: &[Term<'a>],
+        values: &[Term<'a>],
+    ) -> NifResult<Term<'a>> {
+        if keys.len() == values.len() {
+            let keys: Vec<_> = keys.iter().map(|k| k.as_c_arg()).collect();
+            let values: Vec<_> = values.iter().map(|v| v.as_c_arg()).collect();
+
+            unsafe {
+                map::make_map_from_arrays(env.as_c_arg(), &keys, &values)
+                    .map_or_else(|| Err(Error::BadArg), |map| Ok(Term::new(env, map)))
+            }
+        } else {
+            Err(Error::BadArg)
+        }
+    }
+
     /// Construct a new map from pairs of terms
     ///
     /// It is similar to `map_from_arrays` but

rustler_codegen/src/ex_struct.rs

@@ -136,34 +136,31 @@ fn gen_encoder(
     atoms_module_name: &Ident,
     add_exception: bool,
 ) -> TokenStream {
-    let field_defs: Vec<TokenStream> = fields
+    let mut keys = vec![quote! { ::rustler::Encoder::encode(&atom_struct(), env) }];
+    let mut values = vec![quote! { ::rustler::Encoder::encode(&atom_module(), env) }];
+    if add_exception {
+        keys.push(quote! { ::rustler::Encoder::encode(&atom_exception(), env) });
+        values.push(quote! { ::rustler::Encoder::encode(&true, env) });
+    }
+    let (mut data_keys, mut data_values): (Vec<_>, Vec<_>) = fields
         .iter()
         .map(|field| {
             let field_ident = field.ident.as_ref().unwrap();
             let atom_fun = Context::field_to_atom_fun(field);
-            quote_spanned! { field.span() =>
-                map = map.map_put(#atom_fun(), &self.#field_ident).unwrap();
-            }
+            (
+                quote! { ::rustler::Encoder::encode(&#atom_fun(), env) },
+                quote! { ::rustler::Encoder::encode(&self.#field_ident, env) },
+            )
         })
-        .collect();
-
-    let exception_field = if add_exception {
-        quote! {
-            map = map.map_put(atom_exception(), true).unwrap();
-        }
-    } else {
-        quote! {}
-    };
+        .unzip();
+    keys.append(&mut data_keys);
+    values.append(&mut data_values);
 
     super::encode_decode_templates::encoder(
         ctx,
         quote! {
             use #atoms_module_name::*;
-            let mut map = ::rustler::types::map::map_new(env);
-            map = map.map_put(atom_struct(), atom_module()).unwrap();
-            #exception_field
-            #(#field_defs)*
-            map
+            ::rustler::Term::map_from_term_arrays(env, &[#(#keys),*], &[#(#values),*]).unwrap()
         },
     )
 }

rustler_codegen/src/map.rs

@@ -107,26 +107,23 @@ fn gen_decoder(ctx: &Context, fields: &[&Field], atoms_module_name: &Ident) -> T
 }
 
 fn gen_encoder(ctx: &Context, fields: &[&Field], atoms_module_name: &Ident) -> TokenStream {
-    let field_defs: Vec<TokenStream> = fields
+    let (keys, values): (Vec<_>, Vec<_>) = fields
         .iter()
         .map(|field| {
             let field_ident = field.ident.as_ref().unwrap();
             let atom_fun = Context::field_to_atom_fun(field);
-
-            quote_spanned! { field.span() =>
-                map = map.map_put(#atom_fun(), &self.#field_ident).unwrap();
-            }
+            (
+                quote! { ::rustler::Encoder::encode(&#atom_fun(), env) },
+                quote! { ::rustler::Encoder::encode(&self.#field_ident, env) },
+            )
         })
-        .collect();
+        .unzip();
 
     super::encode_decode_templates::encoder(
         ctx,
         quote! {
             use #atoms_module_name::*;
-
-            let mut map = ::rustler::types::map::map_new(env);
-            #(#field_defs)*
-            map
+            ::rustler::Term::map_from_term_arrays(env, &[#(#keys),*], &[#(#values),*]).unwrap()
         },
     )
 }

rustler_codegen/src/tagged_enum.rs

@@ -321,14 +321,17 @@ fn gen_named_encoder(
                 .as_ref()
                 .expect("Named fields must have an ident.");
             let atom_fun = Context::field_to_atom_fun(field);
-            (atom_fun, field_ident)
+            (
+                quote! { ::rustler::Encoder::encode(&#atom_fun(), env) },
+                quote! { ::rustler::Encoder::encode(&#field_ident, env) },
+            )
         })
         .unzip();
     quote! {
         #enum_name :: #variant_ident{
             #(#field_decls)*
         } => {
-            let map = ::rustler::Term::map_from_arrays(env, &[#(#keys()),*], &[#(#values),*])
+            let map = ::rustler::Term::map_from_term_arrays(env, &[#(#keys),*], &[#(#values),*])
                 .expect("Failed to create map");
             ::rustler::types::tuple::make_tuple(env, &[::rustler::Encoder::encode(&#atom_fn(), env), map])
         }

rustler_tests/lib/rustler_test.ex

@@ -95,6 +95,7 @@ defmodule RustlerTest do
   def tagged_enum_1_echo(_), do: err()
   def tagged_enum_2_echo(_), do: err()
   def tagged_enum_3_echo(_), do: err()
+  def tagged_enum_4_echo(_), do: err()
   def untagged_enum_echo(_), do: err()
   def untagged_enum_with_truthy(_), do: err()
   def untagged_enum_for_issue_370(_), do: err()

rustler_tests/native/rustler_test/src/lib.rs

@@ -70,6 +70,7 @@ rustler::init!(
         test_codegen::tagged_enum_1_echo,
         test_codegen::tagged_enum_2_echo,
         test_codegen::tagged_enum_3_echo,
+        test_codegen::tagged_enum_4_echo,
         test_codegen::untagged_enum_echo,
         test_codegen::untagged_enum_with_truthy,
         test_codegen::untagged_enum_for_issue_370,

rustler_tests/native/rustler_test/src/test_codegen.rs

@@ -44,6 +44,7 @@ pub fn record_echo(record: AddRecord) -> AddRecord {
 pub struct AddMap {
     lhs: i32,
     rhs: i32,
+    loc: (u32, u32),
 }
 
 #[rustler::nif]
@@ -57,12 +58,14 @@ pub fn map_echo(map: AddMap) -> AddMap {
 pub struct AddStruct {
     lhs: i32,
     rhs: i32,
+    loc: (u32, u32),
 }
 
 #[derive(Debug, NifException)]
 #[module = "AddException"]
 pub struct AddException {
     message: String,
+    loc: (u32, u32),
 }
 
 #[rustler::nif]
@@ -125,6 +128,32 @@ pub fn tagged_enum_3_echo(tagged_enum: TaggedEnum3) -> TaggedEnum3 {
     tagged_enum
 }
 
+#[derive(NifTaggedEnum)]
+pub enum TaggedEnum4 {
+    Unit,
+    Unnamed(u64, bool),
+    Named {
+        size: u64,
+        filename: String,
+    },
+    Long {
+        f0: bool,
+        f1: u8,
+        f2: u8,
+        f3: u8,
+        f4: u8,
+        f5: Option<i32>,
+        f6: Option<i32>,
+        f7: Option<i32>,
+        f8: Option<i32>,
+    },
+}
+
+#[rustler::nif]
+pub fn tagged_enum_4_echo(tagged_enum: TaggedEnum4) -> TaggedEnum4 {
+    tagged_enum
+}
+
 #[derive(NifUntaggedEnum)]
 pub enum UntaggedEnum {
     Foo(u32),

rustler_tests/test/codegen_test.exs

@@ -1,9 +1,9 @@
 defmodule AddStruct do
-  defstruct lhs: 0, rhs: 0
+  defstruct lhs: 0, rhs: 0, loc: {1, 1}
 end
 
 defmodule AddException do
-  defexception message: ""
+  defexception message: "", loc: {1, 1}
 end
 
 defmodule AddRecord do
@@ -40,13 +40,13 @@ defmodule RustlerTest.CodegenTest do
   end
 
   describe "map" do
-    test "transcoder" do
-      value = %{lhs: 1, rhs: 2}
+    test "transcoder 1" do
+      value = %{lhs: 1, rhs: 2, loc: {52, 15}}
       assert value == RustlerTest.map_echo(value)
     end
 
     test "with invalid map" do
-      value = %{lhs: "invalid", rhs: 2}
+      value = %{lhs: "invalid", rhs: 2, loc: {57, 15}}
 
       assert_raise ErlangError, "Erlang error: \"Could not decode field :lhs on %{}\"", fn ->
         assert value == RustlerTest.map_echo(value)
@@ -56,7 +56,7 @@ defmodule RustlerTest.CodegenTest do
 
   describe "struct" do
     test "transcoder" do
-      value = %AddStruct{lhs: 45, rhs: 123}
+      value = %AddStruct{lhs: 45, rhs: 123, loc: {66, 15}}
       assert value == RustlerTest.struct_echo(value)
 
       assert %ErlangError{original: :invalid_struct} ==
@@ -66,19 +66,27 @@ defmodule RustlerTest.CodegenTest do
     end
 
     test "with invalid struct" do
-      value = %AddStruct{lhs: "lhs", rhs: 123}
+      value = %AddStruct{lhs: "lhs", rhs: 123, loc: {76, 15}}
 
       assert_raise ErlangError,
                    "Erlang error: \"Could not decode field :lhs on %AddStruct{}\"",
                    fn ->
                      RustlerTest.struct_echo(value)
                    end
+
+      value = %AddStruct{lhs: 45, rhs: 123, loc: {-76, -15}}
+
+      assert_raise ErlangError,
+                   "Erlang error: \"Could not decode field :loc on %AddStruct{}\"",
+                   fn ->
+                     RustlerTest.struct_echo(value)
+                   end
     end
   end
 
   describe "exception" do
     test "transcoder" do
-      value = %AddException{message: "testing"}
+      value = %AddException{message: "testing", loc: {96, 15}}
       assert value == RustlerTest.exception_echo(value)
 
       assert %ErlangError{original: :invalid_struct} ==
@@ -88,13 +96,21 @@ defmodule RustlerTest.CodegenTest do
     end
 
     test "with invalid struct" do
-      value = %AddException{message: 'this is a charlist'}
+      value = %AddException{message: 'this is a charlist', loc: {106, 15}}
 
       assert_raise ErlangError,
                    "Erlang error: \"Could not decode field :message on %AddException{}\"",
                    fn ->
                      RustlerTest.exception_echo(value)
                    end
+
+      value = %AddException{message: "testing", loc: %{line: 114, col: 15}}
+
+      assert_raise ErlangError,
+                   "Erlang error: \"Could not decode field :loc on %AddException{}\"",
+                   fn ->
+                     RustlerTest.exception_echo(value)
+                   end
     end
   end
 
@@ -279,6 +295,41 @@ defmodule RustlerTest.CodegenTest do
              end)
   end
 
+  test "tagged enum transcoder 4" do
+    assert :unit == RustlerTest.tagged_enum_4_echo(:unit)
+
+    assert {:unnamed, 10_000_000_000, true} ==
+             RustlerTest.tagged_enum_4_echo({:unnamed, 10_000_000_000, true})
+
+    assert {:named, %{filename: "\u2200", size: 123}} ==
+             RustlerTest.tagged_enum_4_echo({:named, %{filename: "\u2200", size: 123}})
+
+    long_map = %{f0: true, f1: 8, f2: 5, f3: 12, f4: 12, f5: 15, f6: nil, f7: nil, f8: nil}
+
+    assert {:long, long_map} ==
+             RustlerTest.tagged_enum_4_echo({:long, long_map})
+
+    assert %ErlangError{original: :invalid_variant} ==
+             assert_raise(ErlangError, fn ->
+               RustlerTest.tagged_enum_4_echo(:unnamed)
+             end)
+
+    assert %ErlangError{original: :invalid_variant} ==
+             assert_raise(ErlangError, fn ->
+               RustlerTest.tagged_enum_4_echo({:unit, 2, false})
+             end)
+
+    assert %ErlangError{original: :invalid_variant} ==
+             assert_raise(ErlangError, fn ->
+               RustlerTest.tagged_enum_4_echo({:named, "@", 45})
+             end)
+
+    assert %ErlangError{original: :invalid_variant} ==
+             assert_raise(ErlangError, fn ->
+               RustlerTest.tagged_enum_4_echo(nil)
+             end)
+  end
+
   test "untagged enum transcoder" do
     assert 123 == RustlerTest.untagged_enum_echo(123)
     assert "Hello" == RustlerTest.untagged_enum_echo("Hello")