Use IdentitySet and IdentityHashMap in EOG iteration for CF-DFG to improve performance

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/EOGWorklist.kt

@@ -55,11 +55,13 @@ abstract class LatticeElement<T>(open val elements: T) : Comparable<LatticeEleme
  * Implements the [LatticeElement] for a lattice over a set of nodes. The lattice itself is
  * constructed by the powerset.
  */
-class PowersetLattice(override val elements: Set<Node>) : LatticeElement<Set<Node>>(elements) {
+class PowersetLattice(override val elements: IdentitySet<Node>) :
+    LatticeElement<Set<Node>>(elements) {
     override fun lub(other: LatticeElement<Set<Node>>) =
-        PowersetLattice(other.elements.union(this.elements))
+        PowersetLattice(this.elements.union(other.elements))
 
-    override fun duplicate() = PowersetLattice(this.elements.toSet())
+    override fun duplicate(): LatticeElement<Set<Node>> =
+        PowersetLattice(this.elements.toIdentitySet())
 
     override fun compareTo(other: LatticeElement<Set<Node>>): Int {
         return if (this.elements.containsAll(other.elements)) {
@@ -135,7 +137,7 @@ open class State<K, V> : IdentityHashMap<K, LatticeElement<V>>() {
             // upper bound
             this[newNode] = newLatticeElement.lub(current)
         } else {
-            this[newNode] = newLatticeElement.duplicate()
+            this[newNode] = newLatticeElement
         }
         return true
     }
@@ -147,13 +149,15 @@ open class State<K, V> : IdentityHashMap<K, LatticeElement<V>>() {
  */
 class Worklist<K : Any, N : Any, V>() {
     /** A mapping of nodes to the state which is currently available there. */
-    var globalState: MutableMap<K, State<N, V>> = mutableMapOf()
+    var globalState = IdentityHashMap<K, State<N, V>>()
         private set
 
     /** A list of all nodes which have already been visited. */
     private val alreadySeen = IdentitySet<K>()
 
-    constructor(globalState: MutableMap<K, State<N, V>> = mutableMapOf()) : this() {
+    constructor(
+        globalState: IdentityHashMap<K, State<N, V>> = IdentityHashMap<K, State<N, V>>()
+    ) : this() {
         this.globalState = globalState
     }
 
@@ -242,7 +246,9 @@ inline fun <reified K : Node, V> iterateEOG(
     startState: State<K, V>,
     transformation: (K, State<K, V>, Worklist<K, K, V>) -> State<K, V>
 ): State<K, V>? {
-    val worklist = Worklist(mutableMapOf(Pair(startNode, startState)))
+    val initialState = IdentityHashMap<K, State<K, V>>()
+    initialState[startNode] = startState
+    val worklist = Worklist(initialState)
     worklist.push(startNode, startState)
 
     while (worklist.isNotEmpty()) {
@@ -273,7 +279,7 @@ inline fun <reified K : PropertyEdge<Node>, N : Any, V> iterateEOG(
     startState: State<N, V>,
     transformation: (K, State<N, V>, Worklist<K, N, V>) -> State<N, V>
 ): State<N, V>? {
-    val globalState = mutableMapOf<K, State<N, V>>()
+    val globalState = IdentityHashMap<K, State<N, V>>()
     for (startEdge in startEdges) {
         globalState[startEdge] = startState
     }

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/IdentitySet.kt

@@ -146,3 +146,16 @@ fun <T> identitySetOf(vararg elements: T): IdentitySet<T> {
 
     return set
 }
+
+infix fun <T> IdentitySet<T>.union(other: Iterable<T>): IdentitySet<T> {
+    val set = identitySetOf<T>()
+    set += this
+    set += other
+    return set
+}
+
+fun <T> Collection<T>.toIdentitySet(): IdentitySet<T> {
+    val set = identitySetOf<T>()
+    set += this
+    return set
+}

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/ControlFlowSensitiveDFGPass.kt

@@ -91,7 +91,7 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
             clearFlowsOfVariableDeclarations(node)
             val startState = DFGPassState<Set<Node>>()
 
-            startState.declarationsState.push(node, PowersetLattice(setOf()))
+            startState.declarationsState.push(node, PowersetLattice(identitySetOf()))
             val finalState =
                 iterateEOG(node.nextEOGEdges, startState, ::transfer) as? DFGPassState ?: return
 
@@ -157,19 +157,19 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
         if (initializer != null) {
             // A variable declaration with an initializer => The initializer flows to the
             // declaration. This also affects tuples. We split it up later.
-            state.push(currentNode, PowersetLattice(setOf(initializer)))
+            state.push(currentNode, PowersetLattice(identitySetOf(initializer)))
 
             if (currentNode is TupleDeclaration) {
                 // For a tuple declaration, we write the elements in this statement. We do not
                 // really care about the tuple when using the elements subsequently.
                 currentNode.elements.forEach {
-                    doubleState.pushToDeclarationsState(it, PowersetLattice(setOf(it)))
+                    doubleState.pushToDeclarationsState(it, PowersetLattice(identitySetOf(it)))
                 }
             } else {
                 // We also wrote something to this variable declaration here.
                 doubleState.pushToDeclarationsState(
                     currentNode,
-                    PowersetLattice(setOf(currentNode))
+                    PowersetLattice(identitySetOf(currentNode))
                 )
             }
         } else if (isSimpleAssignment(currentNode)) {
@@ -182,7 +182,7 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
                 (assignment.target as? Declaration ?: (assignment.target as? Reference)?.refersTo)
                     ?.let {
                         doubleState.declarationsState[it] =
-                            PowersetLattice(setOf(assignment.target as Node))
+                            PowersetLattice(identitySetOf(assignment.target as Node))
                     }
             }
         } else if (isIncOrDec(currentNode)) {
@@ -194,7 +194,8 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
 
             if (writtenDeclaration != null) {
                 state.push(input, doubleState.declarationsState[writtenDeclaration])
-                doubleState.declarationsState[writtenDeclaration] = PowersetLattice(setOf(input))
+                doubleState.declarationsState[writtenDeclaration] =
+                    PowersetLattice(identitySetOf(input))
             }
         } else if (isCompoundAssignment(currentNode)) {
             // We write to the lhs, but it also serves as an input => We first get all previous
@@ -209,7 +210,8 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
                 state.push(lhs, doubleState.declarationsState[writtenDeclaration])
 
                 // The whole current node is the place of the last update, not (only) the lhs!
-                doubleState.declarationsState[writtenDeclaration] = PowersetLattice(setOf(lhs))
+                doubleState.declarationsState[writtenDeclaration] =
+                    PowersetLattice(identitySetOf(lhs))
             }
         } else if (
             (currentNode as? Reference)?.access == AccessValues.READ &&
@@ -270,19 +272,23 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
 
             iterable?.let {
                 writtenTo?.let {
-                    state.push(writtenTo, PowersetLattice(setOf(iterable)))
+                    state.push(writtenTo, PowersetLattice(identitySetOf(iterable)))
                     // Add the variable declaration (or the reference) to the list of previous
                     // write nodes in this path
-                    state.declarationsState[writtenDeclaration] = PowersetLattice(setOf(writtenTo))
+                    state.declarationsState[writtenDeclaration] =
+                        PowersetLattice(identitySetOf(writtenTo))
                 }
             }
         } else if (currentNode is FunctionDeclaration) {
             // We have to add the parameters
             currentNode.parameters.forEach {
-                doubleState.pushToDeclarationsState(it, PowersetLattice(setOf(it)))
+                doubleState.pushToDeclarationsState(it, PowersetLattice(identitySetOf(it)))
             }
         } else if (currentNode is ReturnStatement) {
-            doubleState.returnStatements.push(currentNode, PowersetLattice(setOf(currentNode)))
+            doubleState.returnStatements.push(
+                currentNode,
+                PowersetLattice(identitySetOf(currentNode))
+            )
         } else {
             doubleState.declarationsState.push(
                 currentNode,