Add Postfix

Author: kyay10

prettify-kotlin-annotation/src/commonMain/kotlin/io/github/kyay10/prettifykotlin/annotations.kt

@@ -1,4 +1,5 @@
 package io.github.kyay10.prettifykotlin
 
 annotation class Pretty(val name: String)
-annotation class Prefix(val prefix: String, val postfix: String)
+annotation class Prefix(val prefix: String, val suffix: String)
+annotation class Postfix(val suffix: String)

src/main/kotlin/io/github/kyay10/prettifykotlin/PrettyFoldingBuilder.kt

@@ -12,17 +12,20 @@ import org.jetbrains.kotlin.idea.util.findAnnotation
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtLabelReferenceExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtSuperExpression
 import org.jetbrains.kotlin.psi.KtThisExpression
 import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
 import org.jetbrains.kotlin.psi.ValueArgument
+import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.uast.UExpression
 import org.jetbrains.uast.toUElementOfType
 
 val PRETTY_FQNAME = FqName("io.github.kyay10.prettifykotlin.Pretty")
 val PREFIX_FQNAME = FqName("io.github.kyay10.prettifykotlin.Prefix")
+val POSTFIX_FQNAME = FqName("io.github.kyay10.prettifykotlin.Postfix")
 
 class PrettyFoldingBuilder : FoldingBuilderEx() {
   override fun buildFoldRegions(root: PsiElement, document: Document, quick: Boolean): Array<FoldingDescriptor> =
@@ -60,6 +63,19 @@ class PrettyFoldingBuilder : FoldingBuilderEx() {
           add(PrettyFoldingDescriptor(leftPar, prefix, group, neverExpands = false))
           add(PrettyFoldingDescriptor(rightPar, suffix, group, neverExpands = false))
         }
+
+        override fun visitDotQualifiedExpression(expression: KtDotQualifiedExpression) = impure {
+          super.visitDotQualifiedExpression(expression)
+          val selector = expression.selectorExpression
+          ensure(selector is KtSimpleNameExpression)
+          val reference = selector.mainReference.bind().resolve()
+          ensure(reference is KtAnnotated)
+
+          val annotation = reference.findAnnotation(POSTFIX_FQNAME).bind()
+          val suffix = annotation.valueArguments.singleOrNull().bind().constValue
+          ensure(suffix is String)
+          add(PrettyFoldingDescriptor(expression.allChildren.first { it.text == "." }, suffix))
+        }
       })
     }.toTypedArray()