kotlin2cpg: set arg names on when-expressions (#3468)

joernio · Aug 10, 2023 · 54dff55 · 54dff55
1 parent 8cccc17
commit 54dff55
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/AstCreator.scala b/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/AstCreator.scala
@@ -201,7 +201,7 @@ class AstCreator(fileWithMeta: KtFileWithMeta, xTypeInfoProvider: TypeInfoProvid
       case typedExpr: KtThisExpression  => Seq(astForThisExpression(typedExpr, argIdxMaybe, argNameMaybe))
       case typedExpr: KtThrowExpression => Seq(astForUnknown(typedExpr, argIdxMaybe))
       case typedExpr: KtTryExpression   => Seq(astForTry(typedExpr, argIdxMaybe))
-      case typedExpr: KtWhenExpression  => Seq(astForWhen(typedExpr, argIdxMaybe))
+      case typedExpr: KtWhenExpression  => Seq(astForWhen(typedExpr, argIdxMaybe, argNameMaybe))
       case typedExpr: KtWhileExpression => Seq(astForWhile(typedExpr))
       case typedExpr: KtNamedFunction if Option(typedExpr.getName).isEmpty =>
         Seq(astForAnonymousFunction(typedExpr, argIdxMaybe))

diff --git a/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/KtPsiToAst.scala b/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/KtPsiToAst.scala
@@ -941,7 +941,7 @@ trait KtPsiToAst {
       } else if (expr.getInitializer.isInstanceOf[KtPostfixExpression]) {
         astForPostfixExpression(expr.getInitializer.asInstanceOf[KtPostfixExpression], None, None)
       } else if (expr.getInitializer.isInstanceOf[KtWhenExpression]) {
-        astForWhenAsExpression(expr.getInitializer.asInstanceOf[KtWhenExpression], None)
+        astForWhenAsExpression(expr.getInitializer.asInstanceOf[KtWhenExpression], None, None)
       } else if (expr.getInitializer.isInstanceOf[KtIfExpression]) {
         astForIfAsExpression(expr.getInitializer.asInstanceOf[KtIfExpression], None, None)
       } else {
@@ -1703,10 +1703,12 @@ trait KtPsiToAst {
     }
   }
 
-  def astForWhenAsExpression(expr: KtWhenExpression, argIdx: Option[Int])(implicit
+  def astForWhenAsExpression(expr: KtWhenExpression, argIdx: Option[Int], argNameMaybe: Option[String])(implicit
     typeInfoProvider: TypeInfoProvider
   ): Ast = {
-    val callNode = withArgumentIndex(operatorCallNode("<operator>.when", "<operator>.when", None), argIdx)
+    val callNode =
+      withArgumentIndex(operatorCallNode("<operator>.when", "<operator>.when", None), argIdx)
+        .argumentName(argNameMaybe)
 
     val subjectExpressionAsts = astsForExpression(expr.getSubjectExpression, None)
     val subjectBlock          = blockNode(expr.getSubjectExpression, "", "")
@@ -1727,9 +1729,11 @@ trait KtPsiToAst {
     callAst(callNode, List(subjectBlockAst) ++ argAsts)
   }
 
-  def astForWhen(expr: KtWhenExpression, argIdx: Option[Int])(implicit typeInfoProvider: TypeInfoProvider): Ast = {
+  def astForWhen(expr: KtWhenExpression, argIdx: Option[Int], argNameMaybe: Option[String])(implicit
+    typeInfoProvider: TypeInfoProvider
+  ): Ast = {
     typeInfoProvider.usedAsExpression(expr) match {
-      case Some(true) => astForWhenAsExpression(expr, argIdx)
+      case Some(true) => astForWhenAsExpression(expr, argIdx, argNameMaybe)
       case _          => astForWhenAsStatement(expr, argIdx)
     }
   }

diff --git a/joern-cli/frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/CallTests.scala b/joern-cli/frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/CallTests.scala
@@ -413,4 +413,17 @@ class CallTests extends KotlinCode2CpgFixture(withOssDataflow = false) {
       c.argument.map(_.argumentName).flatten.l shouldBe List("two", "one")
     }
   }
+
+  "CPG for code with simple call having an when-expression passed in with an argument name" should {
+    val cpg = code("""
+      |package mypkg
+      |fun f1(one: String, two: String)  = println(one + " " + two)
+      |fun f2() = f1(two = "this",  one = when(Random(1).nextBoolean()) { true -> "that" else -> "thatother" })
+      |""".stripMargin)
+
+    "should contain a CALL node with arguments with their ARGUMENT_NAME property set" in {
+      val List(c: Call) = cpg.method.nameExact("f1").callIn.l
+      c.argument.map(_.argumentName).flatten.l shouldBe List("two", "one")
+    }
+  }
 }