MappingsService.java

package knit;

import com.intellij.ide.util.JavaAnonymousClassesHelper;
import com.intellij.lang.jvm.JvmModifier;
import com.intellij.openapi.ui.Messages;
import com.intellij.psi.*;
import com.intellij.psi.util.PsiTreeUtil;
import knit.mapping.*;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class MappingsService {
    private static final String NO_PACKAGE_PACKAGE = "nopackage";
    private File mappingDirectory = null;

    private final Map<PsiClass, ClassMapping> classMappings = new HashMap<>();
    private final Map<PsiField, FieldMapping> fieldMappings = new HashMap<>();
    private final Map<PsiMethod, MethodMapping> methodMappings = new HashMap<>();
    private final Map<PsiParameter, LocalVariableMapping> parameterMappings = new HashMap<>();
    private final Map<PsiClass, Boolean> isInScope = new HashMap<>();
    private final Map<PsiClass, File> mappingFiles = new HashMap<>();

    public void loadMappings(File mappingDirectory) {
        clearMappings();
        this.mappingDirectory = mappingDirectory;
    }

    public void clearMappings() {
        mappingDirectory = null;
        clearCache();
    }

    public void clearCache() {
        classMappings.clear();
        fieldMappings.clear();
        methodMappings.clear();
        isInScope.clear();
        mappingFiles.clear();
    }

    public boolean hasMappings() {
        return mappingDirectory != null;
    }

    public void markChanged(PsiElement element) {
        PsiClass rootClass = element instanceof PsiClass ? (PsiClass) element : PsiTreeUtil.getTopmostParentOfType(element, PsiClass.class);
        ClassMapping mapping = getMapping(rootClass);

        if (mappingFiles.get(rootClass) != null) {
            mappingFiles.get(rootClass).delete();
        }

        File mappingFile = getMappingFile(rootClass.getQualifiedName());
        mappingFiles.put(rootClass, mappingFile);

        try {
            MappingSerializer.writeClass(mappingFile, mapping);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public boolean isInScope(PsiElement element) {
        PsiClass rootClass = element instanceof PsiClass ? (PsiClass) element : PsiTreeUtil.getTopmostParentOfType(element, PsiClass.class);
        return isInScope.computeIfAbsent(rootClass, k -> getMappingFile(rootClass.getQualifiedName()).exists());
    }

    public ClassMapping getMapping(PsiClass clazz) {
        clazz.getQualifiedName();
        return classMappings.computeIfAbsent(clazz, k -> {
            PsiClass containingClass = PsiTreeUtil.getParentOfType(clazz, PsiClass.class);
            if (containingClass != null) {
                ClassMapping containingClassMapping = getMapping(containingClass);
                String relativeName = clazz instanceof PsiAnonymousClass ? JavaAnonymousClassesHelper.getName((PsiAnonymousClass) clazz).substring(1) : clazz.getName();
                String fullName = containingClassMapping.obfuscatedName + "$" + relativeName;

                for (ClassMapping classMapping : containingClassMapping.nestedClasses) {
                    if (classMapping.name.equals(relativeName)) {
                        return classMapping;
                    }
                }

                for (ClassMapping classMapping : containingClassMapping.nestedClasses) {
                    if (classMapping.obfuscatedName.equals(fullName)) {
                        return classMapping;
                    }
                }

                ClassMapping classMapping = new ClassMapping(fullName, fullName);
                containingClassMapping.nestedClasses.add(classMapping);
                return classMapping;
            }

            String mappingName = getMappingName(clazz.getQualifiedName());
            File mappingFile = getMappingFile(clazz.getQualifiedName());

            if (!mappingFile.exists()) {
                return new ClassMapping(mappingName, mappingName);
            }

            try {
                mappingFiles.put(clazz, mappingFile);
                return MappingSerializer.readClass(mappingFile);
            } catch (MappingSerializer.MappingFormatException e) {
                Messages.showErrorDialog(e.getMessage() + " in " + e.file + (e.line == -1 ? "" : " on line " + e.line), "Failed to Read Mappings File");
                return new ClassMapping(mappingName, mappingName);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        });
    }


    public MethodMapping getMapping(PsiMethod method) {
        return methodMappings.computeIfAbsent(method, k -> {
            PsiClass clazz = method.getContainingClass();
            ClassMapping containingClassMapping = method.isConstructor() ? classMappings.computeIfAbsent(clazz, k1 -> {
                PsiClass containingClass = PsiTreeUtil.getParentOfType(clazz, PsiClass.class);
                if (containingClass != null) {
                    ClassMapping containingClassMapping1 = getMapping(containingClass);
                    String actualName1 = clazz instanceof PsiAnonymousClass ? JavaAnonymousClassesHelper.getName((PsiAnonymousClass) clazz) : clazz.getName();

                    for (ClassMapping classMapping : containingClassMapping1.nestedClasses) {
                        if (classMapping.name.equals(actualName1)) {
                            return classMapping;
                        }
                    }

                    for (ClassMapping classMapping : containingClassMapping1.nestedClasses) {
                        if (classMapping.obfuscatedName.equals(actualName1)) {
                            return classMapping;
                        }
                    }

                    ClassMapping classMapping = new ClassMapping(actualName1, actualName1);
                    containingClassMapping1.nestedClasses.add(classMapping);
                    return classMapping;
                }

                String mappingName = getMappingName(clazz.getQualifiedName());
                File mappingFile = getMappingFile(clazz.getQualifiedName());

                if (!mappingFile.exists()) {
                    return new ClassMapping(mappingName, mappingName);
                }

                try {
                    mappingFiles.put(clazz, mappingFile);
                    return MappingSerializer.readClass(mappingFile);
                } catch (MappingSerializer.MappingFormatException e) {
                    Messages.showErrorDialog(e.getMessage() + " in " + e.file + (e.line == -1 ? "" : " on line " + e.line), "Failed to Read Mappings File");
                    return new ClassMapping(mappingName, mappingName);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }) : getMapping(method.getContainingClass());
            String actualName = method.isConstructor() ? "<init>" : method.getName();
            String obfuscatedDescriptor = getObfuscatedDescriptor(method);

            for (MethodMapping methodMapping : containingClassMapping.methods) {
                if (methodMapping.name.equals(actualName) && methodMapping.obfuscatedDescriptor.equals(obfuscatedDescriptor)) {
                    return methodMapping;
                }
            }

            for (MethodMapping methodMapping : containingClassMapping.methods) {
                if (methodMapping.obfuscatedName.equals(actualName) && methodMapping.obfuscatedDescriptor.equals(obfuscatedDescriptor)) {
                    return methodMapping;
                }
            }

            MethodMapping methodMapping = new MethodMapping(actualName, obfuscatedDescriptor, actualName);
            containingClassMapping.methods.add(methodMapping);
            return methodMapping;
        });
    }

    public FieldMapping getMapping(PsiField field) {
        String name = field.getName();
        return fieldMappings.computeIfAbsent(field, k -> {
            ClassMapping containingClassMapping = getMapping(field.getContainingClass());
            String obfuscatedDescriptor = getObfuscatedDescriptor(field.getType());

            for (FieldMapping fieldMapping : containingClassMapping.fields) {
                if (fieldMapping.name.equals(name) && fieldMapping.obfuscatedDescriptor.equals(obfuscatedDescriptor)) {
                    return fieldMapping;
                }
            }

            for (FieldMapping fieldMapping : containingClassMapping.fields) {
                if (fieldMapping.obfuscatedName.equals(name) && fieldMapping.obfuscatedDescriptor.equals(obfuscatedDescriptor)) {
                    return fieldMapping;
                }
            }

            FieldMapping fieldMapping = new FieldMapping(name, obfuscatedDescriptor, name);
            containingClassMapping.fields.add(fieldMapping);
            return fieldMapping;
        });
    }

    public LocalVariableMapping getMapping(PsiParameter parameter) {
        return parameterMappings.computeIfAbsent(parameter, k -> {
            if (parameter.getDeclarationScope() instanceof PsiMethod) {
                PsiMethod method = (PsiMethod) parameter.getDeclarationScope();
                int index = getParameterJvmIndex(parameter);
                MethodMapping enclosingMethodMapping = getMapping(method);

                for (LocalVariableMapping localVariableMapping : enclosingMethodMapping.localVariables) {
                    if (localVariableMapping.index == index) {
                        return localVariableMapping;
                    }
                }

                LocalVariableMapping localVariableMapping = new LocalVariableMapping(index, "arg" + index);
                enclosingMethodMapping.localVariables.add(localVariableMapping);
                return localVariableMapping;
            } else { // not yet supported
                return new LocalVariableMapping(0, "");
            }
        });
    }

    private int getParameterJvmIndex(PsiParameter parameter) {
        PsiMethod method = (PsiMethod) parameter.getDeclarationScope();

        int index = method.hasModifier(JvmModifier.STATIC) ? 0 : 1;
        boolean found = false;
        for (PsiParameter currentParameter : method.getParameterList().getParameters()) {
            if (currentParameter.equals(parameter)) {
                found = true;
                break;
            }
            index += currentParameter.getType().equals(PsiType.LONG) || currentParameter.getType().equals(PsiType.DOUBLE) ? 2 : 1;
        }

        if (!found) {
            throw new AssertionError("parameter not found");
        }

        return index;
    }

    private String getObfuscatedDescriptor(PsiMethod method) {
        StringBuilder descriptor = new StringBuilder();
        descriptor.append('(');

        for (PsiParameter parameter : method.getParameterList().getParameters()) {
            descriptor.append(getObfuscatedDescriptor(parameter.getType()));
        }

        descriptor.append(')');
        descriptor.append(method.isConstructor() ? "V" : getObfuscatedDescriptor(method.getReturnType()));

        return descriptor.toString();
    }

    private String getObfuscatedDescriptor(PsiType type) {
        if (type instanceof PsiPrimitiveType) {
            return ((PsiPrimitiveType) type).getKind().getBinaryName();
        }

        if (type instanceof PsiClassType) {
            PsiClass resolved = ((PsiClassType) type).resolve();

            if (resolved instanceof PsiTypeParameter) {
                PsiClassType[] bounds = ((PsiTypeParameter) resolved).getExtendsList().getReferencedTypes();

                if (bounds.length > 0) {
                    return getObfuscatedDescriptor(bounds[0]);
                }
            }

            return "L" + (resolved != null ? getMapping(resolved).obfuscatedName : ((PsiClassType) type).rawType().getCanonicalText().replace('.', '/')) + ";";
        }

        if (type instanceof PsiArrayType) {
            return "[" + getObfuscatedDescriptor(((PsiArrayType) type).getComponentType());
        }

        throw new IllegalStateException("couldn't find descriptor for " + type);
    }

    private File getMappingFile(String name) {
        return new File(mappingDirectory, getMappingName(name) + ".mapping");
    }

    public String getMappingName(String name) {
        if (name.startsWith(NO_PACKAGE_PACKAGE + ".")) {
            name = name.substring(name.indexOf('.') + 1);
        }

        return name.replace('.', '/');
    }

}