package org.jetbrains.debugger.memory.tracking;

import com.intellij.debugger.engine.SuspendContextImpl;
import com.intellij.debugger.engine.jdi.ThreadReferenceProxy;
import com.sun.jdi.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.stream.Collectors;

class HashBasedTracking extends InstanceTrackingStrategy {
  private final ReferenceType myClassType;
  private Set<Integer> myHashesSet;

  HashBasedTracking(@NotNull SuspendContextImpl suspendContext,
                    @NotNull ReferenceType classType,
                    @NotNull List<ObjectReference> initialInstances) {
    myClassType = classType;
    myHashesSet = toSetOfHashes(evalHashCodes(suspendContext, initialInstances));
  }

  @NotNull
  @Override
  protected List<ObjectReference> updateImpl(@NotNull SuspendContextImpl suspendContext,
                                             @NotNull List<ObjectReference> references) {
    Map<ObjectReference, Optional<Integer>> ref2hash = evalHashCodes(suspendContext, references);
    List<ObjectReference> newInstances = new ArrayList<>();

    for(ObjectReference ref : references) {
      Optional<Integer> hash = ref2hash.get(ref);
      if(hash.isPresent() && !myHashesSet.contains(hash.get())) {
        newInstances.add(ref);
      }
    }

    myHashesSet = toSetOfHashes(ref2hash);
    return newInstances;
  }

  @NotNull
  private Map<ObjectReference, Optional<Integer>> evalHashCodes(@NotNull SuspendContextImpl suspendContext,
                                                                @NotNull List<ObjectReference> references) {
    Method hashCodeMethod = getHashCodeMethod(myClassType);
    Map<ObjectReference, Optional<Integer>> result = new HashMap<>();
    ThreadReferenceProxy threadProxy = suspendContext.getThread();
    ThreadReference thread = threadProxy != null ? threadProxy.getThreadReference() : null;
    if(thread == null) {
      return result;
    }

    for (ObjectReference ref : references) {
      Optional<Integer> hash = Optional.empty();
      try {
        Value value = ref.invokeMethod(thread, hashCodeMethod, Collections.emptyList(), 0);
        hash = Optional.of(((IntegerValue)value).value());
      } catch (InvalidTypeException | ClassNotLoadedException |
          IncompatibleThreadStateException | InvocationException ignored) {
      }

      result.put(ref, hash);
    }

    return result;
  }

  private Set<Integer> toSetOfHashes(@NotNull Map<ObjectReference, Optional<Integer>> ref2hashCode) {
    return ref2hashCode.values().stream()
        .filter(Optional::isPresent)
        .map(Optional::get)
        .collect(Collectors.toCollection(HashSet::new));
  }

  @Nullable
  private Method getHashCodeMethod(@Nullable ReferenceType referenceType) {
    if (referenceType == null) {
      return null;
    }

    return referenceType.methodsByName("hashCode").stream()
        .filter(method -> method.argumentTypeNames().isEmpty())
        .findFirst()
        .orElseGet(null);
  }
}

@@ -1942,16 +1942,42 @@
 hread ==
+ null %7C%7C hashCodeMethod ==
  null) %7B
@@ -1969,32 +1969,32 @@
 thod == null) %7B%0A
-
       return res
@@ -2111,16 +2111,123 @@
   try %7B%0A
+        if(!ref.isCollected()) %7B%0A          ref.disableCollection();%0A          if(!ref.isCollected()) %7B%0A    
         
@@ -2309,13 +2309,65 @@
 t(),
- 0);%0A
+%0A                ObjectReference.INVOKE_NONVIRTUAL);%0A    
     
@@ -2421,16 +2421,72 @@
 lue());%0A
+          %7D%0A          ref.enableCollection();%0A        %7D%0A
       %7D 

Check that object is not collected before hash calculation