Custom conversion rewrites as gradle tasks

Bug: 222647019
Change-Id: Ifdddc41a608ec11044c50923843b126672925372
diff --git a/buildSrc/src/main/java/desugaredlibrary/AsmRewriter.java b/buildSrc/src/main/java/desugaredlibrary/AsmRewriter.java
new file mode 100644
index 0000000..cc1fcf2
--- /dev/null
+++ b/buildSrc/src/main/java/desugaredlibrary/AsmRewriter.java
@@ -0,0 +1,53 @@
+// Copyright (c) 2022, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package desugaredlibrary;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+public abstract class AsmRewriter {
+
+  public static int ASM_VERSION = Opcodes.ASM9;
+
+  public static byte[] transformInvoke(byte[] bytes, MethodTransformer transformer) {
+    ClassReader reader = new ClassReader(bytes);
+    ClassWriter writer = new ClassWriter(reader, 0);
+    ClassVisitor subvisitor = new InvokeTransformer(writer, transformer);
+    reader.accept(subvisitor, 0);
+    return writer.toByteArray();
+  }
+
+  public static class InvokeTransformer extends ClassVisitor {
+
+    private final MethodTransformer transformer;
+
+    InvokeTransformer(ClassWriter writer, MethodTransformer transformer) {
+      super(ASM_VERSION, writer);
+      this.transformer = transformer;
+    }
+
+    @Override
+    public MethodVisitor visitMethod(
+        int access, String name, String descriptor, String signature, String[] exceptions) {
+      MethodVisitor sub = super.visitMethod(access, name, descriptor, signature, exceptions);
+      transformer.setMv(sub);
+      return transformer;
+    }
+  }
+
+  public static class MethodTransformer extends MethodVisitor {
+
+    protected MethodTransformer(int api) {
+      super(api);
+    }
+
+    public void setMv(MethodVisitor visitor) {
+      this.mv = visitor;
+    }
+  }
+}
diff --git a/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriteDescription.java b/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriteDescription.java
new file mode 100644
index 0000000..137ac12
--- /dev/null
+++ b/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriteDescription.java
@@ -0,0 +1,54 @@
+// Copyright (c) 2022, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package desugaredlibrary;
+
+import com.google.common.collect.ImmutableSet;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+public class CustomConversionAsmRewriteDescription {
+
+  static final String WRAP_CONVERT = "wrap_convert";
+  static final String INVERTED_WRAP_CONVERT = "inverted_wrap_convert";
+  static final String CONVERT = "convert";
+
+  private static final Set<String> ENUM_WRAP_CONVERT_OWNER =
+      ImmutableSet.of(
+          "j$/nio/file/StandardOpenOption",
+          "j$/nio/file/LinkOption",
+          "j$/nio/file/attribute/PosixFilePermission",
+          "j$/util/stream/Collector$Characteristics");
+  private static final Set<String> WRAP_CONVERT_OWNER =
+      ImmutableSet.of(
+          "j$/nio/file/spi/FileSystemProvider",
+          "j$/nio/file/spi/FileTypeDetector",
+          "j$/nio/file/Path",
+          "j$/nio/file/WatchEvent",
+          "j$/nio/file/attribute/BasicFileAttributes",
+          "j$/nio/file/attribute/BasicFileAttributeView",
+          "j$/nio/file/attribute/FileOwnerAttributeView",
+          "j$/nio/file/attribute/PosixFileAttributes",
+          "j$/nio/file/attribute/PosixFileAttributeView");
+
+  static Map<String, String> getJavaWrapConvertOwnerMap() {
+    return computeConvertOwnerMap("$VivifiedWrapper");
+  }
+
+  static Map<String, String> getJ$WrapConvertOwnerMap() {
+    return computeConvertOwnerMap("$Wrapper");
+  }
+
+  private static HashMap<String, String> computeConvertOwnerMap(String suffix) {
+    HashMap<String, String> map = new HashMap<>();
+    for (String theEnum : ENUM_WRAP_CONVERT_OWNER) {
+      map.put(theEnum, theEnum + "$EnumConversion");
+    }
+    for (String owner : WRAP_CONVERT_OWNER) {
+      map.put(owner, owner + suffix);
+    }
+    return map;
+  }
+}
diff --git a/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriter.java b/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriter.java
new file mode 100644
index 0000000..6107e67
--- /dev/null
+++ b/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriter.java
@@ -0,0 +1,183 @@
+// Copyright (c) 2022, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package desugaredlibrary;
+
+import static desugaredlibrary.AsmRewriter.ASM_VERSION;
+import static desugaredlibrary.CustomConversionAsmRewriteDescription.CONVERT;
+import static desugaredlibrary.CustomConversionAsmRewriteDescription.INVERTED_WRAP_CONVERT;
+import static desugaredlibrary.CustomConversionAsmRewriteDescription.WRAP_CONVERT;
+import static desugaredlibrary.CustomConversionAsmRewriter.CustomConversionVersion.LEGACY;
+import static org.objectweb.asm.Opcodes.INVOKESTATIC;
+
+import com.google.common.io.ByteStreams;
+import desugaredlibrary.AsmRewriter.MethodTransformer;
+import java.io.BufferedOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.OpenOption;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.Enumeration;
+import java.util.Map;
+import java.util.zip.CRC32;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipFile;
+import java.util.zip.ZipOutputStream;
+
+public class CustomConversionAsmRewriter {
+
+  public enum CustomConversionVersion {
+    LEGACY,
+    LATEST
+  }
+
+  public CustomConversionAsmRewriter(CustomConversionVersion legacy) {
+    this.legacy = legacy;
+  }
+
+  private final CustomConversionVersion legacy;
+  private final Map<String, String> javaWrapConvertOwnerMap =
+      CustomConversionAsmRewriteDescription.getJavaWrapConvertOwnerMap();
+  private final Map<String, String> j$WrapConvertOwnerMap =
+      CustomConversionAsmRewriteDescription.getJ$WrapConvertOwnerMap();
+
+  public static void generateJars(Path jar, Path outputDirectory) throws IOException {
+    for (CustomConversionVersion version : CustomConversionVersion.values()) {
+      new CustomConversionAsmRewriter(version).convert(jar, outputDirectory);
+    }
+  }
+
+  private void convert(Path jar, Path outputDirectory) throws IOException {
+    String fileName = jar.getFileName().toString();
+    String newFileName =
+        fileName.substring(0, fileName.length() - "_raw.jar".length())
+            + (legacy == LEGACY ? "_legacy" : "")
+            + ".jar";
+    Path convertedJar = outputDirectory.resolve(newFileName);
+    internalConvert(jar, convertedJar);
+  }
+
+  private void internalConvert(Path jar, Path convertedJar) throws IOException {
+    OpenOption[] options =
+        new OpenOption[] {StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING};
+    try (ZipOutputStream out =
+        new ZipOutputStream(
+            new BufferedOutputStream(Files.newOutputStream(convertedJar, options)))) {
+      new CustomConversionAsmRewriter(legacy).convert(jar, out, legacy);
+    }
+  }
+
+  private void convert(
+      Path desugaredLibraryFiles, ZipOutputStream out, CustomConversionVersion legacy)
+      throws IOException {
+    try (ZipFile zipFile = new ZipFile(desugaredLibraryFiles.toFile(), StandardCharsets.UTF_8)) {
+      final Enumeration<? extends ZipEntry> entries = zipFile.entries();
+      while (entries.hasMoreElements()) {
+        ZipEntry entry = entries.nextElement();
+        try (InputStream entryStream = zipFile.getInputStream(entry)) {
+          handleFile(entry, entryStream, out, legacy);
+        }
+      }
+    }
+  }
+
+  private void handleFile(
+      ZipEntry entry, InputStream input, ZipOutputStream out, CustomConversionVersion legacy)
+      throws IOException {
+    if (!entry.getName().endsWith(".class")) {
+      return;
+    }
+    if (legacy == LEGACY
+        && (entry.getName().contains("java.nio.file")
+            || entry.getName().contains("ApiFlips")
+            || entry.getName().contains("java.adapter"))) {
+      return;
+    }
+    final byte[] bytes = ByteStreams.toByteArray(input);
+    input.close();
+    final byte[] rewrittenBytes = transformInvoke(bytes);
+    writeToZipStream(out, entry.getName(), rewrittenBytes, ZipEntry.STORED);
+  }
+
+  public static void writeToZipStream(
+      ZipOutputStream stream, String entry, byte[] content, int compressionMethod)
+      throws IOException {
+    int offset = 0;
+    int length = content.length;
+    CRC32 crc = new CRC32();
+    crc.update(content, offset, length);
+    ZipEntry zipEntry = new ZipEntry(entry);
+    zipEntry.setMethod(compressionMethod);
+    zipEntry.setSize(length);
+    zipEntry.setCrc(crc.getValue());
+    zipEntry.setTime(0);
+    stream.putNextEntry(zipEntry);
+    stream.write(content, offset, length);
+    stream.closeEntry();
+  }
+
+  private byte[] transformInvoke(byte[] bytes) {
+    return AsmRewriter.transformInvoke(bytes, new CustomConversionRewriter(ASM_VERSION));
+  }
+
+  class CustomConversionRewriter extends MethodTransformer {
+
+    protected CustomConversionRewriter(int api) {
+      super(api);
+    }
+
+    @Override
+    public void visitMethodInsn(
+        int opcode, String owner, String name, String descriptor, boolean isInterface) {
+      if (opcode == INVOKESTATIC) {
+        if (name.equals(WRAP_CONVERT)) {
+          convertInvoke(
+              opcode,
+              owner,
+              descriptor,
+              isInterface,
+              javaWrapConvertOwnerMap,
+              j$WrapConvertOwnerMap);
+          return;
+        }
+        if (name.equals(INVERTED_WRAP_CONVERT)) {
+          convertInvoke(
+              opcode,
+              owner,
+              descriptor,
+              isInterface,
+              j$WrapConvertOwnerMap,
+              javaWrapConvertOwnerMap);
+          return;
+        }
+      }
+      super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
+    }
+
+    private void convertInvoke(
+        int opcode,
+        String owner,
+        String descriptor,
+        boolean isInterface,
+        Map<String, String> javaMap,
+        Map<String, String> j$Map) {
+      if (!javaMap.containsKey(owner)
+          || !j$Map.containsKey(owner)
+          || !(owner.startsWith("java") || owner.startsWith("j$"))) {
+        throw new RuntimeException("Cannot transform wrap_convert method for " + owner);
+      }
+      if (owner.startsWith("java")) {
+        String newOwner = j$Map.get(owner);
+        super.visitMethodInsn(opcode, newOwner, CONVERT, descriptor, isInterface);
+        return;
+      }
+      assert owner.startsWith("j$");
+      String newOwner = javaMap.get(owner);
+      super.visitMethodInsn(opcode, newOwner, CONVERT, descriptor, isInterface);
+    }
+  }
+}
diff --git a/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriterTask.java b/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriterTask.java
new file mode 100644
index 0000000..bae3c1c
--- /dev/null
+++ b/buildSrc/src/main/java/desugaredlibrary/CustomConversionAsmRewriterTask.java
@@ -0,0 +1,78 @@
+// Copyright (c) 2022, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package desugaredlibrary;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import javax.inject.Inject;
+import org.gradle.api.DefaultTask;
+import org.gradle.api.tasks.InputFile;
+import org.gradle.api.tasks.OutputDirectory;
+import org.gradle.api.tasks.TaskAction;
+import org.gradle.workers.IsolationMode;
+import org.gradle.workers.WorkerExecutor;
+
+public class CustomConversionAsmRewriterTask extends DefaultTask {
+
+  private final WorkerExecutor workerExecutor;
+
+  private File rawJar;
+  private File outputDirectory;
+
+  @Inject
+  public CustomConversionAsmRewriterTask(WorkerExecutor workerExecutor) {
+    this.workerExecutor = workerExecutor;
+  }
+
+  @InputFile
+  public File getRawJar() {
+    return rawJar;
+  }
+
+  public void setRawJar(File rawJar) {
+    this.rawJar = rawJar;
+  }
+
+  @OutputDirectory
+  public File getOutputDirectory() {
+    return outputDirectory;
+  }
+
+  public void setOutputDirectory(File outputDirectory) {
+    this.outputDirectory = outputDirectory;
+  }
+
+  @TaskAction
+  void exec() {
+    workerExecutor.submit(
+        Run.class,
+        config -> {
+          config.setIsolationMode(IsolationMode.NONE);
+          config.params(rawJar, outputDirectory);
+        });
+  }
+
+  public static class Run implements Runnable {
+
+    private final File rawJar;
+    private final File outputDirectory;
+
+    @Inject
+    public Run(File rawJar, File outputDirectory) {
+      this.rawJar = rawJar;
+      this.outputDirectory = outputDirectory;
+    }
+
+    @Override
+    public void run() {
+      try {
+        CustomConversionAsmRewriter.generateJars(rawJar.toPath(), outputDirectory.toPath());
+      } catch (IOException e) {
+        throw new UncheckedIOException(e);
+      }
+    }
+  }
+}