From 234eeb73cb9060364f825a62b73508b1d56cf1d5 Mon Sep 17 00:00:00 2001 From: Colin Snover Date: Thu, 26 Oct 2023 18:04:15 -0500 Subject: [PATCH] Implement `assert` for data enums in BinWrite This directive was being accepted and then silently discarded instead of generating code. --- binrw/tests/derive/write/assert.rs | 34 ++++++++++++++++ .../src/binrw/codegen/write_options.rs | 3 +- .../src/binrw/codegen/write_options/enum.rs | 10 +++-- .../binrw/codegen/write_options/prelude.rs | 40 +++++++++++-------- .../src/binrw/codegen/write_options/struct.rs | 36 ++++++----------- .../src/binrw/parser/top_level_attrs.rs | 5 --- 6 files changed, 78 insertions(+), 50 deletions(-) diff --git a/binrw/tests/derive/write/assert.rs b/binrw/tests/derive/write/assert.rs index fbfd3f63..7884f431 100644 --- a/binrw/tests/derive/write/assert.rs +++ b/binrw/tests/derive/write/assert.rs @@ -32,6 +32,40 @@ fn top_level_assert_fail() { } } +#[test] +fn top_level_assert_self_enum() { + #[binwrite] + #[bw(assert(!matches!(self, Test::A(1))))] + #[derive(PartialEq)] + enum Test { + A(u32), + } + + let mut x = Cursor::new(Vec::new()); + if let Err(err) = x.write_be(&Test::A(1)) { + assert!(matches!(err, binrw::Error::AssertFail { .. })); + } else { + panic!("Assert error expected"); + } +} + +#[test] +fn assert_enum_variant() { + #[binwrite] + #[derive(PartialEq)] + enum Test { + #[bw(assert(self_0 != &1))] + A(u32), + } + + let mut x = Cursor::new(Vec::new()); + if let Err(err) = x.write_be(&Test::A(1)) { + assert!(matches!(err, binrw::Error::AssertFail { .. })); + } else { + panic!("Assert error expected"); + } +} + #[test] fn top_level_assert_self_struct() { #[binwrite] diff --git a/binrw_derive/src/binrw/codegen/write_options.rs b/binrw_derive/src/binrw/codegen/write_options.rs index abbe386e..90717850 100644 --- a/binrw_derive/src/binrw/codegen/write_options.rs +++ b/binrw_derive/src/binrw/codegen/write_options.rs @@ -62,8 +62,9 @@ fn generate_map(input: &Input, name: Option<&Ident>, map: &TokenStream) -> Token let magic = input.magic(); let endian = input.endian(); - prelude::PreludeGenerator::new(write_data, Some(input), name, &writer_var) + prelude::PreludeGenerator::new(write_data, input, name, &writer_var) .prefix_magic(magic) + .prefix_assertions() .prefix_endian(endian) .prefix_imports() .finish() diff --git a/binrw_derive/src/binrw/codegen/write_options/enum.rs b/binrw_derive/src/binrw/codegen/write_options/enum.rs index 1d135a08..b08f9daa 100644 --- a/binrw_derive/src/binrw/codegen/write_options/enum.rs +++ b/binrw_derive/src/binrw/codegen/write_options/enum.rs @@ -17,9 +17,10 @@ pub(crate) fn generate_unit_enum( None => generate_unit_enum_magic(&writer_var, &en.fields), }; - PreludeGenerator::new(write, Some(input), name, &writer_var) + PreludeGenerator::new(write, input, name, &writer_var) .prefix_map_stream() .prefix_magic(&en.magic) + .prefix_assertions() .prefix_endian(&en.endian) .prefix_imports() .finish() @@ -67,7 +68,9 @@ impl<'a> EnumGenerator<'a> { let writer_var = &self.writer_var; let writing = match variant { EnumVariant::Variant { options, .. } => { - StructGenerator::new(None, options, None, &self.writer_var) + let input = Input::Struct(variant.clone().into()); + + StructGenerator::new(&input, options, None, &self.writer_var) .write_fields() .prefix_prelude() .finish() @@ -108,9 +111,10 @@ impl<'a> EnumGenerator<'a> { fn prefix_prelude(mut self) -> Self { let out = self.out; - self.out = PreludeGenerator::new(out, Some(self.input), self.name, &self.writer_var) + self.out = PreludeGenerator::new(out, self.input, self.name, &self.writer_var) .prefix_map_stream() .prefix_magic(&self.en.magic) + .prefix_assertions() .prefix_endian(&self.en.endian) .prefix_imports() .finish(); diff --git a/binrw_derive/src/binrw/codegen/write_options/prelude.rs b/binrw_derive/src/binrw/codegen/write_options/prelude.rs index 8196e497..6731de20 100644 --- a/binrw_derive/src/binrw/codegen/write_options/prelude.rs +++ b/binrw_derive/src/binrw/codegen/write_options/prelude.rs @@ -1,7 +1,7 @@ use crate::{ binrw::{ codegen::{ - get_destructured_imports, get_endian, + get_assertions, get_destructured_imports, get_endian, sanitization::{ARGS, MAP_WRITER_TYPE_HINT, OPT, WRITER, WRITE_METHOD}, }, parser::{CondEndian, Input, Magic}, @@ -14,7 +14,7 @@ use syn::spanned::Spanned; pub(crate) struct PreludeGenerator<'a> { out: TokenStream, - input: Option<&'a Input>, + input: &'a Input, name: Option<&'a Ident>, writer_var: &'a TokenStream, } @@ -22,7 +22,7 @@ pub(crate) struct PreludeGenerator<'a> { impl<'a> PreludeGenerator<'a> { pub(crate) fn new( out: TokenStream, - input: Option<&'a Input>, + input: &'a Input, name: Option<&'a Ident>, writer_var: &'a TokenStream, ) -> Self { @@ -34,11 +34,19 @@ impl<'a> PreludeGenerator<'a> { } } + pub(super) fn prefix_assertions(mut self) -> Self { + let assertions = get_assertions(self.input.assertions()); + let out = self.out; + self.out = quote! { + #(#assertions)* + #out + }; + + self + } + pub(crate) fn prefix_imports(mut self) -> Self { - if let Some(imports) = self - .input - .and_then(|input| get_destructured_imports(input.imports(), self.name, true)) - { + if let Some(imports) = get_destructured_imports(self.input.imports(), self.name, true) { let out = self.out; self.out = quote! { let #imports = #ARGS; @@ -81,16 +89,14 @@ impl<'a> PreludeGenerator<'a> { } pub(crate) fn prefix_map_stream(mut self) -> Self { - if let Some(input) = self.input { - if let Some(map_stream) = input.map_stream() { - let outer_writer = input.stream_ident_or(WRITER); - let inner_writer = &self.writer_var; - let tail = self.out; - self.out = quote_spanned_any! { map_stream.span()=> - let #inner_writer = &mut #MAP_WRITER_TYPE_HINT::(#map_stream)(#outer_writer); - #tail - }; - } + if let Some(map_stream) = self.input.map_stream() { + let outer_writer = self.input.stream_ident_or(WRITER); + let inner_writer = &self.writer_var; + let tail = self.out; + self.out = quote_spanned_any! { map_stream.span()=> + let #inner_writer = &mut #MAP_WRITER_TYPE_HINT::(#map_stream)(#outer_writer); + #tail + }; } self diff --git a/binrw_derive/src/binrw/codegen/write_options/struct.rs b/binrw_derive/src/binrw/codegen/write_options/struct.rs index 1a9673da..137f70e2 100644 --- a/binrw_derive/src/binrw/codegen/write_options/struct.rs +++ b/binrw_derive/src/binrw/codegen/write_options/struct.rs @@ -1,6 +1,6 @@ use super::{prelude::PreludeGenerator, struct_field::write_field}; use crate::binrw::{ - codegen::{get_assertions, sanitization::WRITER}, + codegen::sanitization::{THIS, WRITER}, parser::{Input, Struct}, }; use proc_macro2::TokenStream; @@ -8,16 +8,15 @@ use quote::quote; use syn::Ident; pub(super) fn generate_struct(input: &Input, name: Option<&Ident>, st: &Struct) -> TokenStream { - StructGenerator::new(Some(input), st, name, &input.stream_ident_or(WRITER)) + StructGenerator::new(input, st, name, &input.stream_ident_or(WRITER)) .write_fields() - .prefix_assertions() .prefix_prelude() .prefix_borrow_fields() .finish() } -pub(crate) struct StructGenerator<'input> { - input: Option<&'input Input>, +pub(super) struct StructGenerator<'input> { + input: &'input Input, st: &'input Struct, name: Option<&'input Ident>, writer_var: &'input TokenStream, @@ -25,8 +24,8 @@ pub(crate) struct StructGenerator<'input> { } impl<'input> StructGenerator<'input> { - pub(crate) fn new( - input: Option<&'input Input>, + pub(super) fn new( + input: &'input Input, st: &'input Struct, name: Option<&'input Ident>, writer_var: &'input TokenStream, @@ -40,30 +39,19 @@ impl<'input> StructGenerator<'input> { } } - pub(crate) fn prefix_prelude(mut self) -> Self { + pub(super) fn prefix_prelude(mut self) -> Self { self.out = PreludeGenerator::new(self.out, self.input, self.name, self.writer_var) .prefix_map_stream() .prefix_magic(&self.st.magic) .prefix_endian(&self.st.endian) + .prefix_assertions() .prefix_imports() .finish(); self } - fn prefix_assertions(mut self) -> Self { - let assertions = get_assertions(&self.st.assertions); - - let out = self.out; - self.out = quote! { - #(#assertions)* - #out - }; - - self - } - - pub(crate) fn write_fields(mut self) -> Self { + pub(super) fn write_fields(mut self) -> Self { let write_fields = self .st .fields @@ -77,8 +65,8 @@ impl<'input> StructGenerator<'input> { self } - pub(crate) fn prefix_borrow_fields(mut self) -> Self { - let borrow_fields = self.name.as_ref().map(|name| { + pub(super) fn prefix_borrow_fields(mut self) -> Self { + let borrow_fields = self.name.map(|name| { let pattern = self.st.fields_pattern(); Some(quote! { @@ -96,7 +84,7 @@ impl<'input> StructGenerator<'input> { self } - pub(crate) fn finish(self) -> TokenStream { + pub(super) fn finish(self) -> TokenStream { self.out } } diff --git a/binrw_derive/src/binrw/parser/top_level_attrs.rs b/binrw_derive/src/binrw/parser/top_level_attrs.rs index b6173f0a..4b93502e 100644 --- a/binrw_derive/src/binrw/parser/top_level_attrs.rs +++ b/binrw_derive/src/binrw/parser/top_level_attrs.rs @@ -320,11 +320,6 @@ attr_struct! { pub(crate) magic: Magic, #[from(RW:Import, RW:ImportRaw)] pub(crate) imports: Imports, - // TODO: Does this make sense? It is not known what properties will - // exist in order to construct a valid variant. The assertions all get - // copied and used as if they were applied to each variant in the enum, - // so the only way this ever works is if every variant contains the same - // properties being checked by the assertion. #[from(RW:Assert)] pub(crate) assertions: Vec, #[from(RO:PreAssert)]