pysrc2cpg: Local Node Line & Column Numbers (#2314)

* Extract the earliest identifier node line and column number info and add this to the local node

Author: DavidBakerEffendi

joern-cli/frontends/pysrc2cpg/src/main/scala/io/joern/pysrc2cpg/ContextStack.scala

@@ -1,14 +1,25 @@
 package io.joern.pysrc2cpg
 
-import io.shiftleft.codepropertygraph.generated.nodes
-import io.shiftleft.codepropertygraph.generated.nodes.{NewClosureBinding, NewIdentifier, NewLocal, NewMethod, NewNode}
+import io.joern.pysrc2cpg.ContextStack.transferLineColInfo
 import io.joern.pysrc2cpg.memop._
+import io.shiftleft.codepropertygraph.generated.nodes
+import io.shiftleft.codepropertygraph.generated.nodes._
 import org.slf4j.LoggerFactory
 
 import scala.collection.mutable
 
 object ContextStack {
   private val logger = LoggerFactory.getLogger(getClass)
+
+  def transferLineColInfo(src: NewIdentifier, tgt: NewLocal): Unit = {
+    src.lineNumber match {
+      // If there are multiple occurrences and the local is already set, ignore later updates
+      case Some(srcLineNo) if tgt.lineNumber.isEmpty || !tgt.lineNumber.exists(_ < srcLineNo) =>
+        tgt.lineNumber(src.lineNumber)
+        tgt.columnNumber(src.columnNumber)
+      case _ =>
+    }
+  }
 }
 
 class ContextStack {
@@ -156,6 +167,7 @@ class ContextStack {
         !moduleMethodContext.get.variables.contains(name)
       ) {
         val localNode = createLocal(name, None)
+        transferLineColInfo(identifier, localNode)
         createAstEdge(localNode, moduleMethodContext.get.methodBlockNode.get, moduleMethodContext.get.order.getAndInc)
         moduleMethodContext.get.variables.put(name, localNode)
       }
@@ -202,7 +214,8 @@ class ContextStack {
         } else if (memOp == Store) {
           var variableNode = lookupVariableInMethod(name, contextStack)
           if (variableNode.isEmpty) {
-            val localNode              = createLocal(name, None)
+            val localNode = createLocal(name, None)
+            transferLineColInfo(identifier, localNode)
             val enclosingMethodContext = findEnclosingMethodContext(contextStack)
             createAstEdge(localNode, enclosingMethodContext.methodBlockNode.get, enclosingMethodContext.order.getAndInc)
             enclosingMethodContext.variables.put(name, localNode)
@@ -253,6 +266,7 @@ class ContextStack {
             if (!contextHasVariable) {
               if (context != moduleMethodContext.get) {
                 val localNode = createLocal(name, Some(closureBindingId))
+                transferLineColInfo(identifier, localNode)
                 createAstEdge(localNode, methodContext.methodBlockNode.get, methodContext.order.getAndInc)
                 methodContext.variables.put(name, localNode)
               } else {
@@ -262,6 +276,7 @@ class ContextStack {
                 // For example this happens when there are wildcard imports directly into the
                 // modules namespace.
                 val localNode = createLocal(name, None)
+                transferLineColInfo(identifier, localNode)
                 createAstEdge(localNode, methodContext.methodBlockNode.get, methodContext.order.getAndInc)
                 methodContext.variables.put(name, localNode)
               }

joern-cli/frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/cpg/VariableReferencingCpgTests.scala

@@ -1,8 +1,8 @@
 package io.joern.pysrc2cpg.cpg
 
 import io.joern.pysrc2cpg.Py2CpgTestContext
+import io.shiftleft.codepropertygraph.generated.EvaluationStrategies
 import io.shiftleft.semanticcpg.language._
-import io.shiftleft.codepropertygraph.generated.{EvaluationStrategies, nodes}
 import org.scalatest.freespec.AnyFreeSpec
 import org.scalatest.matchers.should.Matchers
 
@@ -22,6 +22,19 @@ class VariableReferencingCpgTests extends AnyFreeSpec with Matchers {
       localNode.referencingIdentifiers.lineNumber(2).code.head shouldBe "x"
       localNode.referencingIdentifiers.lineNumber(3).code.head shouldBe "x"
     }
+
+    "test local variable line and column numbers" in {
+      val f = cpg.local.nameExact("f").head
+      f.lineNumber shouldBe Some(1)
+      f.columnNumber shouldBe Some(1)
+      val x = cpg.method.name("f").local.nameExact("x").head
+      val y = cpg.method.name("f").local.nameExact("y").head
+      x.lineNumber shouldBe Some(2)
+      x.columnNumber shouldBe Some(3)
+      y.lineNumber shouldBe Some(3)
+      y.columnNumber shouldBe Some(3)
+    }
+
   }
 
   "parameter variable reference" - {
@@ -38,6 +51,20 @@ class VariableReferencingCpgTests extends AnyFreeSpec with Matchers {
       paramNode.referencingIdentifiers.lineNumber(2).code.head shouldBe "x"
       paramNode.referencingIdentifiers.lineNumber(3).code.head shouldBe "x"
     }
+
+    "test local variable line and column numbers" in {
+      val f = cpg.local.nameExact("f").head
+      f.lineNumber shouldBe Some(1)
+      f.columnNumber shouldBe Some(1)
+
+      val x = cpg.method.name("f").parameter.name("x").head
+      val y = cpg.method.name("f").local.name("y").head
+      x.lineNumber shouldBe Some(1)
+      x.columnNumber shouldBe Some(7)
+      y.lineNumber shouldBe Some(3)
+      y.columnNumber shouldBe Some(3)
+    }
+
   }
 
   "comprehension variable reference" - {
@@ -63,6 +90,19 @@ class VariableReferencingCpgTests extends AnyFreeSpec with Matchers {
       val localXNode = cpg.local("x").filter(_.definingBlock.astParent.isMethod.isEmpty).head
       cpg.identifier("x").lineNumber(2).refsTo.foreach(local => local shouldBe localXNode)
     }
+
+    "test local variable line and column numbers" in {
+      val f = cpg.local.nameExact("f").head
+      f.lineNumber shouldBe Some(3)
+      f.columnNumber shouldBe Some(1)
+
+      val x = cpg.local("x").filterNot(_.definingBlock.astParent.isMethod.isEmpty).head
+      val y = cpg.local("y").filterNot(_.definingBlock.astParent.isMethod.isEmpty).head
+      x.lineNumber shouldBe Some(1)
+      x.columnNumber shouldBe Some(1)
+      y.lineNumber shouldBe Some(2)
+      y.columnNumber shouldBe Some(13)
+    }
   }
 
   "comprehension tuple variable reference" - {