Do work on "original" type in call graph extension. (#349)

* Do work on "original" type in call graph extension.

* Also do not work on original for call extension.

For MethodInst extension i did not change the code as it will be removed
soon.

* Fmt.

semanticcpg/src/main/scala/io/shiftleft/semanticcpg/language/callgraphextension/Call.scala

@@ -1,27 +1,24 @@
 package io.shiftleft.semanticcpg.language.callgraphextension
 
+import gremlin.scala.GremlinScala
 import io.shiftleft.codepropertygraph.generated.{EdgeTypes, nodes}
 import io.shiftleft.semanticcpg.language._
-import io.shiftleft.semanticcpg.language.types.expressions.{Call => OriginalCall}
-import io.shiftleft.semanticcpg.language.types.structure.{Method => OriginalMethod, MethodInst => OriginalMethodInst}
 
-class Call(original: OriginalCall) {
+class Call(override val raw: GremlinScala[nodes.Call]) extends NodeSteps[nodes.Call](raw) {
 
   /**
   The callee method
     */
-  def calledMethod(implicit callResolver: ICallResolver): OriginalMethod = {
+  def calledMethod(implicit callResolver: ICallResolver): NodeSteps[nodes.Method] = {
     calledMethodInstance(callResolver).method
   }
 
   /**
   The callee method instance
     */
-  def calledMethodInstance(implicit callResolver: ICallResolver): OriginalMethodInst =
-    new OriginalMethodInst(
-      original
-        .sideEffect(callResolver.resolveDynamicCallSite)
-        .raw
+  def calledMethodInstance(implicit callResolver: ICallResolver): NodeSteps[nodes.MethodInst] =
+    new NodeSteps[nodes.MethodInst](
+      sideEffect(callResolver.resolveDynamicCallSite).raw
         .out(EdgeTypes.CALL)
         .cast[nodes.MethodInst])
 

semanticcpg/src/main/scala/io/shiftleft/semanticcpg/language/callgraphextension/Method.scala

@@ -3,27 +3,26 @@ package io.shiftleft.semanticcpg.language.callgraphextension
 import gremlin.scala._
 import io.shiftleft.codepropertygraph.generated.{EdgeTypes, nodes}
 import io.shiftleft.semanticcpg.language._
-import io.shiftleft.semanticcpg.language.types.expressions.{Call => OriginalCall}
-import io.shiftleft.semanticcpg.language.types.structure.{Method => OriginalMethod, MethodInst => OriginalMethodInst}
+import io.shiftleft.semanticcpg.language.types.structure.{Method => OriginalMethod}
 
-class Method(original: OriginalMethod) {
+class Method(override val raw: GremlinScala[nodes.Method]) extends Steps[nodes.Method](raw) {
 
   /**
     * Intended for internal use!
     * Traverse to direct and transitive callers of the method.
     */
-  def calledByIncludingSink(sourceTrav: OriginalMethod, resolve: Boolean = true)(
-      implicit callResolver: ICallResolver): OriginalMethod = {
+  def calledByIncludingSink(sourceTrav: Steps[nodes.Method], resolve: Boolean = true)(
+      implicit callResolver: ICallResolver): Steps[nodes.Method] = {
     val sourceMethods = sourceTrav.raw.toSet
-    val sinkMethods = original.raw.dedup.toList()
+    val sinkMethods = raw.dedup.toList()
 
     if (sourceMethods.isEmpty || sinkMethods.isEmpty) {
-      new OriginalMethod(original.graph.V(-1).asInstanceOf[GremlinScala[nodes.Method]])
+      new Steps[nodes.Method](graph.V(-1).cast[nodes.Method])
     } else {
       val ids = sinkMethods.map(_.id)
-      val methodTrav = original.graph.V(ids: _*)
+      val methodTrav = graph.V(ids: _*)
 
-      new OriginalMethod(
+      new Steps[nodes.Method](
         methodTrav
           .emit(_.is(P.within(sourceMethods)))
           .repeat(
@@ -37,31 +36,29 @@ class Method(original: OriginalMethod) {
               .dedup
               .simplePath()
           )
-          .asInstanceOf[GremlinScala[nodes.Method]]
+          .cast[nodes.Method]
       )
     }
   }
 
   /**
     * Traverse to direct callers of this method
     * */
-  def caller(implicit callResolver: ICallResolver): OriginalMethod =
+  def caller(implicit callResolver: ICallResolver): Steps[nodes.Method] =
     callIn(callResolver).method
 
   /**
     * Traverse to methods called by this method
     * */
-  def callee(implicit callResolver: ICallResolver): OriginalMethod =
-    new Call(original.callOut).calledMethod(callResolver)
+  def callee(implicit callResolver: ICallResolver): Steps[nodes.Method] =
+    new OriginalMethod(raw).callOut.calledMethod(callResolver)
 
   /**
     * Incoming call sites
     * */
-  def callIn(implicit callResolver: ICallResolver): OriginalCall = {
-    new OriginalCall(
-      original
-        .sideEffect(callResolver.resolveDynamicMethodCallSites)
-        .raw
+  def callIn(implicit callResolver: ICallResolver): Steps[nodes.Call] = {
+    new Steps[nodes.Call](
+      sideEffect(callResolver.resolveDynamicMethodCallSites).raw
         .in(EdgeTypes.REF)
         .in(EdgeTypes.CALL)
         .cast[nodes.Call])
@@ -70,21 +67,23 @@ class Method(original: OriginalMethod) {
   /**
     * Traverse to direct and transitive callers of the method.
     * */
-  def calledBy(sourceTrav: OriginalMethod)(implicit callResolver: ICallResolver): OriginalMethod = {
-    new Method(caller(callResolver)).calledByIncludingSink(sourceTrav)(callResolver)
+  def calledBy(sourceTrav: Steps[nodes.Method])(implicit callResolver: ICallResolver): Steps[nodes.Method] = {
+    caller(callResolver).calledByIncludingSink(sourceTrav)(callResolver)
   }
 
   /**
     * Traverse to direct and transitive callers of the method.
     * */
-  def calledBy(sourceTrav: OriginalMethodInst)(implicit callResolver: ICallResolver): OriginalMethod = {
-    new Method(caller(callResolver)).calledByIncludingSink(sourceTrav.method)(callResolver)
+  def calledBy(sourceTrav: Steps[nodes.MethodInst])(implicit callResolver: ICallResolver,
+                                                    x: DummyImplicit): Steps[nodes.Method] = {
+    caller(callResolver).calledByIncludingSink(sourceTrav.method)(callResolver)
   }
 
   /**
     * Outgoing call sites to methods where fullName matches `regex`.
     * */
-  def callOut(regex: String)(implicit callResolver: ICallResolver): OriginalCall =
-    original.callOut.filter(new Call(_).calledMethod.fullName(regex))
+  def callOutRegex(regex: String)(implicit callResolver: ICallResolver): Steps[nodes.Call] = {
+    new OriginalMethod(raw).callOut.filter(_.calledMethod.fullName(regex))
+  }
 
 }

semanticcpg/src/main/scala/io/shiftleft/semanticcpg/language/callgraphextension/MethodInst.scala

@@ -11,14 +11,14 @@ class MethodInst(original: OriginalMethodInst) {
     * Traverse to direct and transitive callers of the method.
     * */
   def calledBy(sourceTrav: OriginalMethod)(implicit callResolver: ICallResolver): OriginalMethod = {
-    new Method(caller(callResolver)).calledByIncludingSink(sourceTrav)(callResolver)
+    caller(callResolver).calledByIncludingSink(sourceTrav)(callResolver)
   }
 
   /**
     * Traverse to direct and transitive callers of the method.
     * */
   def calledBy(sourceTrav: OriginalMethodInst)(implicit callResolver: ICallResolver): OriginalMethod = {
-    new Method(caller(callResolver)).calledByIncludingSink(sourceTrav.method)(callResolver)
+    caller(callResolver).calledByIncludingSink(sourceTrav.method)(callResolver)
   }
 
   /**
@@ -32,7 +32,7 @@ class MethodInst(original: OriginalMethodInst) {
     * Traverse to methods called by method.
     * */
   def callee(implicit callResolver: ICallResolver): OriginalMethod = {
-    new Method(original.method).callee(callResolver)
+    original.method.callee(callResolver)
   }
 
   /**

semanticcpg/src/main/scala/io/shiftleft/semanticcpg/language/package.scala

@@ -176,11 +176,11 @@ package object language {
 
   // Call graph extension
 
-  implicit def toMethodForCallGraph[X <% OriginalMethod](original: X): Method =
-    new Method(original)
+  implicit def toMethodForCallGraph(steps: Steps[nodes.Method]): Method =
+    new Method(steps.raw)
 
-  implicit def toCallForCallGraph[X <% OriginalCall](original: X): Call =
-    new Call(original)
+  implicit def toCallForCallGraph(steps: Steps[nodes.Call]): Call =
+    new Call(steps.raw)
 
   implicit def toMethodInstForCallGraph[X <% OriginalMethodInst](original: X): MethodInst =
     new MethodInst(original)