Merge addWithoutResize and rehashIfNeeded into one function.

core/src/main/scala/org/apache/spark/util/collection/OpenHashMap.scala

@@ -73,9 +73,8 @@ class OpenHashMap[K : ClassTag, @specialized(Long, Int, Double) V: ClassTag](
       haveNullValue = true
       nullValue = v
     } else {
-      val pos = _keySet.addWithoutResize(k) & OpenHashSet.POSITION_MASK
+      val pos = _keySet.add(k, grow, move) & OpenHashSet.POSITION_MASK
       _values(pos) = v
-      _keySet.rehashIfNeeded(k, grow, move)
       _oldValues = null
     }
   }
@@ -96,11 +95,10 @@ class OpenHashMap[K : ClassTag, @specialized(Long, Int, Double) V: ClassTag](
       }
       nullValue
     } else {
-      val pos = _keySet.addWithoutResize(k)
+      val pos = _keySet.add(k, grow, move)
       if ((pos & OpenHashSet.NONEXISTENCE_MASK) != 0) {
         val newValue = defaultValue
         _values(pos & OpenHashSet.POSITION_MASK) = newValue
-        _keySet.rehashIfNeeded(k, grow, move)
         newValue
       } else {
         _values(pos) = mergeValue(_values(pos))

core/src/main/scala/org/apache/spark/util/collection/OpenHashSet.scala

@@ -106,21 +106,20 @@ class OpenHashSet[@specialized(Long, Int) T: ClassTag](
    * and rehash all elements.
    */
   def add(k: T) {
-    addWithoutResize(k)
-    rehashIfNeeded(k, grow, move)
+    add(k, grow, move)
   }
 
   /**
-   * Add an element to the set. This one differs from add in that it doesn't trigger rehashing.
-   * The caller is responsible for calling rehashIfNeeded.
+   * Add an element to the set. Use the given allocateFunc and moveFunc to rehash the
+   * set if it is overloaded.
    *
    * Use (retval & POSITION_MASK) to get the actual position, and
    * (retval & NONEXISTENCE_MASK) == 0 for prior existence.
    *
    * @return The position where the key is placed, plus the highest order bit is set if the key
    *         does not exists previously.
    */
-  def addWithoutResize(k: T): Int = {
+  def add(k: T, allocateFunc: (Int) => Unit, moveFunc: (Int, Int) => Unit): Int = {
     var pos = hashcode(hasher.hash(k)) & _mask
     var i = 1
     while (true) {
@@ -129,6 +128,11 @@ class OpenHashSet[@specialized(Long, Int) T: ClassTag](
         _data(pos) = k
         _bitset.set(pos)
         _size += 1
+        //Rehash the set if it is overloaded
+        if (_size > _growThreshold) {
+          rehash(k, allocateFunc, moveFunc)
+          pos = getPos(k)
+        }
         return pos | NONEXISTENCE_MASK
       } else if (_data(pos) == k) {
         // Found an existing key.
@@ -144,20 +148,6 @@ class OpenHashSet[@specialized(Long, Int) T: ClassTag](
     INVALID_POS
   }
 
-  /**
-   * Rehash the set if it is overloaded.
-   * @param k A parameter unused in the function, but to force the Scala compiler to specialize
-   *          this method.
-   * @param allocateFunc Callback invoked when we are allocating a new, larger array.
-   * @param moveFunc Callback invoked when we move the key from one position (in the old data array)
-   *                 to a new position (in the new data array).
-   */
-  def rehashIfNeeded(k: T, allocateFunc: (Int) => Unit, moveFunc: (Int, Int) => Unit) {
-    if (_size > _growThreshold) {
-      rehash(k, allocateFunc, moveFunc)
-    }
-  }
-
   /**
    * Return the position of the element in the underlying array, or INVALID_POS if it is not found.
    */

core/src/main/scala/org/apache/spark/util/collection/PrimitiveKeyOpenHashMap.scala

@@ -62,9 +62,8 @@ class PrimitiveKeyOpenHashMap[@specialized(Long, Int) K: ClassTag,
 
   /** Set the value for a key */
   def update(k: K, v: V) {
-    val pos = _keySet.addWithoutResize(k) & OpenHashSet.POSITION_MASK
+    val pos = _keySet.add(k, grow, move) & OpenHashSet.POSITION_MASK
     _values(pos) = v
-    _keySet.rehashIfNeeded(k, grow, move)
     _oldValues = null
   }
 
@@ -75,11 +74,10 @@ class PrimitiveKeyOpenHashMap[@specialized(Long, Int) K: ClassTag,
    * @return the newly updated value.
    */
   def changeValue(k: K, defaultValue: => V, mergeValue: (V) => V): V = {
-    val pos = _keySet.addWithoutResize(k)
+    val pos = _keySet.add(k, grow, move)
     if ((pos & OpenHashSet.NONEXISTENCE_MASK) != 0) {
       val newValue = defaultValue
       _values(pos & OpenHashSet.POSITION_MASK) = newValue
-      _keySet.rehashIfNeeded(k, grow, move)
       newValue
     } else {
       _values(pos) = mergeValue(_values(pos))

graphx/src/main/scala/org/apache/spark/graphx/util/collection/PrimitiveKeyOpenHashMap.scala

@@ -72,23 +72,21 @@ class PrimitiveKeyOpenHashMap[@specialized(Long, Int) K: ClassTag,
 
   /** Set the value for a key */
   def update(k: K, v: V) {
-    val pos = keySet.addWithoutResize(k) & OpenHashSet.POSITION_MASK
+    val pos = keySet.add(k, grow, move) & OpenHashSet.POSITION_MASK
     _values(pos) = v
-    keySet.rehashIfNeeded(k, grow, move)
     _oldValues = null
   }
 
 
   /** Set the value for a key */
   def setMerge(k: K, v: V, mergeF: (V, V) => V) {
-    val pos = keySet.addWithoutResize(k)
+    val pos = keySet.add(k, grow, move)
     val ind = pos & OpenHashSet.POSITION_MASK
     if ((pos & OpenHashSet.NONEXISTENCE_MASK) != 0) { // if first add
       _values(ind) = v
     } else {
       _values(ind) = mergeF(_values(ind), v)
     }
-    keySet.rehashIfNeeded(k, grow, move)
     _oldValues = null
   }
 
@@ -100,11 +98,10 @@ class PrimitiveKeyOpenHashMap[@specialized(Long, Int) K: ClassTag,
    * @return the newly updated value.
    */
   def changeValue(k: K, defaultValue: => V, mergeValue: (V) => V): V = {
-    val pos = keySet.addWithoutResize(k)
+    val pos = keySet.add(k, grow, move)
     if ((pos & OpenHashSet.NONEXISTENCE_MASK) != 0) {
       val newValue = defaultValue
       _values(pos & OpenHashSet.POSITION_MASK) = newValue
-      keySet.rehashIfNeeded(k, grow, move)
       newValue
     } else {
       _values(pos) = mergeValue(_values(pos))