[Retrace] Flatmap ambiguous results

Bug: 170491519
Change-Id: Ic3e4edb23e152b5d40aa6eaaf7907982da08b274
diff --git a/src/main/java/com/android/tools/r8/retrace/Result.java b/src/main/java/com/android/tools/r8/retrace/Result.java
index 12e92d3..692ca9b 100644
--- a/src/main/java/com/android/tools/r8/retrace/Result.java
+++ b/src/main/java/com/android/tools/r8/retrace/Result.java
@@ -14,4 +14,6 @@
   public abstract Stream<R> stream();
 
   public abstract RR forEach(Consumer<R> resultConsumer);
+
+  public abstract boolean isAmbiguous();
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceApi.java b/src/main/java/com/android/tools/r8/retrace/RetraceApi.java
index 7bb9232..d3bd6d3 100644
--- a/src/main/java/com/android/tools/r8/retrace/RetraceApi.java
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceApi.java
@@ -14,11 +14,11 @@
 @Keep
 public interface RetraceApi {
 
+  RetraceClassResult retrace(ClassReference classReference);
+
   RetraceMethodResult retrace(MethodReference methodReference);
 
   RetraceFieldResult retrace(FieldReference fieldReference);
 
-  RetraceClassResult retrace(ClassReference classReference);
-
   RetraceTypeResult retrace(TypeReference typeReference);
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceClassResult.java b/src/main/java/com/android/tools/r8/retrace/RetraceClassResult.java
index 945aa0f..9ea9614 100644
--- a/src/main/java/com/android/tools/r8/retrace/RetraceClassResult.java
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceClassResult.java
@@ -15,7 +15,9 @@
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.retrace.RetraceClassResult.Element;
-import com.android.tools.r8.utils.Box;
+import com.android.tools.r8.utils.Pair;
+import com.google.common.collect.ImmutableList;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -61,7 +63,7 @@
           if (mappedRanges == null || mappedRanges.getMappedRanges().isEmpty()) {
             return null;
           }
-          return mappedRanges;
+          return mappedRanges.getMappedRanges();
         },
         RetraceMethodResult::new);
   }
@@ -70,17 +72,20 @@
       String name,
       BiFunction<ClassNamingForNameMapper, String, T> lookupFunction,
       ResultConstructor<T, R> constructor) {
-    Box<R> elementBox = new Box<>();
+    List<Pair<Element, T>> mappings = new ArrayList<>();
     forEach(
         element -> {
-          assert !elementBox.isSet();
-          T mappedRangesForT = null;
-          if (element.mapper != null) {
-            mappedRangesForT = lookupFunction.apply(element.mapper, name);
+          if (mapper != null) {
+            assert element.mapper != null;
+            T mappedElements = lookupFunction.apply(element.mapper, name);
+            if (mappedElements != null) {
+              mappings.add(new Pair<>(element, mappedElements));
+              return;
+            }
           }
-          elementBox.set(constructor.create(element, mappedRangesForT, name, retracer));
+          mappings.add(new Pair<>(element, null));
         });
-    return elementBox.get();
+    return constructor.create(this, mappings, name, retracer);
   }
 
   boolean hasRetraceResult() {
@@ -106,9 +111,14 @@
   }
 
   private interface ResultConstructor<T, R> {
-    R create(Element element, T mappings, String obfuscatedName, RetraceApi retraceApi);
+    R create(
+        RetraceClassResult classResult,
+        List<Pair<Element, T>> mappings,
+        String obfuscatedName,
+        RetraceApi retraceApi);
   }
 
+  @Override
   public boolean isAmbiguous() {
     // Currently we have no way of producing ambiguous class results.
     return false;
@@ -180,7 +190,7 @@
             if (mappedRanges == null || mappedRanges.getMappedRanges().isEmpty()) {
               return null;
             }
-            return mappedRanges;
+            return mappedRanges.getMappedRanges();
           },
           RetraceMethodResult::new);
     }
@@ -189,11 +199,17 @@
         String name,
         BiFunction<ClassNamingForNameMapper, String, T> lookupFunction,
         ResultConstructor<T, R> constructor) {
-      return constructor.create(
-          this,
-          mapper != null ? lookupFunction.apply(mapper, name) : null,
-          name,
-          classResult.retracer);
+      List<Pair<Element, T>> mappings = ImmutableList.of();
+      if (mapper != null) {
+        T result = lookupFunction.apply(mapper, name);
+        if (result != null) {
+          mappings = ImmutableList.of(new Pair<>(this, result));
+        }
+      }
+      if (mappings.isEmpty()) {
+        mappings = ImmutableList.of(new Pair<>(this, null));
+      }
+      return constructor.create(classResult, mappings, name, classResult.retracer);
     }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceFieldResult.java b/src/main/java/com/android/tools/r8/retrace/RetraceFieldResult.java
index 265b3b2..114f193 100644
--- a/src/main/java/com/android/tools/r8/retrace/RetraceFieldResult.java
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceFieldResult.java
@@ -9,79 +9,83 @@
 import com.android.tools.r8.naming.MemberNaming.FieldSignature;
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.utils.DescriptorUtils;
+import com.android.tools.r8.utils.Pair;
 import java.util.List;
-import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 @Keep
 public class RetraceFieldResult extends Result<RetraceFieldResult.Element, RetraceFieldResult> {
 
-  private final RetraceClassResult.Element classElement;
-  private final List<MemberNaming> memberNamings;
+  private final RetraceClassResult classResult;
+  private final List<Pair<RetraceClassResult.Element, List<MemberNaming>>> memberNamings;
   private final String obfuscatedName;
   private final RetraceApi retracer;
 
   RetraceFieldResult(
-      RetraceClassResult.Element classElement,
-      List<MemberNaming> memberNamings,
+      RetraceClassResult classResult,
+      List<Pair<RetraceClassResult.Element, List<MemberNaming>>> memberNamings,
       String obfuscatedName,
       RetraceApi retracer) {
-    this.classElement = classElement;
+    this.classResult = classResult;
     this.memberNamings = memberNamings;
     this.obfuscatedName = obfuscatedName;
     this.retracer = retracer;
-    assert classElement != null;
-    assert memberNamings == null
-        || (!memberNamings.isEmpty() && memberNamings.stream().allMatch(Objects::nonNull));
+    assert classResult != null;
+    assert !memberNamings.isEmpty();
   }
 
-  private boolean hasRetraceResult() {
-    return memberNamings != null;
-  }
-
+  @Override
   public boolean isAmbiguous() {
-    if (!hasRetraceResult()) {
+    if (memberNamings.size() > 1) {
+      return true;
+    }
+    List<MemberNaming> mappings = memberNamings.get(0).getSecond();
+    if (mappings == null) {
       return false;
     }
-    assert memberNamings != null;
-    return memberNamings.size() > 1;
-  }
-
-  public RetracedField getUnknownReference() {
-    return RetracedField.createUnknown(classElement.getRetracedClass(), obfuscatedName);
+    return mappings.size() > 1;
   }
 
   @Override
   public Stream<Element> stream() {
-    if (!hasRetraceResult()) {
-      return Stream.of(new Element(this, classElement, getUnknownReference()));
-    }
-    assert !memberNamings.isEmpty();
     return memberNamings.stream()
-        .map(
-            memberNaming -> {
-              assert memberNaming.isFieldNaming();
-              FieldSignature fieldSignature =
-                  memberNaming.getOriginalSignature().asFieldSignature();
-              RetracedClass holder =
-                  fieldSignature.isQualified()
-                      ? RetracedClass.create(
-                          Reference.classFromDescriptor(
-                              DescriptorUtils.javaTypeToDescriptor(
-                                  fieldSignature.toHolderFromQualified())))
-                      : classElement.getRetracedClass();
-              return new Element(
-                  this,
-                  classElement,
-                  RetracedField.create(
-                      holder,
-                      Reference.field(
-                          holder.getClassReference(),
-                          fieldSignature.isQualified()
-                              ? fieldSignature.toUnqualifiedName()
-                              : fieldSignature.name,
-                          Reference.typeFromTypeName(fieldSignature.type))));
+        .flatMap(
+            mappedNamePair -> {
+              RetraceClassResult.Element classElement = mappedNamePair.getFirst();
+              List<MemberNaming> memberNamings = mappedNamePair.getSecond();
+              if (memberNamings == null) {
+                return Stream.of(
+                    new RetraceFieldResult.Element(
+                        this,
+                        classElement,
+                        RetracedField.createUnknown(
+                            classElement.getRetracedClass(), obfuscatedName)));
+              }
+              return memberNamings.stream()
+                  .map(
+                      memberNaming -> {
+                        FieldSignature fieldSignature =
+                            memberNaming.getOriginalSignature().asFieldSignature();
+                        RetracedClass holder =
+                            fieldSignature.isQualified()
+                                ? RetracedClass.create(
+                                    Reference.classFromDescriptor(
+                                        DescriptorUtils.javaTypeToDescriptor(
+                                            fieldSignature.toHolderFromQualified())))
+                                : classElement.getRetracedClass();
+                        return new Element(
+                            this,
+                            classElement,
+                            RetracedField.create(
+                                holder,
+                                Reference.field(
+                                    holder.getClassReference(),
+                                    fieldSignature.isQualified()
+                                        ? fieldSignature.toUnqualifiedName()
+                                        : fieldSignature.name,
+                                    Reference.typeFromTypeName(fieldSignature.type))));
+                      });
             });
   }
 
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceMethodResult.java b/src/main/java/com/android/tools/r8/retrace/RetraceMethodResult.java
index 93c87e8..7a41f74 100644
--- a/src/main/java/com/android/tools/r8/retrace/RetraceMethodResult.java
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceMethodResult.java
@@ -8,10 +8,11 @@
 import com.android.tools.r8.naming.ClassNamingForNameMapper.MappedRange;
 import com.android.tools.r8.naming.ClassNamingForNameMapper.MappedRangesOfName;
 import com.android.tools.r8.naming.MemberNaming.MethodSignature;
-import com.android.tools.r8.naming.Range;
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.references.TypeReference;
 import com.android.tools.r8.utils.DescriptorUtils;
+import com.android.tools.r8.utils.Pair;
+import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Consumer;
@@ -21,107 +22,125 @@
 public class RetraceMethodResult extends Result<RetraceMethodResult.Element, RetraceMethodResult> {
 
   private final String obfuscatedName;
-  private final RetraceClassResult.Element classElement;
-  private final MappedRangesOfName mappedRanges;
+  private final RetraceClassResult classResult;
+  private final List<Pair<RetraceClassResult.Element, List<MappedRange>>> mappedRanges;
   private final RetraceApi retracer;
-  private Boolean isAmbiguousCached = null;
 
   RetraceMethodResult(
-      RetraceClassResult.Element classElement,
-      MappedRangesOfName mappedRanges,
+      RetraceClassResult classResult,
+      List<Pair<RetraceClassResult.Element, List<MappedRange>>> mappedRanges,
       String obfuscatedName,
       RetraceApi retracer) {
-    this.classElement = classElement;
+    this.classResult = classResult;
     this.mappedRanges = mappedRanges;
     this.obfuscatedName = obfuscatedName;
     this.retracer = retracer;
-    assert classElement != null;
+    assert classResult != null;
+    // TODO(mkroghj): Enable this when we have frame results.
+    // assert !mappedRanges.isEmpty();
   }
 
-  public RetracedMethod getUnknownReference() {
-    return RetracedMethod.createUnknown(classElement.getRetracedClass(), obfuscatedName);
-  }
-
-  private boolean hasRetraceResult() {
-    return mappedRanges != null && mappedRanges.getMappedRanges().size() > 0;
-  }
-
+  @Override
   public boolean isAmbiguous() {
-    if (isAmbiguousCached != null) {
-      return isAmbiguousCached;
+    if (mappedRanges.size() > 1) {
+      return true;
     }
-    if (!hasRetraceResult()) {
+    List<MappedRange> methodRanges = mappedRanges.get(0).getSecond();
+    if (methodRanges == null || methodRanges.isEmpty()) {
       return false;
     }
-    assert mappedRanges != null;
-    Range minifiedRange = null;
-    boolean seenNull = false;
-    for (MappedRange mappedRange : mappedRanges.getMappedRanges()) {
-      if (minifiedRange != null && !minifiedRange.equals(mappedRange.minifiedRange)) {
-        isAmbiguousCached = true;
+    MappedRange lastRange = methodRanges.get(0);
+    for (MappedRange mappedRange : methodRanges) {
+      if (mappedRange != lastRange
+          && (mappedRange.minifiedRange == null
+              || !mappedRange.minifiedRange.equals(lastRange.minifiedRange))) {
         return true;
-      } else if (minifiedRange == null) {
-        if (seenNull) {
-          isAmbiguousCached = true;
-          return true;
-        }
-        seenNull = true;
       }
-      minifiedRange = mappedRange.minifiedRange;
     }
-    isAmbiguousCached = false;
     return false;
   }
 
   public RetraceMethodResult narrowByLine(int linePosition) {
-    if (!hasRetraceResult()) {
-      return this;
-    }
-    List<MappedRange> narrowedRanges = this.mappedRanges.allRangesForLine(linePosition, false);
-    if (narrowedRanges.isEmpty()) {
-      narrowedRanges = new ArrayList<>();
-      for (MappedRange mappedRange : this.mappedRanges.getMappedRanges()) {
-        if (mappedRange.minifiedRange == null) {
-          narrowedRanges.add(mappedRange);
+    List<Pair<RetraceClassResult.Element, List<MappedRange>>> narrowedRanges = new ArrayList<>();
+    List<Pair<RetraceClassResult.Element, List<MappedRange>>> noMappingRanges = new ArrayList<>();
+    for (Pair<RetraceClassResult.Element, List<MappedRange>> mappedRange : mappedRanges) {
+      if (mappedRange.getSecond() == null) {
+        noMappingRanges.add(new Pair<>(mappedRange.getFirst(), null));
+        continue;
+      }
+      List<MappedRange> ranges =
+          new MappedRangesOfName(mappedRange.getSecond()).allRangesForLine(linePosition, false);
+      boolean hasAddedRanges = false;
+      if (!ranges.isEmpty()) {
+        narrowedRanges.add(new Pair<>(mappedRange.getFirst(), ranges));
+        hasAddedRanges = true;
+      } else {
+        narrowedRanges = new ArrayList<>();
+        for (MappedRange mapped : mappedRange.getSecond()) {
+          if (mapped.minifiedRange == null) {
+            narrowedRanges.add(new Pair<>(mappedRange.getFirst(), ImmutableList.of(mapped)));
+            hasAddedRanges = true;
+          }
         }
       }
+      if (!hasAddedRanges) {
+        narrowedRanges.add(new Pair<>(mappedRange.getFirst(), null));
+      }
     }
     return new RetraceMethodResult(
-        classElement, new MappedRangesOfName(narrowedRanges), obfuscatedName, retracer);
+        classResult,
+        narrowedRanges.isEmpty() ? noMappingRanges : narrowedRanges,
+        obfuscatedName,
+        retracer);
   }
 
   @Override
   public Stream<Element> stream() {
-    if (!hasRetraceResult()) {
-      return Stream.of(new Element(this, classElement, getUnknownReference(), null));
-    }
-    return mappedRanges.getMappedRanges().stream()
-        .map(
-            mappedRange -> {
-              MethodSignature signature = mappedRange.signature;
-              RetracedClass holder =
-                  mappedRange.signature.isQualified()
-                      ? RetracedClass.create(
-                          Reference.classFromDescriptor(
-                              DescriptorUtils.javaTypeToDescriptor(
-                                  mappedRange.signature.toHolderFromQualified())))
-                      : classElement.getRetracedClass();
-              List<TypeReference> formalTypes = new ArrayList<>(signature.parameters.length);
-              for (String parameter : signature.parameters) {
-                formalTypes.add(Reference.typeFromTypeName(parameter));
+    return mappedRanges.stream()
+        .flatMap(
+            mappedRangePair -> {
+              RetraceClassResult.Element classElement = mappedRangePair.getFirst();
+              List<MappedRange> mappedRanges = mappedRangePair.getSecond();
+              if (mappedRanges == null) {
+                return Stream.of(
+                    new Element(
+                        this,
+                        classElement,
+                        RetracedMethod.createUnknown(
+                            classElement.getRetracedClass(), obfuscatedName),
+                        null));
               }
-              TypeReference returnType =
-                  Reference.returnTypeFromDescriptor(
-                      DescriptorUtils.javaTypeToDescriptor(signature.type));
-              RetracedMethod retracedMethod =
-                  RetracedMethod.create(
-                      holder,
-                      Reference.method(
-                          holder.getClassReference(),
-                          signature.isQualified() ? signature.toUnqualifiedName() : signature.name,
-                          formalTypes,
-                          returnType));
-              return new Element(this, classElement, retracedMethod, mappedRange);
+              return mappedRanges.stream()
+                  .map(
+                      mappedRange -> {
+                        MethodSignature signature = mappedRange.signature;
+                        RetracedClass holder =
+                            mappedRange.signature.isQualified()
+                                ? RetracedClass.create(
+                                    Reference.classFromDescriptor(
+                                        DescriptorUtils.javaTypeToDescriptor(
+                                            mappedRange.signature.toHolderFromQualified())))
+                                : classElement.getRetracedClass();
+                        List<TypeReference> formalTypes =
+                            new ArrayList<>(signature.parameters.length);
+                        for (String parameter : signature.parameters) {
+                          formalTypes.add(Reference.typeFromTypeName(parameter));
+                        }
+                        TypeReference returnType =
+                            Reference.returnTypeFromDescriptor(
+                                DescriptorUtils.javaTypeToDescriptor(signature.type));
+                        RetracedMethod retracedMethod =
+                            RetracedMethod.create(
+                                holder,
+                                Reference.method(
+                                    holder.getClassReference(),
+                                    signature.isQualified()
+                                        ? signature.toUnqualifiedName()
+                                        : signature.name,
+                                    formalTypes,
+                                    returnType));
+                        return new Element(this, classElement, retracedMethod, mappedRange);
+                      });
             });
   }
 
@@ -192,6 +211,5 @@
       return RetraceUtils.getSourceFile(
           classElement, methodReference.getHolderClass(), sourceFile, retraceMethodResult.retracer);
     }
-
   }
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceRegularExpression.java b/src/main/java/com/android/tools/r8/retrace/RetraceRegularExpression.java
index 8606324..bccb9dd 100644
--- a/src/main/java/com/android/tools/r8/retrace/RetraceRegularExpression.java
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceRegularExpression.java
@@ -523,7 +523,7 @@
         for (RetraceString retraceString : strings) {
           retraceResult.forEach(
               element -> {
-                final RetraceString newRetraceString =
+                RetraceString newRetraceString =
                     retraceString.updateContext(
                         context -> context.withClassContext(element, element.getRetracedClass()));
                 retraceStrings.add(newRetraceString);
@@ -761,7 +761,7 @@
             }
             return strings;
           }
-          String methodName = matcher.group(captureGroup);
+          String fieldName = matcher.group(captureGroup);
           List<RetraceString> retracedStrings = new ArrayList<>();
           for (RetraceString retraceString : strings) {
             if (retraceString.getClassContext() == null) {
@@ -769,7 +769,7 @@
               return strings;
             }
             final RetraceFieldResult retraceFieldResult =
-                retraceString.getClassContext().lookupField(methodName);
+                retraceString.getClassContext().lookupField(fieldName);
             assert !retraceFieldResult.isAmbiguous();
             retraceFieldResult.forEach(
                 element -> {
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceTypeResult.java b/src/main/java/com/android/tools/r8/retrace/RetraceTypeResult.java
index 4f7ec0a..d546d11 100644
--- a/src/main/java/com/android/tools/r8/retrace/RetraceTypeResult.java
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceTypeResult.java
@@ -14,11 +14,15 @@
   private final TypeReference obfuscatedType;
   private final RetraceApi retracer;
 
-  RetraceTypeResult(TypeReference obfuscatedType, RetraceApi retracer) {
+  private RetraceTypeResult(TypeReference obfuscatedType, RetraceApi retracer) {
     this.obfuscatedType = obfuscatedType;
     this.retracer = retracer;
   }
 
+  static RetraceTypeResult create(TypeReference obfuscatedType, RetraceApi retracer) {
+    return new RetraceTypeResult(obfuscatedType, retracer);
+  }
+
   @Override
   public Stream<Element> stream() {
     // Handle void and primitive types as single element results.
@@ -36,6 +40,7 @@
         .map(classElement -> new Element(classElement.getRetracedClass().getRetracedType()));
   }
 
+  @Override
   public boolean isAmbiguous() {
     return false;
   }
diff --git a/src/main/java/com/android/tools/r8/retrace/Retracer.java b/src/main/java/com/android/tools/r8/retrace/Retracer.java
index 583ca7c..d38e203 100644
--- a/src/main/java/com/android/tools/r8/retrace/Retracer.java
+++ b/src/main/java/com/android/tools/r8/retrace/Retracer.java
@@ -57,6 +57,6 @@
 
   @Override
   public RetraceTypeResult retrace(TypeReference typeReference) {
-    return new RetraceTypeResult(typeReference, this);
+    return RetraceTypeResult.create(typeReference, this);
   }
 }