implement bitfield for enum with a width

If we know the width of an enum type, we don't need 'power of 2' number
of variants.

BUG=None
TEST=cargo test

Change-Id: I8148b28f86bb8e4fd4f67d8a6382fc713dad1439
Reviewed-on: https://chromium-review.googlesource.com/1530455
Commit-Ready: Jingkui Wang <jkwang@google.com>
Tested-by: Jingkui Wang <jkwang@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: David Tolnay <dtolnay@chromium.org>

3 files changed, 225 insertions, 55 deletions

diff --git a/bit_field/bit_field_derive/bit_field_derive.rs b/bit_field/bit_field_derive/bit_field_derive.rs
index c0dfe2a..a0938a1 100644
--- a/bit_field/bit_field_derive/bit_field_derive.rs
+++ b/bit_field/bit_field_derive/bit_field_derive.rs
@@ -65,6 +65,81 @@ fn bitfield_impl(ast: &DeriveInput) -> Result<TokenStream> {
     }
 }
 
+fn bitfield_enum_impl(ast: &DeriveInput, data: &DataEnum) -> Result<TokenStream> {
+    let mut ast = ast.clone();
+    let mut width = None;
+    let mut bits_idx = 0;
+
+    for (i, attr) in ast.attrs.iter().enumerate() {
+        if let Some(w) = try_parse_bits_attr(attr)? {
+            bits_idx = i;
+            width = Some(w);
+        }
+    }
+
+    if width.is_some() {
+        ast.attrs.remove(bits_idx);
+    }
+
+    match width {
+        None => bitfield_enum_without_width_impl(&ast, data),
+        Some(width) => bitfield_enum_with_width_impl(&ast, data, &width),
+    }
+}
+
+fn bitfield_enum_with_width_impl(
+    ast: &DeriveInput,
+    data: &DataEnum,
+    width: &LitInt,
+) -> Result<TokenStream> {
+    if width.value() > 64 {
+        return Err(Error::new(
+            Span::call_site(),
+            "max width of bitfield enum is 64",
+        ));
+    }
+    let bits = width.value() as u8;
+    let declare_discriminants = get_declare_discriminants_for_enum(bits, ast, data);
+
+    let ident = &ast.ident;
+    let type_name = ident.to_string();
+    let variants = &data.variants;
+    let match_discriminants = variants.iter().map(|variant| {
+        let variant = &variant.ident;
+        quote! {
+            discriminant::#variant => Ok(#ident::#variant),
+        }
+    });
+
+    let expanded = quote! {
+        #ast
+
+        impl bit_field::BitFieldSpecifier for #ident {
+            const FIELD_WIDTH: u8 = #bits;
+            type SetterType = Self;
+            type GetterType = Result<Self, bit_field::Error>;
+
+            #[inline]
+            fn from_u64(val: u64) -> Self::GetterType {
+                struct discriminant;
+                impl discriminant {
+                    #(#declare_discriminants)*
+                }
+                match val {
+                    #(#match_discriminants)*
+                    v => Err(bit_field::Error::new(#type_name, v)),
+                }
+            }
+
+            #[inline]
+            fn into_u64(val: Self::SetterType) -> u64 {
+                val as u64
+            }
+        }
+    };
+
+    Ok(expanded)
+}
 // Expand to an impl of BitFieldSpecifier for an enum like:
 //
 //     #[bitfield]
@@ -85,59 +160,20 @@ fn bitfield_impl(ast: &DeriveInput) -> Result<TokenStream> {
 //         suffix: BitField5,
 //     }
 //
-fn bitfield_enum_impl(ast: &DeriveInput, data: &DataEnum) -> Result<TokenStream> {
+fn bitfield_enum_without_width_impl(ast: &DeriveInput, data: &DataEnum) -> Result<TokenStream> {
+    let ident = &ast.ident;
     let variants = &data.variants;
     let len = variants.len();
     if len.count_ones() != 1 {
         return Err(Error::new(
             Span::call_site(),
-            "#[bitfield] expected a number of variants which is a power of 2",
+            "#[bitfield] expected a number of variants which is a power of 2 when bits is not \
+             specified for the enum",
         ));
     }
 
     let bits = len.trailing_zeros() as u8;
-    let upper_bound = 2u64.pow(bits as u32);
-
-    let ident = &ast.ident;
-
-    let declare_discriminants = variants.iter().map(|variant| {
-        let variant = &variant.ident;
-        let span = variant.span();
-
-        let assertion = quote_spanned! {span=>
-            // If IS_IN_BOUNDS is true, this evaluates to 0.
-            //
-            // If IS_IN_BOUNDS is false, this evaluates to `0 - 1` which
-            // triggers a compile error on underflow when referenced below. The
-            // error is not beautiful but does carry the span of the problematic
-            // enum variant so at least it points to the right line.
-            //
-            //     error: any use of this value will cause an error
-            //       --> bit_field/test.rs:10:5
-            //        |
-            //     10 |     OutOfBounds = 0b111111,
-            //        |     ^^^^^^^^^^^ attempt to subtract with overflow
-            //        |
-            //
-            //     error[E0080]: erroneous constant used
-            //      --> bit_field/test.rs:5:1
-            //       |
-            //     5 | #[bitfield]
-            //       | ^^^^^^^^^^^ referenced constant has errors
-            //
-            const ASSERT: u64 = 0 - !IS_IN_BOUNDS as u64;
-        };
-
-        quote! {
-            const #variant: u64 = {
-                const IS_IN_BOUNDS: bool = (#ident::#variant as u64) < #upper_bound;
-
-                #assertion
-
-                #ident::#variant as u64 + ASSERT
-            };
-        }
-    });
+    let declare_discriminants = get_declare_discriminants_for_enum(bits, ast, data);
 
     let match_discriminants = variants.iter().map(|variant| {
         let variant = &variant.ident;
@@ -176,6 +212,58 @@ fn bitfield_enum_impl(ast: &DeriveInput, data: &DataEnum) -> Result<TokenStream>
     Ok(expanded)
 }
 
+fn get_declare_discriminants_for_enum(
+    bits: u8,
+    ast: &DeriveInput,
+    data: &DataEnum,
+) -> Vec<TokenStream> {
+    let variants = &data.variants;
+    let upper_bound = 2u64.pow(bits as u32);
+    let ident = &ast.ident;
+
+    variants
+        .iter()
+        .map(|variant| {
+            let variant = &variant.ident;
+            let span = variant.span();
+
+            let assertion = quote_spanned! {span=>
+                // If IS_IN_BOUNDS is true, this evaluates to 0.
+                //
+                // If IS_IN_BOUNDS is false, this evaluates to `0 - 1` which
+                // triggers a compile error on underflow when referenced below. The
+                // error is not beautiful but does carry the span of the problematic
+                // enum variant so at least it points to the right line.
+                //
+                //     error: any use of this value will cause an error
+                //       --> bit_field/test.rs:10:5
+                //        |
+                //     10 |     OutOfBounds = 0b111111,
+                //        |     ^^^^^^^^^^^ attempt to subtract with overflow
+                //        |
+                //
+                //     error[E0080]: erroneous constant used
+                //      --> bit_field/test.rs:5:1
+                //       |
+                //     5 | #[bitfield]
+                //       | ^^^^^^^^^^^ referenced constant has errors
+                //
+                const ASSERT: u64 = 0 - !IS_IN_BOUNDS as u64;
+            };
+
+            quote! {
+                const #variant: u64 = {
+                    const IS_IN_BOUNDS: bool = (#ident::#variant as u64) < #upper_bound;
+
+                    #assertion
+
+                    #ident::#variant as u64 + ASSERT
+                };
+            }
+        })
+        .collect()
+}
+
 fn bitfield_struct_impl(ast: &DeriveInput, fields: &FieldsNamed) -> Result<TokenStream> {
     let name = &ast.ident;
     let vis = &ast.vis;
@@ -235,14 +323,9 @@ fn parse_bits_attr(attrs: &[Attribute]) -> Result<Option<LitInt>> {
         if attr.path.is_ident("doc") {
             continue;
         }
-
-        if attr.path.is_ident("bits") {
-            if let Meta::NameValue(name_value) = attr.parse_meta()? {
-                if let Lit::Int(int) = name_value.lit {
-                    expected_bits = Some(int);
-                    continue;
-                }
-            }
+        if let Some(v) = try_parse_bits_attr(attr)? {
+            expected_bits = Some(v);
+            continue;
         }
 
         return Err(Error::new_spanned(attr, "unrecognized attribute"));
@@ -251,6 +334,18 @@ fn parse_bits_attr(attrs: &[Attribute]) -> Result<Option<LitInt>> {
     Ok(expected_bits)
 }
 
+// This function will return None if the attribute is not #[bits = *].
+fn try_parse_bits_attr(attr: &Attribute) -> Result<Option<LitInt>> {
+    if attr.path.is_ident("bits") {
+        if let Meta::NameValue(name_value) = attr.parse_meta()? {
+            if let Lit::Int(int) = name_value.lit {
+                return Ok(Some(int));
+            }
+        }
+    }
+    Ok(None)
+}
+
 fn get_struct_def(vis: &Visibility, name: &Ident, fields: &[FieldSpec]) -> TokenStream {
     let mut field_types = Vec::new();
     for spec in fields {
diff --git a/bit_field/src/lib.rs b/bit_field/src/lib.rs
index cafee38..64a309b 100644
--- a/bit_field/src/lib.rs
+++ b/bit_field/src/lib.rs
@@ -101,8 +101,36 @@
 //! }
 //! ```
 //!
-//! Finally, fields may be of user-defined enum types where the enum has a
-//! number of variants which is a power of 2 and the discriminant values
+//! Finally, fields may be of user-defined enum types. The enum must satisfy one of the following
+//! requirements.
+//!
+//! The enum has `#[bits = N]` attributes with it. `N` will be the width of the field. The getter
+//! function of this enum field will return `Result<EnumType, u64>`. Raw value that does not match
+//! any variant will result in an `Err(u64)`.
+//!
+//! ```
+//! extern crate bit_field;
+//!
+//! use bit_field::*;
+//!
+//! #[bitfield]
+//! #[bits = 2]
+//! #[derive(Debug, PartialEq)]
+//! enum TwoBits {
+//!     Zero = 0b00,
+//!     One = 0b01,
+//!     Three = 0b11,
+//! }
+//!
+//! #[bitfield]
+//! struct Struct {
+//!     prefix: BitField1,
+//!     two_bits: TwoBits,
+//!     suffix: BitField5,
+//! }
+//! ```
+//!
+//! The enum has a number of variants which is a power of 2 and the discriminant values
 //! (explicit or implicit) are 0 through (2^n)-1. In this case the generated
 //! getter and setter are defined in terms of the given enum type.
 //!
@@ -235,12 +263,43 @@
 //! }
 //! ```
 
+use std::fmt::{self, Display};
+
 #[allow(unused_imports)]
 #[macro_use]
 extern crate bit_field_derive;
 
 pub use bit_field_derive::bitfield;
 
+/// Error type for bit field get.
+#[derive(Debug)]
+pub struct Error {
+    type_name: &'static str,
+    val: u64,
+}
+
+impl Error {
+    pub fn new(type_name: &'static str, val: u64) -> Error {
+        Error { type_name, val }
+    }
+
+    pub fn raw_val(&self) -> u64 {
+        self.val
+    }
+}
+
+impl Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "enum field type {} has a bad value {}",
+            self.type_name, self.val
+        )
+    }
+}
+
+impl std::error::Error for Error {}
+
 #[doc(hidden)]
 pub trait BitFieldSpecifier {
     // Width of this field in bits.
diff --git a/bit_field/tests/test_enum.rs b/bit_field/tests/test_enum.rs
index d926d9c..8746b06 100644
--- a/bit_field/tests/test_enum.rs
+++ b/bit_field/tests/test_enum.rs
@@ -12,10 +12,21 @@ enum TwoBits {
 }
 
 #[bitfield]
+#[bits = 3]
+#[derive(Debug, PartialEq)]
+enum ThreeBits {
+    Zero = 0b00,
+    One = 0b01,
+    Two = 0b10,
+    Three = 0b111,
+}
+
+#[bitfield]
 struct Struct {
     prefix: BitField1,
     two_bits: TwoBits,
-    suffix: BitField5,
+    three_bits: ThreeBits,
+    suffix: BitField2,
 }
 
 #[test]
@@ -30,5 +41,10 @@ fn test_enum() {
 
     s.set(0, 8, 0b_1010_1010);
     //                   ^^ TwoBits
+    //               ^^_^ Three Bits.
     assert_eq!(s.get_two_bits(), TwoBits::One);
+    assert_eq!(s.get_three_bits().unwrap_err().raw_val(), 0b101);
+
+    s.set_three_bits(ThreeBits::Two);
+    assert_eq!(s.get(0, 8), 0b_1001_0010);
 }