From 00a858d90a880d5a30563d14e58fd64e2f3e241e Mon Sep 17 00:00:00 2001
From: Morten Piibeleht <morten.piibeleht@gmail.com>
Date: Thu, 27 Jun 2019 21:57:27 +1200
Subject: [PATCH] Fix Distributed.head_and_tail (JuliaLang/julia#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.

(cherry picked from commit a7aa3070d8ca35b0a7b32fe41d318e9223c330c8)
---
 src/pmap.jl              | 11 ++++-------
 test/distributed_exec.jl | 26 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/pmap.jl b/src/pmap.jl
index d1839e8..e6d8a20 100644
--- a/src/pmap.jl
+++ b/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
 
 """
diff --git a/test/distributed_exec.jl b/test/distributed_exec.jl
index 62cc8ee..2c7de5d 100644
--- a/test/distributed_exec.jl
+++ b/test/distributed_exec.jl
@@ -1547,6 +1547,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())