From 39420b80ef0a985f96ab77f42abd9dec508f4621 Mon Sep 17 00:00:00 2001
From: Dominik Nakamura <dnaka91@gmail.com>
Date: Sun, 15 Oct 2023 23:11:58 +0900
Subject: [PATCH] fix: create specialized encoding for non-zero types

---
 crates/stef-build/src/encode.rs               | 22 ++++++++++++++++++-
 .../compiler__compile@types-nested.stef.snap  | 11 +++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/crates/stef-build/src/encode.rs b/crates/stef-build/src/encode.rs
index f762070..190c9d6 100644
--- 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
index f739caf..4851967 100644
--- 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);
+                                    },
+                                );
                             },
                         );
                     },