From 6962d3f32c28f679d7e2df75afbca510f40c952e Mon Sep 17 00:00:00 2001
From: Matthew Russo <matthew@matthewclayrusso.com>
Date: Sun, 22 Sep 2024 19:28:58 -0700
Subject: [PATCH] Fix bounds generation for generics in derive(Arbitrary)
 (#511)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The implementation of UseTracker expects that iteration over items of
used_map gives items in insertion order. However, the order of BTreeSet
is based on Ord, not insertion.

I use a Vec of tuples. An alternative would be indexmap crate, but since
the maps are supposed to be very small, it is probably not worthy. There
already was a mention about the crate in the comments worrying about
compile times.

Co-authored-by: Petr Nevyhoštěný <petr.nevyhosteny@gmail.com>
---
 proptest-derive/src/use_tracking.rs  | 27 +++++++++++++++----------
 proptest-derive/tests/use_tracker.rs | 30 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 10 deletions(-)
 create mode 100644 proptest-derive/tests/use_tracker.rs

diff --git a/proptest-derive/src/use_tracking.rs b/proptest-derive/src/use_tracking.rs
index ee1685bc..c81c64a9 100644
--- a/proptest-derive/src/use_tracking.rs
+++ b/proptest-derive/src/use_tracking.rs
@@ -10,11 +10,8 @@
 //! track uses of type variables that need `Arbitrary` bounds in our
 //! impls.
 
-// Perhaps ordermap would be better, but our maps are so small that we care
-// much more about the increased compile times incured by including ordermap.
-// We need to preserve insertion order in any case, so HashMap is not useful.
 use std::borrow::Borrow;
-use std::collections::{BTreeMap, HashSet};
+use std::collections::HashSet;
 
 use syn;
 
@@ -30,10 +27,13 @@ use crate::util;
 /// or similar and thus needs an `Arbitrary` bound added to them.
 pub struct UseTracker {
     /// Tracks 'usage' of a type variable name.
-    /// Allocation of this map will happen at once and no further
+    /// Allocation of this "map" will happen at once and no further
     /// allocation will happen after that. Only potential updates
     /// will happen after initial allocation.
-    used_map: BTreeMap<syn::Ident, bool>,
+    /// We need to preserve insertion order, thus using Vec instead of BTreeMap
+    /// or HashMap. A potential alternative would be indexmap crate, but our
+    /// maps are so small that it would not bring any significant benefit.
+    used_map: Vec<(syn::Ident, bool)>,
     /// Extra types to bound by `Arbitrary` in the `where` clause.
     where_types: HashSet<syn::Type>,
     /// The generics that we are doing this for.
@@ -76,7 +76,11 @@ impl UseTracker {
     /// a type variable and this call has no effect.
     fn use_tyvar(&mut self, tyvar: impl Borrow<syn::Ident>) {
         if self.track {
-            if let Some(used) = self.used_map.get_mut(tyvar.borrow()) {
+            if let Some(used) = self
+                .used_map
+                .iter_mut()
+                .find_map(|(ty, used)| (ty == tyvar.borrow()).then(|| used))
+            {
                 *used = true;
             }
         }
@@ -84,7 +88,7 @@ impl UseTracker {
 
     /// Returns true iff the type variable given exists.
     fn has_tyvar(&self, ty_var: impl Borrow<syn::Ident>) -> bool {
-        self.used_map.contains_key(ty_var.borrow())
+        self.used_map.iter().any(|(ty, _)| ty == ty_var.borrow())
     }
 
     /// Mark the type as used.
@@ -101,8 +105,11 @@ impl UseTracker {
         for_not: Option<syn::TypeParamBound>,
     ) -> DeriveResult<()> {
         {
-            let mut iter =
-                self.used_map.values().zip(self.generics.type_params_mut());
+            let mut iter = self
+                .used_map
+                .iter()
+                .map(|(_, used)| used)
+                .zip(self.generics.type_params_mut());
             if let Some(for_not) = for_not {
                 iter.try_for_each(|(&used, tv)| {
                     // Steal the attributes:
diff --git a/proptest-derive/tests/use_tracker.rs b/proptest-derive/tests/use_tracker.rs
new file mode 100644
index 00000000..d41e6da9
--- /dev/null
+++ b/proptest-derive/tests/use_tracker.rs
@@ -0,0 +1,30 @@
+// Copyright 2018 The proptest developers
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use std::marker::PhantomData;
+
+use proptest::prelude::Arbitrary;
+use proptest_derive::Arbitrary;
+
+#[derive(Debug)]
+struct NotArbitrary;
+
+#[derive(Debug, Arbitrary)]
+// Generic types are not in alphabetical order on purpose.
+struct Foo<V, T, U> {
+    v: V,
+    t: T,
+    u: PhantomData<U>,
+}
+
+#[test]
+fn asserting_arbitrary() {
+    fn assert_arbitrary<T: Arbitrary>() {}
+
+    assert_arbitrary::<Foo<i32, i32, NotArbitrary>>();
+}