Make cong safe in the presence of 0 (#50910)

Co-authored-by: Sukera <11753998+Seelengrab@users.noreply.github.com>
JuliaLang · Aug 16, 2023 · 1c536dd · 1c536dd · nanosoldier · Aug 16, 2023
1 parent cea3d16
commit 1c536dd
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/base/partr.jl b/base/partr.jl
@@ -20,7 +20,7 @@ const heaps = [Vector{taskheap}(undef, 0), Vector{taskheap}(undef, 0)]
 const heaps_lock = [SpinLock(), SpinLock()]
 
 
-cong(max::UInt32) = ccall(:jl_rand_ptls, UInt32, (UInt32,), max) + UInt32(1)
+cong(max::UInt32) = iszero(max) ? UInt32(0) : ccall(:jl_rand_ptls, UInt32, (UInt32,), max) + UInt32(1)
 
 
 function multiq_sift_up(heap::taskheap, idx::Int32)

diff --git a/src/julia_internal.h b/src/julia_internal.h
@@ -1220,6 +1220,8 @@ void jl_push_excstack(jl_excstack_t **stack JL_REQUIRE_ROOTED_SLOT JL_ROOTING_AR
 
 STATIC_INLINE uint64_t cong(uint64_t max, uint64_t *seed) JL_NOTSAFEPOINT
 {
+    if (max == 0)
+        return 0;
     uint64_t mask = ~(uint64_t)0;
     --max;
     mask >>= __builtin_clzll(max|1);

diff --git a/test/threads.jl b/test/threads.jl
@@ -327,3 +327,7 @@ end
     @test_throws ArgumentError @macroexpand(@threads 1) # arg isn't an Expr
     @test_throws ArgumentError @macroexpand(@threads if true 1 end) # arg doesn't start with for
 end
+
+@testset "rand_ptls underflow" begin
+    @test Base.Partr.cong(UInt32(0)) == 0
+end