Caching symbols on symbol resolution

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

@@ -34,6 +34,7 @@ import de.fraunhofer.aisec.cpg.graph.statements.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Block
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.CallExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.ReferenceTag
 import de.fraunhofer.aisec.cpg.graph.types.FunctionPointerType
 import de.fraunhofer.aisec.cpg.graph.types.IncompleteType
 import de.fraunhofer.aisec.cpg.graph.types.Type
@@ -74,6 +75,12 @@ class ScopeManager : ScopeProvider {
     /** Represents an alias with the name [to] for the particular name [from]. */
     data class Alias(var from: Name, var to: Name)
 
+    /**
+     * A cache map of unique tags (computed with [Reference.buildUniqueTag]) and their respective
+     * [ValueDeclaration]. This is used by [resolveReference] as a caching mechanism.
+     */
+    private val symbolTable = mutableMapOf<ReferenceTag, ValueDeclaration>()
+
     /**
      * In some languages, we can define aliases for names. An example is renaming package imports in
      * Go, e.g., to avoid name conflicts.
@@ -607,45 +614,62 @@ class ScopeManager : ScopeProvider {
      * Resolves only references to Values in the current scope, static references to other visible
      * records are not resolved over the ScopeManager.
      *
-     * @param scope
      * @param ref
      * @return
      *
      * TODO: We should merge this function with [.resolveFunction]
      */
-    @JvmOverloads
-    fun resolveReference(ref: Reference, startScope: Scope? = currentScope): ValueDeclaration? {
+    fun resolveReference(ref: Reference): ValueDeclaration? {
+        val startScope = ref.scope
+
+        // Retrieve a unique tag for the particular reference based on the current scope
+        val tag = ref.uniqueTag
+
+        // If we find a match in our symbol table, we can immediately return the declaration
+        var decl = symbolTable[tag]
+        if (decl != null) {
+            return decl
+        }
+
         val (scope, name) = extractScope(ref, startScope)
 
         // Try to resolve value declarations according to our criteria
-        return resolve<ValueDeclaration>(scope) {
-                if (it.name.lastPartsMatch(name)) {
-                    val helper = ref.resolutionHelper
-                    return@resolve when {
-                        // If the reference seems to point to a function (using a function pointer)
-                        // the entire signature is checked for equality
-                        helper?.type is FunctionPointerType && it is FunctionDeclaration -> {
-                            val fptrType = helper.type as FunctionPointerType
-                            // TODO(oxisto): Support multiple return values
-                            val returnType = it.returnTypes.firstOrNull() ?: IncompleteType()
-                            returnType == fptrType.returnType &&
-                                it.hasSignature(fptrType.parameters)
-                        }
-                        // If our language has first-class functions, we can safely return them as a
-                        // reference
-                        ref.language is HasFirstClassFunctions -> {
-                            true
-                        }
-                        // Otherwise, we are not looking for functions here
-                        else -> {
-                            it !is FunctionDeclaration
+        decl =
+            resolve<ValueDeclaration>(scope) {
+                    if (it.name.lastPartsMatch(name)) {
+                        val helper = ref.resolutionHelper
+                        return@resolve when {
+                            // If the reference seems to point to a function (using a function
+                            // pointer) the entire signature is checked for equality
+                            helper?.type is FunctionPointerType && it is FunctionDeclaration -> {
+                                val fptrType = helper.type as FunctionPointerType
+                                // TODO(oxisto): Support multiple return values
+                                val returnType = it.returnTypes.firstOrNull() ?: IncompleteType()
+                                returnType == fptrType.returnType &&
+                                    it.hasSignature(fptrType.parameters)
+                            }
+                            // If our language has first-class functions, we can safely return them
+                            // as a reference
+                            ref.language is HasFirstClassFunctions -> {
+                                true
+                            }
+                            // Otherwise, we are not looking for functions here
+                            else -> {
+                                it !is FunctionDeclaration
+                            }
                         }
                     }
+
+                    return@resolve false
                 }
+                .firstOrNull()
 
-                return@resolve false
-            }
-            .firstOrNull()
+        // Update the symbol cache, if we found a declaration for the tag
+        if (decl != null) {
+            symbolTable[tag] = decl
+        }
+
+        return decl
     }
 
     /**

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/scopes/Scope.kt

@@ -82,6 +82,23 @@ abstract class Scope(
         return this is LoopScope
     }
 
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as Scope
+
+        if (astNode != other.astNode) return false
+        return name == other.name
+    }
+
+    override fun hashCode(): Int {
+        var result = astNode?.hashCode() ?: 0
+        result = 31 * result + (parent?.hashCode() ?: 0)
+        result = 31 * result + (name?.hashCode() ?: 0)
+        return result
+    }
+
     /** Returns the [GlobalScope] of this scope by traversing its parents upwards. */
     val globalScope: Scope?
         get() {

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/statements/expressions/Reference.kt

@@ -153,4 +153,16 @@ open class Reference : Expression(), HasType.TypeObserver {
             prev.registerTypeObserver(this)
         }
     }
+
+    /**
+     * This function builds a unique tag for the particular reference, based on the [startScope].
+     * Its purpose is to cache symbol resolutions, similar to LLVMs system of Unified Symbol
+     * Resolution (USR).
+     */
+    val uniqueTag: ReferenceTag
+        get() {
+            return Objects.hash(this.name, this.resolutionHelper, this.scope)
+        }
 }
+
+typealias ReferenceTag = Int

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

@@ -205,7 +205,7 @@ open class SymbolResolver(ctx: TranslationContext) : ComponentPass(ctx) {
             // Peek into the declaration, and if it is a variable, we can proceed normally, as we
             // are running into the special case explained above. Otherwise, we abort here (for
             // now).
-            wouldResolveTo = scopeManager.resolveReference(current, current.scope)
+            wouldResolveTo = scopeManager.resolveReference(current)
             if (wouldResolveTo !is VariableDeclaration && wouldResolveTo !is ParameterDeclaration) {
                 return
             }
@@ -214,9 +214,7 @@ open class SymbolResolver(ctx: TranslationContext) : ComponentPass(ctx) {
         // Only consider resolving, if the language frontend did not specify a resolution. If we
         // already have populated the wouldResolveTo variable, we can re-use this instead of
         // resolving again
-        var refersTo =
-            current.refersTo
-                ?: wouldResolveTo ?: scopeManager.resolveReference(current, current.scope)
+        var refersTo = current.refersTo ?: wouldResolveTo ?: scopeManager.resolveReference(current)
 
         var recordDeclType: Type? = null
         if (currentClass != null) {

cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/passes/SymbolResolverTest.kt

@@ -31,6 +31,8 @@ import de.fraunhofer.aisec.cpg.TestUtils.assertRefersTo
 import de.fraunhofer.aisec.cpg.graph.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.ConstructExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberCallExpression
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.ReferenceTag
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
@@ -71,4 +73,23 @@ class SymbolResolverTest {
         assertNotNull(construct)
         assertInvokes(construct, constructor)
     }
+
+    @Test
+    fun testUniqueTags() {
+        val result = GraphExamples.getConditionalExpression()
+
+        val map = mutableMapOf<ReferenceTag, MutableList<Reference>>()
+
+        val refs = result.refs
+        refs.forEach {
+            // Build a unique tag based on the scope of the reference is in (since this is usually
+            // the start scope)
+            val list = map.computeIfAbsent(it.uniqueTag) { mutableListOf() }
+            list += it
+
+            // All elements in the list must have the same scope and name
+            assertEquals(1, list.map { ref -> ref.scope }.toSet().size)
+            assertEquals(1, list.map { ref -> ref.name }.toSet().size)
+        }
+    }
 }