content: /*
   * 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 static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  import org.junit.jupiter.api.Test;
  
  /** To test simplified implementation of IdentityHashMap. */
  public class TestLowCostIdentityHashMap {
  
    @Test
    public void testIdentityHashMap() {
      LowCostIdentityHashMap<String, StringWithKey> hashMap = new LowCostIdentityHashMap<>();
  
      StringWithKey obj1 = new StringWithKey("s1key", "s1value");
      StringWithKey obj2 = new StringWithKey("s2key", "s2value");
      StringWithKey obj3 = new StringWithKey("s3key", "s3value");
      StringWithKey obj4 = new StringWithKey("s1key", "s4value");
      StringWithKey obj5 = new StringWithKey("s5key", "s5value");
  
      assertNull(hashMap.put(obj1));
      assertNull(hashMap.put(obj2));
      assertNull(hashMap.put(obj3));
      assertEquals(obj1, hashMap.put(obj4));
      assertNull(hashMap.put(obj5));
  
      assertEquals(4, hashMap.size());
  
      assertEquals(obj4, hashMap.get("s1key"));
  
      assertNull(hashMap.remove("abc"));
  
      assertEquals(obj3, hashMap.remove("s3key"));
  
      assertEquals(3, hashMap.size());
  
      assertFalse(hashMap.isEmpty());
  
      StringWithKey nextValue = hashMap.getNextValue();
  
      assertNotNull(nextValue);
  
      assertTrue(
          (hashMap.get("s1key") == nextValue
              || hashMap.get("s2key") == nextValue
              || hashMap.get("s5key") == nextValue));
  
      assertTrue(hashMap.containsValue(obj4));
      assertTrue(hashMap.containsValue(obj2));
      assertTrue(hashMap.containsValue(obj5));
  
      assertEquals(obj4, hashMap.remove("s1key"));
  
      nextValue = hashMap.getNextValue();
  
      assertNotNull(nextValue);
  
      assertTrue(hashMap.get("s2key") == nextValue || hashMap.get("s5key") == nextValue);
  
      assertEquals(2, hashMap.size());
  
      assertEquals(obj2, hashMap.remove("s2key"));
      assertEquals(obj5, hashMap.remove("s5key"));
  
      assertEquals(0, hashMap.size());
  
      assertTrue(hashMap.isEmpty());
    }
  
    @Test
    public void testLargeMap() {
      LowCostIdentityHashMap<String, StringWithKey> hashMap = new LowCostIdentityHashMap<>();
  
      String[] keys = new String[200];
      for (int i = 0; i < 200; i++) {
        keys[i] = "s" + i + "key";
      }
  
      for (int i = 0; i < 100; i++) {
        if (i % 5 == 0 && i != 0) {
          StringWithKey obj = new StringWithKey(keys[i - 5], "s" + i + "value");
          StringWithKey retObj = hashMap.put(obj);
          assertNotNull(retObj);
          StringWithKey obj1 = new StringWithKey(keys[i], "s" + 2 * i + "value");
          StringWithKey retObj1 = hashMap.put(obj1);
          assertNull(retObj1);
        } else {
          StringWithKey obj = new StringWithKey(keys[i], "s" + i + "value");
          StringWithKey retObj = hashMap.put(obj);
          assertNull(retObj);
        }
      }
      assertEquals(100, hashMap.size());
      for (int i = 0; i < 100; i++) {
        StringWithKey returnObj = hashMap.get(keys[i]);
        assertNotNull(returnObj);
        if (i == 95) {
          assertEquals("s190value", returnObj.getValue());
          continue;
        }
        if (i % 5 == 0) {
          assertEquals("s" + (i + 5) + "value", returnObj.getValue());
        } else {
          assertEquals("s" + i + "value", returnObj.getValue());
        }
      }
  
      for (int i = 0; i < 100; i++) {
        if (i % 4 == 0) {
          StringWithKey returnObj = hashMap.remove(keys[i]);
          assertNotNull(returnObj);
          assertFalse(hashMap.containsKey(keys[i]));
        }
        StringWithKey obj = new StringWithKey(keys[100 + i], "s" + (100 + i) + "value");
        StringWithKey retObj = hashMap.put(obj);
        assertNull(retObj);
        assertTrue(hashMap.containsKey(keys[100 + i]));
      }
      assertEquals(175, hashMap.size());
      for (int i = 0; i < 100; i++) {
        StringWithKey retObj = hashMap.getNextValue();
        assertNotNull(retObj);
        hashMap.remove(retObj.getKey());
      }
      assertFalse(hashMap.isEmpty());
      assertEquals(75, hashMap.size());
      hashMap.clear();
      assertTrue(hashMap.isEmpty());
    }
  
    private static class StringWithKey implements ValueWithKeyIncluded<String> {
  
      private final String myValue;
      private final String myKey;
  
      StringWithKey(String myKey, String myValue) {
        this.myKey = myKey;
        this.myValue = myValue;
      }
  
      @Override
      public String getKey() {
        return myKey;
      }
  
      String getValue() {
        return myValue;
      }
    }
  }
  