Improvements to function and record inference (#1586)

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

@@ -49,6 +49,12 @@ private constructor(
     /** Enables the inference of variables, such as global variables. */
     val inferVariables: Boolean,
 
+    /**
+     * A very EXPERIMENTAL feature. If this is enabled, we will try to infer return types of
+     * functions based on the context of the call it originated out of. This is disabled by default.
+     */
+    val inferReturnTypes: Boolean,
+
     /**
      * Uses heuristics to add DFG edges for call expressions to unresolved functions (i.e.,
      * functions not implemented in the given source code).
@@ -61,6 +67,7 @@ private constructor(
         private var inferRecords: Boolean = true,
         private var inferFunctions: Boolean = true,
         private var inferVariables: Boolean = true,
+        private var inferReturnTypes: Boolean = false,
         private var inferDfgForUnresolvedCalls: Boolean = true
     ) {
         fun enabled(infer: Boolean) = apply { this.enabled = infer }
@@ -73,6 +80,8 @@ private constructor(
 
         fun inferVariables(infer: Boolean) = apply { this.inferVariables = infer }
 
+        fun inferReturnTypes(infer: Boolean) = apply { this.inferReturnTypes = infer }
+
         fun inferDfgForUnresolvedCalls(infer: Boolean) = apply {
             this.inferDfgForUnresolvedCalls = infer
         }
@@ -84,6 +93,7 @@ private constructor(
                 inferRecords,
                 inferFunctions,
                 inferVariables,
+                inferReturnTypes,
                 inferDfgForUnresolvedCalls
             )
     }

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

@@ -111,9 +111,6 @@ class ScopeManager : ScopeProvider {
     val currentRecord: RecordDeclaration?
         get() = this.firstScopeIsInstanceOrNull<RecordScope>()?.astNode as? RecordDeclaration
 
-    val currentTypedefs: Collection<TypedefDeclaration>
-        get() = this.getCurrentTypedefs(currentScope)
-
     val currentNamespace: Name?
         get() {
             val namedScope = this.firstScopeIsInstanceOrNull<NameScope>()
@@ -237,7 +234,7 @@ class ScopeManager : ScopeProvider {
                     is Block -> BlockScope(nodeToScope)
                     is WhileStatement,
                     is DoStatement,
-                    is AssertStatement -> LoopScope(nodeToScope as Statement)
+                    is AssertStatement -> LoopScope(nodeToScope)
                     is ForStatement,
                     is ForEachStatement -> LoopScope(nodeToScope as Statement)
                     is SwitchStatement -> SwitchScope(nodeToScope)
@@ -508,29 +505,6 @@ class ScopeManager : ScopeProvider {
         scope?.addTypedef(typedef)
     }
 
-    private fun getCurrentTypedefs(searchScope: Scope?): Collection<TypedefDeclaration> {
-        val typedefs = mutableMapOf<Name, TypedefDeclaration>()
-
-        val path = mutableListOf<ValueDeclarationScope>()
-        var current = searchScope
-
-        // We need to build a path from the current scope to the top most one
-        while (current != null) {
-            if (current is ValueDeclarationScope) {
-                path += current
-            }
-            current = current.parent
-        }
-
-        // And then follow the path in reverse. This ensures us that a local definition
-        // overwrites / shadows one that was there on a higher scope.
-        for (scope in path.reversed()) {
-            typedefs.putAll(scope.typedefs)
-        }
-
-        return typedefs.values
-    }
-
     /**
      * Resolves only references to Values in the current scope, static references to other visible
      * records are not resolved over the ScopeManager.

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

@@ -204,6 +204,14 @@ interface HasAnonymousIdentifier : LanguageTrait {
  */
 interface HasGlobalVariables : LanguageTrait
 
+/**
+ * A language trait, that specifies that this language has global functions directly in the
+ * [GlobalScope], i.e., not within a namespace, but directly contained in a
+ * [TranslationUnitDeclaration]. For example, C++ has global functions, Java and Go do not (as every
+ * function is either in a class or a namespace).
+ */
+interface HasGlobalFunctions : LanguageTrait
+
 /**
  * A common super-class for all language traits that arise because they are an ambiguity of a
  * function call, e.g., function-style casts. This means that we cannot differentiate between a

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

@@ -25,7 +25,6 @@
  */
 package de.fraunhofer.aisec.cpg.passes.inference
 
-import de.fraunhofer.aisec.cpg.InferenceConfiguration
 import de.fraunhofer.aisec.cpg.ScopeManager
 import de.fraunhofer.aisec.cpg.TranslationContext
 import de.fraunhofer.aisec.cpg.TypeManager
@@ -34,11 +33,14 @@ import de.fraunhofer.aisec.cpg.frontends.Language
 import de.fraunhofer.aisec.cpg.graph.*
 import de.fraunhofer.aisec.cpg.graph.declarations.*
 import de.fraunhofer.aisec.cpg.graph.scopes.Scope
+import de.fraunhofer.aisec.cpg.graph.statements.ReturnStatement
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.BinaryOperator
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.CallExpression
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.ConstructExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Expression
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.TypeExpression
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.UnaryOperator
 import de.fraunhofer.aisec.cpg.graph.types.*
 import de.fraunhofer.aisec.cpg.helpers.Util.debugWithFileLocation
 import de.fraunhofer.aisec.cpg.helpers.Util.errorWithFileLocation
@@ -80,7 +82,7 @@ class Inference internal constructor(val start: Node, override val ctx: Translat
         code: String?,
         isStatic: Boolean,
         signature: List<Type?>,
-        returnType: Type?,
+        incomingReturnType: Type?,
         hint: CallExpression? = null
     ): FunctionDeclaration? {
         if (!ctx.config.inferenceConfiguration.inferFunctions) {
@@ -108,25 +110,43 @@ class Inference internal constructor(val start: Node, override val ctx: Translat
                 }
             inferred.code = code
 
-            debugWithFileLocation(
-                hint,
-                log,
-                "Inferred a new {} declaration {} with parameter types {} in $it",
-                if (inferred is MethodDeclaration) "method" else "function",
-                inferred.name,
-                signature.map { it?.name }
-            )
-
             // Create parameter declarations and receiver (only for methods).
             if (inferred is MethodDeclaration) {
                 createInferredReceiver(inferred, record)
             }
             createInferredParameters(inferred, signature)
 
             // Set the type and return type(s)
-            returnType?.let { inferred.returnTypes = listOf(it) }
+            var returnType =
+                if (
+                    ctx.config.inferenceConfiguration.inferReturnTypes &&
+                        incomingReturnType is UnknownType &&
+                        hint != null
+                ) {
+                    inferReturnType(hint) ?: unknownType()
+                } else {
+                    incomingReturnType
+                }
+
+            if (returnType is TupleType) {
+                inferred.returnTypes = returnType.types
+            } else if (returnType != null) {
+                inferred.returnTypes = listOf(returnType)
+            }
+
             inferred.type = FunctionType.computeType(inferred)
 
+            debugWithFileLocation(
+                hint,
+                log,
+                "Inferred a new {} declaration {} with parameter types {} and return types {} in {}",
+                if (inferred is MethodDeclaration) "method" else "function",
+                inferred.name,
+                signature.map { it?.name },
+                inferred.returnTypes.map { it.name },
+                it
+            )
+
             // Add it to the scope
             scopeManager.addDeclaration(inferred)
 
@@ -528,6 +548,73 @@ class Inference internal constructor(val start: Node, override val ctx: Translat
         this.scopeManager = ctx.scopeManager
         this.typeManager = ctx.typeManager
     }
+
+    /**
+     * This function tries to infer a return type for an inferred [FunctionDeclaration] based the
+     * original [CallExpression] (as the [hint]) parameter that was used to infer the function.
+     */
+    fun inferReturnType(hint: CallExpression): Type? {
+        // Try to find out, if the supplied hint is part of an assignment. If yes, we can use their
+        // type as the return type of the function
+        var targetType =
+            ctx.currentComponent.assignments.singleOrNull { it.value == hint }?.target?.type
+        if (targetType != null && targetType !is UnknownType) {
+            return targetType
+        }
+
+        // Look for an "argument holder". These can be different kind of nodes
+        val holder =
+            ctx.currentComponent.allChildren<ArgumentHolder> { it.hasArgument(hint) }.singleOrNull()
+        when (holder) {
+            is UnaryOperator -> {
+                // If it's a boolean operator, the return type is probably a boolean
+                if (holder.operatorCode == "!") {
+                    return hint.language?.builtInTypes?.values?.firstOrNull { it is BooleanType }
+                }
+                // If it's a numeric operator, return the largest numeric type that we have; we
+                // prefer integers to floats
+                if (holder.operatorCode in listOf("+", "-", "++", "--")) {
+                    val numericTypes =
+                        hint.language
+                            ?.builtInTypes
+                            ?.values
+                            ?.filterIsInstance<NumericType>()
+                            ?.sortedWith(
+                                compareBy<NumericType> { it.bitWidth }
+                                    .then { a, b -> preferIntegerType(a, b) }
+                            )
+
+                    return numericTypes?.lastOrNull()
+                }
+            }
+            is ConstructExpression -> {
+                return holder.type
+            }
+            is BinaryOperator -> {
+                // If it is on the right side, it's probably the same as on the left-side (and
+                // vice versa)
+                if (hint == holder.rhs) {
+                    return holder.lhs.type
+                } else if (hint == holder.lhs) {
+                    return holder.rhs.type
+                }
+            }
+            is ReturnStatement -> {
+                // If this is part of a return statement, we can take the return type
+                val func =
+                    hint.firstParentOrNull { it is FunctionDeclaration } as? FunctionDeclaration
+                val returnTypes = func?.returnTypes
+
+                return if (returnTypes != null && returnTypes.size > 1) {
+                    TupleType(returnTypes)
+                } else {
+                    returnTypes?.singleOrNull()
+                }
+            }
+        }
+
+        return null
+    }
 }
 
 /** Provides information about the inference status of a node. */
@@ -605,3 +692,12 @@ fun RecordDeclaration.inferMethod(
             call
         ) as? MethodDeclaration
 }
+
+/** A small helper function that prefers [IntegerType] when comparing two [NumericType] types. */
+fun preferIntegerType(a: NumericType, b: NumericType): Int {
+    return when {
+        a is IntegerType && b is IntegerType -> 0
+        a is IntegerType && b !is IntegerType -> 1
+        else -> -1
+    }
+}