001    /**
002     * Licensed to the Apache Software Foundation (ASF) under one
003     * or more contributor license agreements.  See the NOTICE file
004     * distributed with this work for additional information
005     * regarding copyright ownership.  The ASF licenses this file
006     * to you under the Apache License, Version 2.0 (the
007     * "License"); you may not use this file except in compliance
008     * with the License.  You may obtain a copy of the License at
009     *
010     *     http://www.apache.org/licenses/LICENSE-2.0
011     *
012     * Unless required by applicable law or agreed to in writing, software
013     * distributed under the License is distributed on an "AS IS" BASIS,
014     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015     * See the License for the specific language governing permissions and
016     * limitations under the License.
017     */
018    
019    package org.apache.hadoop.util;
020    
021    import org.apache.hadoop.classification.InterfaceAudience;
022    import org.apache.hadoop.classification.InterfaceStability;
023    
024    import com.google.common.base.Preconditions;
025    
026    /**
027     * The IdentityHashStore stores (key, value) mappings in an array.
028     * It is similar to java.util.HashTable, but much more lightweight.
029     * Neither inserting nor removing an element ever leads to any garbage
030     * getting created (assuming the array doesn't need to be enlarged).
031     *
032     * Unlike HashTable, it compares keys using
033     * {@link System#identityHashCode(Object)} and the identity operator.
034     * This is useful for types like ByteBuffer which have expensive hashCode
035     * and equals operators.
036     *
037     * We use linear probing to resolve collisions.  This avoids the need for
038     * the overhead of linked list data structures.  It also means that it is
039     * expensive to attempt to remove an element that isn't there, since we
040     * have to look at the entire array to be sure that it doesn't exist.
041     *
042     * @param <K>    The key type to use.
043     * @param <V>    THe value type to use.
044     */
045    @InterfaceAudience.Private
046    @InterfaceStability.Evolving
047    @SuppressWarnings("unchecked")
048    public final class IdentityHashStore<K, V> {
049      /**
050       * Even elements are keys; odd elements are values.
051       * The array has size 1 + Math.pow(2, capacity).
052       */
053      private Object buffer[];
054    
055      private int numInserted = 0;
056    
057      private int capacity;
058    
059      /**
060       * The default maxCapacity value to use.
061       */
062      private static final int DEFAULT_MAX_CAPACITY = 2;
063    
064      public IdentityHashStore(int capacity) {
065        Preconditions.checkArgument(capacity >= 0);
066        if (capacity == 0) {
067          this.capacity = 0;
068          this.buffer = null;
069        } else {
070          // Round the capacity we need up to a power of 2.
071          realloc((int)Math.pow(2,
072              Math.ceil(Math.log(capacity) / Math.log(2))));
073        }
074      }
075    
076      private void realloc(int newCapacity) {
077        Preconditions.checkArgument(newCapacity > 0);
078        Object prevBuffer[] = buffer;
079        this.capacity = newCapacity;
080        // Each element takes two array slots -- one for the key, 
081        // and another for the value.  We also want a load factor 
082        // of 0.50.  Combine those together and you get 4 * newCapacity.
083        this.buffer = new Object[4 * newCapacity];
084        this.numInserted = 0;
085        if (prevBuffer != null) {
086          for (int i = 0; i < prevBuffer.length; i += 2) {
087            if (prevBuffer[i] != null) {
088              putInternal(prevBuffer[i], prevBuffer[i + 1]);
089            }
090          }
091        }
092      }
093    
094      private void putInternal(Object k, Object v) {
095        int hash = System.identityHashCode(k);
096        final int numEntries = buffer.length / 2;
097        int index = hash % numEntries;
098        while (true) {
099          if (buffer[2 * index] == null) {
100            buffer[2 * index] = k;
101            buffer[1 + (2 * index)] = v;
102            numInserted++;
103            return;
104          }
105          index = (index + 1) % numEntries;
106        }
107      }
108    
109      /**
110       * Add a new (key, value) mapping.
111       *
112       * Inserting a new (key, value) never overwrites a previous one.
113       * In other words, you can insert the same key multiple times and it will
114       * lead to multiple entries.
115       */
116      public void put(K k, V v) {
117        Preconditions.checkNotNull(k);
118        if (buffer == null) {
119          realloc(DEFAULT_MAX_CAPACITY);
120        } else if (numInserted + 1 > capacity) {
121          realloc(capacity * 2);
122        }
123        putInternal(k, v);
124      }
125    
126      private int getElementIndex(K k) {
127        if (buffer == null) {
128          return -1;
129        }
130        final int numEntries = buffer.length / 2;
131        int hash = System.identityHashCode(k);
132        int index = hash % numEntries;
133        int firstIndex = index;
134        do {
135          if (buffer[2 * index] == k) {
136            return index;
137          }
138          index = (index + 1) % numEntries;
139        } while (index != firstIndex);
140        return -1;
141      }
142    
143      /**
144       * Retrieve a value associated with a given key.
145       */
146      public V get(K k) {
147        int index = getElementIndex(k);
148        if (index < 0) {
149          return null;
150        }
151        return (V)buffer[1 + (2 * index)];
152      }
153    
154      /**
155       * Retrieve a value associated with a given key, and delete the
156       * relevant entry.
157       */
158      public V remove(K k) {
159        int index = getElementIndex(k);
160        if (index < 0) {
161          return null;
162        }
163        V val = (V)buffer[1 + (2 * index)];
164        buffer[2 * index] = null;
165        buffer[1 + (2 * index)] = null;
166        numInserted--;
167        return val;
168      }
169    
170      public boolean isEmpty() {
171        return numInserted == 0;
172      }
173    
174      public int numElements() {
175        return numInserted;
176      }
177    
178      public int capacity() {
179        return capacity;
180      }
181    
182      public interface Visitor<K, V> {
183        void accept(K k, V v);
184      }
185    
186      /**
187       * Visit all key, value pairs in the IdentityHashStore.
188       */
189      public void visitAll(Visitor<K, V> visitor) {
190        int length = buffer == null ? 0 : buffer.length;
191        for (int i = 0; i < length; i += 2) {
192          if (buffer[i] != null) {
193            visitor.accept((K)buffer[i], (V)buffer[i + 1]);
194          }
195        }
196      }
197    }