From 002b9cb3cb3ee02f60371efee8a2c1d119ee70fe Mon Sep 17 00:00:00 2001
From: Alexey Kudravtsev <cdr@intellij.com>
Date: Thu, 1 Sep 2011 13:25:41 +0400
Subject: [PATCH] sql tests

---
 .../com/intellij/psi/impl/source/text/DiffLog.java  | 13 +++++++++++--
 .../com/intellij/psi/impl/source/tree/TreeUtil.java |  7 ++++++-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/platform/lang-impl/src/com/intellij/psi/impl/source/text/DiffLog.java b/platform/lang-impl/src/com/intellij/psi/impl/source/text/DiffLog.java
index 9d1d2b441ac3..72b7239c9575 100644
--- a/platform/lang-impl/src/com/intellij/psi/impl/source/text/DiffLog.java
+++ b/platform/lang-impl/src/com/intellij/psi/impl/source/text/DiffLog.java
@@ -33,6 +33,7 @@ import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.impl.source.tree.CompositeElement;
 import com.intellij.psi.impl.source.tree.FileElement;
 import com.intellij.psi.impl.source.tree.TreeElement;
+import com.intellij.psi.impl.source.tree.TreeUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.diff.DiffTreeChangeBuilder;
 import org.jetbrains.annotations.NotNull;
@@ -148,7 +149,7 @@ public class DiffLog implements DiffTreeChangeBuilder<ASTNode,ASTNode> {
         event.setParent(psiParent);
         event.setFile(file);
         event.setOldChild(psiOldChild);
-        PsiElement psiNewChild = newNode.getPsi();
+        PsiElement psiNewChild = getPsi(newNode, file);
         event.setNewChild(psiNewChild);
         ((PsiManagerEx)file.getManager()).beforeChildReplacement(event);
       }
@@ -224,7 +225,7 @@ public class DiffLog implements DiffTreeChangeBuilder<ASTNode,ASTNode> {
       }
 
       PsiElement psiParent = myOldParent.getPsi();
-      PsiElement psiChild = file.isPhysical() ? myNewNode.getPsi() : null;
+      PsiElement psiChild = getPsi(myNewNode, file);
       if (psiParent != null && psiChild != null) {
         PsiTreeChangeEventImpl event = new PsiTreeChangeEventImpl(file.getManager());
         event.setParent(psiParent);
@@ -253,6 +254,14 @@ public class DiffLog implements DiffTreeChangeBuilder<ASTNode,ASTNode> {
 
       DebugUtil.checkTreeStructure(myOldParent);
     }
+
+  }
+
+  private static PsiElement getPsi(ASTNode node, PsiFile file) {
+    node.putUserData(TreeUtil.CONTAINING_FILE_KEY_AFTER_REPARSE, ((PsiFileImpl)file).getTreeElement());
+    PsiElement psiChild = file.isPhysical() ? node.getPsi() : null;
+    node.putUserData(TreeUtil.CONTAINING_FILE_KEY_AFTER_REPARSE, null);
+    return psiChild;
   }
 
   private static class ReplaceFileElement extends LogEntry {
diff --git a/platform/lang-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java b/platform/lang-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java
index 95768c9edacd..74148224604e 100644
--- a/platform/lang-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java
+++ b/platform/lang-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java
@@ -228,10 +228,15 @@ public class TreeUtil {
     return nextLeaf((TreeElement)node, null);
   }
 
-  public static FileElement getFileElement(TreeElement parent) {
+  public static Key<FileElement> CONTAINING_FILE_KEY_AFTER_REPARSE = Key.create("CONTAINING_FILE_KEY_AFTER_REPARSE");
+  public static FileElement getFileElement(TreeElement element) {
+    TreeElement parent = element;
     while (parent != null && !(parent instanceof FileElement)) {
       parent = parent.getTreeParent();
     }
+    if (parent == null) {
+      parent = element.getUserData(CONTAINING_FILE_KEY_AFTER_REPARSE);
+    }
     return (FileElement)parent;
   }