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

@@ -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]

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()

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();

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,15 +14,15 @@ 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,
 }
 
 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::<W, _, _>(#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::<W, _, _>(#map_stream)(#outer_writer);
+                #tail
+            };
         }
 
         self

binrw_derive/src/binrw/codegen/write_options/struct.rs

@@ -1,32 +1,31 @@
 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;
 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,
     out: TokenStream,
 }
 
 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
     }
 }

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<Assert>,
         #[from(RO:PreAssert)]