package com.engebretson;

import org.openjdk.jmh.annotations.*;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * Benchmark demonstrating performance impact of polymorphic call sites on HashMap methods.
 * 
 * This test shows that HashMap.hashCode() and equals() suffer from megamorphic
 * call site overhead due to AbstractMap implementations using polymorphic iteration patterns.
 * 
 * The setup poisons call sites with HashMap, TreeMap, and LinkedHashMap to create
 * polymorphic call sites before benchmarking begins.
 */
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
@State(Scope.Benchmark)
@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
@Measurement(iterations = 4, time = 1, timeUnit = TimeUnit.SECONDS)
@Fork(value = 1, jvmArgs = {"-XX:+UseParallelGC", "-Xmx3g"})
public class HashMapMethodsBenchmark {

    private static final int POISON_ITERATIONS = 10000;

    @Param({"0", "1", "100"})
    public int mapSize;

    public HashMap<String, Integer> testHashMap;
    public HashMap<String, Integer> otherHashMap;
    public TreeMap<String, Integer> testTreeMap;
    public LinkedHashMap<String, Integer> testLinkedHashMap;

    @Setup(Level.Trial)
    public void setup() {
        // Create test data with identical contents
        testHashMap = new HashMap<>();
        otherHashMap = new HashMap<>();
        testTreeMap = new TreeMap<>();
        testLinkedHashMap = new LinkedHashMap<>();

        for (int i = 0; i < mapSize; i++) {
            String key = "key" + i;
            Integer value = i;
            testHashMap.put(key, value);
            otherHashMap.put(key, value);
            testTreeMap.put(key, value);
            testLinkedHashMap.put(key, value);
        }

        // Poison hashCode() call sites
        for (int i = 0; i < POISON_ITERATIONS; i++) {
            Map<String, Integer> map = (i % 3 == 0) ? testHashMap : (i % 3 == 1) ? testTreeMap : testLinkedHashMap;
            int hash = map.hashCode();
            if (hash == 0 && mapSize > 0) throw new RuntimeException();
        }

        // Poison equals() call sites
        for (int i = 0; i < POISON_ITERATIONS; i++) {
            Map<String, Integer> map1 = (i % 3 == 0) ? testHashMap : (i % 3 == 1) ? testTreeMap : testLinkedHashMap;
            Map<String, Integer> map2 = (i % 3 == 0) ? otherHashMap : (i % 3 == 1) ? testTreeMap : testLinkedHashMap;
            boolean equal = map1.equals(map2);
            if (!equal && map1 == testHashMap && map2 == otherHashMap) throw new RuntimeException();
        }
    }

    @Benchmark
    public int hashMapHashCode() {
        return testHashMap.hashCode();
    }

    @Benchmark
    public int manualHashCode() {
        int hash = 0;
        for (Map.Entry<String, Integer> entry : testHashMap.entrySet()) {
            hash += entry.hashCode();
        }
        return hash;
    }

    @Benchmark
    public boolean hashMapEquals() {
        return testHashMap.equals(otherHashMap);
    }

    @Benchmark
    public boolean manualEquals() {
        if (testHashMap.size() != otherHashMap.size()) return false;
        
        for (Map.Entry<String, Integer> entry : testHashMap.entrySet()) {
            String key = entry.getKey();
            Integer value = entry.getValue();
            Integer otherValue = otherHashMap.get(key);
            if (!Objects.equals(value, otherValue)) {
                return false;
            }
        }
        return true;
    }
}