[JPS][IC] Register lookups as usages in jps dependency graph

^KT-65429 In Progress


Merge-request: KT-MR-14485
Merged-by: Evgenii Mazhukin <evgenii.mazhukin@jetbrains.com>

jps/jps-platform-api-signatures/build.gradle.kts

@@ -8,6 +8,8 @@ plugins {
 
 dependencies {
     implementation(kotlinStdlib())
+    compileOnly(commonDependency("org.jetbrains.intellij.deps:asm-all"))
+    compileOnly(intellijPlatformUtil())
 }
 
 sourceSets {

jps/jps-platform-api-signatures/src/org/jetbrains/jps/builders/java/dependencyView/Callbacks.java

@@ -0,0 +1,99 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.builders.java.dependencyView;
+
+import com.intellij.util.SmartList;
+import org.jetbrains.jps.dependency.Usage;
+import org.jetbrains.org.objectweb.asm.ClassReader;
+
+import java.io.File;
+import java.nio.file.Path;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.concurrent.Future;
+
+public final class Callbacks {
+
+    public interface ConstantRef {
+        String getOwner();
+        String getName();
+        String getDescriptor();
+    }
+
+    public interface Backend {
+        default void associate(String classFileName, String sourceFileName, ClassReader cr) {
+            associate(classFileName, Collections.singleton(sourceFileName), cr);
+        }
+        default void associate(String classFileName, Collection<String> sources, ClassReader cr) {
+            associate(classFileName, sources, cr, false);
+        }
+        void associate(String classFileName, Collection<String> sources, ClassReader cr, boolean isGenerated);
+        void registerImports(String className, Collection<String> classImports, Collection<String> staticImports);
+        void registerConstantReferences(String className, Collection<ConstantRef> cRefs);
+        default void registerUsage(String className, Usage usage) {
+        }
+        default void registerUsage(Path source, Usage usage) {
+        }
+    }
+
+    public static ConstantRef createConstantReference(String ownerClass, String fieldName, String descriptor) {
+        return new ConstantRef() {
+            @Override
+            public String getOwner() {
+                return ownerClass;
+            }
+
+            @Override
+            public String getName() {
+                return fieldName;
+            }
+
+            @Override
+            public String getDescriptor() {
+                return descriptor;
+            }
+        };
+    }
+
+    public static final class ConstantAffection {
+        public static final ConstantAffection EMPTY = new ConstantAffection();
+        private final boolean myKnown;
+        private final Collection<File> myAffectedFiles;
+
+        public ConstantAffection(final Collection<File> affectedFiles) {
+            myAffectedFiles = affectedFiles;
+            myKnown = true;
+        }
+
+        public ConstantAffection() {
+            myKnown = false;
+            myAffectedFiles = null;
+        }
+
+        public boolean isKnown(){
+            return myKnown;
+        }
+
+        public Collection<File> getAffectedFiles (){
+            return myAffectedFiles;
+        }
+
+        public static ConstantAffection compose(final Collection<? extends ConstantAffection> affections) {
+            if (affections.isEmpty()) {
+                return new ConstantAffection(Collections.emptyList()); // return a 'known' affection here
+            }
+            if (affections.size() == 1) {
+                return affections.iterator().next();
+            }
+            for (ConstantAffection a : affections) {
+                if (!a.isKnown()) {
+                    return EMPTY;
+                }
+            }
+            final Collection<File> affected = new SmartList<>();
+            for (ConstantAffection affection : affections) {
+                affected.addAll(affection.getAffectedFiles());
+            }
+            return new ConstantAffection(affected);
+        }
+    }
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/ExternalizableGraphElement.java

@@ -0,0 +1,9 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency;
+
+import java.io.IOException;
+
+public interface ExternalizableGraphElement {
+
+    void write(GraphDataOutput out) throws IOException;
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/GraphDataInput.java

@@ -0,0 +1,13 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency;
+
+import java.io.DataInput;
+import java.io.IOException;
+import java.util.Collection;
+
+public interface GraphDataInput extends DataInput {
+
+    <T extends ExternalizableGraphElement> T readGraphElement() throws IOException;
+
+    <T extends ExternalizableGraphElement, C extends Collection<? super T>> C readGraphElementCollection(C acc) throws IOException;
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/GraphDataOutput.java

@@ -0,0 +1,14 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency;
+
+import org.jetbrains.annotations.NotNull;
+
+import java.io.DataOutput;
+import java.io.IOException;
+
+public interface GraphDataOutput extends DataOutput {
+
+    <T extends ExternalizableGraphElement> void writeGraphElement(@NotNull T elem) throws IOException;
+
+    <T extends ExternalizableGraphElement> void writeGraphElementCollection(Class<? extends T> elemType, @NotNull Iterable<T> col) throws IOException;
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/ReferenceID.java

@@ -0,0 +1,12 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency;
+
+/**
+ * This is a data property of a Node, used to reference the Node by other Nodes.
+ * For example, is a Node represents a java class, a full qualified name of the class can be used as the Node's ReferenceID
+ */
+public interface ReferenceID extends ExternalizableGraphElement {
+    boolean equals(Object other);
+
+    int hashCode();
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/Usage.java

@@ -0,0 +1,20 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency;
+
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * A base class describing a "usage" of some graph Node or one of the Node's internal elements
+ */
+public interface Usage extends ExternalizableGraphElement {
+
+    /**
+     * @return the ID referring to the Node being used or the Node, to which the used element belongs to
+     */
+    @NotNull
+    ReferenceID getElementOwner();
+
+    boolean equals(Object other);
+
+    int hashCode();
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/java/JvmElementUsage.java

@@ -0,0 +1,55 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency.java;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jps.dependency.GraphDataInput;
+import org.jetbrains.jps.dependency.GraphDataOutput;
+import org.jetbrains.jps.dependency.Usage;
+
+import java.io.IOException;
+
+abstract class JvmElementUsage implements Usage {
+
+    private final @NotNull JvmNodeReferenceID myOwner;
+
+    JvmElementUsage(@NotNull JvmNodeReferenceID owner) {
+        myOwner = owner;
+    }
+
+    JvmElementUsage(GraphDataInput in) throws IOException {
+        myOwner = new JvmNodeReferenceID(in);
+    }
+
+    @Override
+    public void write(GraphDataOutput out) throws IOException {
+        myOwner.write(out);
+    }
+
+    @Override
+    public @NotNull JvmNodeReferenceID getElementOwner() {
+        return myOwner;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        final JvmElementUsage jvmUsage = (JvmElementUsage)o;
+
+        if (!myOwner.equals(jvmUsage.myOwner)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return myOwner.hashCode();
+    }
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/java/JvmNodeReferenceID.java

@@ -0,0 +1,61 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency.java;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jps.dependency.GraphDataInput;
+import org.jetbrains.jps.dependency.GraphDataOutput;
+import org.jetbrains.jps.dependency.ReferenceID;
+
+import java.io.IOException;
+
+public final class JvmNodeReferenceID implements ReferenceID {
+    private final String myName;
+
+    public JvmNodeReferenceID(@NotNull String name) {
+        myName = name;
+    }
+
+    public JvmNodeReferenceID(GraphDataInput in) throws IOException {
+        myName = in.readUTF();
+    }
+
+    @Override
+    public void write(GraphDataOutput out) throws IOException {
+        out.writeUTF(myName);
+    }
+
+    /**
+     * @return either JVM class name (FQ-name) or JVM module name
+     */
+    public String getNodeName() {
+        return myName;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        final JvmNodeReferenceID that = (JvmNodeReferenceID)o;
+
+        if (!myName.equals(that.myName)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return myName.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return "JVM_NODE_ID:" + myName;
+    }
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/java/LookupNameUsage.java

@@ -0,0 +1,26 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency.java;
+
+import org.jetbrains.jps.dependency.GraphDataInput;
+
+import java.io.IOException;
+
+public final class LookupNameUsage extends MemberUsage {
+
+    public LookupNameUsage(String className, String name) {
+        super(className, name);
+    }
+
+    public LookupNameUsage(JvmNodeReferenceID clsId, String name) {
+        super(clsId, name);
+    }
+
+    public LookupNameUsage(GraphDataInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public int hashCode() {
+        return super.hashCode() + 2;
+    }
+}

jps/jps-platform-api-signatures/src/org/jetbrains/jps/dependency/java/MemberUsage.java

@@ -0,0 +1,56 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.dependency.java;
+
+import org.jetbrains.jps.dependency.GraphDataInput;
+import org.jetbrains.jps.dependency.GraphDataOutput;
+
+import java.io.IOException;
+
+public abstract class MemberUsage extends JvmElementUsage {
+
+    private final String myName;
+
+    protected MemberUsage(String className, String name) {
+        this(new JvmNodeReferenceID(className), name);
+    }
+
+    protected MemberUsage(JvmNodeReferenceID clsId, String name) {
+        super(clsId);
+        myName = name;
+    }
+
+    MemberUsage(GraphDataInput in) throws IOException {
+        super(in);
+        myName = in.readUTF();
+    }
+
+    @Override
+    public void write(GraphDataOutput out) throws IOException {
+        super.write(out);
+        out.writeUTF(myName);
+    }
+
+    public String getName() {
+        return myName;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (!super.equals(o)) {
+            return false;
+        }
+
+        final MemberUsage that = (MemberUsage)o;
+
+        if (!myName.equals(that.myName)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return 31 * super.hashCode() + myName.hashCode();
+    }
+}

jps/jps-plugin/src/org/jetbrains/kotlin/jps/targets/KotlinJvmModuleBuildTarget.kt

@@ -13,6 +13,7 @@ import org.jetbrains.jps.builders.java.JavaBuilderUtil
 import org.jetbrains.jps.builders.java.dependencyView.Callbacks
 import org.jetbrains.jps.builders.java.dependencyView.Callbacks.Backend
 import org.jetbrains.jps.builders.storage.BuildDataPaths
+import org.jetbrains.jps.dependency.java.LookupNameUsage
 import org.jetbrains.jps.incremental.*
 import org.jetbrains.jps.model.java.JpsJavaExtensionService
 import org.jetbrains.jps.model.module.JpsSdkDependency
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.build.JvmBuildMetaInfo
 import org.jetbrains.kotlin.build.JvmSourceRoot
 import org.jetbrains.kotlin.cli.common.arguments.CommonCompilerArguments
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.compilerRunner.JpsCompilerEnvironment
 import org.jetbrains.kotlin.compilerRunner.JpsKotlinCompilerRunner
 import org.jetbrains.kotlin.config.IncrementalCompilation
@@ -37,6 +39,7 @@ import org.jetbrains.kotlin.jps.incremental.JpsIncrementalJvmCache
 import org.jetbrains.kotlin.jps.model.k2JvmCompilerArguments
 import org.jetbrains.kotlin.jps.model.kotlinCompilerSettings
 import org.jetbrains.kotlin.jps.statistic.JpsBuilderMetricReporter
+import org.jetbrains.kotlin.jps.targets.impl.LookupUsageRegistrar
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
 import org.jetbrains.kotlin.modules.KotlinModuleXmlBuilder
@@ -385,7 +388,7 @@ class KotlinJvmModuleBuildTarget(kotlinContext: KotlinCompileContext, jpsModuleB
             if (!cache.isMultifileFacade(className)) return emptySet()
 
             // In case of graph implementation of JPS
-            if (previousMappings == null) return emptySet()
+            if (KotlinBuilder.useDependencyGraph || previousMappings == null) return emptySet()
 
             val name = previousMappings.getName(className.internalName)
             return previousMappings.getClassSources(name).toSet()
@@ -413,6 +416,13 @@ class KotlinJvmModuleBuildTarget(kotlinContext: KotlinCompileContext, jpsModuleB
                 )
             }
         }
+        if (KotlinBuilder.useDependencyGraph) {
+            LookupUsageRegistrar().processLookupTracker(
+                environment.services[LookupTracker::class.java],
+                callback,
+                environment.messageCollector
+            )
+        }
 
         val allCompiled = dirtyFilesHolder.allDirtyFiles
         JavaBuilderUtil.registerFilesToCompile(localContext, allCompiled)

jps/jps-plugin/src/org/jetbrains/kotlin/jps/targets/impl/LookupUsageRegistrar.kt

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.jps.targets.impl
+
+import org.jetbrains.jps.builders.java.dependencyView.Callbacks.Backend
+import org.jetbrains.jps.dependency.java.LookupNameUsage
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.incremental.LookupTrackerImpl
+import org.jetbrains.kotlin.incremental.components.LookupTracker
+import java.nio.file.Paths
+
+class LookupUsageRegistrar {
+    fun processLookupTracker(lookupTracker: LookupTracker?, callback: Backend, messageCollector: MessageCollector) {
+        if (!checkRequiredJpsBuildApi()) {
+            messageCollector.report(
+                CompilerMessageSeverity.WARNING,
+                "Can't register lookup usages with this version of JPS. $HOW_TO_FIX"
+            )
+            return
+        }
+
+        when (lookupTracker) {
+            is LookupTrackerImpl -> registerLookupTrackerImplEntries(lookupTracker, callback)
+            else -> {
+                // could be DO_NOTHING tracker, TestLookupTracker, RemoteLookupTrackerClient - the last two are not visible, and the first one is irrelevant
+                messageCollector.report(
+                    CompilerMessageSeverity.WARNING,
+                    "Can't register lookup usages with this compilation setup. lookupTracker is $lookupTracker. $HOW_TO_FIX"
+                )
+            }
+        }
+    }
+
+    // Kotlin plugin can be used with older versions of jps-build, so we check for the availability of APIs.
+    private fun checkRequiredJpsBuildApi(): Boolean {
+        try {
+            Class.forName("org.jetbrains.jps.dependency.java.LookupNameUsage")
+        } catch (_: Throwable) {
+            return false
+        }
+        return true
+    }
+
+    private fun registerLookupTrackerImplEntries(lookupTracker: LookupTrackerImpl, callback: Backend) {
+        for ((lookupKey, fileList) in lookupTracker.lookups.entrySet()) {
+            val symbolOwner = lookupKey.scope.replace('.', '/')
+            val symbolName = lookupKey.name
+            val usage = LookupNameUsage(symbolOwner, symbolName)
+            for (file in fileList) {
+                callback.registerUsage(Paths.get(file), usage)
+            }
+        }
+    }
+
+    private companion object {
+        // these branches should only be reachable if the jps is in the dependency graph mode
+        private const val HOW_TO_FIX =
+            "Kotlin incremental compilation might be incorrect. Consider using build option -Djps.use.dependency.graph=false"
+    }
+}