Add ability to limit the number of generated atoms

src/proper.erl

@@ -297,6 +297,9 @@
 %%% <dt>`{stop_nodes, boolean()}'</dt>
 %%% <dd> Specifies whether parallel PropEr should stop the nodes after running a property
 %%% or not. Defaults to true.</dd>
+%%% <dt>`{atom_limit, infinity | <Positive_integer>}'</dt>
+%%% <dd>The overall number of atoms generated by {@link proper_types:atom/0} will not
+%%% exceed the given number, per size. Default is infinity.</dd>
 %%% </dl>
 %%%
 %%% == Spec testing ==
@@ -539,6 +542,7 @@
 		  | 'quiet'
 		  | 'verbose'
 		  | pos_integer()
+		  | {'atom_limit', 'infinity' | pos_integer()}
 		  | {'constraint_tries',pos_integer()}
 		  | {'false_positive_mfas',false_positive_mfas()}
 		  | {'max_shrinks',non_neg_integer()}
@@ -568,6 +572,7 @@
 	       noshrink         = false           :: boolean(),
 	       constraint_tries = 50              :: pos_integer(),
 	       expect_fail      = false           :: boolean(),
+	       atom_limit       = infinity        :: pos_integer() | 'infinity',
 	       any_type	        :: {'type', proper_types:type()} | 'undefined',
 	       spec_timeout     = infinity        :: timeout(),
 	       skip_mfas        = []              :: [mfa()],
@@ -744,7 +749,8 @@ global_state_init_size_seed(Size, Seed) ->
 -spec global_state_init(opts()) -> 'ok'.
 global_state_init(#opts{start_size = StartSize, constraint_tries = CTries,
 			search_strategy = Strategy, search_steps = SearchSteps,
-			any_type = AnyType, seed = Seed, numworkers = NumWorkers} = Opts) ->
+			any_type = AnyType, seed = Seed, numworkers = NumWorkers,
+			atom_limit = AtomLimit} = Opts) ->
     clean_garbage(),
     put('$size', StartSize - 1),
     put('$left', 0),
@@ -753,6 +759,7 @@ global_state_init(#opts{start_size = StartSize, constraint_tries = CTries,
     grow_size(Opts),
     put('$constraint_tries', CTries),
     put('$any_type', AnyType),
+    put('$atom_limit', AtomLimit),
     put('$property_id', erlang:unique_integer()),
     put('$proper_test_incr', NumWorkers),
     {_, _, _} = Seed, % just an assertion
@@ -773,6 +780,7 @@ global_state_erase() ->
     proper_typeserver:stop(),
     proper_arith:rand_stop(),
     erase('$any_type'),
+    erase('$atom_limit'),
     erase('$constraint_tries'),
     erase('$left'),
     erase('$size'),

src/proper_gen.erl

@@ -406,6 +406,40 @@ float_gen(_Size, Low, High) ->
 %% We make sure we never clash with internal atoms by checking that the first
 %% character is not '$'.
 atom_gen(Size) ->
+    atom_gen(get('$atom_limit'), Size).
+
+atom_gen(infinity, Size) ->
+    atom_gen1(Size);
+atom_gen(Limit, Size) ->
+    ?LAZY(
+	?LET(Atom,
+	     try persistent_term:get({'$proper', ?MODULE, ?FUNCTION_NAME, Size}) of
+		 {N, Atoms} when N =< Limit ->
+		     proper_types:oneof(Atoms);
+		 {N, CurAtoms} ->
+		     ?LET(MoreAtoms,
+			  proper_types:vector(Limit-N, atom_gen1(Size)),
+			  begin
+			      Atoms = CurAtoms ++ MoreAtoms,
+			      persistent_term:put({'$proper', ?MODULE, ?FUNCTION_NAME, Size},
+						  {Limit, Atoms}),
+			      proper_types:oneof(Atoms)
+			  end
+		     )
+	     catch
+		 error:badarg ->
+		     ?LET(Atoms,
+			  proper_types:vector(Limit, atom_gen1(Size)),
+			  begin
+			      persistent_term:put({'$proper', ?MODULE, ?FUNCTION_NAME, Size},
+						  {Limit, Atoms}),
+			      proper_types:oneof(Atoms)
+			  end
+		     )
+	     end,
+	Atom)).
+
+atom_gen1(Size) ->
     ?LET(Str,
 	 ?SUCHTHAT(X,
 		   proper_types:resize(Size,