feat(header-params): Add more test scenarios

Replace full duplicates in Header & QueryVisitors
tabilzad · Oct 10, 2024 · 64d5011 · 64d5011
1 parent 073611c
commit 64d5011
Show file tree

Hide file tree

Showing 6 changed files with 146 additions and 212 deletions.
diff --git a/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/ExpressionsVisitorK2.kt b/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/ExpressionsVisitorK2.kt
@@ -81,6 +81,7 @@ internal class ExpressionsVisitorK2(
         if (parent is EndPoint && parent.body == null) {
 
             val receiveCall = block.statements.findReceiveCallExpression()
+
             val queryParam = block.statements.findQueryParameterExpression()
             if (queryParam.isNotEmpty()) parent.parameters = parent.parameters merge queryParam.toSet()
 

diff --git a/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/K2Utils.kt b/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/K2Utils.kt
@@ -9,9 +9,12 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
+import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.expressions.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.fqName
 import org.jetbrains.kotlin.fir.resolve.toClassSymbol
@@ -211,5 +214,10 @@ fun FirRegularClassSymbol.findEnumParamValue(value: String): List<String> {
     return declarationSymbols.filterIsInstance<FirEnumEntrySymbol>().map { it.name.asString() }
 }
 
+fun FirPropertyAccessExpression.isEnum(session: FirSession): Boolean = this.dispatchReceiver
+    ?.toResolvedCallableSymbol(session)
+    ?.resolvedReturnType
+    ?.toRegularClassSymbol(session)?.isEnumClass == true
+
 
 
diff --git a/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/visitors/HeaderParamsVisitor.kt b/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/visitors/HeaderParamsVisitor.kt
@@ -2,22 +2,11 @@ package io.github.tabilzad.ktor.k2.visitors
 
 import io.github.tabilzad.ktor.k2.ClassIds
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirEvaluatorResult
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkers.declaration.primaryConstructorSymbol
-import org.jetbrains.kotlin.fir.declarations.EnumValueArgumentInfo
 import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.collectEnumEntries
-import org.jetbrains.kotlin.fir.declarations.extractEnumValueArgumentInfo
-import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
-import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
-import org.jetbrains.kotlin.fir.resolve.toClassSymbol
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
-import org.jetbrains.kotlin.fir.types.toLookupTag
-import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 
 class HeaderParamsVisitor(private val session: FirSession) : FirDefaultVisitor<Unit, MutableList<String>>() {
@@ -37,7 +26,6 @@ class HeaderParamsVisitor(private val session: FirSession) : FirDefaultVisitor<U
         }.joinToString(""))
     }
 
-
     override fun visitFunctionCall(functionCall: FirFunctionCall, data: MutableList<String>) {
         val functionFqName = functionCall.dispatchReceiver?.toResolvedCallableSymbol(session)?.callableId?.asSingleFqName()
         if (functionFqName == ClassIds.KTOR_HEADER_PARAM) {
@@ -47,7 +35,6 @@ class HeaderParamsVisitor(private val session: FirSession) : FirDefaultVisitor<U
         }
     }
 
-
     override fun visitLiteralExpression(literalExpression: FirLiteralExpression, data: MutableList<String>) {
         val element = literalExpression.value
         element?.let { data.add(it.toString()) }
@@ -69,68 +56,15 @@ class HeaderParamsVisitor(private val session: FirSession) : FirDefaultVisitor<U
     }
 
     @OptIn(PrivateConstantEvaluatorAPI::class)
-    // TODO(Look into evaluatePropertyInitializer instead of evaluateExpression)
     override fun visitArgumentList(argumentList: FirArgumentList, data: MutableList<String>) {
-
-        if (argumentList is FirResolvedArgumentList) {
-            val g = argumentList.mapping.keys
-                .filterIsInstance<FirFunctionCall>()
-                .map {
-                    FirExpressionEvaluator.evaluateExpression(it, session)
-                }.filterIsInstance<FirEvaluatorResult.Evaluated>().map {
-                    it.result
-                }.filterIsInstance<FirLiteralExpression>()
-
-            g.forEach { it.accept(this, data) }
-
-        }
-
-        argumentList.acceptChildren(this, data)
+        visitArgumentList(session, argumentList, data)
     }
 
     @OptIn(SymbolInternals::class)
     override fun visitPropertyAccessExpression(
         propertyAccessExpression: FirPropertyAccessExpression,
         data: MutableList<String>
     ) {
-
-        val isEnum =
-            propertyAccessExpression.dispatchReceiver?.toResolvedCallableSymbol(session)?.resolvedReturnType?.toRegularClassSymbol(
-                session
-            )?.isEnumClass == true
-        val enumInfo: EnumValueArgumentInfo? = propertyAccessExpression.dispatchReceiver?.extractEnumValueArgumentInfo()
-        val enumEntryAccessor = propertyAccessExpression.calleeReference.toResolvedCallableSymbol()?.name
-
-        if (isEnum) {
-
-            val entries = enumInfo?.enumClassId?.toLookupTag()?.toClassSymbol(session)?.collectEnumEntries()
-            val v = entries?.find { it.name.asString() == enumInfo.enumEntryName.asString() }
-                ?.initializerObjectSymbol
-                ?.primaryConstructorSymbol(session)
-                ?.fir?.delegatedConstructor
-
-            val paramName =
-                v?.resolvedArgumentMapping?.values?.find { it.name.asString() == enumEntryAccessor?.asString() }
-            val paramLiteral = v?.resolvedArgumentMapping?.entries?.find { it.value == paramName }?.key
-
-            val headerParam = (paramLiteral as? FirLiteralExpression)?.value
-            headerParam?.let {
-                data.add(headerParam.toString())
-            }
-        } else {
-            val calleeReference = propertyAccessExpression.calleeReference
-            if (calleeReference is FirResolvedNamedReference) {
-                val fir = calleeReference.resolvedSymbol.fir
-                if (fir is FirProperty) {
-                    val init = fir.initializer
-
-                    if (init is FirLiteralExpression) {
-                        init.accept(this, data)
-                    }
-                }
-            }
-
-        }
-
+        visitPropertyAccessExpression(session, propertyAccessExpression, data)
     }
 }
diff --git a/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/visitors/QueryParamsVisitor.kt b/create-plugin/src/main/kotlin/io/github/tabilzad/ktor/k2/visitors/QueryParamsVisitor.kt
@@ -1,6 +1,7 @@
 package io.github.tabilzad.ktor.k2.visitors
 
 import io.github.tabilzad.ktor.k2.ClassIds
+import io.github.tabilzad.ktor.k2.isEnum
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirEvaluatorResult
 import org.jetbrains.kotlin.fir.FirSession
@@ -9,15 +10,13 @@ import org.jetbrains.kotlin.fir.declarations.EnumValueArgumentInfo
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.collectEnumEntries
 import org.jetbrains.kotlin.fir.declarations.extractEnumValueArgumentInfo
-import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.resolve.toClassSymbol
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.types.toLookupTag
-import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 
 class QueryParamsVisitor(private val session: FirSession) : FirDefaultVisitor<Unit, MutableList<String>>() {
@@ -71,66 +70,75 @@ class QueryParamsVisitor(private val session: FirSession) : FirDefaultVisitor<Un
     @OptIn(PrivateConstantEvaluatorAPI::class)
     // TODO(Look into evaluatePropertyInitializer instead of evaluateExpression)
     override fun visitArgumentList(argumentList: FirArgumentList, data: MutableList<String>) {
-
-        if (argumentList is FirResolvedArgumentList) {
-            val g = argumentList.mapping.keys
-                .filterIsInstance<FirFunctionCall>()
-                .map {
-                    FirExpressionEvaluator.evaluateExpression(it, session)
-                }.filterIsInstance<FirEvaluatorResult.Evaluated>().map {
-                    it.result
-                }.filterIsInstance<FirLiteralExpression>()
-
-            g.forEach { it.accept(this, data) }
-
-        }
-
-        argumentList.acceptChildren(this, data)
+        visitArgumentList(session, argumentList, data)
     }
 
     @OptIn(SymbolInternals::class)
     override fun visitPropertyAccessExpression(
         propertyAccessExpression: FirPropertyAccessExpression,
         data: MutableList<String>
     ) {
+        visitPropertyAccessExpression(session, propertyAccessExpression, data)
+    }
+}
 
-        val isEnum =
-            propertyAccessExpression.dispatchReceiver?.toResolvedCallableSymbol(session)?.resolvedReturnType?.toRegularClassSymbol(
-                session
-            )?.isEnumClass == true
-        val enumInfo: EnumValueArgumentInfo? = propertyAccessExpression.dispatchReceiver?.extractEnumValueArgumentInfo()
-        val enumEntryAccessor = propertyAccessExpression.calleeReference.toResolvedCallableSymbol()?.name
-
-        if (isEnum) {
-
-            val entries = enumInfo?.enumClassId?.toLookupTag()?.toClassSymbol(session)?.collectEnumEntries()
-            val v = entries?.find { it.name.asString() == enumInfo.enumEntryName.asString() }
-                ?.initializerObjectSymbol
-                ?.primaryConstructorSymbol(session)
-                ?.fir?.delegatedConstructor
+@OptIn(PrivateConstantEvaluatorAPI::class)
+fun FirDefaultVisitor<Unit, MutableList<String>>.visitArgumentList(
+    session: FirSession,
+    argumentList: FirArgumentList,
+    data: MutableList<String>,
+) {
+    if (argumentList is FirResolvedArgumentList) {
+        val g = argumentList.mapping.keys
+            .filterIsInstance<FirFunctionCall>()
+            .map {
+                FirExpressionEvaluator.evaluateExpression(it, session)
+            }.filterIsInstance<FirEvaluatorResult.Evaluated>().map {
+                it.result
+            }.filterIsInstance<FirLiteralExpression>()
+
+        g.forEach { it.accept(this, data) }
+    }
 
-            val paramName =
-                v?.resolvedArgumentMapping?.values?.find { it.name.asString() == enumEntryAccessor?.asString() }
-            val paramLiteral = v?.resolvedArgumentMapping?.entries?.find { it.value == paramName }?.key
+    argumentList.acceptChildren(this, data)
+}
 
-            val queryParam = (paramLiteral as? FirLiteralExpression)?.value
-            queryParam?.let {
-                data.add(queryParam.toString())
-            }
-        } else {
-            val calleeReference = propertyAccessExpression.calleeReference
-            if (calleeReference is FirResolvedNamedReference) {
-                val fir = calleeReference.resolvedSymbol.fir
-                if (fir is FirProperty) {
-                    val init = fir.initializer
-
-                    if (init is FirLiteralExpression) {
-                        init.accept(this, data)
-                    }
+@OptIn(SymbolInternals::class)
+fun FirDefaultVisitor<Unit, MutableList<String>>.visitPropertyAccessExpression(
+    session: FirSession,
+    propertyAccessExpression: FirPropertyAccessExpression,
+    data: MutableList<String>
+) {
+    val enumInfo: EnumValueArgumentInfo? = propertyAccessExpression.dispatchReceiver?.extractEnumValueArgumentInfo()
+    val enumEntryAccessor = propertyAccessExpression.calleeReference.toResolvedCallableSymbol()?.name
+
+    if (propertyAccessExpression.isEnum(session)) {
+
+        val entries = enumInfo?.enumClassId?.toLookupTag()?.toClassSymbol(session)?.collectEnumEntries()
+        val v = entries?.find { it.name.asString() == enumInfo.enumEntryName.asString() }
+            ?.initializerObjectSymbol
+            ?.primaryConstructorSymbol(session)
+            ?.fir?.delegatedConstructor
+
+        val paramName =
+            v?.resolvedArgumentMapping?.values?.find { it.name.asString() == enumEntryAccessor?.asString() }
+        val paramLiteral = v?.resolvedArgumentMapping?.entries?.find { it.value == paramName }?.key
+
+        val queryParam = (paramLiteral as? FirLiteralExpression)?.value
+        queryParam?.let {
+            data.add(queryParam.toString())
+        }
+    } else {
+        val calleeReference = propertyAccessExpression.calleeReference
+        if (calleeReference is FirResolvedNamedReference) {
+            val fir = calleeReference.resolvedSymbol.fir
+            if (fir is FirProperty) {
+                val init = fir.initializer
+
+                if (init is FirLiteralExpression) {
+                    init.accept(this, data)
                 }
             }
-
         }
-
     }
 }