Extract local functions

java/kotlin-extractor/src/main/kotlin/KotlinFileExtractor.kt

@@ -35,8 +35,6 @@ open class KotlinFileExtractor(
         }
     }
 
-
-
     fun getLabel(element: IrElement) : String? {
         when (element) {
             is IrFile -> return "@\"${element.path};sourcefile\"" // todo: remove copy-pasted code
@@ -125,35 +123,10 @@ open class KotlinFileExtractor(
         return id
     }
 
-    private val anonymousTypeMap: MutableMap<IrClass, TypeResults> = mutableMapOf()
-
-    override fun useAnonymousClass(c: IrClass): TypeResults {
-        var res = anonymousTypeMap[c]
-        if (res == null) {
-            val javaResult = TypeResult(tw.getFreshIdLabel<DbClass>(), "", "")
-            val kotlinResult = TypeResult(tw.getFreshIdLabel<DbKt_notnull_type>(), "", "")
-            tw.writeKt_notnull_types(kotlinResult.id, javaResult.id)
-            res = TypeResults(javaResult, kotlinResult)
-            anonymousTypeMap[c] = res
-        }
-
-        return res
-    }
-
-    private fun extractAnonymousClassStmt(c: IrClass, callable: Label<out DbCallable>, parent: Label<out DbStmtparent>, idx: Int) {
-        @Suppress("UNCHECKED_CAST")
-        val id = extractClassSource(c) as Label<out DbClass>
-        val stmtId = tw.getFreshIdLabel<DbAnonymousclassdeclstmt>()
-        tw.writeStmts_anonymousclassdeclstmt(stmtId, parent, idx, callable)
-        tw.writeKtAnonymousClassDeclarationStmts(stmtId, id)
-        val locId = tw.getLocation(c)
-        tw.writeHasLocation(stmtId, locId)
-    }
-
     fun extractClassSource(c: IrClass): Label<out DbClassorinterface> {
         val id = if (c.isAnonymousObject) {
             @Suppress("UNCHECKED_CAST")
-            useAnonymousClass(c).javaResult.id as Label<out DbClass>
+            withSourceFile(c.fileOrNull!!).useAnonymousClass(c).javaResult.id as Label<out DbClass>
         } else {
             useClassSource(c)
         }
@@ -183,7 +156,7 @@ open class KotlinFileExtractor(
                 val parentId =
                     if (parent.isAnonymousObject) {
                         @Suppress("UNCHECKED_CAST")
-                        useAnonymousClass(c).javaResult.id as Label<out DbClass>
+                        withSourceFile(c.fileOrNull!!).useAnonymousClass(c).javaResult.id as Label<out DbClass>
                     } else {
                         useClassInstance(parent, listOf()).typeResult.id
                     }
@@ -280,7 +253,7 @@ open class KotlinFileExtractor(
         }
 
         // add method:
-        val obinitLabel = getFunctionLabel(c, "<obinit>", listOf(), pluginContext.irBuiltIns.unitType)
+        val obinitLabel = getFunctionLabel(c, "<obinit>", listOf(), pluginContext.irBuiltIns.unitType, extensionReceiverParameter = null)
         val obinitId = tw.getLabelFor<DbMethod>(obinitLabel)
         val returnType = useType(pluginContext.irBuiltIns.unitType)
         tw.writeMethods(obinitId, "<obinit>", "<obinit>()", returnType.javaResult.id, returnType.kotlinResult.id, parentId, obinitId)
@@ -341,22 +314,34 @@ open class KotlinFileExtractor(
         }
     }
 
-    fun extractFunction(f: IrFunction, parentId: Label<out DbReftype>): Label<out DbCallable> {
+    fun extractFunction(f: IrFunction, parentId: Label<out DbReftype>, label: Label<DbMethod>? = null): Label<out DbCallable> {
         currentFunction = f
 
         f.typeParameters.map { extractTypeParameter(it) }
 
         val locId = tw.getLocation(f)
 
-        val id = useFunction<DbCallable>(f)
+        val id = label ?: useFunction<DbCallable>(f)
 
         val extReceiver = f.extensionReceiverParameter
-        val isExtension = extReceiver != null
-        val idxOffset = if (isExtension) 1 else 0
+        val idxOffset = if (extReceiver != null) 1 else 0
         val paramTypes = f.valueParameters.mapIndexed { i, vp ->
             extractValueParameter(vp, id, i + idxOffset)
         }
-        val paramsSignature = paramTypes.joinToString(separator = ",", prefix = "(", postfix = ")") { it.javaResult.signature!! }
+        val allParamTypes = if (extReceiver != null) {
+            val extendedType = useType(extReceiver.type)
+            @Suppress("UNCHECKED_CAST")
+            tw.writeKtExtensionFunctions(id as Label<DbMethod>, extendedType.javaResult.id, extendedType.kotlinResult.id)
+
+            val t = extractValueParameter(extReceiver, id, 0)
+            val l = mutableListOf(t)
+            l.addAll(paramTypes)
+            l
+        } else {
+            paramTypes
+        }
+
+        val paramsSignature = allParamTypes.joinToString(separator = ",", prefix = "(", postfix = ")") { it.javaResult.signature!! }
 
         if (f.symbol is IrConstructorSymbol) {
             val returnType = useType(erase(f.returnType), TypeContext.RETURN)
@@ -368,12 +353,6 @@ open class KotlinFileExtractor(
             val shortName = f.name.asString()
             @Suppress("UNCHECKED_CAST")
             tw.writeMethods(id as Label<DbMethod>, shortName, "$shortName$paramsSignature", returnType.javaResult.id, returnType.kotlinResult.id, parentId, id)
-            
-            if (extReceiver != null) {
-                val extendedType = useType(extReceiver.type)
-                tw.writeKtExtensionFunctions(id, extendedType.javaResult.id, extendedType.kotlinResult.id)
-                extractValueParameter(extReceiver, id, 0)
-            }
         }
 
         tw.writeHasLocation(id, locId)
@@ -524,11 +503,20 @@ open class KotlinFileExtractor(
             }
             is IrClass -> {
                 if (s.isAnonymousObject) {
-                    extractAnonymousClassStmt(s, callable, parent, idx)
+                    withSourceFile(s.fileOrNull!!).extractAnonymousClassStmt(s, callable, parent, idx)
                 } else {
                     logger.warnElement(Severity.ErrorSevere, "Found non anonymous IrClass as IrStatement: " + s.javaClass, s)
                 }
             }
+            is IrFunction -> {
+                if (s.isLocalFunction()) {
+                    val extractor = withSourceFile(s.fileOrNull!!)
+                    val classId =  extractor.extractGeneratedClass(s, listOf(pluginContext.irBuiltIns.anyType))
+                    extractor.extractAnonymousClassStmt(classId, s, callable, parent, idx)
+                } else {
+                    logger.warnElement(Severity.ErrorSevere, "Expected to find local function", s)
+                }
+            }
             else -> {
                 logger.warnElement(Severity.ErrorSevere, "Unrecognised IrStatement: " + s.javaClass, s)
             }
@@ -680,11 +668,10 @@ open class KotlinFileExtractor(
             val id = tw.getFreshIdLabel<DbMethodaccess>()
             val type = useType(c.type)
             val locId = tw.getLocation(c)
-            val methodId = useFunction<DbMethod>(callTarget)
+
             tw.writeExprs_methodaccess(id, type.javaResult.id, type.kotlinResult.id, parent, idx)
             tw.writeHasLocation(id, locId)
             tw.writeCallableEnclosingExpr(id, callable)
-            tw.writeCallableBinding(id, methodId)
             tw.writeStatementEnclosingExpr(id, enclosingStmt)
 
             if (extractTypeArguments) {
@@ -692,9 +679,27 @@ open class KotlinFileExtractor(
                 extractTypeArguments(c, id, callable, enclosingStmt, -2, true)
             }
 
-            val dr = c.dispatchReceiver
-            if (dr != null) {
-                extractExpressionExpr(dr, callable, id, -1, enclosingStmt)
+            if (callTarget.isLocalFunction()) {
+                val ids = withSourceFile(callTarget.fileOrNull!!).useGeneratedLocalFunctionClass(callTarget)
+
+                val methodId = ids.function
+                tw.writeCallableBinding(id, methodId)
+
+                val idNewexpr = tw.getFreshIdLabel<DbNewexpr>()
+                tw.writeExprs_newexpr(idNewexpr, ids.type.javaResult.id, ids.type.kotlinResult.id, id, -1)
+                tw.writeHasLocation(idNewexpr, locId)
+                tw.writeCallableEnclosingExpr(idNewexpr, callable)
+                tw.writeStatementEnclosingExpr(idNewexpr, enclosingStmt)
+                tw.writeCallableBinding(idNewexpr, ids.constructor)
+
+            } else {
+                val methodId = useFunction<DbMethod>(callTarget)
+                tw.writeCallableBinding(id, methodId)
+
+                val dr = c.dispatchReceiver
+                if (dr != null) {
+                    extractExpressionExpr(dr, callable, id, -1, enclosingStmt)
+                }
             }
 
             val er = c.extensionReceiver
@@ -1054,7 +1059,7 @@ open class KotlinFileExtractor(
 
             val c = (e.type as IrSimpleType).classifier.owner as IrClass
 
-            type = useAnonymousClass(c)
+            type = withSourceFile(c.fileOrNull!!).useAnonymousClass(c)
 
             @Suppress("UNCHECKED_CAST")
             tw.writeIsAnonymClass(type.javaResult.id as Label<DbClass>, id)
@@ -1290,7 +1295,7 @@ open class KotlinFileExtractor(
                 val id = tw.getFreshIdLabel<DbMethodaccess>()
                 val type = useType(e.type)
                 val locId = tw.getLocation(e)
-                val methodLabel = getFunctionLabel(irCallable.parent, "<obinit>", listOf(), e.type)
+                val methodLabel = getFunctionLabel(irCallable.parent, "<obinit>", listOf(), e.type, null)
                 val methodId = tw.getLabelFor<DbMethod>(methodLabel)
                 tw.writeExprs_methodaccess(id, type.javaResult.id, type.kotlinResult.id, exprParent.parent, exprParent.idx)
                 tw.writeHasLocation(id, locId)

java/kotlin-extractor/src/main/kotlin/KotlinSourceFileExtractor.kt

@@ -2,11 +2,17 @@ package com.github.codeql
 
 import com.github.codeql.comments.CommentExtractor
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
+import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.util.fileOrNull
 import org.jetbrains.kotlin.ir.util.packageFqName
 import org.jetbrains.kotlin.ir.util.IdSignature
 
@@ -21,7 +27,24 @@ class KotlinSourceFileExtractor(
 ) :
   KotlinFileExtractor(logger, tw, null, externalClassExtractor, primitiveTypeMapping, pluginContext) {
 
-    val fileClass by lazy {
+    init {
+        if (!stateCache.containsKey(file)){
+            stateCache[file] = SourceFileExtractionState()
+        }
+    }
+
+    data class SourceFileExtractionState(val anonymousTypeMap: MutableMap<IrClass, TypeResults> = mutableMapOf(),
+                                         val generatedLocalFunctionTypeMap: MutableMap<IrFunction, LocalFunctionLabels> = mutableMapOf())
+
+    companion object {
+        private val stateCache: MutableMap<IrFile, SourceFileExtractionState> = mutableMapOf()
+    }
+
+    private val fileExtractionState by lazy {
+        stateCache[file]!!
+    }
+
+    private val fileClass by lazy {
         extractFileClass(file)
     }
 
@@ -70,4 +93,94 @@ class KotlinSourceFileExtractor(
       return id
   }
 
+    fun useAnonymousClass(c: IrClass): TypeResults {
+        var res = fileExtractionState.anonymousTypeMap[c]
+        if (res == null) {
+            val javaResult = TypeResult(tw.getFreshIdLabel<DbClass>(), "", "")
+            val kotlinResult = TypeResult(tw.getFreshIdLabel<DbKt_notnull_type>(), "", "")
+            tw.writeKt_notnull_types(kotlinResult.id, javaResult.id)
+            res = TypeResults(javaResult, kotlinResult)
+            fileExtractionState.anonymousTypeMap[c] = res
+        }
+
+        return res
+    }
+
+    data class LocalFunctionLabels(val type: TypeResults, val constructor: Label<DbConstructor>, val function: Label<DbMethod>)
+
+    fun useGeneratedLocalFunctionClass(f: IrFunction): LocalFunctionLabels {
+        if (!f.isLocalFunction()){
+            logger.warnElement(Severity.ErrorSevere, "Extracting a non-local function as a local one", f)
+        }
+
+        var res = fileExtractionState.generatedLocalFunctionTypeMap[f]
+        if (res == null) {
+            val javaResult = TypeResult(tw.getFreshIdLabel<DbClass>(), "", "")
+            val kotlinResult = TypeResult(tw.getFreshIdLabel<DbKt_notnull_type>(), "", "")
+            tw.writeKt_notnull_types(kotlinResult.id, javaResult.id)
+            res = LocalFunctionLabels(
+                TypeResults(javaResult, kotlinResult),
+                tw.getFreshIdLabel(),
+                tw.getFreshIdLabel())
+            fileExtractionState.generatedLocalFunctionTypeMap[f] = res
+        }
+
+        return res
+    }
+
+    fun extractGeneratedClass(localFunction: IrFunction, superTypes: List<IrType>) : Label<out DbClass> {
+        val ids = withSourceFile(localFunction.fileOrNull!!).useGeneratedLocalFunctionClass(localFunction)
+
+        // Write class
+        @Suppress("UNCHECKED_CAST")
+        val id = ids.type.javaResult.id as Label<out DbClass>
+        val pkgId = extractPackage("")
+        tw.writeClasses(id, "", pkgId, id)
+        val locId = tw.getLocation(localFunction)
+        tw.writeHasLocation(id, locId)
+
+        // Extract local function as a member
+        extractFunction(localFunction, id, ids.function)
+
+        // Extract constructor
+        tw.writeConstrs(ids.constructor, "", "", ids.type.javaResult.id, ids.type.kotlinResult.id, id, ids.constructor)
+        tw.writeHasLocation(ids.constructor, locId)
+
+        // Constructor body
+        val constructorBlockId = tw.getFreshIdLabel<DbBlock>()
+        tw.writeStmts_block(constructorBlockId, ids.constructor, 0, ids.constructor)
+        tw.writeHasLocation(constructorBlockId, locId)
+
+        // Super call
+        val superCallId = tw.getFreshIdLabel<DbSuperconstructorinvocationstmt>()
+        tw.writeStmts_superconstructorinvocationstmt(superCallId, constructorBlockId, 0, ids.function)
+
+        val baseConstructor = superTypes.first().classOrNull!!.owner.declarations.find { it is IrFunction && it.symbol is IrConstructorSymbol }
+        val baseConstructorId = useFunction<DbConstructor>(baseConstructor as IrFunction)
+
+        tw.writeHasLocation(superCallId, locId)
+        @Suppress("UNCHECKED_CAST")
+        tw.writeCallableBinding(superCallId as Label<DbCaller>, baseConstructorId)
+
+        // TODO: We might need to add an `<obinit>` function, and a call to it to match other classes
+
+        addModifiers(id, "public", "static", "final")
+        extractClassSupertypes(superTypes, listOf(), id)
+
+        return id
+    }
+
+    fun extractAnonymousClassStmt(c: IrClass, callable: Label<out DbCallable>, parent: Label<out DbStmtparent>, idx: Int) {
+        @Suppress("UNCHECKED_CAST")
+        val id = extractClassSource(c) as Label<out DbClass>
+        extractAnonymousClassStmt(id, c, callable, parent, idx)
+    }
+
+    fun extractAnonymousClassStmt(id: Label<out DbClass>, locElement: IrElement, callable: Label<out DbCallable>, parent: Label<out DbStmtparent>, idx: Int) {
+        val stmtId = tw.getFreshIdLabel<DbAnonymousclassdeclstmt>()
+        tw.writeStmts_anonymousclassdeclstmt(stmtId, parent, idx, callable)
+        tw.writeKtAnonymousClassDeclarationStmts(stmtId, id)
+        val locId = tw.getLocation(locElement)
+        tw.writeHasLocation(stmtId, locId)
+    }
 }

java/kotlin-extractor/src/main/kotlin/KotlinUsesExtractor.kt

@@ -4,6 +4,7 @@ import com.github.codeql.utils.substituteTypeArguments
 import com.semmle.extractor.java.OdasaOutput
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
@@ -65,22 +66,25 @@ open class KotlinUsesExtractor(
             ?.let { pluginContext.referenceClass(it.asSingleFqName()) }
             ?.owner
 
+    fun withSourceFile(clsFile: IrFile): KotlinSourceFileExtractor {
+        val newTrapWriter = tw.makeSourceFileTrapWriter(clsFile, false)
+        val newLogger = FileLogger(logger.logCounter, newTrapWriter)
+        return KotlinSourceFileExtractor(newLogger, newTrapWriter, clsFile, externalClassExtractor, primitiveTypeMapping, pluginContext)
+    }
+
     /**
      * Gets a KotlinFileExtractor based on this one, except it attributes locations to the file that declares the given class.
      */
     fun withSourceFileOfClass(cls: IrClass): KotlinFileExtractor {
         val clsFile = cls.fileOrNull
 
-        val newTrapWriter =
-            if (isExternalDeclaration(cls) || clsFile == null)
-                tw.makeFileTrapWriter(getIrClassBinaryPath(cls))
-            else
-                tw.makeSourceFileTrapWriter(clsFile, false)
-
-        val newLogger = FileLogger(logger.logCounter, newTrapWriter)
-
-        // TODO: Conditionally KotlinSourceFileExtractor?
-        return KotlinFileExtractor(newLogger, newTrapWriter, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext)
+        if (isExternalDeclaration(cls) || clsFile == null) {
+            val newTrapWriter = tw.makeFileTrapWriter(getIrClassBinaryPath(cls))
+            val newLogger = FileLogger(logger.logCounter, newTrapWriter)
+            return KotlinFileExtractor(newLogger, newTrapWriter, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext)
+        } else {
+            return withSourceFile(clsFile)
+        }
     }
 
     fun useClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>): UseClassInstanceResult {
@@ -150,10 +154,6 @@ open class KotlinUsesExtractor(
             classLabelResult.shortName)
     }
 
-    open fun useAnonymousClass(c: IrClass): TypeResults {
-        throw Exception("Anonymous classes can only be accessed through source file extraction")
-    }
-
     fun useSimpleTypeClass(c: IrClass, args: List<IrTypeArgument>, hasQuestionMark: Boolean): TypeResults {
         if (c.isAnonymousObject) {
             if (args.isNotEmpty()) {
@@ -163,7 +163,7 @@ open class KotlinUsesExtractor(
                 logger.warn(Severity.ErrorHigh, "Unexpected nullable anonymous class")
             }
 
-            return useAnonymousClass(c)
+            return withSourceFile(c.fileOrNull!!).useAnonymousClass(c)
         }
 
         val classInstanceResult = useClassInstance(c, args)
@@ -421,22 +421,40 @@ class X {
         }
 
     fun getFunctionLabel(f: IrFunction) : String {
-        return getFunctionLabel(f.parent, f.name.asString(), f.valueParameters, f.returnType)
+        return getFunctionLabel(f.parent, f.name.asString(), f.valueParameters, f.returnType, f.extensionReceiverParameter)
     }
 
     fun getFunctionLabel(
         parent: IrDeclarationParent,
         name: String,
         parameters: List<IrValueParameter>,
-        returnType: IrType
+        returnType: IrType,
+        extensionReceiverParameter: IrValueParameter?
     ): String {
-        val paramTypeIds = parameters.joinToString { "{${useType(erase(it.type)).javaResult.id}}" }
+        val allParams = if (extensionReceiverParameter == null) {
+                            parameters
+                        } else {
+                            val params = mutableListOf(extensionReceiverParameter)
+                            params.addAll(parameters)
+                            params
+                        }
+        val paramTypeIds = allParams.joinToString { "{${useType(erase(it.type)).javaResult.id}}" }
         val returnTypeId = useType(erase(returnType)).javaResult.id
         val parentId = useDeclarationParent(parent)
         return "@\"callable;{$parentId}.$name($paramTypeIds){$returnTypeId}\""
     }
 
+    protected fun IrFunction.isLocalFunction(): Boolean {
+        return this.visibility == DescriptorVisibilities.LOCAL
+    }
+
     fun <T: DbCallable> useFunction(f: IrFunction): Label<out T> {
+        if (f.isLocalFunction()) {
+            val ids = withSourceFile(f.fileOrNull!!).useGeneratedLocalFunctionClass(f)
+            @Suppress("UNCHECKED_CAST")
+            return ids.function as Label<out T>
+        }
+
         val label = getFunctionLabel(f)
         val id: Label<T> = tw.getLabelFor(label)
         if(isExternalDeclaration(f)) {
@@ -536,7 +554,7 @@ class X {
     fun useClassSource(c: IrClass): Label<out DbClassorinterface> {
         if (c.isAnonymousObject) {
             @Suppress("UNCHECKED_CAST")
-            return useAnonymousClass(c).javaResult.id as Label<DbClass>
+            return withSourceFile(c.fileOrNull!!).useAnonymousClass(c).javaResult.id as Label<DbClass>
         }
 
         // For source classes, the label doesn't include and type arguments
@@ -586,13 +604,18 @@ class X {
      * where `E` is the type variable declared as `List<E>`.
      */
     fun extractClassSupertypes(c: IrClass, id: Label<out DbReftype>, typeArgsQ: List<IrTypeArgument>? = null) {
+        extractClassSupertypes(c.superTypes, c.typeParameters, id, typeArgsQ)
+    }
+
+
+    fun extractClassSupertypes(superTypes: List<IrType>, typeParameters: List<IrTypeParameter>, id: Label<out DbReftype>, typeArgsQ: List<IrTypeArgument>? = null) {
         // Note we only need to substitute type args here because it is illegal to directly extend a type variable.
         // (For example, we can't have `class A<E> : E`, but can have `class A<E> : Comparable<E>`)
         val subbedSupertypes = typeArgsQ?.let { typeArgs ->
-            c.superTypes.map {
-                it.substituteTypeArguments(c.typeParameters, typeArgs)
+            superTypes.map {
+                it.substituteTypeArguments(typeParameters, typeArgs)
             }
-        } ?: c.superTypes
+        } ?: superTypes
 
         for(t in subbedSupertypes) {
             when(t) {