Decode with mem limit

.github/workflows/ci.yml

@@ -22,10 +22,10 @@ jobs:
       - id: set_image
         run: echo "IMAGE=${{ env.IMAGE }}" >> $GITHUB_OUTPUT
 
-# Checks
+  # Checks
   fmt:
     runs-on: ubuntu-latest
-    needs: [set-image]
+    needs: [ set-image ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code
@@ -39,7 +39,7 @@ jobs:
 
   clippy:
     runs-on: ubuntu-latest
-    needs: [set-image]
+    needs: [ set-image ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code/.
@@ -64,7 +64,7 @@ jobs:
 
   checks:
     runs-on: ubuntu-latest
-    needs: [set-image]
+    needs: [ set-image ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code
@@ -99,10 +99,10 @@ jobs:
           export RUSTFLAGS='-Cdebug-assertions=y -Dwarnings'
           time cargo +stable check --verbose --features max-encoded-len
 
-# Tests
+  # Tests
   tests:
     runs-on: ubuntu-latest
-    needs: [set-image]
+    needs: [ set-image ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code
@@ -139,13 +139,18 @@ jobs:
           export RUSTFLAGS='-Cdebug-assertions=y -Dwarnings'
           time cargo +stable test --verbose --features max-encoded-len,std --no-default-features
 
-# Benches
+      - name: Run Nightly Tests
+        run: |
+          export RUSTFLAGS='-Cdebug-assertions=y -Dwarnings'
+          time cargo +nightly test --verbose --lib btree_utils
+
+  # Benches
   bench-rust-nightly:
     runs-on: ubuntu-latest
-    needs: [set-image]
+    needs: [ set-image ]
     strategy:
       matrix:
-        feature: [bit-vec,bytes,generic-array,derive,max-encoded-len]
+        feature: [ bit-vec,bytes,generic-array,derive,max-encoded-len ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code
@@ -163,10 +168,10 @@ jobs:
 
   miri:
     runs-on: ubuntu-latest
-    needs: [set-image]
+    needs: [ set-image ]
     strategy:
       matrix:
-        feature: [bit-vec,bytes,generic-array,arbitrary]
+        feature: [ bit-vec,bytes,generic-array,arbitrary ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code
@@ -184,11 +189,11 @@ jobs:
           export MIRIFLAGS='-Zmiri-disable-isolation'
           time cargo +nightly miri test --features ${{ matrix.feature }} --release
 
-# Build
+  # Build
 
   build-linux-ubuntu-amd64:
     runs-on: ubuntu-latest
-    needs: [set-image, clippy, checks, tests]
+    needs: [ set-image, clippy, checks, tests ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code
@@ -204,7 +209,7 @@ jobs:
 
   publish-dry-run:
     runs-on: ubuntu-latest
-    needs: [set-image, build-linux-ubuntu-amd64]
+    needs: [ set-image, build-linux-ubuntu-amd64 ]
     container: ${{ needs.set-image.outputs.IMAGE }}
     steps:
       - name: Checkout code

Cargo.toml

@@ -1,24 +1,26 @@
 [package]
 name = "parity-scale-codec"
 description = "SCALE - Simple Concatenating Aggregated Little Endians"
-version = "3.6.8"
+version = "3.7.0"
 authors = ["Parity Technologies <[email protected]>"]
 license = "Apache-2.0"
 repository = "https://github.com/paritytech/parity-scale-codec"
 categories = ["encoding"]
 edition = "2021"
+build = "build.rs"
 rust-version = "1.60.0"
 
 [dependencies]
 arrayvec = { version = "0.7", default-features = false }
 serde = { version = "1.0.204", default-features = false, optional = true }
 parity-scale-codec-derive = { path = "derive", version = ">= 3.6.8", default-features = false, optional = true }
-bitvec = { version = "1", default-features = false, features = [ "alloc" ], optional = true }
+bitvec = { version = "1", default-features = false, features = ["alloc"], optional = true }
 bytes = { version = "1", default-features = false, optional = true }
 byte-slice-cast = { version = "1.2.2", default-features = false }
 generic-array = { version = "0.14.7", optional = true }
 arbitrary = { version = "1.3.2", features = ["derive"], optional = true }
 impl-trait-for-tuples = "0.2.2"
+rustversion = "1.0.17"
 
 [dev-dependencies]
 criterion = "0.4.0"
@@ -28,6 +30,10 @@ quickcheck = "1.0"
 proptest = "1.5.0"
 trybuild = "1.0.97"
 paste = "1"
+rustversion = "1"
+
+[build-dependencies]
+rustversion = "1"
 
 [[bench]]
 name = "benches"

build.rs

@@ -0,0 +1,6 @@
+fn main() {
+	println!("cargo:rustc-check-cfg=cfg(nightly)");
+	if rustversion::cfg!(nightly) {
+		println!("cargo:rustc-cfg=nightly");
+	}
+}

derive/src/decode.rs

@@ -12,11 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::utils;
 use proc_macro2::{Ident, Span, TokenStream};
+use quote::ToTokens;
+use std::iter;
 use syn::{spanned::Spanned, Data, Error, Field, Fields};
 
-use crate::utils;
-
 /// Generate function block for function `Decode::decode`.
 ///
 /// * data: data info of the type,
@@ -31,33 +32,20 @@ pub fn quote(
 	crate_path: &syn::Path,
 ) -> TokenStream {
 	match *data {
-		Data::Struct(ref data) => match data.fields {
-			Fields::Named(_) | Fields::Unnamed(_) => create_instance(
-				quote! { #type_name #type_generics },
-				&type_name.to_string(),
-				input,
-				&data.fields,
-				crate_path,
-			),
-			Fields::Unit => {
-				quote_spanned! { data.fields.span() =>
-					::core::result::Result::Ok(#type_name)
-				}
-			},
-		},
+		Data::Struct(ref data) => create_instance(
+			quote! { #type_name #type_generics },
+			&type_name.to_string(),
+			input,
+			&data.fields,
+			crate_path,
+		),
 		Data::Enum(ref data) => {
-			let data_variants =
-				|| data.variants.iter().filter(|variant| !utils::should_skip(&variant.attrs));
-
-			if data_variants().count() > 256 {
-				return Error::new(
-					data.variants.span(),
-					"Currently only enums with at most 256 variants are encodable.",
-				)
-				.to_compile_error();
-			}
+			let variants = match utils::try_get_variants(data) {
+				Ok(variants) => variants,
+				Err(e) => return e.to_compile_error(),
+			};
 
-			let recurse = data_variants().enumerate().map(|(i, v)| {
+			let recurse = variants.iter().enumerate().map(|(i, v)| {
 				let name = &v.ident;
 				let index = utils::variant_index(v, i);
 
@@ -198,12 +186,12 @@ fn create_decode_expr(
 	crate_path: &syn::Path,
 ) -> TokenStream {
 	let encoded_as = utils::get_encoded_as_type(field);
-	let compact = utils::is_compact(field);
+	let compact = utils::get_compact_type(field, crate_path);
 	let skip = utils::should_skip(&field.attrs);
 
 	let res = quote!(__codec_res_edqy);
 
-	if encoded_as.is_some() as u8 + compact as u8 + skip as u8 > 1 {
+	if encoded_as.is_some() as u8 + compact.is_some() as u8 + skip as u8 > 1 {
 		return Error::new(
 			field.span(),
 			"`encoded_as`, `compact` and `skip` can only be used one at a time!",
@@ -213,13 +201,10 @@ fn create_decode_expr(
 
 	let err_msg = format!("Could not decode `{}`", name);
 
-	if compact {
-		let field_type = &field.ty;
+	if let Some(compact) = compact {
 		quote_spanned! { field.span() =>
 			{
-				let #res = <
-					<#field_type as #crate_path::HasCompact>::Type as #crate_path::Decode
-				>::decode(#input);
+				let #res = <#compact as #crate_path::Decode>::decode(#input);
 				match #res {
 					::core::result::Result::Err(e) => return ::core::result::Result::Err(e.chain(#err_msg)),
 					::core::result::Result::Ok(#res) => #res.into(),
@@ -300,3 +285,48 @@ fn create_instance(
 		},
 	}
 }
+
+pub fn quote_decode_with_mem_tracking_checks(data: &Data, crate_path: &syn::Path) -> TokenStream {
+	let fields: Box<dyn Iterator<Item = &Field>> = match data {
+		Data::Struct(data) => Box::new(data.fields.iter()),
+		Data::Enum(ref data) => {
+			let variants = match utils::try_get_variants(data) {
+				Ok(variants) => variants,
+				Err(e) => return e.to_compile_error(),
+			};
+
+			let mut fields: Box<dyn Iterator<Item = &Field>> = Box::new(iter::empty());
+			for variant in variants {
+				fields = Box::new(fields.chain(variant.fields.iter()));
+			}
+			fields
+		},
+		Data::Union(_) => {
+			return Error::new(Span::call_site(), "Union types are not supported.")
+				.to_compile_error();
+		},
+	};
+
+	let processed_fields = fields.filter_map(|field| {
+		if utils::should_skip(&field.attrs) {
+			return None;
+		}
+
+		let field_type = if let Some(compact) = utils::get_compact_type(field, crate_path) {
+			compact
+		} else if let Some(encoded_as) = utils::get_encoded_as_type(field) {
+			encoded_as
+		} else {
+			field.ty.to_token_stream()
+		};
+		Some(quote_spanned! {field.span() => #field_type})
+	});
+
+	quote! {
+		fn check_field<T: #crate_path::DecodeWithMemTracking>() {}
+
+		#(
+			check_field::<#processed_fields>();
+		)*
+	}
+}

derive/src/encode.rs

@@ -28,7 +28,7 @@ fn encode_single_field(
 	crate_path: &syn::Path,
 ) -> TokenStream {
 	let encoded_as = utils::get_encoded_as_type(field);
-	let compact = utils::is_compact(field);
+	let compact = utils::get_compact_type(field, crate_path);
 
 	if utils::should_skip(&field.attrs) {
 		return Error::new(
@@ -38,20 +38,19 @@ fn encode_single_field(
 		.to_compile_error();
 	}
 
-	if encoded_as.is_some() && compact {
+	if encoded_as.is_some() && compact.is_some() {
 		return Error::new(
 			Span::call_site(),
 			"`encoded_as` and `compact` can not be used at the same time!",
 		)
 		.to_compile_error();
 	}
 
-	let final_field_variable = if compact {
+	let final_field_variable = if let Some(compact) = compact {
 		let field_type = &field.ty;
 		quote_spanned! {
 			field.span() => {
-				<<#field_type as #crate_path::HasCompact>::Type as
-				#crate_path::EncodeAsRef<'_, #field_type>>::RefType::from(#field_name)
+				<#compact as #crate_path::EncodeAsRef<'_, #field_type>>::RefType::from(#field_name)
 			}
 		}
 	} else if let Some(encoded_as) = encoded_as {
@@ -298,23 +297,17 @@ fn impl_encode(data: &Data, type_name: &Ident, crate_path: &syn::Path) -> TokenS
 			Fields::Unit => [quote! { 0_usize }, quote!()],
 		},
 		Data::Enum(ref data) => {
-			let data_variants =
-				|| data.variants.iter().filter(|variant| !utils::should_skip(&variant.attrs));
-
-			if data_variants().count() > 256 {
-				return Error::new(
-					data.variants.span(),
-					"Currently only enums with at most 256 variants are encodable.",
-				)
-				.to_compile_error();
-			}
+			let variants = match utils::try_get_variants(data) {
+				Ok(variants) => variants,
+				Err(e) => return e.to_compile_error(),
+			};
 
 			// If the enum has no variants, we don't need to encode anything.
-			if data_variants().count() == 0 {
+			if variants.is_empty() {
 				return quote!();
 			}
 
-			let recurse = data_variants().enumerate().map(|(i, f)| {
+			let recurse = variants.iter().enumerate().map(|(i, f)| {
 				let name = &f.ident;
 				let index = utils::variant_index(f, i);