Fix Distributed.head_and_tail (#32431)

The last return statement is currently referring to a non-existent s
variable. Also fix the related doctest, even though it does not get
tested at the moment.

Author: mortenpi

src/pmap.jl

@@ -238,14 +238,11 @@ Return `head`: the first `n` elements of `c`;
 and `tail`: an iterator over the remaining elements.
 
 ```jldoctest
-julia> a = 1:10
-1:10
-
-julia> b, c = Base.head_and_tail(a, 3)
-([1,2,3],Base.Iterators.Rest{UnitRange{Int64},Int64}(1:10,4))
+julia> b, c = Distributed.head_and_tail(1:10, 3)
+([1, 2, 3], Base.Iterators.Rest{UnitRange{Int64},Int64}(1:10, 3))
 
 julia> collect(c)
-7-element Array{Any,1}:
+7-element Array{Int64,1}:
   4
   5
   6
@@ -268,7 +265,7 @@ function head_and_tail(c, n)
         i += 1
         head[i] = y[1]
     end
-    return head, Iterators.rest(c, s)
+    return head, Iterators.rest(c, y[2])
 end
 
 """

test/distributed_exec.jl

@@ -1626,6 +1626,32 @@ let code = """
     @test success(`$(Base.julia_cmd()) --startup-file=no -e $code`)
 end
 
+# PR 32431: tests for internal Distributed.head_and_tail
+let (h, t) = Distributed.head_and_tail(1:10, 3)
+    @test h == 1:3
+    @test collect(t) == 4:10
+end
+let (h, t) = Distributed.head_and_tail(1:10, 0)
+    @test h == []
+    @test collect(t) == 1:10
+end
+let (h, t) = Distributed.head_and_tail(1:3, 5)
+    @test h == 1:3
+    @test collect(t) == []
+end
+let (h, t) = Distributed.head_and_tail(1:3, 3)
+    @test h == 1:3
+    @test collect(t) == []
+end
+let (h, t) = Distributed.head_and_tail(Int[], 3)
+    @test h == []
+    @test collect(t) == []
+end
+let (h, t) = Distributed.head_and_tail(Int[], 0)
+    @test h == []
+    @test collect(t) == []
+end
+
 # Run topology tests last after removing all workers, since a given
 # cluster at any time only supports a single topology.
 rmprocs(workers())