/*
 * Copyright 2000-2016 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.memory.component;

import com.intellij.debugger.memory.utils.StackFrameItem;
import com.intellij.openapi.components.AbstractProjectComponent;
import com.intellij.openapi.project.Project;
import com.intellij.xdebugger.XDebugSession;
import com.intellij.xdebugger.XDebugSessionListener;
import com.sun.jdi.ObjectReference;
import com.sun.jdi.ReferenceType;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

public class CreationPositionTracker extends AbstractProjectComponent {
  /**
   * Stores all tracked instance for each debug session.
   */
  private final ConcurrentHashMap<XDebugSession, Map<ObjectReference, List<StackFrameItem>>>
      mySession2Reference2Stack = new ConcurrentHashMap<>();

  private final ConcurrentHashMap<XDebugSession, Map<ObjectReference, List<StackFrameItem>>>
      myPinnedSession2Reference2Stack = new ConcurrentHashMap<>();

  public CreationPositionTracker(Project project) {
    super(project);
  }

  @Nullable
  public static CreationPositionTracker getInstance(@NotNull Project project) {
    return project.isDisposed() ? null : project.getComponent(CreationPositionTracker.class);
  }

  @Nullable
  public List<StackFrameItem> getStack(@NotNull XDebugSession session, @NotNull ObjectReference ref) {
    List<StackFrameItem> stack = extract(mySession2Reference2Stack, session, ref);
    return stack != null ? stack : extract(myPinnedSession2Reference2Stack, session, ref);
  }

  public void addStack(@NotNull XDebugSession session, @NotNull ObjectReference ref,
                       @NotNull List<StackFrameItem> stack) {
    if (!mySession2Reference2Stack.containsKey(session)) {
      mySession2Reference2Stack.put(session, new ConcurrentHashMap<>());
      session.addSessionListener(new XDebugSessionListener() {
        @Override
        public void sessionStopped() {
          mySession2Reference2Stack.remove(session);
        }
      });
    }

    mySession2Reference2Stack.get(session).put(ref, stack);
  }

  public void releaseBySession(@NotNull XDebugSession session) {
    if(mySession2Reference2Stack.containsKey(session)) {
      mySession2Reference2Stack.put(session, new ConcurrentHashMap<>());
    }
  }

  public void unpinStacks(@NotNull XDebugSession session, @NotNull ReferenceType ref) {
    Map<ObjectReference, List<StackFrameItem>> ref2Stack = myPinnedSession2Reference2Stack.getOrDefault(session, null);
    if (ref2Stack != null) {
      Iterator<ObjectReference> iterator = ref2Stack.keySet().iterator();
      while (iterator.hasNext()) {
        ObjectReference reference = iterator.next();
        if (ref.equals(reference.referenceType())) {
          iterator.remove();
        }
      }

      if (ref2Stack.isEmpty()) {
        myPinnedSession2Reference2Stack.remove(session);
      }
    }
  }

  public void pinStacks(@NotNull XDebugSession session, @NotNull ReferenceType ref) {
    Map<ObjectReference, List<StackFrameItem>> ref2Stack = mySession2Reference2Stack.getOrDefault(session, null);
    if (ref2Stack != null) {
      Map<ObjectReference, List<StackFrameItem>> ref2StacksByReferenceType = ref2Stack.entrySet().stream()
          .filter(entry -> ref.equals(entry.getKey().referenceType()))
          .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
      myPinnedSession2Reference2Stack.put(session, ref2StacksByReferenceType);
    }
  }

  @Nullable
  private static<T> T extract(@NotNull Map<XDebugSession, Map<ObjectReference, T>> map,
                              @NotNull XDebugSession session, @NotNull ObjectReference ref) {
    Map<ObjectReference, T> ref2something = map.getOrDefault(session, null);
    if(ref2something != null) {
      return ref2something.getOrDefault(ref, null);
    }

    return null;
  }
}

@@ -1455,26 +1455,24 @@
 Item%3E%3E%3E%0A    
-  
 mySession2Re
@@ -1608,26 +1608,24 @@
 Item%3E%3E%3E%0A    
-  
 myPinnedSess
@@ -1707,16 +1707,25 @@
 Tracker(
+@NotNull 
 Project 
@@ -2070,16 +2070,22 @@
 ) %7B%0A    
+final 
 List%3CSta
@@ -2866,16 +2866,17 @@
 %7B%0A    if
+ 
 (mySessi
@@ -3086,32 +3086,38 @@
 Type ref) %7B%0A    
+final 
 Map%3CObjectRefere
@@ -3190,33 +3190,24 @@
 tack.get
-OrDefault
 (session
 , null);
@@ -3190,38 +3190,32 @@
 tack.get(session
-, null
 );%0A    if (ref2S
@@ -3236,16 +3236,22 @@
 %7B%0A      
+final 
 Iterator
@@ -3353,16 +3353,22 @@
         
+final 
 ObjectRe
@@ -3696,32 +3696,38 @@
 Type ref) %7B%0A    
+final 
 Map%3CObjectRefere
@@ -3794,33 +3794,24 @@
 tack.get
-OrDefault
 (session
 , null);
@@ -3794,38 +3794,32 @@
 tack.get(session
-, null
 );%0A    if (ref2S
@@ -3836,24 +3836,30 @@
 ll) %7B%0A      
+final 
 Map%3CObjectRe
@@ -3955,18 +3955,16 @@
         
-  
 .filter(
@@ -4016,18 +4016,16 @@
 ype()))%0A
-  
         
@@ -4210,16 +4210,17 @@
 e static
+ 
 %3CT%3E T ex
@@ -4305,32 +4305,33 @@
                 
+ 
 @NotNull XDebugS
@@ -4382,16 +4382,22 @@
 ) %7B%0A    
+final 
 Map%3CObje
@@ -4439,25 +4439,16 @@
 .get
-OrDefault
 (session
 , nu
@@ -4443,22 +4443,16 @@
 (session
-, null
 );%0A    i
@@ -4452,16 +4452,17 @@
 ;%0A    if
+ 
 (ref2som
@@ -4513,27 +4513,12 @@
 .get
-OrDefault(ref, null
+(ref
 );%0A 

Fix typos and spacing, add notnull annotation, simplify map.get calls