EmbedRClassesBytecodeTransformer.java

package com.example;

import com.android.build.api.transform.DirectoryInput;
import com.android.build.api.transform.Format;
import com.android.build.api.transform.QualifiedContent;
import com.android.build.api.transform.Status;
import com.android.build.api.transform.Transform;
import com.android.build.api.transform.TransformException;
import com.android.build.api.transform.TransformInput;
import com.android.build.api.transform.TransformInvocation;
import com.android.build.api.transform.TransformOutputProvider;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.NotFoundException;
import javassist.bytecode.ClassFile;
import javassist.bytecode.ConstPool;
import kotlin.io.FilesKt;

import static com.android.build.api.transform.QualifiedContent.DefaultContentType.CLASSES;
import static com.android.build.api.transform.QualifiedContent.Scope.PROJECT;


/**
 * Transforms bytecode of classes, that are embedded with fat-aar.
 *
 * Because we embed multiple modules into one fat-aar, aapt wont generate
 * R class for embedded modules. So we just replace references of their R
 * classes with current module R class.
 */
public class EmbedRClassesBytecodeTransformer extends Transform {

    private final Map<String, String> classTransformTable;
    private final Logger logger = Logging.getLogger(this.getClass());

    public EmbedRClassesBytecodeTransformer(final List<String> libraryPackages, final String targetPackage) {
        // We need jvm internal representation to use ConstPool.renameClass
        // Example transform: com.example.lib -> com/example/lib/R$style
        List<String> resourceTypes = Arrays.asList("anim", "animator", "array", "attr", "bool", "color", "dimen",
                "drawable", "font", "fraction", "id", "integer", "interpolator", "layout", "menu", "mipmap", "plurals",
                "raw", "string", "style", "styleable", "transition", "xml");

        HashMap<String, String> map = new HashMap<>();
        for (String resource : resourceTypes) {
            String targetClass = targetPackage.replace(".", "/") + "/R$" + resource;
            for (String libraryPackage : libraryPackages) {
                String fromClass = libraryPackage.replace(".", "/") + "/R$" + resource;
                map.put(fromClass, targetClass);
            }
        }
        classTransformTable = map;
    }

    @Override
    public String getName() {
        return "fatAarRTransform";
    }

    @Override
    public Set<QualifiedContent.ContentType> getInputTypes() {
        HashSet<QualifiedContent.ContentType> set = new HashSet<>();
        set.add(CLASSES);
        return set;
    }

    @Override
    public Set<? super QualifiedContent.Scope> getScopes() {
        HashSet<QualifiedContent.Scope> set = new HashSet<>();
        set.add(PROJECT);
        return set;
    }

    @Override
    public boolean isIncremental() {
        return true;
    }

    @Override
    public void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException {
        final TransformOutputProvider outputProvider = transformInvocation.getOutputProvider();
        final boolean isIncremental = transformInvocation.isIncremental();
        if (!isIncremental) {
            outputProvider.deleteAll();
        }

        final File outputDir = outputProvider.getContentLocation("classes", getOutputTypes(), getScopes(), Format.DIRECTORY);

        if (isIncremental) {
            logger.debug("Incremental bytecode transform: " + outputDir);
            logger.debug("Count of files in this dir: " + countFilesInDir(outputDir));
        } else {
            logger.debug("Incremental bytecode transform: not incremental");
        }

        try {
            for (final TransformInput input : transformInvocation.getInputs()) {
                for (final DirectoryInput directoryInput : input.getDirectoryInputs()) {
                    final File directoryFile = directoryInput.getFile();

                    final ClassPool classPool = new ClassPool();
                    classPool.insertClassPath(directoryFile.getAbsolutePath());

                    for (final File originalClassFile : getChangedClassesList(isIncremental, directoryInput)) {
                        if (!originalClassFile.getPath().endsWith(".class")) {
                            continue; // ignore anything that is not class file
                        }
                        File relative = FilesKt.relativeTo(originalClassFile, directoryFile);
                        String className = filePathToClassname(relative);
                        final CtClass ctClass = classPool.get(className);
                        transformClass(ctClass);
                        ctClass.writeFile(outputDir.getAbsolutePath());
                    }

                    logger.debug("Count of files after transform: " + countFilesInDir(directoryFile));
                }
                if (!input.getJarInputs().isEmpty()) {
                    String paths = input.getJarInputs().stream()
                            .map(QualifiedContent::getFile)
                            .map(File::getPath)
                            .collect(Collectors.joining());
                    throw new RuntimeException("Expected only directory inputs, but found jar inputs: " + paths);
                }
            }
        } catch (NotFoundException | CannotCompileException e) {
            throw new RuntimeException(e);
        }
    }

    private List<File> getChangedClassesList(final boolean isIncremental, final DirectoryInput directoryInput)
            throws IOException {
        if (!isIncremental) {
            return Files.walk(directoryInput.getFile().toPath())
                    .filter(Files::isRegularFile)
                    .map(Path::toFile)
                    .collect(Collectors.toList());
        } else {
            final Map<File, Status> changedFiles = directoryInput.getChangedFiles();
            changedFiles.entrySet().stream()
                    .filter(it -> it.getValue() == Status.REMOVED)
                    .forEach(it -> it.getKey().delete());

            return changedFiles.entrySet().stream()
                    .filter(it -> it.getValue() == Status.ADDED || it.getValue() == Status.CHANGED)
                    .map(Map.Entry::getKey)
                    .filter(File::isFile)
                    .collect(Collectors.toList());
        }
    }

    // Imports substitution happens here
    private void transformClass(final CtClass ctClass) {
        ClassFile classFile = ctClass.getClassFile();
        ConstPool constPool = classFile.getConstPool();
        constPool.renameClass(classTransformTable);
    }

    private String filePathToClassname(File file) {
        return file.getPath().replace("/", ".")
                .replace("\\", ".")
                .replace(".class", "");
    }

    private long countFilesInDir(File dir) throws IOException {
        return Files.walk(dir.toPath())
                .parallel()
                .filter(p -> !p.toFile().isDirectory())
                .count();
    }
}