/* Copyright (c) 2008-2023, Nathan Sweet
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following
 * conditions are met:
 *
 * - Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 * - Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
 * disclaimer in the documentation and/or other materials provided with the distribution.
 * - Neither the name of Esoteric Software nor the names of its contributors may be used to endorse or promote products derived
 * from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING,
 * BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */

package org.apache.fory.collection;

import java.util.Iterator;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiConsumer;
import org.apache.fory.util.Preconditions;

// Derived from
// https://github.com/EsotericSoftware/kryo/blob/135df69526615bb3f6b34846e58ba3fec3b631c3/src/com/esotericsoftware/kryo/util/ObjectMap.java.

/**
 * An unordered map where the keys and values are objects. Null keys are not allowed. No allocation
 * is done except when growing the table size.
 *
 * <p>This class performs fast contains and remove (typically O(1), worst case O(n) but that is rare
 * in practice). Add may be slightly slower, depending on hash collisions. Hashcodes are rehashed to
 * reduce collisions and the need to resize. Load factors greater than 0.91 greatly increase the
 * chances to resize to the next higher POT size.
 *
 * <p>Unordered sets and maps are not designed to provide especially fast iteration.
 *
 * <p>This implementation uses linear probing with the backward shift algorithm for removal.
 * Hashcodes are rehashed using Fibonacci hashing, instead of the more common power-of-two mask, to
 * better distribute poor hashCodes (see <a href=
 * "https://probablydance.com/2018/06/16/fibonacci-hashing-the-optimization-that-the-world-forgot-or-a-better-alternative-to-integer-modulo/">Malte
 * Skarupke's blog post</a>). Linear probing continues to work even when all hashCodes collide, just
 * more slowly.
 *
 * @author Nathan Sweet
 * @author Tommy Ettinger
 */
@SuppressWarnings("unchecked")
public class ForyObjectMap<K, V> {
  static final long MASK_NUMBER = 0x9E3779B97F4A7C15L;
  static final Object dummy = new Object();

  public int size;

  K[] keyTable;
  V[] valueTable;

  float loadFactor;
  int threshold;

  /**
   * Used by {@link #place(Object)} to bit shift the upper bits of a {@code long} into a usable
   * range (&gt;= 0 and &lt;= {@link #mask}). The shift can be negative, which is convenient to
   * match the number of bits in mask: if mask is a 7-bit number, a shift of -7 shifts the upper 7
   * bits into the lowest 7 positions. This class sets the shift &gt; 32 and &lt; 64, which if used
   * with an int will still move the upper bits of an int to the lower bits due to Java's implicit
   * modulus on shifts.
   *
   * <p>{@link #mask} can also be used to mask the low bits of a number, which may be faster for
   * some hashcodes, if {@link #place(Object)} is overridden.
   */
  protected int shift;

  /**
   * A bitmask used to confine hashcodes to the size of the table. Must be all 1 bits in its low
   * positions, ie a power of two minus 1. If {@link #place(Object)} is overridden, this can be used
   * instead of {@link #shift} to isolate usable bits of a hash.
   */
  protected int mask;

  /** Creates a new map with an initial capacity of 1 and a load factor of 0.6. */
  public ForyObjectMap() {
    this(1, 0.6f);
  }

  /**
   * Creates a new map with a load factor of 0.6.
   *
   * @param initialCapacity If not a power of two, it is increased to the next nearest power of two.
   */
  public ForyObjectMap(int initialCapacity) {
    this(initialCapacity, 0.6f);
  }

  /**
   * Creates a new map with the specified initial capacity and load factor. This map will hold
   * initialCapacity items before growing the backing table.
   *
   * @param initialCapacity If not a power of two, it is increased to the next nearest power of two.
   */
  public ForyObjectMap(int initialCapacity, float loadFactor) {
    if (loadFactor <= 0f || loadFactor >= 1f) {
      throw new IllegalArgumentException("loadFactor must be > 0 and < 1: " + loadFactor);
    }
    this.loadFactor = loadFactor;

    int tableSize = tableSize(initialCapacity, loadFactor);
    threshold = (int) (tableSize * loadFactor);
    mask = tableSize - 1;
    shift = Long.numberOfLeadingZeros(mask);

    keyTable = (K[]) new Object[tableSize];
    valueTable = (V[]) new Object[tableSize];
  }

  /**
   * Returns an {@code index >= 0 and <= mask} for the specified {@code item}.
   *
   * <p>The default implementation uses Fibonacci hashing on the item's {@link Object#hashCode()}:
   * the hashcode is multiplied by a long constant (2 to the 64th, divided by the golden ratio) then
   * the uppermost bits are shifted into the lowest positions to obtain an index in the desired
   * range. Multiplication by a long may be slower than int (eg on GWT) but greatly improves
   * rehashing, allowing even very poor hashcodes, such as those that only differ in their upper
   * bits, to be used without high collision rates. Fibonacci hashing has increased collision rates
   * when all or most hashcodes are multiples of larger Fibonacci numbers (see <a href=
   * "https://probablydance.com/2018/06/16/fibonacci-hashing-the-optimization-that-the-world-forgot-or-a-better-alternative-to-integer-modulo/">Malte
   * Skarupke's blog post</a>).
   *
   * <p>This method can be overridden to customizing hashing. This may be useful eg in the unlikely
   * event that most hashcodes are Fibonacci numbers, if keys provide poor or incorrect hashcodes,
   * or to simplify hashing if keys provide high quality hashcodes and don't need Fibonacci hashing:
   * {@code return item.hashCode() & mask;}
   *
   * @see #mask
   */
  protected int place(K item) {
    return (int) (item.hashCode() * MASK_NUMBER >>> shift);
  }

  /**
   * Returns the index of the key if already present, else -(index + 1) for the next empty index.
   * This can be overridden in this package to compare for equality differently than {@link
   * Object#equals(Object)}.
   */
  int locateKey(K key) {
    if (key == null) {
      throw new IllegalArgumentException("key cannot be null.");
    }
    K[] keyTable = this.keyTable;
    int mask = this.mask;
    for (int i = place(key); ; i = i + 1 & mask) {
      K other = keyTable[i];
      if (other == null) {
        return -(i + 1); // Empty space is available.
      }
      if (other.equals(key)) {
        return i; // Same key was found.
      }
    }
  }

  /** Returns the old value associated with the specified key, or null. */
  public V put(K key, V value) {
    int i = locateKey(key);
    if (i >= 0) { // Existing key was found.
      V oldValue = valueTable[i];
      valueTable[i] = value;
      return oldValue;
    }
    i = -(i + 1); // Empty space was found.
    keyTable[i] = key;
    valueTable[i] = value;
    if (++size >= threshold) {
      resize(keyTable.length << 1);
    }
    return null;
  }

  /** Skips checks for existing keys, doesn't increment size. */
  private void putResize(K key, V value) {
    K[] keyTable = this.keyTable;
    int mask = this.mask;
    for (int i = place(key); ; i = (i + 1) & mask) {
      if (keyTable[i] == null) {
        keyTable[i] = key;
        valueTable[i] = value;
        return;
      }
    }
  }

  /** Returns the value for the specified key, or null if the key is not in the map. */
  public <T extends K> V get(T key) {
    int mask = this.mask;
    K[] keyTable = this.keyTable;
    for (int i = place(key); ; i = i + 1 & mask) {
      K other = keyTable[i];
      if (other == null) {
        return null;
      }
      if (other.equals(key)) {
        return valueTable[i];
      }
    }
  }

  /** Returns the value for the specified key, or the default value if the key is not in the map. */
  public V get(K key, V defaultValue) {
    int mask = this.mask;
    K[] keyTable = this.keyTable;
    for (int i = place(key); ; i = i + 1 & mask) {
      K other = keyTable[i];
      if (other == null) {
        return defaultValue;
      }
      if (other.equals(key)) {
        return valueTable[i];
      }
    }
  }

  public V remove(K key) {
    int i = locateKey(key);
    if (i < 0) {
      return null;
    }
    K[] keyTable = this.keyTable;
    V[] valueTable = this.valueTable;
    V oldValue = valueTable[i];
    int mask = this.mask, next = i + 1 & mask;
    while ((key = keyTable[next]) != null) {
      int placement = place(key);
      if ((next - placement & mask) > (i - placement & mask)) {
        keyTable[i] = key;
        valueTable[i] = valueTable[next];
        i = next;
      }
      next = next + 1 & mask;
    }
    keyTable[i] = null;
    valueTable[i] = null;
    size--;
    return oldValue;
  }

  /** Returns true if the map is empty. */
  public boolean isEmpty() {
    return size == 0;
  }

  /**
   * Clears the map and reduces the size of the backing arrays to be the specified capacity /
   * loadFactor, if they are larger.
   */
  public void clear(int maximumCapacity) {
    int tableSize = tableSize(maximumCapacity, loadFactor);
    if (keyTable.length <= tableSize) {
      clear();
      return;
    }
    size = 0;
    resize(tableSize);
  }

  public void clear() {
    if (size == 0) {
      return;
    }
    size = 0;
    ObjectArray.clearObjectArray(keyTable, 0, keyTable.length);
    ObjectArray.clearObjectArray(valueTable, 0, valueTable.length);
  }

  public boolean containsKey(K key) {
    return locateKey(key) >= 0;
  }

  public boolean containsValue(Object value) {
    Preconditions.checkNotNull(value);
    Object[] tab = valueTable;
    for (int i = 1; i < tab.length; i += 2) {
      if (tab[i] == value) {
        return true;
      }
    }
    return false;
  }

  public void forEach(BiConsumer<? super K, ? super V> action) {
    Objects.requireNonNull(action);
    K[] keyTable = this.keyTable;
    V[] valueTable = this.valueTable;
    int i = keyTable.length;
    while (i-- > 0) {
      K key = keyTable[i];
      if (key == null) {
        continue;
      }
      V value = valueTable[i];
      action.accept(key, value);
    }
  }

  /** Returns an Iterable for the entries in the map. Remove isn't supported. */
  public Iterable<Map.Entry<K, V>> iterable() {
    return MapIterator::new;
  }

  private class MapIterator implements Iterator<Map.Entry<K, V>> {
    private int nextIndex;

    @Override
    public boolean hasNext() {
      K[] keyTable = ForyObjectMap.this.keyTable;
      for (int i = nextIndex; i < keyTable.length; i++) {
        if (keyTable[i] != null) {
          nextIndex = i;
          return true;
        }
      }
      return false;
    }

    @Override
    public Map.Entry<K, V> next() {
      return new MapEntry<>(keyTable[nextIndex], valueTable[nextIndex++]);
    }
  }

  final void resize(int newSize) {
    int oldCapacity = keyTable.length;
    threshold = (int) (newSize * loadFactor);
    mask = newSize - 1;
    shift = Long.numberOfLeadingZeros(mask);

    K[] oldKeyTable = keyTable;
    V[] oldValueTable = valueTable;

    keyTable = (K[]) new Object[newSize];
    valueTable = (V[]) new Object[newSize];

    if (size > 0) {
      for (int i = 0; i < oldCapacity; i++) {
        K key = oldKeyTable[i];
        if (key != null) {
          putResize(key, oldValueTable[i]);
        }
      }
    }
  }

  @Override
  public int hashCode() {
    int h = size;
    K[] keyTable = this.keyTable;
    V[] valueTable = this.valueTable;
    for (int i = 0, n = keyTable.length; i < n; i++) {
      K key = keyTable[i];
      if (key != null) {
        h += key.hashCode();
        V value = valueTable[i];
        if (value != null) {
          h += value.hashCode();
        }
      }
    }
    return h;
  }

  @Override
  public boolean equals(Object obj) {
    if (obj == this) {
      return true;
    }
    if (!(obj instanceof ForyObjectMap)) {
      return false;
    }
    ForyObjectMap other = (ForyObjectMap) obj;
    if (other.size != size) {
      return false;
    }
    K[] keyTable = this.keyTable;
    V[] valueTable = this.valueTable;
    for (int i = 0, n = keyTable.length; i < n; i++) {
      K key = keyTable[i];
      if (key != null) {
        V value = valueTable[i];
        if (value == null) {
          if (other.get(key, dummy) != null) {
            return false;
          }
        } else {
          if (!value.equals(other.get(key))) {
            return false;
          }
        }
      }
    }
    return true;
  }

  public String toString(String separator) {
    return toString(separator, false);
  }

  @Override
  public String toString() {
    return toString(", ", true);
  }

  private String toString(String separator, boolean braces) {
    if (size == 0) {
      return braces ? "{}" : "";
    }
    StringBuilder buffer = new StringBuilder(32);
    if (braces) {
      buffer.append('{');
    }
    K[] keyTable = this.keyTable;
    V[] valueTable = this.valueTable;
    int i = keyTable.length;
    while (i-- > 0) {
      K key = keyTable[i];
      if (key == null) {
        continue;
      }
      buffer.append(separator);
      buffer.append(key == this ? "(this)" : key);
      buffer.append('=');
      V value = valueTable[i];
      buffer.append(value == this ? "(this)" : value);
    }
    if (braces) {
      buffer.append('}');
    }
    return buffer.toString();
  }

  public static int tableSize(int capacity, float loadFactor) {
    if (capacity < 0) {
      throw new IllegalArgumentException("capacity must be >= 0: " + capacity);
    }
    int tableSize = nextPowerOfTwo(Math.max(2, (int) Math.ceil(capacity / loadFactor)));
    if (tableSize > 1 << 30) {
      throw new IllegalArgumentException("The required capacity is too large: " + capacity);
    }
    return tableSize;
  }

  public static int nextPowerOfTwo(int value) {
    if (value == 0) {
      return 1;
    }
    value--;
    value |= value >> 1;
    value |= value >> 2;
    value |= value >> 4;
    value |= value >> 8;
    value |= value >> 16;
    return value + 1;
  }
}
