Replace the Scala threading instrumenter with IntelliJ's TMH Java sources

README.md

@@ -437,6 +437,10 @@ Generate JVM bytecode to assert that a method is called on the correct IDEA thre
 4. `com.intellij.util.concurrency.annotations.RequiresReadLockAbsence`
 5. `com.intellij.util.concurrency.annotations.RequiresWriteLock`
 
+This is the same instrumentation that IntelliJ IDEA applies to its own codebase. Note that, matching IntelliJ IDEA's
+behavior, `@RequiresReadLock` generates a call to `ThreadingAssertions.softAssertReadAccess`, which logs an error
+instead of throwing an exception when read access is missing.
+
 See: [IntelliJ IDEA ThreadingAssertions.java](https://github.com/JetBrains/intellij-community/blob/5758eb99b4a1971ebe75cda755693cc930949465/platform/core-api/src/com/intellij/util/concurrency/ThreadingAssertions.java)
 
 #### `instrumentNotNullAnnotations :: SettingKey[Boolean]`

build.sbt

@@ -134,8 +134,9 @@ lazy val ideaSupport = (project in file("ideaSupport"))
       "io.get-coursier" %% "coursier" % "2.1.24",
       "commons-io" % "commons-io" % "2.22.0",
 
-      // Used to compile the test fixtures of the @NotNull instrumentation tests
-      "org.jetbrains" % "annotations" % "26.1.0" % Test
+      // Used by the sources copied from IntelliJ IDEA in the threadingModelHelper package
+      // and to compile the test fixtures of the @NotNull instrumentation tests
+      "org.jetbrains" % "annotations" % "26.1.0"
     ),
   )
 

ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/threadingModelHelper/TMHAssertionGenerator.java

@@ -0,0 +1,17 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package org.jetbrains.jps.devkit.threadingModelHelper.
+// The only changes are the package name and the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm).
+package org.jetbrains.sbtidea.instrumentation.threadingModelHelper;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.MethodVisitor;
+
+@ApiStatus.Internal
+public interface TMHAssertionGenerator {
+  boolean isMyAnnotation(String annotationDescriptor);
+
+  AnnotationVisitor getAnnotationChecker(int api, Runnable onShouldGenerateAssertion);
+
+  void generateAssertion(MethodVisitor writer, int methodStartLineNumber);
+}

ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/threadingModelHelper/TMHAssertionGenerator2.java

@@ -0,0 +1,112 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package org.jetbrains.jps.devkit.threadingModelHelper.
+// The only changes are the package name and the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm).
+package org.jetbrains.sbtidea.instrumentation.threadingModelHelper;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.NotNull;
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+
+import java.util.Set;
+
+@ApiStatus.Internal
+public final class TMHAssertionGenerator2 implements TMHAssertionGenerator {
+
+  private static final String THREAD_ASSERTIONS_CLASS_NAME = "com/intellij/util/concurrency/ThreadingAssertions";
+  private static final String GENERATE_ASSERTION_PARAMETER = "generateAssertion";
+
+  private final String myThreadAssertionsClassName;
+  private final Type myAnnotationClass;
+  private final String myAssertionMethodName;
+
+  TMHAssertionGenerator2(String threadAssertionsClassName, Type annotationClass, String assertionMethodName) {
+    myThreadAssertionsClassName = threadAssertionsClassName;
+    myAnnotationClass = annotationClass;
+    myAssertionMethodName = assertionMethodName;
+  }
+
+  @Override
+  public boolean isMyAnnotation(String annotationDescriptor) {
+    return myAnnotationClass.getDescriptor().equals(annotationDescriptor);
+  }
+
+  @Override
+  public AnnotationVisitor getAnnotationChecker(int api, Runnable onShouldGenerateAssertion) {
+    return new AnnotationChecker(api, onShouldGenerateAssertion);
+  }
+
+  @Override
+  public void generateAssertion(MethodVisitor writer, int methodStartLineNumber) {
+    if (methodStartLineNumber != -1) {
+      Label generatedCodeStart = new Label();
+      writer.visitLabel(generatedCodeStart);
+      writer.visitLineNumber(methodStartLineNumber, generatedCodeStart);
+    }
+    writer.visitMethodInsn(
+      Opcodes.INVOKESTATIC,
+      myThreadAssertionsClassName,
+      myAssertionMethodName,
+      "()V",
+      false
+    );
+  }
+
+  static class AnnotationChecker extends AnnotationVisitor {
+    private boolean myShouldGenerateAssertion = true;
+    private final Runnable myOnShouldGenerateAssertion;
+
+    private AnnotationChecker(int api, Runnable onShouldGenerateAssertion) {
+      super(api);
+      myOnShouldGenerateAssertion = onShouldGenerateAssertion;
+    }
+
+    @Override
+    public void visit(String annotationParameterName, Object value) {
+      if (GENERATE_ASSERTION_PARAMETER.equals(annotationParameterName) && Boolean.FALSE.equals(value)) {
+        myShouldGenerateAssertion = false;
+      }
+    }
+
+    @Override
+    public void visitEnd() {
+      if (myShouldGenerateAssertion) {
+        myOnShouldGenerateAssertion.run();
+      }
+    }
+  }
+
+  // TODO avoid hardcoding annotation names
+  static @NotNull Set<? extends TMHAssertionGenerator> generators() {
+    return GENERATORS;
+  }
+
+  private static final Set<? extends TMHAssertionGenerator> GENERATORS = generators(
+    THREAD_ASSERTIONS_CLASS_NAME,
+    "com/intellij/util/concurrency/annotations"
+  );
+
+  public static @NotNull Set<? extends TMHAssertionGenerator> generators(
+    @NotNull String threadAssertionsClassName,
+    @NotNull String packageString
+  ) {
+    return Set.of(
+      generator(threadAssertionsClassName, packageString + "/RequiresEdt", "assertEventDispatchThread"),
+      generator(threadAssertionsClassName, packageString + "/RequiresBackgroundThread", "assertBackgroundThread"),
+      generator(threadAssertionsClassName, packageString + "/RequiresReadLock", "softAssertReadAccess"),
+      generator(threadAssertionsClassName, packageString + "/RequiresReadLockAbsence", "assertNoReadAccess"),
+      generator(threadAssertionsClassName, packageString + "/RequiresWriteLock", "assertWriteAccess")
+    );
+  }
+
+  private static @NotNull TMHAssertionGenerator generator(
+    @NotNull String threadAssertionsClassName,
+    @NotNull String annotationClassName,
+    @NotNull String assertionMethodName
+  ) {
+    return new TMHAssertionGenerator2(threadAssertionsClassName, Type.getType("L" + annotationClassName + ";"), assertionMethodName);
+  }
+}

ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/threadingModelHelper/TMHInstrumenter.java

@@ -0,0 +1,134 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package org.jetbrains.jps.devkit.threadingModelHelper.
+// The only changes are the package name, the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm),
+// Opcodes.API_VERSION (a JetBrains ASM addition) -> Opcodes.ASM9 and the import of FailSafeMethodVisitor
+// (originally com.intellij.compiler.instrumentation.FailSafeMethodVisitor).
+package org.jetbrains.sbtidea.instrumentation.threadingModelHelper;
+
+import org.jetbrains.sbtidea.instrumentation.notNullVerification.FailSafeMethodVisitor;
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+public final class TMHInstrumenter {
+  public static boolean instrument(ClassReader classReader,
+                                   ClassVisitor classWriter,
+                                   Set<? extends TMHAssertionGenerator> generators,
+                                   boolean generateLineNumbers) {
+    AnnotatedMethodsCollector collector = new AnnotatedMethodsCollector(generators);
+    int options = ClassReader.SKIP_FRAMES;
+    if (!generateLineNumbers) {
+      options |= ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG;
+    }
+    classReader.accept(collector, options);
+    if (collector.annotatedMethods.isEmpty()) {
+      return false;
+    }
+    Instrumenter instrumenter = new Instrumenter(classWriter, collector.annotatedMethods);
+    classReader.accept(instrumenter, 0);
+    return true;
+  }
+
+  private static final class AnnotatedMethodsCollector extends ClassVisitor {
+    final Set<? extends TMHAssertionGenerator> assertionGenerators;
+    final Map<MethodKey, InstrumentationInfo> annotatedMethods = new HashMap<>();
+
+    AnnotatedMethodsCollector(Set<? extends TMHAssertionGenerator> assertionGenerators) {
+      super(Opcodes.ASM9);
+      this.assertionGenerators = assertionGenerators;
+    }
+
+    @Override
+    public MethodVisitor visitMethod(int access, final String name, final String methodDescriptor, String signature, String[] exceptions) {
+      return new MethodVisitor(Opcodes.ASM9) {
+        private final MethodKey methodKey = new MethodKey(name, methodDescriptor);
+        private boolean annotated = false;
+        private boolean firstLineNumberVisited = false;
+
+        @Override
+        public AnnotationVisitor visitAnnotation(String annotationDescriptor, boolean visible) {
+          for (TMHAssertionGenerator assertionGenerator : assertionGenerators) {
+            if (assertionGenerator.isMyAnnotation(annotationDescriptor)) {
+              return assertionGenerator.getAnnotationChecker(Opcodes.ASM9, () -> {
+                annotatedMethods.put(methodKey, new InstrumentationInfo(assertionGenerator));
+                annotated = true;
+              });
+            }
+          }
+          return super.visitAnnotation(annotationDescriptor, visible);
+        }
+
+        @Override
+        public void visitLineNumber(int line, Label start) {
+          super.visitLineNumber(line, start);
+          if (annotated && !firstLineNumberVisited) {
+            annotatedMethods.get(methodKey).methodStartLineNumber = line;
+            firstLineNumberVisited = true;
+          }
+        }
+      };
+    }
+  }
+
+  private static final class Instrumenter extends ClassVisitor {
+    private final Map<MethodKey, InstrumentationInfo> myAnnotatedMethods;
+
+    Instrumenter(ClassVisitor writer, Map<MethodKey, InstrumentationInfo> annotatedMethods) {
+      super(Opcodes.ASM9, writer);
+      myAnnotatedMethods = annotatedMethods;
+    }
+
+    @Override
+    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+      InstrumentationInfo instrumentationInfo = myAnnotatedMethods.get(new MethodKey(name, descriptor));
+      if (instrumentationInfo == null) {
+        return super.visitMethod(access, name, descriptor, signature, exceptions);
+      }
+      return new FailSafeMethodVisitor(Opcodes.ASM9, super.visitMethod(access, name, descriptor, signature, exceptions)) {
+        @Override
+        public void visitCode() {
+          instrumentationInfo.assertionGenerator.generateAssertion(mv, instrumentationInfo.methodStartLineNumber);
+          super.visitCode();
+        }
+      };
+    }
+  }
+
+  private static final class MethodKey {
+    final String name;
+    final String descriptor;
+
+    private MethodKey(String name, String descriptor) {
+      this.name = name;
+      this.descriptor = descriptor;
+    }
+
+    @Override
+    public int hashCode() {
+      int result = 1;
+      result = 31 * result + name.hashCode();
+      result = 31 * result + descriptor.hashCode();
+      return result;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      return obj == this ||
+             obj instanceof MethodKey && ((MethodKey)obj).name.equals(name) && ((MethodKey)obj).descriptor.equals(descriptor);
+    }
+  }
+
+  private static final class InstrumentationInfo {
+    final TMHAssertionGenerator assertionGenerator;
+    int methodStartLineNumber = -1;
+
+    private InstrumentationInfo(TMHAssertionGenerator generator) {assertionGenerator = generator;}
+  }
+}