fix: create specialized encoding for non-zero types

dnaka91 · Oct 15, 2023 · 39420b8 · 39420b8
1 parent ba90911
commit 39420b8
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/crates/stef-build/src/encode.rs b/crates/stef-build/src/encode.rs
@@ -303,6 +303,26 @@ fn compile_data_type(ty: &DataType<'_>, name: TokenStream) -> TokenStream {
             );
             quote! { ::stef::buf::encode_option(w, &#name, |w, v| { #ty; }) }
         }
+        DataType::NonZero(ty) => {
+            if matches!(
+                **ty,
+                DataType::U8
+                    | DataType::U16
+                    | DataType::U32
+                    | DataType::U64
+                    | DataType::U128
+                    | DataType::I8
+                    | DataType::I16
+                    | DataType::I32
+                    | DataType::I64
+                    | DataType::I128
+            ) {
+                quote! { (#name).encode(w) }
+            } else {
+                compile_data_type(ty, name)
+            }
+        }
+
         DataType::BoxString => quote! { ::stef::buf::encode_string(w, &*#name) },
         DataType::BoxBytes => quote! { ::stef::buf::encode_bytes(w, &*#name) },
         DataType::Tuple(types) => match types.len() {
@@ -335,7 +355,7 @@ fn compile_data_type(ty: &DataType<'_>, name: TokenStream) -> TokenStream {
             );
             quote! { ::stef::buf::encode_array(w, &#name, |w, v| { #ty; }) }
         }
-        DataType::NonZero(_) | DataType::External(_) => {
+        DataType::External(_) => {
             quote! { (#name).encode(w) }
         }
     }

diff --git a/crates/stef-build/tests/snapshots/compiler__compile@types-nested.stef.snap b/crates/stef-build/tests/snapshots/compiler__compile@types-nested.stef.snap
@@ -25,7 +25,16 @@ impl ::stef::Encode for Sample {
                             w,
                             &v,
                             |w, v| {
-                                (v).encode(w);
+                                ::stef::buf::encode_hash_map(
+                                    w,
+                                    &v,
+                                    |w, k| {
+                                        ::stef::buf::encode_i64(w, *k);
+                                    },
+                                    |w, v| {
+                                        ::stef::buf::encode_string(w, &*v);
+                                    },
+                                );
                             },
                         );
                     },