Skip to content

Commit

Permalink
Cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@macladson
macladson committed May 29, 2024
1 parent a7f3ada commit d767785
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 45 deletions.
81 changes: 38 additions & 43 deletions ssz_derive/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -587,7 +587,6 @@ fn ssz_encode_derive_stable_container(
let mut active_fields = BitVector::<#max_fields>::new();

let mut working_field: usize = 0;

#(
if self.#struct_fields_vec.is_some() {
active_fields.set(working_field, true).expect("Should not be out of bounds");
Expand Down Expand Up @@ -1358,6 +1357,8 @@ fn ssz_decode_derive_stable_container(
let mut is_fixed_lens = vec![];
let mut fixed_lens = vec![];

let mut working_index: usize = 0;

for (ty, ident, field_opts) in parse_ssz_fields(struct_data) {
let ident = match ident {
Some(ref ident) => ident,
Expand Down Expand Up @@ -1405,14 +1406,21 @@ fn ssz_decode_derive_stable_container(
from_ssz_bytes = quote! { <#ty as ssz::Decode>::from_ssz_bytes(slice) };

register_types.push(quote! {
builder.register_type::<#ty>()?
if bitvector.get(#working_index).unwrap() {
builder.register_type::<#ty>()?;
}
});
decodes.push(quote! {
decoder.decode_next()?
let #ident = if bitvector.get(#working_index).unwrap() {
decoder.decode_next()?
} else {
None
};
});
}

fixed_decodes.push(quote! {
let #ident = if bitvector.get(#working_index).unwrap() {
start = end;
end = end
.checked_add(#ssz_fixed_len)
Expand All @@ -1425,9 +1433,14 @@ fn ssz_decode_derive_stable_container(
expected: end
})?;
#from_ssz_bytes?
} else {
None
};
});
is_fixed_lens.push(is_ssz_fixed_len);
fixed_lens.push(ssz_fixed_len);

working_index += 1;
}

let output = quote! {
Expand Down Expand Up @@ -1465,14 +1478,8 @@ fn ssz_decode_derive_stable_container(
let mut start: usize = 0;
let mut end = start;

let mut working_field: usize = 0;
#(
let #field_names = if bitvector.get(working_field).unwrap() {
#fixed_decodes
} else {
None
};
working_field += 1;
#fixed_decodes
)*

Ok(Self {
Expand All @@ -1483,27 +1490,16 @@ fn ssz_decode_derive_stable_container(
} else {
let mut builder = ssz::SszDecoderBuilder::new(bytes);

let mut working_field: usize = 0;
#(
if bitvector.get(working_field).unwrap() {
#register_types
}
working_field += 1;
#register_types
)*

let mut decoder = builder.build()?;

let mut working_field: usize = 0;
#(
let #field_names = if bitvector.get(working_field).unwrap() {
#decodes
} else {
None
};
working_field += 1;
#decodes
)*


Ok(Self {
#(
#field_names,
Expand All @@ -1529,15 +1525,12 @@ fn ssz_decode_derive_profile_container(
let mut decodes = vec![];
let mut is_fixed_lens = vec![];
let mut fixed_lens = vec![];
let mut optional_field_indices: Vec<usize> = vec![];
let mut optional_field_names: Vec<&Ident> = vec![];
// Since we use a truncated bitvector, we need to keep track of which optional field we are up
// to.
let mut working_optional_index: usize = 0;

for (field_index, (ty, ident, field_opts)) in
parse_ssz_fields(struct_data).into_iter().enumerate()
{
for (ty, ident, field_opts) in parse_ssz_fields(struct_data) {
let mut is_optional = false;
let ident = match ident {
Some(ref ident) => ident,
Expand Down Expand Up @@ -1565,7 +1558,6 @@ fn ssz_decode_derive_profile_container(

// Check if field is optional.
if ty_inner_type("Option", ty).is_some() {
optional_field_indices.push(field_index);
optional_field_names.push(ident);
is_optional = true;
}
Expand All @@ -1583,18 +1575,9 @@ fn ssz_decode_derive_profile_container(
register_types.push(quote! {
builder.register_type_parameterized(#is_ssz_fixed_len, #ssz_fixed_len)?;
});
// If the field is optional, we need to check the bitvector before decoding.
if is_optional {
decodes.push(quote! {
if bitvector.get(#working_optional_index).unwrap()
if self.#ident.is_some() {
let #ident = decoder.decode_next_with(|slice| #module::from_ssz_bytes(slice))?;
} else {
// Value is None so just decode an Option::default();
let #ident = <_>::default();
}
});
}
decodes.push(quote! {
let #ident = decoder.decode_next_with(|slice| #module::from_ssz_bytes(slice))?;
});
} else {
is_ssz_fixed_len = quote! { <#ty as ssz::Decode>::is_ssz_fixed_len() };
ssz_fixed_len = quote! { <#ty as ssz::Decode>::ssz_fixed_len() };
Expand All @@ -1603,9 +1586,19 @@ fn ssz_decode_derive_profile_container(
register_types.push(quote! {
builder.register_type::<#ty>()?;
});
decodes.push(quote! {
let #ident = decoder.decode_next()?;
});
if is_optional {
decodes.push(quote! {
let #ident = if bitvector.get(#working_optional_index).unwrap() {
decoder.decode_next()?
} else {
<_>::default()
};
});
} else {
decodes.push(quote! {
let #ident = decoder.decode_next()?;
});
}
}

// If the field is optional, we need to check the bitvector before decoding.
Expand Down Expand Up @@ -1653,6 +1646,8 @@ fn ssz_decode_derive_profile_container(
is_fixed_lens.push(is_ssz_fixed_len);
fixed_lens.push(ssz_fixed_len);
}

// Increment the working index so we check the next field of the bitvector.
if is_optional {
working_optional_index += 1
};
Expand Down
13 changes: 11 additions & 2 deletions ssz_derive/tests/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -268,6 +268,8 @@ struct Shape {
#[ssz(struct_behaviour = "profile")]
struct Square {
side: u16,
#[ssz(skip_serializing, skip_deserializing)]
skip: Vec<u8>,
color: u8,
}

Expand All @@ -284,7 +286,9 @@ struct Circle {
struct ShapeVec {
side: Option<u16>,
color: Option<u8>,
radius: Option<ssz_types::VariableList<u8, typenum::U8>>,
#[ssz(skip_serializing, skip_deserializing)]
skip: Vec<u8>,
radius: Option<ssz_types::VariableList<u8, typenum::U4>>,
}

#[test]
Expand Down Expand Up @@ -317,7 +321,11 @@ fn shape_2() {
/// Square(side=0x42, color=1)
/// 420001
fn square() {
let square = Square { side: 42, color: 1 };
let square = Square {
side: 42,
skip: vec![],
color: 1,
};

assert_encode_decode(&square, &vec![42, 0, 1]);
}
Expand All @@ -339,6 +347,7 @@ fn shape_3() {
let shape = ShapeVec {
side: None,
color: Some(1),
skip: vec![],
radius: Some(vec![1, 2, 3, 4].into()),
};

Expand Down

0 comments on commit d767785

Please sign in to comment.