diff --git a/platform/platform-impl/src/com/intellij/ui/tree/MapBasedTree.java b/platform/platform-impl/src/com/intellij/ui/tree/MapBasedTree.java
index 4dcdebf431cb..64a8876c4b9a 100644
--- a/platform/platform-impl/src/com/intellij/ui/tree/MapBasedTree.java
+++ b/platform/platform-impl/src/com/intellij/ui/tree/MapBasedTree.java
@@ -79,17 +79,15 @@ public final class MapBasedTree<K, N> {
   public boolean updateRoot(N node) {
     if (node == (root == null ? null : root.node)) return false;
 
-    K key = getKey(node);
-    Entry<N> entry = null;
-    if (key != null) {
-      entry = new Entry<>(null, node);
-    }
-    else if (root == null) {
-      return false;
-    }
-    root = entry;
     map.clear();
-    if (key != null) map.put(key, entry);
+    if (node == null) {
+      root = null;
+    }
+    else {
+      root = new Entry<>(null, node);
+      K key = keyFunction.apply(node);
+      if (key != null) map.put(key, root);
+    }
     return true;
   }