initial support for class instrumentation

GitOrigin-RevId: 4b4308c792b8252e11c46db3242860a330cff0e2

jps/jps-builders/src/org/jetbrains/jps/bazel/StorageManager.java

@@ -1,6 +1,7 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.jps.bazel;
 
+import com.intellij.compiler.instrumentation.InstrumentationClassFinder;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jps.bazel.impl.AbiJarBuilder;
@@ -9,13 +10,21 @@ import org.jetbrains.jps.dependency.DependencyGraph;
 import org.jetbrains.jps.dependency.GraphConfiguration;
 import org.jetbrains.jps.dependency.impl.DependencyGraphImpl;
 import org.jetbrains.jps.dependency.impl.PersistentMVStoreMapletFactory;
+import org.jetbrains.jps.javac.Iterators;
 
+import java.io.ByteArrayInputStream;
 import java.io.Closeable;
 import java.io.IOException;
+import java.io.InputStream;
+import java.net.MalformedURLException;
+import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.jetbrains.jps.javac.Iterators.collect;
@@ -26,6 +35,7 @@ public class StorageManager implements Closeable {
   private GraphConfiguration myGraphConfig;
   private ZipOutputBuilderImpl myOutputBuilder;
   private AbiJarBuilder myAbiOutputBuilder;
+  private InstrumentationClassFinder myInstrumentationClassFinder;
 
   public StorageManager(BuildContext context) {
     myContext = context;
@@ -35,7 +45,7 @@ public class StorageManager implements Closeable {
     close();
     Path output = myContext.getOutputZip();
     Path abiOutput = myContext.getAbiOutputZip();
-    Path srcSnapshot = DataPaths.getConfigStateStoreFile(myContext);
+    Path srcSnapshotStore = DataPaths.getConfigStateStoreFile(myContext);
 
     BuildProcessLogger logger = myContext.getBuildLogger();
     if (logger.isEnabled() && !myContext.isRebuild()) {
@@ -53,7 +63,7 @@ public class StorageManager implements Closeable {
     if (abiOutput != null) {
       Files.deleteIfExists(abiOutput);
     }
-    Files.deleteIfExists(srcSnapshot);
+    Files.deleteIfExists(srcSnapshotStore);
     Files.deleteIfExists(DataPaths.getDepGraphStoreFile(myContext));
 
     cleanDependenciesBackupDir(myContext);
@@ -98,12 +108,27 @@ public class StorageManager implements Closeable {
     if (builder == null) {
       Path abiOutputPath = myContext.getAbiOutputZip();
       if (abiOutputPath != null) {
-        myAbiOutputBuilder = builder = new AbiJarBuilder(abiOutputPath);
+        myAbiOutputBuilder = builder = new AbiJarBuilder(abiOutputPath, getInstrumentationClassFinder());
       }
     }
     return builder;
   }
 
+  public @NotNull InstrumentationClassFinder getInstrumentationClassFinder() throws MalformedURLException {
+    InstrumentationClassFinder finder = myInstrumentationClassFinder;
+    if (finder == null) {
+      myInstrumentationClassFinder = finder = createInstrumentationClassFinder(path -> {
+        try {
+          return getOutputBuilder().getContent(path);
+        }
+        catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      });
+    }
+    return finder;
+  }
+
   public DependencyGraph getGraph() throws IOException {
     return getGraphConfiguration().getGraph();
   }
@@ -134,4 +159,34 @@ public class StorageManager implements Closeable {
     }
   }
 
+  private InstrumentationClassFinder createInstrumentationClassFinder(Function<String, byte[]> outputContentLookup) throws MalformedURLException {
+    final URL jrt = tryGetJrtURL();
+    List<URL> platformCp = jrt != null? List.of(jrt) : List.of();
+    final List<URL> urls = new ArrayList<>();
+    for (Path path : Iterators.map(myContext.getBinaryDependencies().getElements(), myContext.getPathMapper()::toPath)) {
+      urls.add(path.toUri().toURL());
+    }
+    return new InstrumentationClassFinder(platformCp.toArray(URL[]::new), urls.toArray(URL[]::new)) {
+      @Override
+      protected InputStream lookupClassBeforeClasspath(String internalClassName) {
+        final byte[] content = outputContentLookup.apply(internalClassName);
+        return content != null? new ByteArrayInputStream(content) : null;
+      }
+    };
+  }
+
+  private static URL tryGetJrtURL() {
+    final String home = System.getProperty("java.home");
+    Path jrtFsPath = Path.of(home).normalize().resolve("lib").resolve("jrt-fs.jar");
+    if (Files.isRegularFile(jrtFsPath)) {
+      // this is a modular jdk where platform classes are stored in a jrt-fs image
+      try {
+        return InstrumentationClassFinder.createJDKPlatformUrl(home);
+      }
+      catch (MalformedURLException ignored) {
+      }
+    }
+    return null;
+  }
+
 }

jps/jps-builders/src/org/jetbrains/jps/bazel/impl/AbiJarBuilder.java

@@ -3,6 +3,7 @@ package org.jetbrains.jps.bazel.impl;
 
 import com.dynatrace.hash4j.hashing.HashStream64;
 import com.dynatrace.hash4j.hashing.Hashing;
+import com.intellij.compiler.instrumentation.InstrumentationClassFinder;
 import org.jetbrains.annotations.Nullable;
 
 import java.io.*;
@@ -20,9 +21,16 @@ public class AbiJarBuilder extends ZipOutputBuilderImpl {
   private final Map<String, Long> myPackageIndex = new TreeMap<>(); // directoryEntryName -> digestOf(content entries)
   private boolean myPackageIndexChanged;
   private boolean myShouldStoreIndex;
+  @Nullable
+  private final InstrumentationClassFinder myClassFinder;
 
   public AbiJarBuilder(Path outputZip) throws IOException {
+    this(outputZip, null);
+  }
+
+  public AbiJarBuilder(Path outputZip, @Nullable InstrumentationClassFinder classFinder) throws IOException {
     super(outputZip);
+    myClassFinder = classFinder;
     byte[] content = getContent(PACKAGE_INDEX_STORAGE_ENTRY_NAME);
     if (content != null) {
       readPackageIndex(content, myPackageIndex);
@@ -86,9 +94,12 @@ public class AbiJarBuilder extends ZipOutputBuilderImpl {
   }
 
   private byte @Nullable [] filterAbiJarContent(byte[] content) {
+    if (myClassFinder == null) {
+      return content; // no instrumentation, if class finder is not specified
+    }
     // todo: check content and instrument it before adding
     // todo: for java use JavaAbiClassVisitor, for kotlin-generated classes use KotlinAnnotationVisitor, abiMetadataProcessor
-    return content;
+    return JavaAbiClassFilter.filter(content, myClassFinder);
   }
 
   @Override

jps/jps-builders/src/org/jetbrains/jps/bazel/impl/JavaAbiClassFilter.java

@@ -0,0 +1,106 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.jps.bazel.impl;
+
+import com.intellij.compiler.instrumentation.FailSafeClassReader;
+import com.intellij.compiler.instrumentation.InstrumentationClassFinder;
+import com.intellij.compiler.instrumentation.InstrumenterClassWriter;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.org.objectweb.asm.*;
+import org.jetbrains.org.objectweb.asm.tree.FieldNode;
+import org.jetbrains.org.objectweb.asm.tree.MethodNode;
+
+import java.util.*;
+
+public class JavaAbiClassFilter extends ClassVisitor {
+
+  private boolean isAbiClass;
+  private Set<String> myExcludedClasses = new HashSet<>();
+  private List<FieldNode> myFields = new ArrayList<>();
+  private List<MethodNode> myMethods = new ArrayList<>();
+
+  private JavaAbiClassFilter(ClassVisitor delegate) {
+    super(Opcodes.API_VERSION, delegate);
+  }
+
+  public static byte @Nullable [] filter(byte[] classBytes, InstrumentationClassFinder finder) {
+    ClassReader reader = new FailSafeClassReader(classBytes);
+    int version = InstrumenterClassWriter.getClassFileVersion(reader);
+    ClassWriter writer = new InstrumenterClassWriter(reader, InstrumenterClassWriter.getAsmClassWriterFlags(version), finder);
+    JavaAbiClassFilter abiVisitor = new JavaAbiClassFilter(writer);
+    reader.accept(
+      abiVisitor, ClassReader.SKIP_FRAMES | ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG
+    );
+    return abiVisitor.isAbiClass? writer.toByteArray() : null;
+  }
+
+  @Override
+  public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
+    isAbiClass = isAbiVisible(access);
+    if (isAbiClass) {
+      super.visit(version, access, name, signature, superName, interfaces);
+    }
+    else {
+      myExcludedClasses.add(name);
+    }
+  }
+
+  private static boolean isAbiVisible(int access) {
+    return (access & (Opcodes.ACC_PUBLIC | Opcodes.ACC_PROTECTED)) != 0;
+  }
+
+  @Override
+  public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
+    if (isAbiVisible(access)) {
+      FieldNode field = new FieldNode(Opcodes.API_VERSION, access, name, descriptor, signature, value);
+      myFields.add(field);
+      return field;
+    }
+    return null;
+  }
+
+  @Override
+  public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+    if (isAbiVisible(access)) {
+      MethodNode method = new MethodNode(Opcodes.API_VERSION, access, name, descriptor, signature, exceptions);
+      myMethods.add(method);
+      return method;
+    }
+    return null;
+  }
+
+  @Override
+  public void visitEnd() {
+    Collections.sort(myFields, Comparator.comparing(f -> f.name));
+    for (FieldNode field : myFields) {
+      field.accept(cv);
+    }
+    Collections.sort(myMethods, Comparator.comparing(m -> m.name));
+    for (MethodNode method : myMethods) {
+      method.accept(cv);
+    }
+    super.visitEnd();
+  }
+
+  @Override
+  public void visitNestMember(String nestMember) {
+    if (nestMember == null || !myExcludedClasses.contains(nestMember)) {
+      super.visitNestMember(nestMember);
+    }
+  }
+
+  @Override
+  public void visitPermittedSubclass(String permittedSubclass) {
+    if (permittedSubclass == null || !myExcludedClasses.contains(permittedSubclass)) {
+      super.visitPermittedSubclass(permittedSubclass);
+    }
+  }
+
+  @Override
+  public void visitInnerClass(String name, String outerName, String innerName, int access) {
+    // innerName == null for anonymous classes
+    if (isAbiVisible(access) && innerName != null && !myExcludedClasses.contains(name)) {
+      super.visitInnerClass(name, outerName, innerName, access);
+    }
+  }
+
+}