diff --git a/beacon_chain/types/src/attestation_record.rs b/beacon_chain/types/src/attestation_record.rs index ed38664a9..e4ddffd62 100644 --- a/beacon_chain/types/src/attestation_record.rs +++ b/beacon_chain/types/src/attestation_record.rs @@ -31,7 +31,7 @@ impl Encodable for AttestationRecord { s.append(&self.shard_id); s.append_vec(&self.oblique_parent_hashes); s.append(&self.shard_block_hash); - s.append_vec(&self.attester_bitfield.to_bytes()); + s.append(&self.attester_bitfield); s.append(&self.justified_slot); s.append(&self.justified_block_hash); s.append_vec(&self.aggregate_sig.as_bytes()); diff --git a/beacon_chain/utils/boolean-bitfield/src/lib.rs b/beacon_chain/utils/boolean-bitfield/src/lib.rs index 2a08930fe..b90e4949e 100644 --- a/beacon_chain/utils/boolean-bitfield/src/lib.rs +++ b/beacon_chain/utils/boolean-bitfield/src/lib.rs @@ -86,6 +86,7 @@ impl BooleanBitfield { } /// Returns a vector of bytes representing the bitfield + /// Note that this returns the bit layout of the underlying implementation in the `bit-vec` crate. pub fn to_bytes(&self) -> Vec { self.0.to_bytes() } @@ -99,6 +100,28 @@ impl default::Default for BooleanBitfield { } } +// borrowed from bit_vec crate +fn reverse_bits(byte: u8) -> u8 { + let mut result = 0; + for i in 0..8 { + result = result | ((byte >> i) & 1) << (7 - i); + } + result +} + +impl ssz::Encodable for BooleanBitfield { + // ssz_append encodes Self according to the `ssz` spec. + // Note that we have to flip the endianness of the encoding with `reverse_bits` to account for an implementation detail of `bit-vec` crate. + fn ssz_append(&self, s: &mut ssz::SszStream) { + let bytes: Vec = self + .to_bytes() + .iter() + .map(|&byte| reverse_bits(byte)) + .collect(); + s.append_vec(&bytes); + } +} + impl ssz::Decodable for BooleanBitfield { fn ssz_decode(bytes: &[u8], index: usize) -> Result<(Self, usize), ssz::DecodeError> { let len = ssz::decode::decode_length(bytes, index, ssz::LENGTH_BYTES)?; @@ -109,7 +132,18 @@ impl ssz::Decodable for BooleanBitfield { if len == 0 { Ok((BooleanBitfield::new(), index + ssz::LENGTH_BYTES)) } else { - let field = BooleanBitfield::from_bytes(&bytes[(index + 4)..(index + len + 4)]); + let bytes = &bytes[(index + 4)..(index + len + 4)]; + + let mut field = BooleanBitfield::from_elem(0, false); + for (byte_index, byte) in bytes.iter().enumerate() { + for i in 0..8 { + let bit = byte & (1 << i); + if bit != 0 { + field.set(8 * byte_index + i, true); + } + } + } + let index = index + ssz::LENGTH_BYTES + len; Ok((field, index)) } @@ -119,6 +153,7 @@ impl ssz::Decodable for BooleanBitfield { #[cfg(test)] mod tests { use super::*; + use ssz::SszStream; #[test] fn test_empty_bitfield() { @@ -231,4 +266,31 @@ mod tests { let field = BooleanBitfield::from_elem(13, true); assert_eq!(field.num_bytes(), 2); } + + #[test] + fn test_ssz_encode() { + let field = BooleanBitfield::from_elem(5, true); + + let mut stream = SszStream::new(); + stream.append(&field); + assert_eq!(stream.drain(), vec![0, 0, 0, 1, 31]); + + let field = BooleanBitfield::from_elem(18, true); + let mut stream = SszStream::new(); + stream.append(&field); + assert_eq!(stream.drain(), vec![0, 0, 0, 3, 255, 255, 3]); + } + + #[test] + fn test_ssz_decode() { + let encoded = vec![0, 0, 0, 1, 31]; + let (field, _): (BooleanBitfield, usize) = ssz::decode_ssz(&encoded, 0).unwrap(); + let expected = BooleanBitfield::from_elem(5, true); + assert_eq!(field, expected); + + let encoded = vec![0, 0, 0, 3, 255, 255, 3]; + let (field, _): (BooleanBitfield, usize) = ssz::decode_ssz(&encoded, 0).unwrap(); + let expected = BooleanBitfield::from_elem(18, true); + assert_eq!(field, expected); + } }