Fix NBTT reformatting for read-only files

Author: DenWav

build.gradle.kts

@@ -86,6 +86,10 @@ dependencies {
 
     compile(files(gradleToolingExtensionJar))
 
+    // gradle-intellij-plugin doesn't attach sources properly for Kotlin unless it's manually defined here
+    // compileOnly since it's only here for reference
+    compileOnly(kotlin("stdlib-jdk8"))
+
     "jflex"("org.jetbrains.idea:jflex:1.7.0-b7f882a")
     "jflex-skeleton"("org.jetbrains.idea:jflex:1.7.0-c1fdf11:idea@skeleton")
     "grammar-kit"("org.jetbrains.idea:grammar-kit:1.5.1")

src/main/kotlin/com/demonwav/mcdev/nbt/NbtVirtualFile.kt

@@ -13,12 +13,19 @@ package com.demonwav.mcdev.nbt
 import com.demonwav.mcdev.nbt.editor.CompressionSelection
 import com.demonwav.mcdev.nbt.editor.NbtToolbar
 import com.demonwav.mcdev.nbt.lang.NbttFile
+import com.demonwav.mcdev.nbt.lang.NbttLanguage
+import com.demonwav.mcdev.util.invokeLaterAny
+import com.demonwav.mcdev.util.runWriteAction
 import com.intellij.notification.Notification
 import com.intellij.notification.NotificationType
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.PsiManager
+import com.intellij.psi.codeStyle.CodeStyleManager
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
 import java.io.DataOutputStream
@@ -33,20 +40,37 @@ class NbtVirtualFile(private val backingFile: VirtualFile, private val project:
     val parseSuccessful: Boolean
 
     init {
+        this.bytes = byteArrayOf()
+        var text: String
+
         var tempCompressed: Boolean
         var tempParseSuccessful: Boolean
         try {
             val (rootCompound, isCompressed) = Nbt.buildTagTree(backingFile.inputStream, TimeUnit.SECONDS.toMillis(1))
-            this.bytes = rootCompound.toString().toByteArray()
+            text = rootCompound.toString()
             tempCompressed = isCompressed
             tempParseSuccessful = true
         } catch (e: MalformedNbtFileException) {
-            this.bytes = "Malformed NBT file:\n${e.message}".toByteArray()
+            text = "Malformed NBT file:\n${e.message}"
             tempCompressed = false
             tempParseSuccessful = false
         }
+
         this.isCompressed = tempCompressed
         this.parseSuccessful = tempParseSuccessful
+
+        if (this.parseSuccessful) {
+            val psiFile = PsiFileFactory.getInstance(project).createFileFromText(NbttLanguage, text)
+            invokeLaterAny {
+                psiFile.runWriteAction {
+                    this.bytes = PsiDocumentManager.getInstance(project).getDocument(
+                        CodeStyleManager.getInstance(project).reformat(psiFile, true) as PsiFile
+                    )?.immutableCharSequence?.toString()?.toByteArray() ?: byteArrayOf()
+                }
+            }
+        } else {
+            this.bytes = text.toByteArray()
+        }
     }
 
     override fun refresh(asynchronous: Boolean, recursive: Boolean, postRunnable: Runnable?) {
@@ -69,7 +93,7 @@ class NbtVirtualFile(private val backingFile: VirtualFile, private val project:
     override fun getOutputStream(requestor: Any, newModificationStamp: Long, newTimeStamp: Long) =
         VfsUtilCore.outputStreamAddingBOM(NbtOutputStream(this, requestor), this)
 
-    fun writeFile(requestor: Any) {
+    fun writeFile(requester: Any) {
         val nbttFile = PsiManager.getInstance(project).findFile(this) as NbttFile
         val rootTag = nbttFile.getRootCompound()?.getRootCompoundTag()
 
@@ -83,11 +107,14 @@ class NbtVirtualFile(private val backingFile: VirtualFile, private val project:
             return
         }
 
+        this.bytes = PsiDocumentManager.getInstance(project).getDocument(nbttFile)?.immutableCharSequence?.toString()?.toByteArray()
+            ?: byteArrayOf()
+
         // just to be safe
         this.parent.bom = null
         val filteredStream = when (toolbar.selection) {
-            CompressionSelection.GZIP -> GZIPOutputStream(this.parent.getOutputStream(requestor))
-            CompressionSelection.UNCOMPRESSED -> this.parent.getOutputStream(requestor)
+            CompressionSelection.GZIP -> GZIPOutputStream(this.parent.getOutputStream(requester))
+            CompressionSelection.UNCOMPRESSED -> this.parent.getOutputStream(requester)
         }
 
         DataOutputStream(filteredStream).use { stream ->

src/main/kotlin/com/demonwav/mcdev/nbt/editor/NbtFileEditorProvider.kt

@@ -12,7 +12,6 @@ package com.demonwav.mcdev.nbt.editor
 
 import com.demonwav.mcdev.nbt.NbtVirtualFile
 import com.demonwav.mcdev.nbt.filetype.NbtFileType
-import com.intellij.codeInsight.actions.ReformatCodeProcessor
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorPolicy
 import com.intellij.openapi.fileEditor.FileEditorState
@@ -23,7 +22,6 @@ import com.intellij.openapi.project.DumbAware
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.psi.PsiManager
 import java.awt.BorderLayout
 import java.beans.PropertyChangeListener
 import javax.swing.JPanel
@@ -40,11 +38,6 @@ class NbtFileEditorProvider : PsiAwareTextEditorProvider(), DumbAware {
             NonProjectFileWritingAccessProvider.allowWriting(nbtFile)
         }
 
-        // Format to their format preferences
-        PsiManager.getInstance(project).findFile(nbtFile)?.let {
-            ReformatCodeProcessor(it, false).run()
-        }
-
         nbtFile.toolbar = fileEditor.toolbar
         return fileEditor
     }
@@ -54,7 +47,7 @@ class NbtFileEditorProvider : PsiAwareTextEditorProvider(), DumbAware {
     }
 }
 
-private class NbtFileEditor(private val editor: FileEditor, nbtFile: NbtVirtualFile) : FileEditor {
+private class NbtFileEditor(private val editor: FileEditor, private val nbtFile: NbtVirtualFile) : FileEditor {
 
     val toolbar = NbtToolbar(nbtFile)
     private val component = JPanel(BorderLayout())

src/main/kotlin/com/demonwav/mcdev/nbt/editor/NbtToolbar.kt

@@ -19,7 +19,7 @@ import javax.swing.JPanel
 class NbtToolbar(nbtFile: NbtVirtualFile) {
     lateinit var panel: JPanel
     private lateinit var compressionBox: JComboBox<CompressionSelection>
-    private lateinit var saveButton: JButton
+    lateinit var saveButton: JButton
 
     private var lastSelection: CompressionSelection
 
@@ -35,7 +35,7 @@ class NbtToolbar(nbtFile: NbtVirtualFile) {
             compressionBox.isEnabled = false
         } else {
             compressionBox.addActionListener {
-                saveButton.isVisible = lastSelection != selection
+                checkModified()
             }
         }
 
@@ -44,14 +44,19 @@ class NbtToolbar(nbtFile: NbtVirtualFile) {
         }
 
         saveButton.addActionListener {
-            saveButton.isVisible = false
+            lastSelection = selection
+            checkModified()
 
             runWriteTaskLater {
                 nbtFile.writeFile(this)
             }
         }
     }
 
+    private fun checkModified() {
+        saveButton.isVisible = lastSelection != selection
+    }
+
     val selection
         get() = compressionBox.selectedItem as CompressionSelection
 }

src/main/kotlin/com/demonwav/mcdev/util/utils.kt

@@ -19,21 +19,22 @@ import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
+import com.intellij.openapi.util.Computable
+import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import org.jetbrains.annotations.Contract
 
-inline fun runWriteTask(crossinline func: () -> Unit) {
+inline fun <T : Any?> runWriteTask(crossinline func: () -> T): T =
     if (ApplicationManager.getApplication().isWriteAccessAllowed) {
         func()
     } else {
         invokeAndWait {
-            ApplicationManager.getApplication().runWriteAction {
+            ApplicationManager.getApplication().runWriteAction(Computable {
                 func()
-            }
+            })
         }
     }
-}
 
 inline fun runWriteTaskLater(crossinline func: () -> Unit) {
     if (ApplicationManager.getApplication().isWriteAccessAllowed) {
@@ -47,13 +48,14 @@ inline fun runWriteTaskLater(crossinline func: () -> Unit) {
     }
 }
 
-inline fun invokeAndWait(crossinline func: () -> Unit) {
+inline fun <T : Any?> invokeAndWait(crossinline func: () -> T): T =
     if (ApplicationManager.getApplication().isDispatchThread) {
         func()
     } else {
-        ApplicationManager.getApplication().invokeAndWait({ func() }, ModalityState.defaultModalityState())
+        val ref = Ref<T>()
+        ApplicationManager.getApplication().invokeAndWait({ ref.set(func()) }, ModalityState.defaultModalityState())
+        ref.get()
     }
-}
 
 inline fun invokeLater(crossinline func: () -> Unit) {
     if (ApplicationManager.getApplication().isDispatchThread) {
@@ -76,7 +78,8 @@ inline fun <T : Any?> PsiFile.runWriteAction(crossinline func: () -> T) =
 
 inline fun <T : Any?> PsiFile.applyWriteAction(crossinline func: PsiFile.() -> T): T {
     val result = WriteCommandAction.writeCommandAction(this).withGlobalUndo().compute<T, Throwable> { func() }
-    PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(FileDocumentManager.getInstance().getDocument(this.virtualFile) ?: return result)
+    PsiDocumentManager.getInstance(project)
+        .doPostponedOperationsAndUnblockDocument(FileDocumentManager.getInstance().getDocument(this.virtualFile) ?: return result)
     return result
 }
 
@@ -96,13 +99,13 @@ inline fun <T : Collection<*>> T.ifEmpty(func: () -> Unit): T {
 
 @Contract(pure = true)
 inline fun <T, R> Iterable<T>.mapFirstNotNull(transform: (T) -> R?): R? {
-    forEach { transform(it)?.let { return it } }
+    forEach { element -> transform(element)?.let { return it } }
     return null
 }
 
 @Contract(pure = true)
 inline fun <T, R> Array<T>.mapFirstNotNull(transform: (T) -> R?): R? {
-    forEach { transform(it)?.let { return it } }
+    forEach { element -> transform(element)?.let { return it } }
     return null
 }