[Intellij IDEA Plugin Development] What's correct way to mark a PsiElement? Follow

Hello! I'm writing a codegen IDEA Plugin.

I have no idea how to put a custom data into a PsiElement.

By putting custom data into PsiElement, I can know that which PsiFile is generated by my written plugin.

I have try to use UserDataHolder.putUserData and UserDataHolder.getUserData but seemingly it did't works.

What's correct way to mark a PsiElement? Thank you! :)

Attachment

my code that writing user data into a KtFile:

val ktFile = generateKtFile(sourceKtFile)
ktFile.putUserData(markGeneratedKey, "Generated By IDEA Plugin")
(ktFile.parent as PsiDirectory).add(ktFile)

my code that reading user data into KtFile:

private fun checkGeneratedFileIfExistInDirectory(psiDirectory: PsiDirectory, fileName: String): Boolean {
    psiDirectory.children.filterIsInstance<PsiFile>().forEach loop@{ psiFile ->
        if (psiFile.name != fileName) {
            return@loop
        }
        val data = psiFile.getUserData(Constants.markGeneratedKey)
        if (data == "Generated By IDEA Plugin") {
            return true
        }
    }
    return false
}