/*
 * Copyright 2000-2017 JetBrains s.r.o.
 *
 * Licensed 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 com.intellij.debugger.streams.resolve;

import com.intellij.debugger.streams.trace.TraceElement;
import com.intellij.debugger.streams.trace.TraceInfo;
import com.intellij.debugger.streams.wrapper.TraceUtil;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NotNull;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * @author Vitaliy.Bibaev
 */
public class IdentityResolver implements ValuesOrderResolver {
  @NotNull
  @Override
  public Result resolve(@NotNull TraceInfo info) {
    final Map<Integer, TraceElement> before = info.getValuesOrderBefore();
    final Map<Integer, TraceElement> after = info.getValuesOrderAfter();

    final Map<TraceElement, List<TraceElement>> direct = new HashMap<>();
    final Map<TraceElement, List<TraceElement>> reverse = new HashMap<>();

    final Map<Object, List<TraceElement>> grouped = StreamEx
      .of(after.keySet())
      .sorted()
      .map(after::get)
      .groupingBy(TraceUtil::extractKey);

    for (final TraceElement element : before.values()) {
      final Object value = TraceUtil.extractKey(element);

      final List<TraceElement> elements = grouped.get(value);
      if (elements == null || elements.isEmpty()) {
        direct.put(element, Collections.emptyList());
        continue;
      }

      final TraceElement afterItem = elements.get(0);

      direct.put(element, Collections.singletonList(afterItem));
      reverse.put(afterItem, Collections.singletonList(element));

      grouped.put(value, elements.isEmpty() ? Collections.emptyList() : elements.subList(1, elements.size()));
    }

    return Result.of(direct, reverse);
  }
}

@@ -1,13 +1,10 @@
 /
-*%0A *
+/
  Copyrig
@@ -36,566 +36,109 @@
 r.o.
-%0A *%0A * Licensed under the Apache License, Version 2.0 (the %22License%22);%0A * you may not use this file except in compliance with the License.%0A * You may obtain a copy of the License at%0A *%0A * http://www.apache.org/licenses/LICENSE-2.0%0A *%0A * Unless required by applicable law or agreed to in writing, software%0A * distributed under the License is distributed on an %22AS IS%22 BASIS,%0A * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.%0A * See the License for the specific language governing permissions and%0A * limitations under the License.%0A */
+ Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 %0Apac
@@ -627,16 +627,75 @@
 olver %7B%0A
+  private static final Object NULL_MARKER = new Object();%0A%0A
   @NotNu
@@ -1203,25 +1203,32 @@
 upingBy(
-TraceUtil
+IdentityResolver
 ::extrac
@@ -1323,18 +1323,8 @@
 e = 
-TraceUtil.
 extr
@@ -1837,24 +1837,24 @@
 )));%0A    %7D%0A%0A
-
     return R
@@ -1883,11 +1883,193 @@
 se);%0A  %7D
+%0A%0A  @NotNull%0A  private static Object extractKey(@NotNull TraceElement element) %7B%0A    final Object key = TraceUtil.extractKey(element);%0A    return key == null ? NULL_MARKER : key;%0A  %7D
 %0A%7D%0A

Fix possible NPE in IdentityResolver when some value in chain are null