IDEA-89342 Folding: Restore custom fold regions on file re-open

platform/lang-impl/src/com/intellij/codeInsight/folding/impl/DocumentFoldingInfo.java

@@ -84,9 +84,10 @@ public class DocumentFoldingInfo implements JDOMExternalizable, CodeFoldingState
           myPsiElementsOrRangeMarkers.add(element);
         }
         else if (region.isValid()) {
-          myPsiElementsOrRangeMarkers.add(region);
+          RangeMarker marker = editor.getDocument().createRangeMarker(region.getStartOffset(), region.getEndOffset());
+          myPsiElementsOrRangeMarkers.add(marker);
           String placeholderText = region.getPlaceholderText();
-          myPlaceholderTexts.put(region, placeholderText);
+          myPlaceholderTexts.put(marker, placeholderText);
         }
         myExpandedStates.add(expanded ? Boolean.TRUE : Boolean.FALSE);
       }
@@ -269,4 +270,47 @@ public class DocumentFoldingInfo implements JDOMExternalizable, CodeFoldingState
     if (!myFile.isValid()) return "";
     return Long.toString(myFile.getTimeStamp());
   }
+
+  @Override
+  public int hashCode() {
+    int result = myProject != null ? myProject.hashCode() : 0;
+    result = 31 * result + (myFile != null ? myFile.hashCode() : 0);
+    result = 31 * result + (myPsiElementsOrRangeMarkers != null ? myPsiElementsOrRangeMarkers.hashCode() : 0);
+    result = 31 * result + (myExpandedStates != null ? myExpandedStates.hashCode() : 0);
+    result = 31 * result + (myPlaceholderTexts != null ? myPlaceholderTexts.hashCode() : 0);
+    return result;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+
+    DocumentFoldingInfo info = (DocumentFoldingInfo)o;
+
+    if (myExpandedStates != null ? !myExpandedStates.equals(info.myExpandedStates) : info.myExpandedStates != null) {
+      return false;
+    }
+    if (myFile != null ? !myFile.equals(info.myFile) : info.myFile != null) {
+      return false;
+    }
+    if (myPlaceholderTexts != null ? !myPlaceholderTexts.equals(info.myPlaceholderTexts) : info.myPlaceholderTexts != null) {
+      return false;
+    }
+    if (myProject != null ? !myProject.equals(info.myProject) : info.myProject != null) {
+      return false;
+    }
+    if (myPsiElementsOrRangeMarkers != null
+        ? !myPsiElementsOrRangeMarkers.equals(info.myPsiElementsOrRangeMarkers)
+        : info.myPsiElementsOrRangeMarkers != null)
+    {
+      return false;
+    }
+
+    return true;
+  }
 }

platform/lang-impl/src/com/intellij/openapi/fileEditor/impl/text/PsiAwareTextEditorProvider.java

@@ -33,6 +33,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.WriteExternalException;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiDocumentManager;
+import com.intellij.util.Producer;
 import org.jdom.Element;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
@@ -57,14 +58,22 @@ public class PsiAwareTextEditorProvider extends TextEditorProvider {
     // Foldings
     Element child = element.getChild(FOLDING_ELEMENT);
     Document document = FileDocumentManager.getInstance().getCachedDocument(file);
-    if (child != null && document != null) {
-      //PsiDocumentManager.getInstance(project).commitDocument(document);
-      state.FOLDING_STATE = CodeFoldingManager.getInstance(project).readFoldingState(child, document);
+    if (child != null) {
+      if (document == null) {
+        final Element detachedStateCopy = (Element)child.clone();
+        state.setDelayedFoldState(new Producer<CodeFoldingState>() {
+          @Override
+          public CodeFoldingState produce() {
+            Document document = FileDocumentManager.getInstance().getCachedDocument(file);
+            return document == null ? null : CodeFoldingManager.getInstance(project).readFoldingState(detachedStateCopy, document);
+          }
+        });
+      }
+      else {
+        //PsiDocumentManager.getInstance(project).commitDocument(document);
+        state.setFoldingState(CodeFoldingManager.getInstance(project).readFoldingState(child, document));
+      }
     }
-    else {
-      state.FOLDING_STATE = null;
-    }
-
     return state;
   }
 
@@ -74,10 +83,11 @@ public class PsiAwareTextEditorProvider extends TextEditorProvider {
     TextEditorState state = (TextEditorState)_state;
 
     // Foldings
-    if (state.FOLDING_STATE != null) {
+    CodeFoldingState foldingState = state.getFoldingState();
+    if (foldingState != null) {
       Element e = new Element(FOLDING_ELEMENT);
       try {
-        CodeFoldingManager.getInstance(project).writeFoldingState(state.FOLDING_STATE, e);
+        CodeFoldingManager.getInstance(project).writeFoldingState(foldingState, e);
       }
       catch (WriteExternalException e1) {
         //ignore
@@ -94,10 +104,10 @@ public class PsiAwareTextEditorProvider extends TextEditorProvider {
       // Folding
       if (project != null && !editor.isDisposed()) {
         PsiDocumentManager.getInstance(project).commitDocument(editor.getDocument());
-        state.FOLDING_STATE = CodeFoldingManager.getInstance(project).saveFoldingState(editor);
+        state.setFoldingState(CodeFoldingManager.getInstance(project).saveFoldingState(editor));
       }
       else {
-        state.FOLDING_STATE = null;
+        state.setFoldingState(null);
       }
     }
 
@@ -107,12 +117,13 @@ public class PsiAwareTextEditorProvider extends TextEditorProvider {
   protected void setStateImpl(final Project project, final Editor editor, final TextEditorState state) {
     super.setStateImpl(project, editor, state);
     // Folding
-    if (project != null && state.FOLDING_STATE != null){
+    final CodeFoldingState foldState = state.getFoldingState();
+    if (project != null && foldState != null){
       PsiDocumentManager.getInstance(project).commitDocument(editor.getDocument());
       editor.getFoldingModel().runBatchFoldingOperation(
         new Runnable() {
           public void run() {
-            CodeFoldingManager.getInstance(project).restoreFoldingState(editor, state.FOLDING_STATE);
+            CodeFoldingManager.getInstance(project).restoreFoldingState(editor, foldState);
           }
         }
       );

platform/platform-impl/src/com/intellij/openapi/fileEditor/impl/text/CodeFoldingState.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2009 JetBrains s.r.o.
+ * Copyright 2000-2012 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.
@@ -15,5 +15,8 @@
  */
 package com.intellij.openapi.fileEditor.impl.text;
 
+/**
+ * Implementations of this interface are expected to provide correct {@link #equals(Object)} & {@link #hashCode()} implementations.
+ */
 public interface CodeFoldingState {
 }

platform/platform-impl/src/com/intellij/openapi/fileEditor/impl/text/TextEditorState.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2009 JetBrains s.r.o.
+ * Copyright 2000-2012 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.
@@ -15,30 +15,65 @@
  */
 package com.intellij.openapi.fileEditor.impl.text;
 
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.fileEditor.FileEditorState;
 import com.intellij.openapi.fileEditor.FileEditorStateLevel;
+import com.intellij.util.Producer;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * @author Vladimir Kondratyev
  */
 public final class TextEditorState implements FileEditorState {
 
-  public int LINE;
-  public int COLUMN;
-  public float VERTICAL_SCROLL_PROPORTION;
-  public int SELECTION_START;
-  public int SELECTION_END;
+  public int              LINE;
+  public int              COLUMN;
+  public float            VERTICAL_SCROLL_PROPORTION;
+  public int              SELECTION_START;
+  public int              SELECTION_END;
   /**
    * State which describes how editor is folded.
    * This field can be <code>null</code>.
    */
-  public CodeFoldingState FOLDING_STATE;
+  private CodeFoldingState myFoldingState;
+  @Nullable private Producer<CodeFoldingState> myDelayedFoldInfoProducer;
 
   private static final int MIN_CHANGE_DISTANCE = 4;
 
   public TextEditorState() {
   }
 
+  /**
+   * Folding state is more complex than, say, line/column number, that's why it's deserialization can be performed only when
+   * necessary pre-requisites are met (e.g. corresponding {@link Document} is created).
+   * <p/>
+   * However, we can't be sure that those conditions are met on IDE startup (when editor states are read). Current method allows
+   * to register a closure within the current state object which returns folding info if possible.
+   *
+   * @param producer  delayed folding info producer
+   */
+  public void setDelayedFoldState(@NotNull Producer<CodeFoldingState> producer) {
+    myDelayedFoldInfoProducer = producer;
+  }
+
+  @Nullable
+  public CodeFoldingState getFoldingState() {
+    // Assuming single-thread access here.
+    if (myFoldingState == null && myDelayedFoldInfoProducer != null) {
+      myFoldingState = myDelayedFoldInfoProducer.produce();
+      if (myFoldingState != null) {
+        myDelayedFoldInfoProducer = null;
+      }
+    }
+    return myFoldingState;
+  }
+
+  public void setFoldingState(@Nullable CodeFoldingState foldingState) {
+    myFoldingState = foldingState;
+    myDelayedFoldInfoProducer = null;
+  }
+
   public boolean equals(Object o) {
     if (!(o instanceof TextEditorState)) {
       return false;
@@ -51,6 +86,9 @@ public final class TextEditorState implements FileEditorState {
     if (VERTICAL_SCROLL_PROPORTION != textEditorState.VERTICAL_SCROLL_PROPORTION) return false;
     if (SELECTION_START != textEditorState.SELECTION_START) return false;
     if (SELECTION_END != textEditorState.SELECTION_END) return false;
+    CodeFoldingState localFoldingState = getFoldingState();
+    CodeFoldingState theirFoldingState = textEditorState.getFoldingState();
+    if (localFoldingState == null ? theirFoldingState != null : !localFoldingState.equals(theirFoldingState)) return false;
 
     return true;
   }
@@ -68,5 +106,4 @@ public final class TextEditorState implements FileEditorState {
   public String toString() {
     return "[" + LINE + "," + COLUMN + "]";
   }
-
 }