[kotlin] keepTypeArguments Flag (#4544)

* Similar to #4502, but for Kotlin. * Made `TypeRenderer` an object for easier global setting of `keepTypeArgument` * Improved null safety of `expr.getCalleeExpression` under `AstForExpressionsCreator.astsForNonCtorCall`
joernio · May 7, 2024 · b76b2f1 · b76b2f1
1 parent 7e92698
commit b76b2f1
Show file tree

Hide file tree

Showing 8 changed files with 170 additions and 81 deletions.
diff --git a/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/Kotlin2Cpg.scala b/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/Kotlin2Cpg.scala
@@ -7,8 +7,7 @@ import io.joern.kotlin2cpg.files.SourceFilesPicker
 import io.joern.kotlin2cpg.interop.JavasrcInterop
 import io.joern.kotlin2cpg.jar4import.UsesService
 import io.joern.kotlin2cpg.passes.*
-import io.joern.kotlin2cpg.types.ContentSourcesPicker
-import io.joern.kotlin2cpg.types.DefaultTypeInfoProvider
+import io.joern.kotlin2cpg.types.{ContentSourcesPicker, DefaultTypeInfoProvider, TypeRenderer}
 import io.joern.kotlin2cpg.utils.PathUtils
 import io.joern.x2cpg.SourceFiles
 import io.joern.x2cpg.X2CpgFrontend
@@ -228,8 +227,11 @@ class Kotlin2Cpg extends X2CpgFrontend[Config] with UsesService {
 
       new MetaDataPass(cpg, Languages.KOTLIN, config.inputPath).createAndApply()
 
+      val typeRenderer = new TypeRenderer(config.keepTypeArguments)
       val astCreator =
-        new AstCreationPass(sourceFiles, new DefaultTypeInfoProvider(environment), cpg)(config.schemaValidation)
+        new AstCreationPass(sourceFiles, new DefaultTypeInfoProvider(environment, typeRenderer), cpg)(
+          config.schemaValidation
+        )
       astCreator.createAndApply()
 
       val kotlinAstCreatorTypes = astCreator.global.usedTypes.keys().asScala.toList

diff --git a/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/Main.scala b/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/Main.scala
@@ -14,7 +14,8 @@ final case class Config(
   jar4importServiceUrl: Option[String] = None,
   includeJavaSourceFiles: Boolean = false,
   generateNodesForDependencies: Boolean = false,
-  downloadDependencies: Boolean = false
+  downloadDependencies: Boolean = false,
+  keepTypeArguments: Boolean = false
 ) extends X2CpgConfig[Config]
     with DependencyDownloadConfig[Config] {
 
@@ -49,6 +50,10 @@ final case class Config(
   override def withDownloadDependencies(value: Boolean): Config = {
     this.copy(downloadDependencies = value).withInheritedFields(this)
   }
+
+  def withKeepTypeArguments(value: Boolean): Config = {
+    copy(keepTypeArguments = value).withInheritedFields(this)
+  }
 }
 
 private object Frontend {
@@ -82,7 +87,11 @@ private object Frontend {
       opt[Unit]("generate-nodes-for-dependencies")
         .text("Generate nodes for the dependencies of the target project")
         .action((_, c) => c.withGenerateNodesForDependencies(true)),
-      DependencyDownloadConfig.parserOptions
+      DependencyDownloadConfig.parserOptions,
+      opt[Unit]("keep-type-arguments")
+        .hidden()
+        .action((_, c) => c.withKeepTypeArguments(true))
+        .text("Type full names of variables keep their type arguments.")
     )
   }
 }

diff --git a/...rontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/AstForExpressionsCreator.scala b/...rontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/AstForExpressionsCreator.scala
@@ -448,14 +448,19 @@ trait AstForExpressionsCreator(implicit withSchemaValidation: ValidationMode) {
 
     val methodFqName = if (importedNames.isDefinedAt(referencedName)) {
       importedNames(referencedName).getImportedFqName.toString
-    } else if (nameToClass.contains(expr.getCalleeExpression.getText)) {
+    } else if (Option(expr.getCalleeExpression).map(_.getText).exists(nameToClass.contains)) {
       val klass = nameToClass(expr.getCalleeExpression.getText)
       s"${klass.getContainingKtFile.getPackageFqName.toString}.$referencedName"
     } else {
       s"${expr.getContainingKtFile.getPackageFqName.toString}.$referencedName"
     }
-    val explicitSignature     = s"${TypeConstants.any}(${argAsts.map { _ => TypeConstants.any }.mkString(",")})"
-    val explicitFullName      = s"$methodFqName:$explicitSignature"
+    lazy val typeArgs =
+      expr.getTypeArguments.asScala.map(x => typeInfoProvider.typeFullName(x.getTypeReference, TypeConstants.any))
+    val explicitSignature = s"${TypeConstants.any}(${argAsts.map { _ => TypeConstants.any }.mkString(",")})"
+    val explicitFullName =
+      if (typeInfoProvider.typeRenderer.keepTypeArguments && typeArgs.nonEmpty)
+        s"$methodFqName<${typeArgs.mkString(",")}>:$explicitSignature"
+      else s"$methodFqName:$explicitSignature"
     val (fullName, signature) = typeInfoProvider.fullNameWithSignature(expr, (explicitFullName, explicitSignature))
 
     // TODO: add test case to confirm whether the ANY fallback makes sense (could be void)

diff --git a/...ontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/DefaultTypeInfoProvider.scala b/...ontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/DefaultTypeInfoProvider.scala
diff --git a/...-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/TypeInfoProvider.scala b/...-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/TypeInfoProvider.scala
@@ -26,7 +26,7 @@ import org.jetbrains.kotlin.psi.{
 
 case class AnonymousObjectContext(declaration: KtElement)
 
-trait TypeInfoProvider {
+trait TypeInfoProvider(val typeRenderer: TypeRenderer = new TypeRenderer()) {
   def isExtensionFn(fn: KtNamedFunction): Boolean
 
   def usedAsExpression(expr: KtExpression): Option[Boolean]

diff --git a/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/TypeRenderer.scala b/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/TypeRenderer.scala
@@ -12,6 +12,8 @@ import org.jetbrains.kotlin.renderer.{DescriptorRenderer, DescriptorRendererImpl
 import org.jetbrains.kotlin.types.typeUtil.TypeUtilsKt
 import org.jetbrains.kotlin.resolve.jvm.JvmPrimitiveType
 
+import scala.jdk.CollectionConverters.*
+
 object TypeRenderer {
 
   private val cpgUnresolvedType =
@@ -28,6 +30,12 @@ object TypeRenderer {
     "kotlin.ShortArray"   -> "short[]"
   )
 
+}
+
+class TypeRenderer(val keepTypeArguments: Boolean = false) {
+
+  import TypeRenderer.*
+
   private def descriptorRenderer(): DescriptorRenderer = {
     val opts = new DescriptorRendererOptionsImpl
     opts.setParameterNamesInFunctionalTypes(false)
@@ -130,9 +138,20 @@ object TypeRenderer {
             val relevantT = Option(TypeUtilsKt.getImmediateSuperclassNotAny(t)).getOrElse(t)
             stripped(renderer.renderType(relevantT))
           }
-    if (shouldMapPrimitiveArrayTypes && primitiveArrayMappings.contains(rendered)) primitiveArrayMappings(rendered)
-    else if (rendered == TypeConstants.kotlinUnit) TypeConstants.void
-    else rendered
+    val renderedType =
+      if (shouldMapPrimitiveArrayTypes && primitiveArrayMappings.contains(rendered)) primitiveArrayMappings(rendered)
+      else if (rendered == TypeConstants.kotlinUnit) TypeConstants.void
+      else rendered
+
+    if (keepTypeArguments && !t.getArguments.isEmpty) {
+      val typeArgs = t.getArguments.asScala
+        .map(_.getType)
+        .map(render(_, shouldMapPrimitiveArrayTypes, unwrapPrimitives))
+        .mkString(",")
+      s"$renderedType<$typeArgs>"
+    } else {
+      renderedType
+    }
   }
 
   private def isFunctionXType(t: KotlinType): Boolean = {

diff --git a/joern-cli/frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/TypeTests.scala b/joern-cli/frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/TypeTests.scala
@@ -1,7 +1,9 @@
 package io.joern.kotlin2cpg.querying
 
+import io.joern.kotlin2cpg.Config
 import io.joern.kotlin2cpg.testfixtures.KotlinCode2CpgFixture
-import io.shiftleft.semanticcpg.language._
+import io.shiftleft.codepropertygraph.generated.Operators
+import io.shiftleft.semanticcpg.language.*
 
 class TypeTests extends KotlinCode2CpgFixture(withOssDataflow = false) {
 
@@ -64,4 +66,48 @@ class TypeTests extends KotlinCode2CpgFixture(withOssDataflow = false) {
       x.name shouldBe "l"
     }
   }
+
+  "generics with 'keep type arguments' config" should {
+
+    "show the fully qualified type arguments for stdlib `List and `Map` objects" in {
+      val cpg = code("""
+                       |import java.util.ArrayList
+                       |import java.util.HashMap
+                       |
+                       |fun foo() {
+                       |  val stringList = ArrayList<String>()
+                       |  val stringIntMap = HashMap<String, Integer>()
+                       |}
+                       |""".stripMargin)
+        .withConfig(Config().withKeepTypeArguments(true))
+
+      cpg.identifier("stringList").typeFullName.head shouldBe "java.util.ArrayList<java.lang.String>"
+      cpg.identifier("stringIntMap").typeFullName.head shouldBe "java.util.HashMap<java.lang.String,java.lang.Integer>"
+    }
+
+    "show the fully qualified names of external types" in {
+      val cpg = code("""
+                       |import org.apache.flink.streaming.api.datastream.DataStream
+                       |import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment
+                       |import org.apache.flink.streaming.connectors.kafka.FlinkKafkaProducer
+                       |import org.apache.flink.streaming.util.serialization.SimpleStringSchema
+                       |
+                       |import java.util.Properties;
+                       |
+                       |class FlinkKafkaExample {
+                       |    fun main() {
+                       |        val kafkaProducer = FlinkKafkaProducer<String>("kafka-topic")
+                       |    }
+                       |}
+                       |""".stripMargin).withConfig(Config().withKeepTypeArguments(true))
+
+      cpg.call
+        .codeExact("FlinkKafkaProducer<String>(\"kafka-topic\")")
+        .filterNot(_.name == Operators.alloc)
+        .map(_.methodFullName)
+        .head shouldBe "org.apache.flink.streaming.connectors.kafka.FlinkKafkaProducer<java.lang.String>:ANY(ANY)"
+    }
+
+  }
+
 }
diff --git a/...s/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/testfixtures/KotlinCodeToCpgFixture.scala b/...s/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/testfixtures/KotlinCodeToCpgFixture.scala
@@ -19,15 +19,22 @@ trait KotlinFrontend extends LanguageFrontend {
   protected val withTestResourcePaths: Boolean
 
   override val fileSuffix: String = ".kt"
+  private lazy val defaultContentRoot =
+    BFile(ProjectRoot.relativise("joern-cli/frontends/kotlin2cpg/src/test/resources/jars/"))
+  private lazy val defaultConfig: Config =
+    Config(
+      classpath = if (withTestResourcePaths) Set(defaultContentRoot.path.toAbsolutePath.toString) else Set(),
+      includeJavaSourceFiles = true
+    )
 
   override def execute(sourceCodeFile: File): Cpg = {
-    val defaultContentRoot =
-      BFile(ProjectRoot.relativise("joern-cli/frontends/kotlin2cpg/src/test/resources/jars/"))
-    implicit val defaultConfig: Config =
-      Config(
-        classpath = if (withTestResourcePaths) Set(defaultContentRoot.path.toAbsolutePath.toString) else Set(),
-        includeJavaSourceFiles = true
-      )
+    implicit val config: Config = getConfig() match {
+      case Some(config: Config) => config
+      case _ =>
+        setConfig(defaultConfig)
+        defaultConfig
+    }
+
     new Kotlin2Cpg().createCpg(sourceCodeFile.getAbsolutePath).get
   }
 }