Implement Iterator::size_hint() in some cases.

CHANGELOG.md

@@ -13,7 +13,10 @@
 
 ### Changed
 
-* `derive(Exhaust)` generates smaller iterators for structs or enum variants with exactly one field.
+* `derive(Exhaust)` iterators now implement `size_hint()` exactly in some cases;
+  most notably, fieldless enums. (They still do not implement `ExactSizeIterator`.)
+* `derive(Exhaust)` iterators now take up less memory when a struct, or a variant of an enum,
+  has exactly one field.
 
 ## 0.2.4 (2025-08-25)
 

exhaust-macros/src/common.rs

@@ -152,11 +152,14 @@ impl ExhaustContext {
 
     /// Generate the parts of the trait implementations for the iterator and factory
     /// that do not depend on whether the type being exhausted is an enum or struct.
+    ///
+    /// If `iterator_size_hint_body` is `None`, a `size_hint()` method is not generated.
     pub fn impl_iterator_and_factory_traits(
         &self,
         iterator_next_body: TokenStream2,
         iterator_default_body: TokenStream2,
         iterator_clone_body: TokenStream2,
+        iterator_size_hint_body: Option<TokenStream2>,
     ) -> TokenStream2 {
         let exhaust_crate_path = &self.exhaust_crate_path;
         let helpers = self.helpers();
@@ -187,6 +190,17 @@ impl ExhaustContext {
             },
         };
 
+        let iterator_size_hint_decl = if let Some(iterator_size_hint_body) = iterator_size_hint_body
+        {
+            quote! {
+                fn size_hint(&self) -> (usize, #helpers::Option<usize>) {
+                    #iterator_size_hint_body
+                }
+            }
+        } else {
+            TokenStream2::new()
+        };
+
         quote! {
             impl #impl_generics #helpers::Iterator for #iterator_type_name #ty_generics
             where #augmented_where_predicates {
@@ -196,6 +210,7 @@ impl ExhaustContext {
                     #![allow(unreachable_code)] // an iterator or factory might be uninhabited
                     #iterator_next_body
                 }
+                #iterator_size_hint_decl
             }
 
             impl #impl_generics #helpers::FusedIterator for #iterator_type_name #ty_generics

exhaust-macros/src/fields.rs

@@ -19,6 +19,10 @@ pub(crate) struct ExhaustFields {
     pub field_pats: TokenStream2,
     /// Code to implement advancing the iterator. [`Self::field_pats`] should be in scope.
     pub advance: TokenStream2,
+    /// Code to implement `Iterator::size_hint()`,
+    /// or `None` if this is infeasible.
+    /// [`Self::field_pats`] should be in scope.
+    pub iter_size_hint: Option<TokenStream2>,
 }
 
 /// Given a set of fields to exhaust, generate fields and code for the iterator to
@@ -137,6 +141,10 @@ pub(crate) fn exhaustion_of_fields(
                         #helpers::None => #helpers::None,
                     }
                 },
+                // If we have exactly one field, then our size hint is equal to that field’s.
+                iter_size_hint: Some(quote! {
+                    #helpers::size_hint(#field_iter_var)
+                }),
             };
         }
 
@@ -321,5 +329,9 @@ pub(crate) fn exhaustion_of_fields(
             #( #iter_field_names , )*
         },
         advance,
+        // No size hint because in any situation including carrying, we need to be able to ask
+        // what the size hint of a *newly created* iterator will be, which will require expanding
+        // the `Exhaust` trait itself.
+        iter_size_hint: None,
     }
 }

exhaust-macros/src/lib.rs

@@ -174,6 +174,7 @@ fn derive_exhaust_for_struct(
             cloners,
             field_pats,
             advance,
+            iter_size_hint,
         },
     ) = if s.fields.is_empty() {
         // If there are no fields, then
@@ -200,6 +201,13 @@ fn derive_exhaust_for_struct(
                     #helpers::Some(#factory_ctor_expr)
                 }
             },
+            iter_size_hint: Some(quote! {
+                if *done {
+                    const { (0, #helpers::Some(0)) }
+                } else {
+                    const { (1, #helpers::Some(1)) }
+                }
+            }),
         };
 
         let output_type = ctx.item_type.path()?;
@@ -318,6 +326,12 @@ fn derive_exhaust_for_struct(
             let Self { #field_pats } = self;
             Self { #cloners }
         },
+        iter_size_hint.map(|size_hint_body| {
+            quote! {
+                let Self { #field_pats } = self;
+                #size_hint_body
+            }
+        }),
     );
 
     // Struct that is exposed as the `<Self as Exhaust>::Iter` type.
@@ -452,6 +466,7 @@ fn derive_exhaust_for_enum(
                 cloners: state_fields_clone,
                 field_pats,
                 advance,
+                iter_size_hint: _, // TODO: propagate size hint if remainder is fieldless
             } = if target_variant.fields.is_empty() {
                 // TODO: don't even construct this dummy value (needs refactoring)
                 fields::ExhaustFields {
@@ -463,6 +478,7 @@ fn derive_exhaust_for_enum(
                     advance: quote! {
                         compile_error!("can't happen: fieldless ExhaustFields not used")
                     },
+                    iter_size_hint: None,
                 }
             } else {
                 fields::exhaustion_of_fields(
@@ -556,6 +572,42 @@ fn derive_exhaust_for_enum(
         },
     );
 
+    let iter_size_hint_body = {
+        let mut size_hint_arms: Vec<TokenStream2> =
+            Vec::with_capacity(state_enum_progress_variants.len() + 1);
+        let mut remaining_count_so_far: usize = 0;
+        for (original_enum_variant, progress_variant_name) in
+            izip!(&e.variants, &state_enum_progress_variants).rev()
+        {
+            if original_enum_variant.fields.is_empty() {
+                // If the variant is fieldless, we can predict that it has exactly 1 value to exhaust.
+                remaining_count_so_far += 1;
+                size_hint_arms.push(quote! {
+                    #iter_state_enum_type :: #progress_variant_name =>
+                        (#remaining_count_so_far, #helpers::Some(#remaining_count_so_far))
+                });
+            } else {
+                // If the variant has fields, prediction is difficult.
+                // TODO: We should use at least the size hint from this variant's ExhaustFields
+
+                size_hint_arms.push(quote! {
+                    // Note: We can't even increment by 1 because a field might be uninhabited
+                    // (have zero values).
+                    _ => (#remaining_count_so_far, #helpers::None)
+                });
+                // Stop because any further variants would be wrong
+                break;
+            }
+        }
+
+        Some(quote! {
+            match &self.0 {
+                #iter_state_enum_type :: #done_variant => (0, #helpers::Some(0)),
+                #(#size_hint_arms),*
+            }
+        })
+    };
+
     let impls = ctx.impl_iterator_and_factory_traits(
         quote! {
             'variants: loop {
@@ -573,6 +625,7 @@ fn derive_exhaust_for_enum(
                 #( #state_enum_variant_cloners , )*
             })
         },
+        iter_size_hint_body,
     );
 
     let factory_struct_decl_and_impls = match &ctx.factory_type {
@@ -729,6 +782,7 @@ fn derive_exhaust_for_primitive_tuple(size: u64) -> Result<TokenStream2, syn::Er
         cloners,
         field_pats,
         advance,
+        iter_size_hint,
     } = exhaustion_of_fields(&ctx, &synthetic_fields, None, &ConstructorSyntax::Tuple);
     assert!(
         !state_field_decls.is_empty(),
@@ -748,6 +802,12 @@ fn derive_exhaust_for_primitive_tuple(size: u64) -> Result<TokenStream2, syn::Er
             let Self { #field_pats } = self;
             Self { #cloners }
         },
+        iter_size_hint.map(|size_hint_body| {
+            quote! {
+                let Self { #field_pats } = self;
+                #size_hint_body
+            }
+        }),
     );
 
     let iterator_doc = ctx.iterator_doc();

src/mh_.rs

@@ -3,6 +3,7 @@
 
 pub use core::fmt;
 pub use core::iter::{FusedIterator, Iterator, Peekable};
+pub use core::primitive::usize;
 pub use {Clone, Default, None, Option, Some};
 
 /// Convenience trait-alias for helping the derive macro be simpler and generate simpler code.
@@ -23,6 +24,9 @@ pub fn default<T: Default>() -> T {
 pub fn next<I: Iterator>(iterator: &mut I) -> Option<I::Item> {
     iterator.next()
 }
+pub fn size_hint<I: Iterator>(iterator: &I) -> (usize, Option<usize>) {
+    iterator.size_hint()
+}
 pub fn clone<T: Clone>(original: &T) -> T {
     original.clone()
 }

tests/deriving.rs

@@ -10,6 +10,10 @@ mod helper;
 // Don’t glob import the std prelude, so that we check the macro doesn't depend on it.
 use std::prelude::rust_2021 as p;
 
+fn ex<T: exhaust::Exhaust>() -> exhaust::Iter<T> {
+    T::exhaust()
+}
+
 fn c<T: std::fmt::Debug + exhaust::Exhaust>() -> std::vec::Vec<T>
 where
     <T as exhaust::Exhaust>::Iter: std::fmt::Debug,
@@ -32,6 +36,9 @@ where
 
     helper::assert_size_hint_valid(size_hint, result.len());
 
+    // Check the final size hint is not nonzero
+    helper::assert_size_hint_valid(p::Iterator::size_hint(&iterator), 0);
+
     result
 }
 
@@ -53,12 +60,24 @@ struct UnitStructFis;
 #[test]
 fn struct_unit() {
     std::assert_eq!(c::<UnitStruct>(), std::vec![UnitStruct]);
+
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<UnitStruct>()),
+        (1, p::Some(1)),
+        "size hint should be exact"
+    );
 }
 #[test]
 fn struct_unit_fis() {
     std::assert_eq!(c::<UnitStructFis>(), std::vec![UnitStructFis]);
 
     assert_factory_is_self::<UnitStructFis>();
+
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<UnitStructFis>()),
+        (1, p::Some(1)),
+        "size hint should be exact"
+    );
 }
 
 #[derive(Debug, exhaust::Exhaust, PartialEq)]
@@ -203,7 +222,12 @@ fn struct_uninhabited_generic() {
 
 #[test]
 fn struct_uninhabited_nongeneric() {
-    std::assert_eq!(c::<UninhabitedStruct>(), std::vec![])
+    std::assert_eq!(c::<UninhabitedStruct>(), std::vec![]);
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<UninhabitedStruct>()),
+        (0, p::Some(0)),
+        "size hint should be exact"
+    );
 }
 
 #[derive(Debug, exhaust::Exhaust, PartialEq)]
@@ -212,6 +236,12 @@ enum EmptyEnum {}
 #[test]
 fn enum_empty() {
     std::assert_eq!(c::<EmptyEnum>(), std::vec![]);
+
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<EmptyEnum>()),
+        (0, p::Some(0)),
+        "size hint should be exact"
+    );
 }
 
 #[derive(Debug, exhaust::Exhaust, PartialEq)]
@@ -222,6 +252,11 @@ enum OneValueEnum {
 #[test]
 fn enum_one_value() {
     std::assert_eq!(c::<OneValueEnum>(), std::vec![OneValueEnum::Foo]);
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<OneValueEnum>()),
+        (1, p::Some(1)),
+        "size hint should be exact"
+    );
 }
 
 #[derive(Debug, exhaust::Exhaust, PartialEq)]
@@ -237,6 +272,11 @@ fn enum_fieldless_multi() {
         c::<FieldlessEnum>(),
         std::vec![FieldlessEnum::Foo, FieldlessEnum::Bar, FieldlessEnum::Baz]
     );
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<FieldlessEnum>()),
+        (3, p::Some(3)),
+        "size hint should be exact"
+    );
 }
 
 #[derive(Clone, Debug, exhaust::Exhaust, PartialEq)]
@@ -320,6 +360,10 @@ fn enum_generic() {
             EnumWithGeneric::After,
         ]
     );
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<EnumWithGeneric<'static, bool>>()),
+        (1, p::None),
+    );
 }
 
 #[derive(Debug, exhaust::Exhaust, PartialEq)]
@@ -336,6 +380,10 @@ fn enum_with_uninhabited_nongeneric() {
         c::<EnumWithUninhabited>(),
         [EnumWithUninhabited::Before, EnumWithUninhabited::After]
     );
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<EnumWithUninhabited>()),
+        (1, p::None),
+    );
 }
 #[test]
 fn enum_with_uninhabited_generic() {
@@ -346,6 +394,10 @@ fn enum_with_uninhabited_generic() {
             EnumWithGeneric::After,
         ]
     );
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<EnumWithGeneric<std::convert::Infallible>>()),
+        (1, p::None),
+    );
 }
 
 // Test specifically newtypes (structs with exactly one field).
@@ -359,6 +411,12 @@ fn newtype_struct() {
     // using FieldlessEnum as a non-factory_is_self implementation to use in our generic newtype
     c::<NewtypeStruct<FieldlessEnum>>();
 
+    std::assert_eq!(
+        p::Iterator::size_hint(&ex::<NewtypeStruct<bool>>()),
+        (2, p::Some(2)),
+        "size hint should be exact"
+    );
+
     // Check that the newtype's iterator is not bigger (because it is itself a newtype).
     // (This is not technically guaranteed by Rust unless we add a `repr(transparent)`, though.)
     std::assert_eq!(