/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.arrow.memory;

import org.apache.arrow.util.Preconditions;
import org.apache.arrow.util.VisibleForTesting;
import org.checkerframework.checker.initialization.qual.Initialized;
import org.checkerframework.checker.initialization.qual.UnderInitialization;
import org.checkerframework.checker.nullness.qual.Nullable;

/**
 * Highly specialized IdentityHashMap that implements only partial Map APIs. It incurs low initial
 * cost (just two elements by default). It assumes Value includes the Key -
 * Implements @ValueWithKeyIncluded iface that provides "getKey" method.
 *
 * @param <K> Key type
 * @param <V> Value type
 */
public class LowCostIdentityHashMap<K, V extends ValueWithKeyIncluded<K>> {

  /*
   * The internal data structure to hold values.
   */
  private @Nullable Object[] elementData; // elementData[index] = null;

  /* Actual number of values. */
  private int size;

  /*
   * maximum number of elements that can be put in this map before having to
   * rehash.
   */
  private int threshold;

  private static final int DEFAULT_MIN_SIZE = 1;

  /* Default load factor of 0.75; */
  private static final int LOAD_FACTOR = 7500;

  /** Creates a Map with default expected maximum size. */
  public LowCostIdentityHashMap() {
    this(DEFAULT_MIN_SIZE);
  }

  /**
   * Creates a Map with the specified maximum size parameter.
   *
   * @param maxSize The estimated maximum number of entries that will be put in this map.
   */
  public LowCostIdentityHashMap(int maxSize) {
    if (maxSize >= 0) {
      this.size = 0;
      threshold = getThreshold(maxSize);
      elementData = newElementArrayUnderInitialized(computeElementArraySize());
    } else {
      throw new IllegalArgumentException();
    }
  }

  private int getThreshold(@UnderInitialization LowCostIdentityHashMap<K, V> this, int maxSize) {
    // assign the threshold to maxSize initially, this will change to a
    // higher value if rehashing occurs.
    return maxSize > 2 ? maxSize : 2;
  }

  private int computeElementArraySize(@UnderInitialization LowCostIdentityHashMap<K, V> this) {
    int arraySize = (int) (((long) threshold * 10000) / LOAD_FACTOR);
    // ensure arraySize is positive, the above cast from long to int type
    // leads to overflow and negative arraySize if threshold is too big
    return arraySize < 0 ? -arraySize : arraySize;
  }

  /**
   * Create a new element array.
   *
   * @param s the number of elements
   * @return Reference to the element array
   */
  private Object[] newElementArrayInitialized(
      @Initialized LowCostIdentityHashMap<K, V> this, int s) {
    return new Object[s];
  }

  /**
   * Create a new element array.
   *
   * @param s the number of elements
   * @return Reference to the element array
   */
  private Object[] newElementArrayUnderInitialized(
      @UnderInitialization LowCostIdentityHashMap<K, V> this, int s) {
    return new Object[s];
  }

  /**
   * Removes all elements from this map, leaving it empty.
   *
   * @see #isEmpty()
   * @see #size()
   */
  public void clear() {
    size = 0;
    for (int i = 0; i < elementData.length; i++) {
      elementData[i] = null;
    }
  }

  /**
   * Returns whether this map contains the specified key.
   *
   * @param key the key to search for.
   * @return {@code true} if this map contains the specified key, {@code false} otherwise.
   */
  public boolean containsKey(K key) {
    Preconditions.checkNotNull(key);

    int index = findIndex(key, elementData);
    return (elementData[index] == null) ? false : ((V) elementData[index]).getKey() == key;
  }

  /**
   * Returns whether this map contains the specified value.
   *
   * @param value the value to search for.
   * @return {@code true} if this map contains the specified value, {@code false} otherwise.
   */
  public boolean containsValue(V value) {
    Preconditions.checkNotNull(value);

    for (int i = 0; i < elementData.length; i++) {
      if (elementData[i] == value) {
        return true;
      }
    }
    return false;
  }

  /**
   * Returns the value of the mapping with the specified key.
   *
   * @param key the key.
   * @return the value of the mapping with the specified key.
   */
  public @Nullable V get(K key) {
    Preconditions.checkNotNull(key);

    int index = findIndex(key, elementData);

    return (elementData[index] == null)
        ? null
        : (((V) elementData[index]).getKey() == key) ? (V) elementData[index] : null;
  }

  /**
   * Returns the index where the key is found at, or the index of the next empty spot if the key is
   * not found in this table.
   */
  @VisibleForTesting
  int findIndex(@Nullable Object key, @Nullable Object[] array) {
    int length = array.length;
    int index = getModuloHash(key, length);
    int last = (index + length - 1) % length;
    while (index != last) {
      if ((array[index] == null) || ((V) array[index]).getKey() == key) {
        /*
         * Found the key, or the next empty spot (which means key is not
         * in the table)
         */
        break;
      }
      index = (index + 1) % length;
    }
    return index;
  }

  @VisibleForTesting
  static int getModuloHash(@Nullable Object key, int length) {
    return ((System.identityHashCode(key) & 0x7FFFFFFF) % length);
  }

  /**
   * Maps the specified key to the specified value.
   *
   * @param value the value.
   * @return the value of any previous mapping with the specified key or {@code null} if there was
   *     no such mapping.
   */
  public V put(V value) {
    Preconditions.checkNotNull(value);
    K key = value.getKey();
    Preconditions.checkNotNull(key);

    int index = findIndex(key, elementData);

    // if the key doesn't exist in the table
    if (elementData[index] == null || ((V) elementData[index]).getKey() != key) {
      if (++size > threshold) {
        rehash();
        index = findIndex(key, elementData);
      }

      // insert the key and assign the value to null initially
      elementData[index] = null;
    }

    // insert value to where it needs to go, return the old value
    Object result = elementData[index];
    elementData[index] = value;

    return (V) result;
  }

  @VisibleForTesting
  void rehash() {
    int newlength = elementData.length * 15 / 10;
    if (newlength == 0) {
      newlength = 1;
    }
    @Nullable Object[] newData = newElementArrayInitialized(newlength);
    for (int i = 0; i < elementData.length; i++) {
      Object key = (elementData[i] == null) ? null : ((V) elementData[i]).getKey();
      if (key != null) {
        // if not empty
        int index = findIndex(key, newData);
        newData[index] = elementData[i];
      }
    }
    elementData = newData;
    computeMaxSize();
  }

  private void computeMaxSize() {
    threshold = (int) ((long) elementData.length * LOAD_FACTOR / 10000);
  }

  /**
   * Removes the mapping with the specified key from this map.
   *
   * @param key the key of the mapping to remove.
   * @return the value of the removed mapping, or {@code null} if no mapping for the specified key
   *     was found.
   */
  public @Nullable V remove(K key) {
    Preconditions.checkNotNull(key);

    boolean hashedOk;
    int index;
    int next;
    int hash;
    Object result;
    Object object;
    index = next = findIndex(key, elementData);

    if (elementData[index] == null || ((V) elementData[index]).getKey() != key) {
      return null;
    }

    // store the value for this key
    result = elementData[index];
    // clear value to allow movement of the rest of the elements
    elementData[index] = null;
    size--;

    // shift the following elements up if needed
    // until we reach an empty spot
    int length = elementData.length;
    while (true) {
      next = (next + 1) % length;
      object = elementData[next];
      if (object == null) {
        break;
      }

      hash = getModuloHash(((V) object).getKey(), length);
      hashedOk = hash > index;
      if (next < index) {
        hashedOk = hashedOk || (hash <= next);
      } else {
        hashedOk = hashedOk && (hash <= next);
      }
      if (!hashedOk) {
        elementData[index] = object;
        index = next;
        elementData[index] = null;
      }
    }
    return (V) result;
  }

  /**
   * Returns whether this Map has no elements.
   *
   * @return {@code true} if this Map has no elements, {@code false} otherwise.
   * @see #size()
   */
  public boolean isEmpty() {
    return size == 0;
  }

  /**
   * Returns the number of mappings in this Map.
   *
   * @return the number of mappings in this Map.
   */
  public int size() {
    return size;
  }

  /**
   * Special API to return next value - substitute of regular Map.values.iterator().next().
   *
   * @return next available value or null if none available
   */
  public @Nullable V getNextValue() {
    for (int i = 0; i < elementData.length; i++) {
      if (elementData[i] != null) {
        return (V) elementData[i];
      }
    }
    return null;
  }
}
