Provide a way for sets to zipper with more things

modules/internal/ChapelHashtable.chpl

@@ -144,7 +144,7 @@ module ChapelHashtable {
     const minSizePerTask = dataParMinGranularity;
 
     // We are simply slicing up the table here.  Trying to do something
-    //  more intelligent (like evenly dividing up the full slots, led
+    //  more intelligent (like evenly dividing up the full slots), led
     //  to poor speed ups.
 
     if debugAssocDataPar {
@@ -644,5 +644,103 @@ module ChapelHashtable {
 
       rehash(newSize);
     }
+
+    proc _determineEvenChunks(numChunks: int, numFull: int): [] int {
+      var numFullPerTask = numFull / numChunks;
+      var rem = numFull % numChunks;
+
+      var chunkEnds: [0..#(numChunks-1)] int;
+      var chunkIdx = 0;
+      var curNumFull = 0;
+      for i in 0..#(this.tableSize) {
+        if this.isSlotFull(i) {
+          curNumFull += 1;
+        }
+        /* If we've found the last full slot or we've filled our list of chunk
+           end points, exit the loop */
+        if (curNumFull == numFull) ||
+          (chunkIdx >= chunkEnds.size) {
+          break;
+        }
+
+        // Allows us to evenly distribute when the chunks don't divide the
+        // total number of elements evenly.
+        if chunkIdx < rem {
+          if (curNumFull == (numFullPerTask * (chunkIdx+1)) + (chunkIdx+1)) {
+            chunkEnds[chunkIdx] = i;
+            chunkIdx += 1;
+          }
+        } else {
+          if (curNumFull == (numFullPerTask * (chunkIdx+1)) + rem) {
+            chunkEnds[chunkIdx] = i;
+            chunkIdx += 1;
+          }
+        }
+      }
+      return chunkEnds;
+    }
+
+    iter _evenSlots(size: int, param tag) where tag == iterKind.leader {
+      var numChunks = _computeNumChunks(size);
+      var numFullPerTask = size / numChunks;
+      var rem = size % numChunks;
+
+      forall i in 0..#numChunks {
+        if (i < rem) {
+          yield (((numFullPerTask*i)+i)..<((numFullPerTask*(i+1))+(i+1)),);
+        } else {
+          yield (((numFullPerTask*i)+rem)..<((numFullPerTask*(i+1))+rem),);
+        }
+      }
+    }
+
+    proc _guessSlots(followThis) {
+      var iterSpace: range;
+      var intendedSpace = followThis;
+      var chunkSize = intendedSpace.size;
+      var numChunksGuess = this.tableNumFullSlots / chunkSize;
+      var intendedChunkGuess = intendedSpace.low / chunkSize;
+
+      if (intendedChunkGuess >= numChunksGuess) {
+        __primitive("chpl_error",
+                    "zippered iterations have non-equal lengths".c_str());
+      }
+
+      // Determine corresponding chunk to use.
+      var chunkEnds = _determineEvenChunks(numChunksGuess,
+                                           this.tableNumFullSlots);
+      if (intendedChunkGuess == 0) {
+        iterSpace = 0..chunkEnds[intendedChunkGuess];
+      } else if (intendedChunkGuess == chunkEnds.size) {
+        iterSpace = (chunkEnds[intendedChunkGuess-1]+1)..(this.tableSize - 1);
+      } else {
+        iterSpace = (chunkEnds[intendedChunkGuess-1]+1)..chunkEnds[intendedChunkGuess];
+      }
+      return iterSpace;
+    }
+
+    iter _evenSlots(size: int, followThis, param tag) const ref
+      where tag == iterKind.follower {
+      use Types;
+
+      if (size != this.tableNumFullSlots) {
+        __primitive("chpl_error",
+                    "zippered iterations have non-equal lengths".c_str());
+      }
+
+      var iterSpace: range;
+
+      if (isTuple(followThis)) {
+        iterSpace = _guessSlots(followThis(0));
+      } else {
+        iterSpace = _guessSlots(followThis);
+      }
+
+      foreach slot in iterSpace {
+        if this.isSlotFull(slot) then {
+          yield this.table[slot].key;
+        }
+      }
+    }
   }
 }

modules/standard/Set.chpl

@@ -576,6 +576,71 @@ module Set {
         if _htb.isSlotFull(idx) then yield _htb.table[idx].key;
     }
 
+    /*
+      Iterate over the elements of this set. Yields constant references
+      that cannot be modified.
+
+      .. warning::
+
+        Modifying this set while iterating over it may invalidate the
+        references returned by an iterator and is considered undefined
+        behavior.
+
+      :yields: A constant reference to an element in this set.
+
+      :arg size: the number of elements in the leader, when participating in
+                 zippered parallel iteration.  Defaults to the number of
+                 elements in the set.  For standalone and serial iteration, this
+                 must match the number of elements in the set.
+    */
+    @unstable("The contents iterator is unstable because it is new")
+    iter const contents(size: int = this.size) const ref {
+      if (size != _htb.tableNumFullSlots) {
+        iterHalt("serial 'contents' iterator needs to be provided the set's " +
+                 "size as its argument");
+      }
+
+      foreach idx in 0..#_htb.tableSize {
+        if _htb.isSlotFull(idx) then yield _htb.table[idx].key;
+      }
+    }
+
+    @chpldoc.nodoc
+    iter const contents(size: int = this.size, param tag) const ref
+      where tag == iterKind.standalone {
+      if (size != _htb.tableNumFullSlots) {
+        iterHalt("standalone 'contents' iterator needs to be provided the " +
+                 "set's size as its argument");
+      }
+
+      var space = 0..#_htb.tableSize;
+      foreach idx in space.these(tag) do
+        if _htb.isSlotFull(idx) then yield _htb.table[idx].key;
+    }
+
+    @chpldoc.nodoc
+    iter const contents(size: int = this.size, param tag)
+      where tag == iterKind.leader {
+      if (size != _htb.tableNumFullSlots) {
+        iterHalt("leader 'contents' iterator needs to be provided the " +
+                 "set's size as its argument (follower should be provided the" +
+                 " size of the leader)");
+      }
+
+      for followThis in _htb._evenSlots(size, tag) {
+        yield followThis;
+      }
+    }
+
+    @chpldoc.nodoc
+    iter const contents(size: int = this.size, param tag,
+                        followThis) const ref
+    where tag == iterKind.follower {
+      foreach val in _htb._evenSlots(size, followThis, tag) {
+        yield val;
+      }
+    }
+
     @chpldoc.nodoc
     proc const _defaultWriteHelper(ch: fileWriter) throws {
       on this {

test/library/standard/Set/these/setZipDifferentSize.chpl

@@ -7,7 +7,7 @@ proc test() {
   var s2 = s1;
   s2.add(9);
 
-  forall (x, y) in zip(s1, s2) do
+  forall (x, y) in zip(s1.contents(s1.size), s2.contents(s1.size)) do
     var tmp = x + y;
 
   return;

test/library/standard/Set/these/setZipDifferentSize2.chpl

@@ -0,0 +1,16 @@
+use Set;
+
+proc test() {
+  var s1: set(int);
+  for i in 1..8 do s1.add(i);
+
+  var s2 = s1;
+  s2.add(9);
+
+  forall (x, y) in zip(s2.contents(s2.size), s1.contents(s2.size)) do
+    var tmp = x + y;
+
+  return;
+}
+test();
+

test/library/standard/Set/these/setZipDifferentSize2.good

@@ -0,0 +1 @@
+setZipDifferentSize2.chpl:10: error: zippered iterations have non-equal lengths

test/library/standard/Set/these/setZipDifferentSize2.skipif

@@ -0,0 +1,2 @@
+# Follower checks are skipped when compiling with --fast
+COMPOPTS <= --fast

test/library/standard/Set/zippering/arrayShorter.chpl

@@ -0,0 +1,20 @@
+use Set;
+
+config const verbose = false;
+
+var LocalSet= new set(int,parSafe = true);
+LocalSet.add(1);
+LocalSet.add(2);
+LocalSet.add(3);
+LocalSet.add(4);
+LocalSet.add(5);
+
+var A: [0..3] int;
+writeln(A.size, " ", LocalSet.size);
+
+// Expect this to fail to zip, the set is longer than the array
+forall (a,b) in zip(A,LocalSet.contents(A.size)) {
+    a=b;
+    if verbose then writeln(b);
+}
+writeln(A);

test/library/standard/Set/zippering/arrayShorter.good

@@ -0,0 +1,2 @@
+4 5
+arrayShorter.chpl:16: error: zippered iterations have non-equal lengths

test/library/standard/Set/zippering/arrayShorter2.chpl

@@ -0,0 +1,20 @@
+use Set;
+
+config const verbose = false;
+
+var LocalSet= new set(int,parSafe = true);
+LocalSet.add(1);
+LocalSet.add(2);
+LocalSet.add(3);
+LocalSet.add(4);
+LocalSet.add(5);
+
+var A: [0..3] int;
+writeln(A.size, " ", LocalSet.size);
+
+// Expect this to fail to zip, the set is longer than the array
+forall (b, a) in zip(LocalSet.contents(LocalSet.size), A) {
+    a=b;
+    if verbose then writeln(b);
+}
+writeln(A);

test/library/standard/Set/zippering/arrayShorter2.good

@@ -0,0 +1,2 @@
+4 5
+arrayShorter2.chpl:16: error: halt reached - size mismatch in zippered iteration (dimension 0)

test/library/standard/Set/zippering/badContentsCall.chpl

@@ -0,0 +1,23 @@
+// Test to verify that we yell when sending in a bad length
+
+use Set;
+
+config const verbose = false;
+
+var LocalSet= new set(int,parSafe = true);
+LocalSet.add(1);
+LocalSet.add(2);
+LocalSet.add(3);
+LocalSet.add(4);
+LocalSet.add(5);
+
+var A: [0..3] int;
+writeln(A.size, " ", LocalSet.size);
+
+// Expect this to fail with our error message, the leader needs to be sent its
+// own size (I know, I know, I'm open to alternative implementations)
+forall (b, a) in zip(LocalSet.contents(A.size), A) {
+    a=b;
+    if verbose then writeln(b);
+}
+writeln(A);

test/library/standard/Set/zippering/badContentsCall.good

@@ -0,0 +1,2 @@
+4 5
+badContentsCall.chpl:19: error: halt reached - leader 'contents' iterator needs to be provided the set's size as its argument (follower should be provided the size of the leader)

test/library/standard/Set/zippering/sameLength.chpl

@@ -0,0 +1,13 @@
+use Set;
+
+var s1, s2: set(int);
+
+for i in 1..10 { s1.add(i); s2.add(i+10); }
+
+writeln("s1 = ", s1);
+writeln("s2 = ", s2);
+
+forall (i,j) in zip(s1.contents(s1.size), s2.contents(s1.size)) do
+  writeln(i, " ", j);
+
+writeln("all done");

test/library/standard/Set/zippering/sameLength.good

@@ -0,0 +1,13 @@
+all done
+s1 = {5, 7, 8, 1, 4, 6, 2, 10, 9, 3}
+s2 = {16, 20, 12, 17, 19, 13, 15, 11, 18, 14}
+1 17
+2 15
+3 14
+4 19
+5 16
+6 13
+7 20
+8 12
+9 18
+10 11

test/library/standard/Set/zippering/sameLength.prediff

@@ -0,0 +1,3 @@
+#!/bin/sh
+sort -n $1.exec.out.tmp > $1.prediff.tmp
+mv $1.prediff.tmp $1.exec.out.tmp

test/library/standard/Set/zippering/serialVersion.chpl

@@ -0,0 +1,19 @@
+use Set;
+
+var s = new set(string);
+
+s.add("my");
+s.add("great");
+s.add("aunt");
+s.add("sally");
+s.add("told");
+s.add("me");
+s.add("to");
+s.add("sleep");
+s.add("in");
+s.add("tomorrow");
+
+// Check the behavior of the contents iterator when used in a serial context
+for val in s.contents() {
+  writeln(val);
+}

test/library/standard/Set/zippering/serialVersion.good

@@ -0,0 +1,10 @@
+great
+in
+tomorrow
+sleep
+aunt
+me
+sally
+my
+told
+to

test/library/standard/Set/zippering/standaloneVersion.chpl

@@ -0,0 +1,20 @@
+use Set;
+
+var s = new set(string);
+
+s.add("my");
+s.add("great");
+s.add("aunt");
+s.add("sally");
+s.add("told");
+s.add("me");
+s.add("to");
+s.add("sleep");
+s.add("in");
+s.add("tomorrow");
+
+// Check the behavior of the contents iterator when used as a standalone
+// iterator
+forall val in s.contents() {
+  writeln(val);
+}

test/library/standard/Set/zippering/standaloneVersion.good

@@ -0,0 +1,10 @@
+aunt
+great
+in
+me
+my
+sally
+sleep
+to
+told
+tomorrow

test/library/standard/Set/zippering/standaloneVersion.prediff

@@ -0,0 +1,3 @@
+#!/bin/sh
+sort -n $1.exec.out.tmp > $1.prediff.tmp
+mv $1.prediff.tmp $1.exec.out.tmp