How do I run a write action in FormattingService#formatElement/Ranges?

Hi,

Please describe what exactly you are trying to implement and what APIs/extension points you use. Also, please share your code.

I'm trying to directly implement FormattingService with the extension formattingService(or com.intellij.formattingService) to customize javadoc block tag description alignment:

import com.intellij.formatting.FormattingRangesInfo;
import com.intellij.formatting.service.CoreFormattingService;
import com.intellij.formatting.service.FormattingService;
import com.intellij.ide.highlighter.JavaFileType;
import com.intellij.lang.ImportOptimizer;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.command.WriteCommandAction;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Set;


public class ReformatEx implements FormattingService {
	private static final Logger logger = Logger.getInstance(ReformatEx.class);

	@Override
	@NotNull
	public Set<Feature> getFeatures() {
		return Set.of();
	}

	@Override
	public boolean canFormat(@NotNull PsiFile file) {
		return file.getFileType().equals(JavaFileType.INSTANCE);
	}

	@Override
	@NotNull
	public PsiElement formatElement(@NotNull PsiElement element, boolean canChangeWhiteSpaceOnly) {
		// This log is infinitely repeated until I manually cancel the formatting
		logger.info("Reformatting " + element.getClass().getName() + ":\n" + element.getText());
		ApplicationManager.getApplication().invokeLater(() -> {
			WriteCommandAction.runWriteCommandAction(element.getProject(), null, null, () -> {
				// Formatting work here
			});
		});
		return element;
	}

	@Override
	@NotNull
	public PsiElement formatElement(@NotNull PsiElement element, @NotNull TextRange range, boolean canChangeWhiteSpaceOnly) {
		logger.info("Reformatting " + element.getClass().getName() + " with range:\n" + element.getText());
		return element;
	}

	@Override
	public void formatRanges(@NotNull PsiFile file, FormattingRangesInfo rangesInfo, boolean canChangeWhiteSpaceOnly, boolean quickFormat) {
		logger.info("Reformatting " + file.getName() + ":\n" + file.getText());
	}

	@Override
	@NotNull
	public Set<ImportOptimizer> getImportOptimizers(@NotNull PsiFile file) {
		return Set.of();
	}

	@Override
	@Nullable
	public Class<? extends FormattingService> runAfter() {
		return CoreFormattingService.class;
	}
}

Most sources say that to run write actions, I should use Application#invokeLater with WriteCommandAction#runWriteCommandAction, so I used them as above in formatElement(PsiElement, boolean). However, this causes an infinite loop which IntelliJ tries to format the same file over and over, when ‘Optimize imports’ option is enabled.

I want to know if this is the right way to run write actions, and whether this behavior is a bug. Or should I always use AsyncDocumentFormattingService, which seems that I don't need to care about running write actions individually?

Below is my enviornment, if needed.

IntelliJ IDEA 2023.3.4 (Community Edition)
Build #IC-233.14475.28, built on February 13, 2024
Runtime version: 17.0.10+1-b1087.17 amd64
VM: OpenJDK 64-Bit Server VM by JetBrains s.r.o.
Windows 11.0
GC: G1 Young Generation, G1 Old Generation
Memory: 2048M
Cores: 16
Registry:
  debugger.new.tool.window.layout=true
  ide.experimental.ui=true
Non-Bundled Plugins:
  DevKit (233.14475.56)
  io.intellij-sdk-thread-access (2.2.1)
Kotlin: 233.14475.28-IJ

I don't see what can be the reason for retriggering formatting. Maybe the reason is in the formatting logic itself. Is it possible that you can share it? Also, is the formatElement triggered for the same element each time? What element is it? A file? Could you please also share the stacktrace of your formatElement method call (when it is in the loop, not only the first call)?

This went too long, so I created a file collection here: https://gist.github.com/spacedvoid/f0d657193c91d64e502c4134af20b88b

The prefixed number and underscore is just to maintain the order; you can ignore them.

JavadocFormatter.java is the code that i used to format the javadoc; it was mostly made with trial and error, and doesn't work. I was trying to traverse the javadoc PSI tree and format it.

ReformatExtension.java is the code that I used when testing Intellij, replacing the actual formatting process with a log. It still behaves the same.

1st iteration, when I just clicked ‘Reformat Code’ and the specific options popup did not appear because of the breakpoint: 1st_iteration.stacktrace

And idea.log: 1st_iteration.idea.log

Weirdly, PsiElement#getText() returned an empty string.

2nd and later, 2nd_iteration.stacktrace

And idea.log at the 2nd iteration: 2nd_iteration.idea.log

…and it had some issues while showing the whole file's text.

But strangly, after the 3rd iteration, idea.log started to only show this:

2024-03-12 18:41:31,797 [ 392320]   WARN - #com.github.spacedvoid.refex.ReformatEx - Reformatting com.intellij.psi.impl.source.PsiJavaFileImpl:
2024-03-12 18:41:31,895 [ 392418]   WARN - #com.github.spacedvoid.refex.ReformatEx - Write action in progess
2024-03-12 18:41:33,649 [ 394172]   WARN - #com.github.spacedvoid.refex.ReformatEx - Reformatting com.intellij.psi.impl.source.PsiJavaFileImpl:
2024-03-12 18:41:33,672 [ 394195]   WARN - #com.github.spacedvoid.refex.ReformatEx - Write action in progess

And after resuming from the breakpoint, it sometimes went normal, without spamming the log anymore.

But it was not the case when I applied the plugin to my own IDE: main.idea.log

And the last log entry kept repeating until I manually canceled the formatting. I just guess that it was a problem with debugging.

Here are some top of my head tips after reading through the IntelliJ Platform code. Have you questioned, if you really need a separate write action? From what I see in CodeStyleManagerImpl.java, the element has already been checked as writable. Additionally, have you looked at AbstractDocumentFormattingService.java? In lines 46-52 you see that basically only the text offset of the element is extracted, the range is reformatted and the changed element in this range is returned. And if you look at this test case, you'll see that it is a simple string replacement on the document level.

I'd say that you test gradually:

Remove the `invokeLater()` and leave only the log messages, but remove the logging of `element.getText()`. Then you can try out on which elements the methods are called (e.g. in comparison when you have selected text).

If that works without running into any infinite loop, you could try the same approach as in AbstractDocumentFormattingService. You look for the right PSI element and use its text range and the document to replace the string with the formatted version. Then committing the document and returning the element at this range.

I hope this helps investigate.

From what I see in CodeStyleManagerImpl.java, the element has already been checked as writable.

I don't know the entire behavior of the Intellij source code, but if you were talking about CheckUtil#checkWritable, it explicitly states that it checks whether the file/directory is not read-only, separate from whether I can write at the file with/without the lock.

Additionally, have you looked at AbstractDocumentFormattingService.java? In lines 46-52 you see that basically only the text offset of the element is extracted, the range is reformatted and the changed element in this range is returned.

If you look at line 48, and follow into PsiDocumentManagerBase#commitDocument, go to line 343, go to #doCommit, and go into the branches of if(ApplicationManager.getApplication().isDispatchThread()), it runs the commit in write actions.

And if you look at this test case, you'll see that it is a simple string replacement on the document level.

Since Document explicitly states that its content is a text file loaded into memory, AbstractDocumentFormattingService#formatDocument is called to modify the text in memory, and AbstractDocumentFormattingService calls PsiDocumentManagerBase#commitDocument to update the actual file, if I'm correct. And since I'm not going to make a lexer and parse the entire file, I do need the PSI tree, but I can't find anything to get a PsiFile from a Document except PsiDocumentManager#getPsiFile, and it does not clearly state how I should modify it.

And I think I made some mistakes while uploading the log files before, so I'll make an update about it.

Yes, you are right. The checkWritable was not what I thought it was after a quick glance. Could you try reworking your logic using my small debugging test example below and the example given here? My example works and I don't run into the issues you have. However, I'm using JavaPsiFacade and its PsiElementFactory to create the new PSI element for replacement.

package de.halirutan.formattingservice

import com.intellij.formatting.FormattingRangesInfo
import com.intellij.formatting.service.FormattingService
import com.intellij.lang.ImportOptimizer
import com.intellij.lang.java.JavaLanguage
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.util.TextRange
import com.intellij.psi.JavaDocTokenType.DOC_COMMENT_DATA
import com.intellij.psi.JavaPsiFacade
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.javadoc.PsiDocComment
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.elementType

class MyFormattingService: FormattingService {
    private val logger = Logger.getInstance(MyFormattingService::class.java)

    override fun getFeatures(): MutableSet<FormattingService.Feature> {
        return mutableSetOf(FormattingService.Feature.AD_HOC_FORMATTING, FormattingService.Feature.FORMAT_FRAGMENTS)
    }

    override fun canFormat(file: PsiFile): Boolean {
        return file.language == JavaLanguage.INSTANCE
    }

    override fun formatElement(element: PsiElement, canChangeWhiteSpaceOnly: Boolean): PsiElement {
        logger.warn("formatElement: $element")
        return element
    }

    override fun formatElement(element: PsiElement, range: TextRange, canChangeWhiteSpaceOnly: Boolean): PsiElement {
        logger.warn("formatElement with range: $element")
        return element
    }

    override fun formatRanges(
        file: PsiFile,
        rangesInfo: FormattingRangesInfo?,
        canChangeWhiteSpaceOnly: Boolean,
        quickFormat: Boolean
    ) {
        if (canChangeWhiteSpaceOnly) return
        if (rangesInfo == null) return
        for (range in rangesInfo.textRanges) {
            if (range == null) continue
            val docElements = PsiTreeUtil.findChildrenOfType(file, PsiDocComment::class.java)
            for (element in docElements) {
                if (element != null) {
                    if(ApplicationManager.getApplication().isDispatchThread) {
                        logger.warn("Make changes in write action")
                        ApplicationManager.getApplication().runWriteAction {
                            formatDocCommentNode(element)
                        }
                    } else {
                        logger.warn("Make changes without write action")
                        formatDocCommentNode(element)
                    }
                }
            }
        }
        logger.warn("formatRanges: ${rangesInfo.textRanges}")
    }

    private fun formatDocCommentNode(docElement: PsiDocComment) {
        val builder = StringBuilder()
        for (child in docElement.children) {
            if (child.elementType == DOC_COMMENT_DATA) {
                builder.append(child.text + " Hello there")
            } else {
                builder.append(child.text)
            }
        }
        val factory = JavaPsiFacade.getInstance(docElement.project).elementFactory
        val node = factory.createDocCommentFromText(builder.toString())
        docElement.replace(node)
    }

    override fun getImportOptimizers(file: PsiFile): MutableSet<ImportOptimizer> {
        logger.warn("getImportOptimizers: ${file.name}")
        return mutableSetOf()
    }
}

Sorry for the late update, was busy for real-life problems.

I'm highly suspecting that this entire formatting system is messed up. I can't reproduce the behavior, even with my previous code. But your code does work with the following modifications:

import com.intellij.formatting.FormattingRangesInfo;
import com.intellij.formatting.service.CoreFormattingService;
import com.intellij.formatting.service.FormattingService;
import com.intellij.ide.highlighter.JavaFileType;
import com.intellij.lang.ImportOptimizer;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.util.TextRange;
import com.intellij.psi.JavaDocTokenType;
import com.intellij.psi.JavaPsiFacade;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiElementFactory;
import com.intellij.psi.PsiFile;
import com.intellij.psi.javadoc.PsiDocComment;
import com.intellij.psi.javadoc.PsiDocToken;
import com.intellij.psi.util.PsiTreeUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Collection;
import java.util.Set;

/**
* Hello.
*/
public class ReformatEx implements FormattingService {
    private static final Logger logger = Logger.getInstance(ReformatEx.class);

    @Override
    @NotNull
    public Set<Feature> getFeatures() {
        return Set.of(Feature.AD_HOC_FORMATTING, Feature.FORMAT_FRAGMENTS);
    }

    @Override
    public boolean canFormat(@NotNull PsiFile file) {
        return file.getFileType().equals(JavaFileType.INSTANCE);
    }

    @Override
    @NotNull
    public PsiElement formatElement(@NotNull PsiElement element, boolean canChangeWhiteSpaceOnly) {
        logger.warn("Reformatting " + element.getClass().getName() + " in " + element.getContainingFile().getName() + ":\n" + element.getText());
        return element;
    }

    @Override
    @NotNull
    public PsiElement formatElement(@NotNull PsiElement element, @NotNull TextRange range, boolean canChangeWhiteSpaceOnly) {
        logger.warn("Reformatting " + element.getClass().getName() + " in " + element.getContainingFile().getName() + " with range " + range + ":\n" + element.getText());
        return element;
    }

    @Override
    public void formatRanges(@NotNull PsiFile file, FormattingRangesInfo rangesInfo, boolean canChangeWhiteSpaceOnly, boolean quickFormat) {
        logger.warn("Reformatting file " + file.getName() + " with ranges " + rangesInfo.getTextRanges() + ":\n" + file.getText());
        if(!canChangeWhiteSpaceOnly) return;
        for(TextRange range : rangesInfo.getTextRanges()) {
            if(range == null) continue;
            Collection<PsiDocComment> docElements = PsiTreeUtil.findChildrenOfType(file, PsiDocComment.class);
            for(PsiDocComment docElement : docElements) {
                if(docElement != null) {
                    if(ApplicationManager.getApplication().isDispatchThread()) {
                        logger.warn("Make changes in write action");
                        ApplicationManager.getApplication().runWriteAction(() -> formatDocCommentNode(docElement));
                    }
                    else {
                        logger.warn("Make changes without write action");
                        formatDocCommentNode(docElement);
                    }
                }
            }
        }
    }

    private void formatDocCommentNode(PsiDocComment docElement) {
        StringBuilder builder = new StringBuilder();
        for(PsiElement child : docElement.getChildren()) {
            if(child instanceof PsiDocToken token && token.getTokenType() == JavaDocTokenType.DOC_COMMENT_DATA) {
                builder.append(child.getText()).append(" Hello there");
            }
            else {
                builder.append(child.getText());
            }
        }
        PsiElementFactory factory = JavaPsiFacade.getInstance(docElement.getProject()).getElementFactory();
        PsiDocComment node = factory.createDocCommentFromText(builder.toString());
        docElement.replace(node);
    }

    @Override
    @NotNull
    public Set<ImportOptimizer> getImportOptimizers(@NotNull PsiFile file) {
        return Set.of();
    }

    @Override
    @Nullable
    public Class<? extends FormattingService> runAfter() {
        return CoreFormattingService.class;
    }
}

But not in all situations, like the comment at the class itself. Also, the formatter always calls #formatRanges with the entire file range, even if I state that no features are available.
Don't know what they've changed between 2023.3.5 and 2024.1, release notes were too long, didn't read, but I'll just use it for now. I'll make an update if this still happens. Thanks for answering my questions.

Updated to 2024.1, in case:

IntelliJ IDEA 2024.1 (Community Edition)
Build #IC-241.14494.240, built on March 28, 2024
Runtime version: 17.0.10+8-b1207.12 amd64
VM: OpenJDK 64-Bit Server VM by JetBrains s.r.o.
Windows 11.0
Kotlin analyzer version: 2.0.0-ij241-276
GC: G1 Young Generation, G1 Old Generation
Memory: 1964M
Cores: 16
Registry:
 debugger.new.tool.window.layout=true
 ide.experimental.ui=true
Non-Bundled Plugins:
 spacedvoid.ReformatExtension (0.1.0)
 DevKit (241.14494.247)
Kotlin: 241.14494.240-IJ