IntelliJPlatformAnnotationIndexBuilder.java

// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package org.jetbrains.jps.devkit.builder;

import com.google.protobuf.CodedOutputStream;
import com.intellij.compiler.instrumentation.InstrumentationClassFinder;
import com.intellij.compiler.instrumentation.InstrumentationClassFinder.PseudoAnnotation;
import com.intellij.compiler.instrumentation.InstrumentationClassFinder.PseudoClass;
import com.intellij.components.IJPAnnotationIndexModel;
import com.intellij.components.IJPAnnotationIndexModel.ServiceDescriptor;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.util.SystemProperties;
import com.intellij.util.containers.ContainerUtil;
import gnu.trove.THashSet;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.jps.ModuleChunk;
import org.jetbrains.jps.api.GlobalOptions;
import org.jetbrains.jps.builders.DirtyFilesHolder;
import org.jetbrains.jps.builders.java.JavaBuilderUtil;
import org.jetbrains.jps.builders.java.JavaSourceRootDescriptor;
import org.jetbrains.jps.incremental.*;
import org.jetbrains.jps.incremental.instrumentation.ClassProcessingBuilder;
import org.jetbrains.jps.incremental.messages.BuildMessage;
import org.jetbrains.jps.incremental.messages.CompilerMessage;
import org.jetbrains.jps.incremental.messages.FileDeletedEvent;
import org.jetbrains.jps.model.java.JpsJavaExtensionService;
import org.jetbrains.jps.util.JpsPathUtil;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SeekableByteChannel;
import java.nio.file.*;
import java.util.*;

/**
 * Protobuf is used as format for index file to ensure that regardless of any format change, class name list will be possible to read,
 * to avoid rebuilding the whole module - instead, just re-process corresponding class files on disk.
 */
final class IntelliJPlatformAnnotationIndexBuilder extends ClassProcessingBuilder {
  public static final int INDEX_FORMAT_VERSION = 0;
  public static final String INDEX_FILE_NAME = "ij-services";

  private static final Logger LOG = Logger.getInstance(IntelliJPlatformAnnotationIndexBuilder.class);

  private static final Set<OpenOption> WRITE_FILE_OPTIONS = new THashSet<>(Arrays.asList(StandardOpenOption.WRITE,
                                                                                         StandardOpenOption.CREATE,
                                                                                         StandardOpenOption.TRUNCATE_EXISTING));

  private final Set<String> deletedFiles = ContainerUtil.newConcurrentSet();

  private boolean isEnabled;

  IntelliJPlatformAnnotationIndexBuilder() {
    super(BuilderCategory.CLASS_POST_PROCESSOR);
  }

  @Override
  public void buildStarted(CompileContext context) {
    super.buildStarted(context);

    deletedFiles.clear();
    isEnabled = SystemProperties.is(GlobalOptions.PROCESS_IJ_PLATFORM_ANNOTATION_OPTION);

    // no need to listen for rebuild
    if (!JavaBuilderUtil.isForcedRecompilationAllJavaModules(context)) {
      context.addBuildListener(new BuildListener() {
        @Override
        public void filesDeleted(@NotNull FileDeletedEvent event) {
          for (String path : event.getFilePaths()) {
            if (path.endsWith(".class")) {
              deletedFiles.add(path.replace('\\', '/'));
            }
          }
        }
      });
    }
  }

  @Override
  public void buildFinished(CompileContext context) {
    deletedFiles.clear();
    isEnabled = false;

    super.buildFinished(context);
  }

  @Override
  protected boolean isEnabled(CompileContext context, ModuleChunk chunk) {
    return isEnabled;
  }

  @NotNull
  @Override
  public String getPresentableName() {
    return "IJ Platform annotation processor";
  }

  @Override
  protected String getProgressMessage() {
    return null;
  }

  @Override
  public ExitCode build(CompileContext context,
                        ModuleChunk chunk,
                        DirtyFilesHolder<JavaSourceRootDescriptor, ModuleBuildTarget> dirtyFilesHolder,
                        OutputConsumer outputConsumer) throws IOException {
    if (!isEnabled) {
      return ExitCode.NOTHING_DONE;
    }

    Map<String, CompiledClass> compiledClasses = outputConsumer.getCompiledClasses();

    if (LOG.isDebugEnabled()) {
      LOG.debug("build (module=" + chunk.getName() + ", compiledClasses=" + compiledClasses + ")");
    }

    if (compiledClasses.isEmpty()) {
      return deletedFiles.isEmpty() ? ExitCode.NOTHING_DONE : updateIfOnlyFilesRemoved(chunk);
    }
    else {
      return super.build(context, chunk, dirtyFilesHolder, outputConsumer);
    }
  }

  @NotNull
  private ExitCode updateIfOnlyFilesRemoved(@NotNull ModuleChunk chunk) throws IOException {
    ExitCode exitCode = ExitCode.NOTHING_DONE;
    for (ModuleBuildTarget target : chunk.getTargets()) {
      String moduleOutDir = JpsPathUtil.urlToPath(JpsJavaExtensionService.getInstance().getOutputUrl(target.getModule(), target.isTests()));
      if (moduleOutDir == null) {
        continue;
      }

      Path indexFile = Paths.get(moduleOutDir, "META-INF", INDEX_FILE_NAME);
      List<ServiceDescriptor> list = readListOrNull(indexFile, true);
      if (list == null) {
        continue;
      }

      if (saveFilteredList(list, indexFile, moduleOutDir)) {
        exitCode = ExitCode.OK;
      }
    }

    return exitCode;
  }

  @Nullable
  private static List<ServiceDescriptor> readListOrNull(@NotNull Path indexFile, boolean deleteOnFail) throws IOException {
    byte[] bytes;
    try {
      bytes = Files.readAllBytes(indexFile);
    }
    catch (NoSuchFileException e) {
      return null;
    }

    try {
      // todo check version
      return IJPAnnotationIndexModel.Index.parseFrom(bytes).getDescriptorsList();
    }
    catch (IOException e) {
      if (deleteOnFail) {
        Files.delete(indexFile);
      }
      LOG.error("Cannot read index file", e);
      return null;
    }
  }

  @Override
  protected ExitCode performBuild(CompileContext context, ModuleChunk chunk, InstrumentationClassFinder finder, OutputConsumer outputConsumer) {
    ExitCode exitCode = ExitCode.NOTHING_DONE;
    for (ModuleBuildTarget target : chunk.getTargets()) {
      String moduleOutDir = JpsPathUtil.urlToPath(JpsJavaExtensionService.getInstance().getOutputUrl(target.getModule(), target.isTests()));
      if (moduleOutDir == null) {
        context.processMessage(new CompilerMessage(getPresentableName(), BuildMessage.Kind.WARNING, "Module output is null"));
        continue;
      }

      Path indexFile = Paths.get(moduleOutDir, "META-INF", "ij-services");
      List<ServiceDescriptor> descriptors = null;
      for (CompiledClass compiledClass : outputConsumer.getTargetCompiledClasses(target)) {
        try {
          String className = compiledClass.getClassName();
          if (className == null) {
            continue;
          }

          PseudoClass pseudoClass = finder.loadClass(className);
          ServiceDescriptor descriptor = getServiceDescriptor(pseudoClass);
          if (descriptor == null) {
            continue;
          }

          if (descriptors == null) {
            descriptors = new ArrayList<>();
          }
          descriptors.add(descriptor);
          deletedFiles.remove(compiledClass.getOutputFile().getPath().replace('\\', '/'));
        }
        catch (Exception e) {
          context.processMessage(CompilerMessage.createInternalCompilationError("Cannot process class " + compiledClass, e));
        }
      }

      if (descriptors == null) {
        continue;
      }

      try {
        saveNewClasses(descriptors, indexFile, moduleOutDir);

        outputConsumer.registerOutputFile(target, indexFile.toFile(), Collections.emptyList());
        exitCode = ExitCode.OK;
      }
      catch (Exception e) {
        context.processMessage(CompilerMessage.createInternalCompilationError("Cannot save " + indexFile, e));
      }
    }
    return exitCode;
  }

  private void saveNewClasses(@NotNull List<ServiceDescriptor> descriptors, @NotNull Path indexFile, @NotNull String moduleOutDir) throws IOException {
    List<ServiceDescriptor> list = readListOrNull(indexFile, false);
    List<ServiceDescriptor> effectiveList;
    boolean saveEvenIfNotFiltered;
    if (list == null) {
      saveEvenIfNotFiltered = true;
      effectiveList = descriptors;
      Files.createDirectories(indexFile.getParent());
    }
    else {
      saveEvenIfNotFiltered = false;
      effectiveList = list;
      Set<String> existingClassNameSet = new THashSet<>(list.size());
      for (ServiceDescriptor descriptor : list) {
        existingClassNameSet.add(descriptor.getClassName());
      }

      for (ServiceDescriptor descriptor : descriptors) {
        if (!existingClassNameSet.contains(descriptor.getClassName())) {
          list.add(descriptor);
          saveEvenIfNotFiltered = true;
        }
      }
    }

    if (!saveFilteredList(effectiveList, indexFile, moduleOutDir) && saveEvenIfNotFiltered) {
      saveIndex(indexFile, effectiveList);
    }
  }

  private boolean saveFilteredList(@NotNull List<ServiceDescriptor> list, @NotNull Path indexFile, @NotNull String moduleOutDir) throws IOException {
    List<ServiceDescriptor> newList = null;
    for (int i = 0, size = list.size(); i < size; i++) {
      ServiceDescriptor descriptor = list.get(i);
      if (deletedFiles.contains(moduleOutDir + '/' + descriptor.getClassName().replace('.', '/') + ".class")) {
        if (newList == null) {
          newList = new ArrayList<>(size - 1);
          for (int j = 0; j < i; j++) {
            newList.add(list.get(j));
          }
        }
      }
      else if (newList != null) {
        newList.add(descriptor);
      }
    }

    if (newList == null) {
      return false;
    }

    if (newList.isEmpty()) {
      // yes, empty META-INF dir can be left - it is ok
      Files.delete(indexFile);
    }
    else {
      saveIndex(indexFile, newList);
    }
    return true;
  }

  private static void saveIndex(@NotNull Path indexFile, @NotNull List<ServiceDescriptor> descriptors) throws IOException {
    if (LOG.isDebugEnabled()) {
      LOG.debug("Save index to " + indexFile);
    }
    IJPAnnotationIndexModel.Index.Builder builder = IJPAnnotationIndexModel.Index.newBuilder();
    builder.setVersion(INDEX_FORMAT_VERSION);
    builder.addAllDescriptors(descriptors);
    IJPAnnotationIndexModel.Index result = builder.build();

    // protobuf internally in any case calls `getSerializedSize` (see AbstractMessageLite.writeTo), so, no need to use BufferExposingByteArrayOutputStream
    ByteBuffer buffer = ByteBuffer.allocate(result.getSerializedSize());
    result.writeTo(CodedOutputStream.newInstance(buffer));
    try (SeekableByteChannel channel = Files.newByteChannel(indexFile, WRITE_FILE_OPTIONS)) {
      channel.write(buffer);
    }
  }

  @Nullable
  private static ServiceDescriptor getServiceDescriptor(@NotNull PseudoClass pseudoClass) {
    @SuppressWarnings("SpellCheckingInspection")
    PseudoAnnotation serviceAnnotation = pseudoClass.findAnnotation("Lcom/intellij/openapi/components/Service;");
    if (serviceAnnotation == null) {
      return null;
    }

    ServiceDescriptor.Builder result = ServiceDescriptor.newBuilder();
    result.setClassName(pseudoClass.getName());
    @SuppressWarnings("SpellCheckingInspection")
    PseudoAnnotation stateAnnotation = pseudoClass.findAnnotation("Lcom/intellij/openapi/components/State");
    if (stateAnnotation != null) {
      result.setPersistentStateComponent(true);
    }
    return result.build();
  }
}