ringbuf: add #[count(skip)] & #[count(children)] (#1630)

Currently, the `#[derive(ringbuf::Count)]` attribute is pretty simple:
it generates a single counter for every variant of the annotated `enum`
type. For some of the ringbuf events we would like to generate counters
for, this is sufficient. However, there are a few things that would be
nice to have that we can't currently do with the derive attribute:

- Most ringbuf entry enum types have an "Empty"/"None" variant that's
  used for initializing the ringbuffer, but are never actually recorded
  at runtime. We could save a few bytes by not generating a counter for
  these variants.
- Many ringbuf entry enum variants contain a nested enum variant. such
  as a variant representing an error which contains an error kind enum,
  or, the motivating example, `Log::MgsMessage(MgsMessage)` in
  control-plane-agent. In this case, we would like to be able to
  generate a counter of each MGS message variant, which is not currently
  possible with the current `#[derive(ringbuf::Count)]` attribute.

This commit adds new helper attributes to the derive macro:
`#[count(skip)]` and `#[count(children)]`. The `#[count(skip)]`
attribute can be placed on an enum variant to indicate that a counter
*shouldn't* be generated for it. The `#[count(children)]` attribute can
be placed on an enum variant that has a single field, to indicate that a
_nested_ set of counters should be generated for that variant's inner
field. When the `#[count(children)]` attribute is used, the inner field
must also implement the `Count` trait. If it does, the field on the
counter type for that variant will be the child type's `Count::Counters`
type, rather than an `AtomicU32`, and we will increment the nested
counter. My Humility PR #449 adds nice support for interpreting and
displaying these nested counters.

Author: hawkw

lib/ringbuf/examples/count_children.rs

@@ -0,0 +1,43 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#[derive(ringbuf::Count, Debug, Copy, Clone, PartialEq, Eq)]
+pub enum MainEvent {
+    #[count(skip)]
+    None,
+
+    /// This generates a counter for each variant of `Person`.
+    #[count(children)]
+    Person(Person),
+
+    /// This generates a counter for each variant of `Place`.
+    #[count(children)]
+    Place(Place),
+
+    /// This generates a single counter for `Number`.
+    Number(u32),
+}
+
+#[derive(ringbuf::Count, Debug, Copy, Clone, PartialEq, Eq)]
+pub enum Person {
+    Cliff,
+    Matt,
+    Laura,
+    Bryan,
+    John,
+    Steve,
+}
+
+#[derive(ringbuf::Count, Debug, Copy, Clone, PartialEq, Eq)]
+pub enum Place {
+    Emeryville,
+    Cambridge,
+    Austin,
+    SanFrancisco,
+    ZipCode(u16),
+}
+
+ringbuf::counted_ringbuf!(MainEvent, 16, MainEvent::None);
+
+fn main() {}

lib/ringbuf/examples/counts.rs

@@ -10,10 +10,13 @@ use ringbuf::*;
 
 #[derive(ringbuf::Count, Debug, Copy, Clone, PartialEq, Eq)]
 pub enum Event {
+    #[count(skip)]
     NothingHappened,
     SomethingHappened,
     SomethingElse(u32),
-    SecretThirdThing { secret: () },
+    SecretThirdThing {
+        secret: (),
+    },
 }
 
 counted_ringbuf!(Event, 16, Event::NothingHappened);

lib/ringbuf/macros/src/lib.rs

@@ -6,13 +6,32 @@ extern crate proc_macro;
 use proc_macro::TokenStream;
 use proc_macro2::{Ident, Span};
 use quote::{quote, ToTokens};
-use syn::{parse_macro_input, DeriveInput};
-
+use syn::{parse_macro_input, DeriveInput, parse::{Parse, ParseStream}};
+ 
 /// Derives an implementation of the `ringbuf::Count` trait for the annotated
 /// `enum` type.
 ///
 /// Note that this macro can currently only be used on `enum` types.
-#[proc_macro_derive(Count)]
+///
+/// # Variant Attributes
+///
+/// The following attributes may be added on one or more of the variants of the
+/// `enum` type deriving `Count`:
+///
+/// - `#[count(skip)]`: Skip counting this variant. Enums used as ringbuf
+///   entries often have a `None` or `Empty` variant which is used to initialize
+///   the ring buffer but not recorded as an entry at runtime. The
+///   `#[count(skip)]` attribute avoids generating a counter for such variants,
+///   reducing the memory used by the counter struct a little bit.
+///
+/// - `#[count(children)]`: Count variants of a nested enum. Typically, when a
+///   variant of a type deriving `Count` has fields, all instances of that
+///   variant increment the same counter, regardless of the value of the fields.
+///   When a variant has a single field of a type which also implements the
+///   `Count` trait, however, the `#[count(children)]` attribute can be used to
+///   generate an instance of the field type's counter struct, and implement
+///   those counters instead.
+#[proc_macro_derive(Count, attributes(count))]
 pub fn derive_count(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
     match gen_count_impl(input) {
@@ -23,6 +42,7 @@ pub fn derive_count(input: TokenStream) -> TokenStream {
 
 fn gen_count_impl(input: DeriveInput) -> Result<impl ToTokens, syn::Error> {
     let name = &input.ident;
+    let vis = &input.vis;
     let data_enum = match input.data {
         syn::Data::Enum(ref data_enum) => data_enum,
         _ => {
@@ -34,35 +54,98 @@ fn gen_count_impl(input: DeriveInput) -> Result<impl ToTokens, syn::Error> {
     };
     let variants = &data_enum.variants;
     let len = variants.len();
-    let mut variant_names = Vec::with_capacity(len);
+    let mut field_defs = Vec::with_capacity(len);
+    let mut field_inits = Vec::with_capacity(len);
     let mut variant_patterns = Vec::with_capacity(len);
-    for variant in variants {
+    let mut any_skipped = false;
+    'variants: for variant in variants {
         let ident = &variant.ident;
-        variant_patterns.push(match variant.fields {
-            syn::Fields::Unit => quote! { #name::#ident => &counters.#ident },
-            syn::Fields::Named(_) => {
-                quote! { #name::#ident { .. } => &counters.#ident }
+        let mut count_children = None;
+        for attr in &variant.attrs {
+            if !attr.path().is_ident("count") {
+                continue;
             }
-            syn::Fields::Unnamed(_) => {
-                quote! { #name::#ident(..) => &counters.#ident }
+            match attr.parse_args_with(VariantAttr::parse)? {
+                VariantAttr::Skip => {
+                    any_skipped = true;
+                    continue 'variants;
+                },
+                VariantAttr::Children => {
+                    count_children = Some(attr);
+                },
             }
-        });
-        variant_names.push(ident.clone());
-    }
-    let counts_ty = counts_ty(name);
-    let code = quote! {
-        #[doc = concat!(" Ringbuf entry total counts for [`", stringify!(#name), "`].")]
-        #[allow(nonstandard_style)]
-        pub struct #counts_ty {
-            #(
+        }
+
+        if let Some(count_children) = count_children {
+            match variant.fields {
+                syn::Fields::Unit => return Err(syn::Error::new_spanned(
+                    count_children,
+                    "the `count_children` attribute may not be used on unit variants",
+                )),
+                syn::Fields::Named(_) => return Err(syn::Error::new_spanned(
+                    count_children,
+                    "the `count_children` attribute does not currently support variants with named fields",
+                )),
+
+                syn::Fields::Unnamed(_) if variant.fields.len() > 1 => return Err(syn::Error::new_spanned(
+                    count_children,
+                    "the `count_children` attribute does not currently support variants with multiple fields",
+                )),
+                syn::Fields::Unnamed(ref u) => {
+                    let field = u.unnamed.first().unwrap();
+                    let ty = &field.ty;
+                    variant_patterns.push(quote! { #name::#ident(c) => c.count(&counters.#ident) });
+                    field_defs.push(quote! {
+                        #[doc = concat!(
+                            " The total number of times a [`",
+                            stringify!(#name), "::", stringify!(#ident),
+                            "`] entry"
+                        )]
+                        #[doc = " has been recorded by this ringbuf."]
+                        pub #ident: <#ty as ringbuf::Count>::Counters
+                    });
+                    field_inits.push(quote! {
+                        #ident: <#ty as ringbuf::Count>::NEW_COUNTERS
+                    });
+                }
+            }
+        } else {
+            let incr = quote! {
+                counters.#ident.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+            };
+            variant_patterns.push(match variant.fields {
+                syn::Fields::Unit => quote! { #name::#ident => { #incr } },
+                syn::Fields::Named(_) => {
+                    quote! { #name::#ident { .. } => { #incr } }
+                }
+                syn::Fields::Unnamed(_) => {
+                    quote! { #name::#ident(..) => { #incr } }
+                }
+            });
+            field_defs.push(quote! {
                 #[doc = concat!(
                     " The total number of times a [`",
-                    stringify!(#name), "::", stringify!(#variant_names),
+                    stringify!(#name), "::", stringify!(#ident),
                     "`] entry"
                 )]
                 #[doc = " has been recorded by this ringbuf."]
-                pub #variant_names: core::sync::atomic::AtomicU32
-            ),*
+                pub #ident: core::sync::atomic::AtomicU32
+            });
+            field_inits.push(quote! { #ident: core::sync::atomic::AtomicU32::new(0) });
+        }
+    }
+
+    // If we skipped any variants, generate a catchall case.
+    if any_skipped {
+        variant_patterns.push(quote! { _ => {} });
+    }
+
+    let counts_ty = counts_ty(name);
+    let code = quote! {
+        #[doc = concat!(" Ringbuf entry total counts for [`", stringify!(#name), "`].")]
+        #[allow(nonstandard_style)]
+        #vis struct #counts_ty {
+            #(#field_defs),*
         }
 
         #[automatically_derived]
@@ -77,23 +160,48 @@ fn gen_count_impl(input: DeriveInput) -> Result<impl ToTokens, syn::Error> {
             // Lint...
             #[allow(clippy::declare_interior_mutable_const)]
             const NEW_COUNTERS: #counts_ty = #counts_ty {
-                #(#variant_names: core::sync::atomic::AtomicU32::new(0)),*
+                #(#field_inits),*
             };
 
             fn count(&self, counters: &Self::Counters) {
                 #[cfg(all(target_arch = "arm", armv6m))]
                 use ringbuf::rmv6m_atomic_hack::AtomicU32Ext;
 
-                let counter = match self {
+                match self {
                     #(#variant_patterns),*
                 };
-                counter.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
             }
         }
     };
     Ok(code)
 }
 
+#[derive(Copy, Clone, PartialEq, Eq)]
+enum VariantAttr {
+    Skip,
+    Children,
+}
+
+impl Parse for VariantAttr {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let ident = input.fork().parse::<syn::Ident>()?;
+        if ident == "skip" {
+            // consume the token
+            let _: syn::Ident = input.parse()?;
+            Ok(VariantAttr::Skip)
+        } else if ident == "children" {
+            let _: syn::Ident = input.parse()?;
+            Ok(VariantAttr::Children)
+        } else {
+            Err(syn::Error::new(
+                ident.span(),
+                "unrecognized count option, expected `skip` or `children`",
+            ))
+        }
+    }
+}
+
+
 fn counts_ty(ident: &Ident) -> Ident {
     Ident::new(&format!("{ident}Counts"), Span::call_site())
 }

task/control-plane-agent/src/main.rs

@@ -46,7 +46,7 @@ task_slot!(NET, net);
 task_slot!(SYS, sys);
 
 #[allow(dead_code)] // Not all cases are used by all variants
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy, PartialEq, ringbuf::Count)]
 enum Log {
     Empty,
     BarcodeParseError(BarcodeParseError),
@@ -153,9 +153,9 @@ enum MgsMessage {
     ReadRotPage,
 }
 
-ringbuf!(Log, 16, Log::Empty);
+counted_ringbuf!(Log, 16, Log::Empty);
 
-#[derive(Copy, Clone, Debug, PartialEq)]
+#[derive(Copy, Clone, Debug, PartialEq, ringbuf::Count)]
 enum CriticalEvent {
     Empty,
     /// We have received a network request to change power states. This record
@@ -171,7 +171,7 @@ enum CriticalEvent {
 
 // This ringbuf exists to record critical events _only_ and thus not get
 // overwritten by chatter in the debug/trace-style messages.
-ringbuf!(CRITICAL, CriticalEvent, 16, CriticalEvent::Empty);
+counted_ringbuf!(CRITICAL, CriticalEvent, 16, CriticalEvent::Empty);
 
 const SOCKET: SocketName = SocketName::control_plane_agent;