use erased extensions to implement operatorExtension API

schema/build.sbt

@@ -1,6 +1,6 @@
 name := "codepropertygraph-schema"
 
-libraryDependencies += "io.shiftleft" %% "overflowdb-codegen" % "1.86"
+libraryDependencies += "io.shiftleft" %% "overflowdb-codegen" % "1.86+1-abfeddaf"
 
 val generateDomainClasses = taskKey[Seq[File]]("generate overflowdb domain classes for our schema")
 generateDomainClasses := Def.taskDyn {

semanticcpg/src/main/scala/io/shiftleft/semanticcpg/language/operatorextension/Implicits.scala

@@ -1,21 +1,30 @@
 package io.shiftleft.semanticcpg.language.operatorextension
 
 import io.shiftleft.codepropertygraph.Cpg
-import io.shiftleft.codepropertygraph.generated.nodes.{AstNode, Expression}
-import io.shiftleft.semanticcpg.language.operatorextension.nodemethods._
+import io.shiftleft.codepropertygraph.generated.nodes
+import io.shiftleft.semanticcpg.language.operatorextension
 import overflowdb.traversal.Traversal
 
 trait Implicits {
-  implicit def toNodeTypeStartersOperatorExtension(cpg: Cpg): NodeTypeStarters = new NodeTypeStarters(cpg)
-  implicit def toArrayAccessExt(arrayAccess: opnodes.ArrayAccess): ArrayAccessMethods =
-    new ArrayAccessMethods(arrayAccess)
-  implicit def toArrayAccessTrav(steps: Traversal[opnodes.ArrayAccess]): ArrayAccessTraversal =
-    new ArrayAccessTraversal(steps)
-  implicit def toAssignmentExt(assignment: opnodes.Assignment): AssignmentMethods = new AssignmentMethods(assignment)
-  implicit def toAssignmentTrav(steps: Traversal[opnodes.Assignment]): AssignmentTraversal =
-    new AssignmentTraversal(steps)
-  implicit def toTargetExt(call: Expression): TargetMethods = new TargetMethods(call)
-  implicit def toTargetTrav(steps: Traversal[Expression]): TargetTraversal = new TargetTraversal(steps)
-  implicit def toOpAstNodeExt[A <: AstNode](node: A): OpAstNodeMethods[A] = new OpAstNodeMethods(node)
-  implicit def toOpAstNodeTrav[A <: AstNode](steps: Traversal[A]): OpAstNode[A] = new OpAstNode(steps)
+  implicit def toAssignmentOperatorExtension(node: operatorextension.Assignment): operatorextension.AssignmentExt =
+    new operatorextension.AssignmentExt(node)
+  implicit def toArrayAccessOperatorExtension(node: operatorextension.ArrayAccess): operatorextension.ArrayAccessExt =
+    new operatorextension.ArrayAccessExt(node)
+  implicit def toAstNodeOperatorExtension(node: nodes.AstNode): operatorextension.AstNodeExt =
+    new operatorextension.AstNodeExt(node)
+  implicit def toExpressionNodeOperatorExtension(node: nodes.Expression): operatorextension.ExpressionNodeExt =
+    new operatorextension.ExpressionNodeExt(node)
+
+  implicit def toAssignmentTraversalOperatorExtension[T <: operatorextension.Assignment](
+      traversal: Traversal[T]): operatorextension.AssignmentTraversal[T] =
+    new operatorextension.AssignmentTraversal(traversal)
+  implicit def toArrayAccessTraversalOperatorExtension[T <: operatorextension.ArrayAccess](
+      traversal: Traversal[T]): operatorextension.AstTraversal[T] = new operatorextension.AstTraversal(traversal)
+  implicit def toAstTraversalOperatorExtension[T <: nodes.AstNode](
+      traversal: Traversal[T]): operatorextension.AstTraversal[T] = new operatorextension.AstTraversal(traversal)
+  implicit def toExpressionTraversalOperatorExtension[T <: nodes.Expression](
+      traversal: Traversal[T]): operatorextension.ExpressionTraversal[T] =
+    new operatorextension.ExpressionTraversal(traversal)
+  implicit def toNodeStartersOperatorExtension(cpg: Cpg): operatorextension.NodeTypeStarters =
+    new operatorextension.NodeTypeStarters(cpg)
 }

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

@@ -0,0 +1,124 @@
+package io.shiftleft.semanticcpg.language
+
+import io.shiftleft.codepropertygraph.Cpg
+import io.shiftleft.codepropertygraph.generated.{Operators, nodes}
+import overflowdb.traversal.help.{Doc, TraversalSource}
+import overflowdb.traversal.{NodeOps, Traversal}
+
+package object operatorextension {
+  trait AssignmentT
+  type Assignment = nodes.Call_Ref[AssignmentT]
+  trait ArithmeticT
+  type Arithmetic = nodes.Call_Ref[ArithmeticT]
+  trait ArrayAccessT
+  type ArrayAccess = nodes.Call_Ref[ArrayAccessT]
+
+  def asAssignment(e: nodes.StoredNode): Option[Assignment] = {
+    e match {
+      case c: nodes.Call =>
+        c.methodFullName match {
+          case Operators.assignment | Operators.assignmentAnd | Operators.assignmentOr | Operators.assignmentXor |
+              Operators.assignmentMinus | Operators.assignmentPlus | Operators.assignmentArithmeticShiftRight |
+              Operators.assignmentLogicalShiftRight | Operators.assignmentShiftLeft | Operators.assignmentDivision |
+              Operators.assignmentModulo | Operators.assignmentMultiplication | Operators.assignmentExponentiation =>
+            Some(c.asInstanceOf[Assignment])
+          case _ => None
+        }
+      case _ => None
+    }
+  }
+  def asArithmetic(e: nodes.StoredNode): Option[Arithmetic] = {
+    e match {
+      case c: nodes.Call =>
+        c.methodFullName match {
+          case Operators.addition | Operators.minus | Operators.modulo | Operators.multiplication | Operators.division |
+              Operators.exponentiation =>
+            Some(c.asInstanceOf[Arithmetic])
+          case _ => None
+        }
+      case _ => None
+    }
+  }
+
+  def asArrayAccess(e: nodes.StoredNode): Option[ArrayAccess] = {
+    e match {
+      case c: nodes.Call =>
+        c.methodFullName match {
+          case Operators.computedMemberAccess | Operators.indirectComputedMemberAccess | Operators.indexAccess |
+              Operators.indirectIndexAccess =>
+            Some(c.asInstanceOf[ArrayAccess])
+          case _ => None
+        }
+      case _ => None
+    }
+  }
+
+  def inAssignment(a: nodes.AstNode): Option[Assignment] = {
+    var iter = a._astIn
+    while (iter.hasNext) {
+      val next = iter.next()
+      val resOption = asAssignment(next)
+      if (resOption.isDefined) return resOption
+      iter = next._astIn
+    }
+    None
+  }
+
+  class AssignmentExt(val assignment: Assignment) extends AnyVal {
+    def target: nodes.Expression = assignment.argument(1)
+    def source: nodes.Expression = assignment.argument(2)
+  }
+
+  class ArrayAccessExt(val arrayAccess: ArrayAccess) extends AnyVal {
+    def array: nodes.Expression =
+      arrayAccess.argument(1)
+
+    def subscripts: Traversal[nodes.Identifier] =
+      arrayAccess.argument(2).start.ast.isIdentifier
+  }
+
+  class AstNodeExt(val node: nodes.AstNode) extends AnyVal {
+    def inAssignment: Traversal[Assignment] = Traversal.from(operatorextension.inAssignment(node))
+    def assignments: Traversal[Assignment] = node.start.ast.flatMap(operatorextension.asAssignment)
+    def arithmetics: Traversal[Arithmetic] = node.start.ast.flatMap(operatorextension.asArithmetic)
+  }
+
+  //this looks seriously wrong, but it is what the old code did
+  class ExpressionNodeExt(val node: nodes.Expression) extends AnyVal {
+    def isArrayAccess: Traversal[ArrayAccess] = node.start.ast.flatMap(operatorextension.asArrayAccess)
+  }
+
+  class AssignmentTraversal[T <: Assignment](val traversal: Traversal[T]) extends AnyVal {
+    def target: Traversal[nodes.Expression] = traversal.map(_.target)
+    def source: Traversal[nodes.Expression] = traversal.map(_.source)
+  }
+  class ArrayAccessTraversal[T <: ArrayAccess](val traversal: Traversal[T]) extends AnyVal {
+    def array: Traversal[nodes.Expression] = traversal.map(_.array)
+    def subscripts: Traversal[nodes.Expression] = traversal.flatMap(_.subscripts)
+  }
+
+  class AstTraversal[T <: nodes.AstNode](val traversal: Traversal[T]) extends AnyVal {
+    def isAssignment: Traversal[Assignment] = traversal.flatMap(asAssignment)
+    def assignments: Traversal[Assignment] = traversal.ast.flatMap(asAssignment)
+    def inAssignment: Traversal[Assignment] = traversal.flatMap(operatorextension.inAssignment)
+
+    def isArithmetic: Traversal[Arithmetic] = traversal.flatMap(asArithmetic)
+    def arithmetics: Traversal[Arithmetic] = traversal.ast.isArithmetic
+  }
+  //this looks seriously wrong, but it is what the old code did
+  class ExpressionTraversal[T <: nodes.Expression](val traversal: Traversal[T]) extends AnyVal {
+    def isArrayAccess: Traversal[ArrayAccess] = traversal.ast.flatMap(asArrayAccess)
+  }
+
+  @TraversalSource
+  implicit class NodeTypeStarters(val cpg: Cpg) extends AnyVal {
+
+    @Doc("All assignments")
+    def assignment: Traversal[Assignment] =
+      cpg.call.flatMap(operatorextension.asAssignment)
+
+    @Doc("All arithmetic operations")
+    def arithmetic: Traversal[Arithmetic] =
+      cpg.call.flatMap(operatorextension.asArithmetic)
+  }
+}