Update cfr and fixes a few related issues (#414)

* Update cfr and fixes a few related issues

fixes #368 context for token references, missing some field and method
references

* Prevent crash when cfr fails to decompile a lambda and creates a fallback

Signed-off-by: liach <liach@users.noreply.github.com>

* Fix issue with parameters incorrectly spilling into other methods

Signed-off-by: liach <liach@users.noreply.github.com>

* try improve performance of cfr source

Co-authored-by: liach <liach@users.noreply.github.com>

4 files changed, 51 insertions, 30 deletions

diff --git a/enigma/build.gradle b/enigma/build.gradle
index ec1c878c..1f53f461 100644
--- a/enigma/build.gradle
+++ b/enigma/build.gradle
@@ -9,7 +9,7 @@ dependencies {
     implementation 'org.ow2.asm:asm-util:9.2'
 
     implementation 'net.fabricmc:procyon-fabric-compilertools:0.5.35.13'
-    implementation 'net.fabricmc:cfr:0.0.4'
+    implementation 'net.fabricmc:cfr:0.0.6'
 
     proGuard 'com.guardsquare:proguard-base:7.1.0-beta5'
 }
diff --git a/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrDecompiler.java b/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrDecompiler.java
index ffc4788c..91fd5a27 100644
--- a/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrDecompiler.java
+++ b/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrDecompiler.java
@@ -24,7 +24,6 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 import org.objectweb.asm.tree.ClassNode;
 
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.Map;
 
 public class CfrDecompiler implements Decompiler {
@@ -33,7 +32,7 @@ public class CfrDecompiler implements Decompiler {
     private final SourceSettings settings;
 
     public CfrDecompiler(ClassProvider classProvider, SourceSettings sourceSettings) {
-        Map<String, String> options = new HashMap<>();
+        Map<String, String> options = Map.of("trackbytecodeloc", "true");
 
         state = new DCCommonState(OptionsImpl.getFactory().create(options), new ClassFileSource2() {
             @Override
diff --git a/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrSource.java b/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrSource.java
index ba831d76..7c63afcc 100644
--- a/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrSource.java
+++ b/enigma/src/main/java/cuchaz/enigma/source/cfr/CfrSource.java
@@ -13,11 +13,12 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 public class CfrSource implements Source {
     private final SourceSettings settings;
     private final ClassFile tree;
-    private final SourceIndex index;
-    private final String string;
     private final DCCommonState state;
     private final TypeUsageInformation typeUsage;
     private final Options options;
+    private final EntryRemapper mapper;
+
+    private SourceIndex index;
 
     public CfrSource(SourceSettings settings, ClassFile tree, DCCommonState state, TypeUsageInformation typeUsage, Options options, @Nullable EntryRemapper mapper) {
         this.settings = settings;
@@ -25,16 +26,7 @@ public class CfrSource implements Source {
         this.state = state;
         this.typeUsage = typeUsage;
         this.options = options;
-
-        EnigmaDumper dumper = new EnigmaDumper(new StringBuilder(), settings, typeUsage, options, mapper);
-        tree.dump(state.getObfuscationMapping().wrap(dumper));
-        index = dumper.getIndex();
-        string = dumper.getString();
-    }
-
-    @Override
-    public String asString() {
-        return string;
+        this.mapper = mapper;
     }
 
     @Override
@@ -44,6 +36,23 @@ public class CfrSource implements Source {
 
     @Override
     public SourceIndex index() {
+        ensureDecompiled();
         return index;
     }
+
+    @Override
+    public String asString() {
+        ensureDecompiled();
+        return index.getSource();
+    }
+
+    private synchronized void ensureDecompiled() {
+        if (index != null) {
+            return;
+        }
+
+        EnigmaDumper dumper = new EnigmaDumper(new StringBuilder(), settings, typeUsage, options, mapper);
+        tree.dump(state.getObfuscationMapping().wrap(dumper));
+        index = dumper.getIndex();
+    }
 }
diff --git a/enigma/src/main/java/cuchaz/enigma/source/cfr/EnigmaDumper.java b/enigma/src/main/java/cuchaz/enigma/source/cfr/EnigmaDumper.java
index d0734c5e..b85851fa 100644
--- a/enigma/src/main/java/cuchaz/enigma/source/cfr/EnigmaDumper.java
+++ b/enigma/src/main/java/cuchaz/enigma/source/cfr/EnigmaDumper.java
@@ -12,17 +12,16 @@ import cuchaz.enigma.translation.representation.entry.Entry;
 import cuchaz.enigma.translation.representation.entry.FieldEntry;
 import cuchaz.enigma.translation.representation.entry.LocalVariableEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
-import org.benf.cfr.reader.bytecode.analysis.types.JavaArrayTypeInstance;
-import org.benf.cfr.reader.bytecode.analysis.types.JavaGenericBaseInstance;
+import org.benf.cfr.reader.bytecode.analysis.loc.HasByteCodeLoc;
 import org.benf.cfr.reader.bytecode.analysis.types.JavaRefTypeInstance;
 import org.benf.cfr.reader.bytecode.analysis.types.JavaTypeInstance;
 import org.benf.cfr.reader.bytecode.analysis.types.MethodPrototype;
-import org.benf.cfr.reader.bytecode.analysis.types.RawJavaType;
 import org.benf.cfr.reader.bytecode.analysis.variables.NamedVariable;
 import org.benf.cfr.reader.entities.AccessFlag;
 import org.benf.cfr.reader.entities.ClassFile;
 import org.benf.cfr.reader.entities.ClassFileField;
 import org.benf.cfr.reader.entities.Field;
+import org.benf.cfr.reader.entities.Method;
 import org.benf.cfr.reader.state.TypeUsageInformation;
 import org.benf.cfr.reader.util.getopt.Options;
 import org.benf.cfr.reader.util.output.Dumper;
@@ -49,6 +48,7 @@ public class EnigmaDumper extends StringStreamDumper {
     private final TypeUsageInformation typeUsage;
     private final MovableDumperContext dumperContext;
     private boolean muteLine = false;
+    private MethodEntry contextMethod = null;
 
     public EnigmaDumper(StringBuilder sb, SourceSettings sourceSettings, TypeUsageInformation typeUsage, Options options,
             @Nullable EntryRemapper mapper) {
@@ -68,18 +68,19 @@ public class EnigmaDumper extends StringStreamDumper {
     }
 
     private MethodEntry getMethodEntry(MethodPrototype method) {
-        if (method == null || method.getClassType() == null) {
+        if (method == null || method.getOwner() == null) {
             return null;
         }
 
         MethodDescriptor desc = new MethodDescriptor(method.getOriginalDescriptor());
 
-        return new MethodEntry(getClassEntry(method.getClassType()), method.getName(), desc);
+        return new MethodEntry(getClassEntry(method.getOwner()), method.getName(), desc);
     }
 
     private LocalVariableEntry getParameterEntry(MethodPrototype method, int parameterIndex, String name) {
         MethodEntry owner = getMethodEntry(method);
-        if (owner == null) {
+        // params may be not computed if cfr creates a lambda expression fallback, e.g. in PointOfInterestSet
+        if (owner == null || !method.parametersComputed()) {
             return null;
         }
 
@@ -253,7 +254,7 @@ public class EnigmaDumper extends StringStreamDumper {
             if (defines) {
                 index.addDeclaration(token, entry); // override as cfr reuses local vars
             } else {
-                index.addReference(token, entry, null);
+                index.addReference(token, entry, contextMethod);
             }
         }
 
@@ -276,7 +277,7 @@ public class EnigmaDumper extends StringStreamDumper {
             if (defines) {
                 this.index.addDeclaration(token, entry);
             } else {
-                this.index.addReference(token, entry, null);
+                this.index.addReference(token, entry, contextMethod);
             }
         }
 
@@ -293,27 +294,31 @@ public class EnigmaDumper extends StringStreamDumper {
     public Dumper identifier(String name, Object ref, boolean defines) {
         super.identifier(name, ref, defines);
         Entry<?> entry;
+        if (defines) {
+            refs.remove(ref);
+            return this;
+        }
         if ((entry = refs.get(ref)) == null) {
             return this;
         }
         int now = sb.length();
         Token token = new Token(now - name.length(), now, name);
-        index.addReference(token, entry, null);
+        index.addReference(token, entry, contextMethod);
         return this;
     }
 
     @Override
-    public Dumper fieldName(String name, Field field, JavaTypeInstance owner, boolean hiddenDeclaration, boolean defines) {
-        super.fieldName(name, field, owner, hiddenDeclaration, defines);
+    public Dumper fieldName(String name, String descriptor, JavaTypeInstance owner, boolean hiddenDeclaration, boolean isStatic, boolean defines) {
+        super.fieldName(name, descriptor, owner, hiddenDeclaration, isStatic, defines);
         int now = sb.length();
         Token token = new Token(now - name.length(), now, name);
-        Entry<?> entry = field == null ? null : getFieldEntry(owner, name, field.getDescriptor());
+        if (descriptor != null) {
+            Entry<?> entry = getFieldEntry(owner, name, descriptor);
 
-        if (entry != null) {
             if (defines) {
                 index.addDeclaration(token, entry);
             } else {
-                index.addReference(token, entry, null);
+                index.addReference(token, entry, contextMethod);
             }
         }
 
@@ -348,7 +353,7 @@ public class EnigmaDumper extends StringStreamDumper {
             if (defines) {
                 index.addDeclaration(token, getClassEntry(type));
             } else {
-                index.addReference(token, getClassEntry(type), null);
+                index.addReference(token, getClassEntry(type), contextMethod);
             }
             return;
         }
@@ -367,6 +372,14 @@ public class EnigmaDumper extends StringStreamDumper {
         return new EnigmaDumper(this.sb, sourceSettings, innerclassTypeUsageInformation, options, mapper, index, dumperContext);
     }
 
+    @Override
+    public void informBytecodeLoc(HasByteCodeLoc loc) {
+        Collection<Method> methods = loc.getLoc().getMethods();
+        if (!methods.isEmpty()) {
+            this.contextMethod = getMethodEntry(methods.iterator().next().getMethodPrototype());
+        }
+    }
+
     public SourceIndex getIndex() {
         index.setSource(getString());
         return index;