LongToObjectHashMap.java

// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.
 
package org.wpilib.units.collections;
 
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;

public class LongToObjectHashMap<V> {
  private static final int INITIAL_SIZE = 0;
  private static final int INITIAL_CAPACITY = 8; // NOTE: must be a power of two

  private static final double LOAD_FACTOR = 75.00 / 100;

  private int m_size = INITIAL_SIZE;

  private int m_capacity = INITIAL_CAPACITY;

  private long[] m_keys = new long[m_capacity];

  private boolean[] m_uses = new boolean[m_capacity];

  @SuppressWarnings("unchecked")
  private V[] m_values = (V[]) new Object[m_capacity];

  public LongToObjectHashMap() {}

  public V put(long key, V value) {
    int bucket = bucket(key);
 
    // Increment the bucket until we hit an open space (there's always going to be at least one)
    while (m_uses[bucket]) {
      if (m_keys[bucket] == key) {
        // replace the existing value
        var oldValue = m_values[bucket];
        m_values[bucket] = value;
        return oldValue;
      }
      bucket = safeIncrement(bucket);
    }
 
    m_uses[bucket] = true;
    m_keys[bucket] = key;
    m_values[bucket] = value;
    m_size++;
 
    if (m_size > maxSize()) {
      grow();
    }
    return null;
  }

  public V get(long key) {
    int bucket = bucket(key);
    while (m_uses[bucket]) {
      if (m_keys[bucket] == key) {
        // found it
        return m_values[bucket];
      }
      bucket = safeIncrement(bucket);
    }
    return null;
  }

  public V remove(long key) {
    int bucket = bucket(key);
    while (m_uses[bucket]) {
      if (m_keys[bucket] == key) {
        // found it
        // TODO: Shrink the map when below a certain load factor
        //       Current use cases don't remove elements from the map, so there's not much use
        //       for shrinking at the moment.
        m_size--;
        m_keys[bucket] = 0L;
        m_uses[bucket] = false;
 
        var oldValue = m_values[bucket];
        m_values[bucket] = null;
        return oldValue;
      }
      bucket = safeIncrement(bucket);
    }
 
    return null;
  }

  public boolean containsKey(long key) {
    int bucket = bucket(key);
    while (m_uses[bucket]) {
      if (m_keys[bucket] == key) {
        return true;
      }
      bucket = safeIncrement(bucket);
    }
    return false;
  }

  public void clear() {
    if (m_size == 0) {
      // Nothing to do
      return;
    }
    m_size = 0;
 
    Arrays.fill(m_uses, false);
    Arrays.fill(m_keys, 0L);
    Arrays.fill(m_values, null);
  }

  public int size() {
    return m_size;
  }
 
  // package-private for tests
  int capacity() {
    return m_capacity;
  }

  public boolean isEmpty() {
    return m_size == 0;
  }

  public ReadOnlyPrimitiveLongSet keySet() {
    // copy the sparse key array into a compact array
    final long[] keys = new long[m_size];
    int i = 0;
    for (int bucket = 0; bucket < m_capacity; bucket++) {
      if (m_uses[bucket]) {
        keys[i] = m_keys[bucket];
        i++;
      }
    }
 
    return new ReadOnlyPrimitiveLongSet(keys);
  }

  public Collection<V> values() {
    Collection<V> values = new ArrayList<>();
    for (int bucket = 0; bucket < m_capacity; bucket++) {
      if (m_uses[bucket]) {
        values.add(m_values[bucket]);
      }
    }
    return List.copyOf(values); // return a readonly copy
  }

  @FunctionalInterface
  public interface IteratorFunction<V> {
    void accept(long key, V value);
  }

  public void forEach(IteratorFunction<? super V> function) {
    for (int bucket = 0; bucket < m_capacity; bucket++) {
      if (m_uses[bucket]) {
        function.accept(m_keys[bucket], m_values[bucket]);
      }
    }
  }
 
  private void grow() {
    final int currentSize = m_size;
    final int oldCapacity = m_capacity;
    if (oldCapacity * LOAD_FACTOR >= currentSize) {
      // We're below the maximum allowed size for the current capacity
      // Nothing to do
      return;
    }
 
    final int newCapacity = oldCapacity * 2;
    final int newMask = newCapacity - 1;
 
    final boolean[] oldUses = m_uses;
    final long[] oldKeys = m_keys;
    final V[] oldValues = m_values;
 
    final boolean[] newUses = new boolean[newCapacity];
    final long[] newKeys = new long[newCapacity];
    @SuppressWarnings("unchecked")
    final V[] newValues = (V[]) new Object[newCapacity];
 
    for (int oldBucket = 0; oldBucket < oldCapacity; oldBucket++) {
      if (!oldUses[oldBucket]) {
        // Bucket is empty, skip
        continue;
      }
      final long key = oldKeys[oldBucket];
      final V value = oldValues[oldBucket];
 
      int newBucket = (int) (hash(key) & newMask);
      while (newUses[newBucket]) {
        newBucket = (newBucket + 1) & newMask;
      }
 
      newUses[newBucket] = true;
      newKeys[newBucket] = key;
      newValues[newBucket] = value;
    }
 
    m_capacity = newCapacity;
    m_uses = newUses;
    m_keys = newKeys;
    m_values = newValues;
  }
 
  private int maxSize() {
    return (int) (m_capacity * LOAD_FACTOR);
  }

  private long hash(long key) {
    return 31 + (key ^ (key >>> 15) ^ (key >>> 31) ^ (key << 31));
  }

  private int mask() {
    return m_capacity - 1;
  }

  private int bucket(long key) {
    var hash = hash(key);
    return (int) (hash & mask());
  }

  private int safeIncrement(int bucket) {
    return (bucket + 1) & mask();
  }
}