Extend package pattern with wildcard

Bug: b/155618698
Change-Id: Idc8de62d098a1037cb1dd04dbd3f5ba042332fe7
diff --git a/build.gradle b/build.gradle
index 4abc2d5..1467498 100644
--- a/build.gradle
+++ b/build.gradle
@@ -817,25 +817,25 @@
                 "--output",
                 output,
                 "--map",
-                "com.google.common->com.android.tools.r8.com.google.common",
+                "com.google.common.**->com.android.tools.r8.com.google.common",
                 "--map",
-                "com.google.gson->com.android.tools.r8.com.google.gson",
+                "com.google.gson.**->com.android.tools.r8.com.google.gson",
                 "--map",
-                "com.google.thirdparty->com.android.tools.r8.com.google.thirdparty",
+                "com.google.thirdparty.**->com.android.tools.r8.com.google.thirdparty",
                 "--map",
-                "org.objectweb.asm->com.android.tools.r8.org.objectweb.asm",
+                "org.objectweb.asm.**->com.android.tools.r8.org.objectweb.asm",
                 "--map",
-                "it.unimi.dsi.fastutil->com.android.tools.r8.it.unimi.dsi.fastutil",
+                "it.unimi.dsi.fastutil.**->com.android.tools.r8.it.unimi.dsi.fastutil",
                 "--map",
-                "kotlin->com.android.tools.r8.jetbrains.kotlin",
+                "kotlin.**->com.android.tools.r8.jetbrains.kotlin",
                 "--map",
-                "kotlinx->com.android.tools.r8.jetbrains.kotlinx",
+                "kotlinx.**->com.android.tools.r8.jetbrains.kotlinx",
                 "--map",
-                "org.jetbrains->com.android.tools.r8.org.jetbrains",
+                "org.jetbrains.**->com.android.tools.r8.org.jetbrains",
                 "--map",
-                "org.intellij->com.android.tools.r8.org.intellij",
+                "org.intellij.**->com.android.tools.r8.org.intellij",
                 "--map",
-                "org.checkerframework->com.android.tools.r8.org.checkerframework"
+                "org.checkerframework.**->com.android.tools.r8.org.checkerframework"
         ])
     }
 }
@@ -1041,7 +1041,7 @@
             "--output",
             output,
             "--map",
-            "kotlinx.metadata->com.android.tools.r8.jetbrains.kotlinx.metadata"
+            "kotlinx.metadata.**->com.android.tools.r8.jetbrains.kotlinx.metadata"
     ])
 }
 
diff --git a/d8_r8/main/build.gradle.kts b/d8_r8/main/build.gradle.kts
index a262dc5..4ba82d6 100644
--- a/d8_r8/main/build.gradle.kts
+++ b/d8_r8/main/build.gradle.kts
@@ -128,38 +128,40 @@
              "$deps",
              "--output",
              "$output",
+             // Add identity mapping to enforce no relocation of things already in package
+             // com.android.tools.r8.
              "--map",
-             "com.android.resources->com.android.tools.r8.com.android.resources",
+             "com.android.tools.r8.**->com.android.tools.r8",
              "--map",
-             "com.android.support->com.android.tools.r8.com.android.support",
+             "com.android.**->com.android.tools.r8.com.android",
              "--map",
-             "com.google.common->com.android.tools.r8.com.google.common",
+             "com.android.build.shrinker.**->com.android.tools.r8.resourceshrinker",
              "--map",
-             "com.android.build.shrinker->com.android.tools.r8.resourceshrinker",
+             "com.google.common.**->com.android.tools.r8.com.google.common",
              "--map",
-             "com.google.gson->com.android.tools.r8.com.google.gson",
+             "com.google.gson.**->com.android.tools.r8.com.google.gson",
              "--map",
-             "com.google.thirdparty->com.android.tools.r8.com.google.thirdparty",
+             "com.google.thirdparty.**->com.android.tools.r8.com.google.thirdparty",
              "--map",
-             "org.objectweb.asm->com.android.tools.r8.org.objectweb.asm",
+             "org.objectweb.asm.**->com.android.tools.r8.org.objectweb.asm",
              "--map",
-             "it.unimi.dsi.fastutil->com.android.tools.r8.it.unimi.dsi.fastutil",
+             "it.unimi.dsi.fastutil.**->com.android.tools.r8.it.unimi.dsi.fastutil",
              "--map",
-             "kotlin->com.android.tools.r8.jetbrains.kotlin",
+             "kotlin.**->com.android.tools.r8.jetbrains.kotlin",
              "--map",
-             "kotlinx->com.android.tools.r8.jetbrains.kotlinx",
+             "kotlinx.**->com.android.tools.r8.jetbrains.kotlinx",
              "--map",
-             "org.jetbrains->com.android.tools.r8.org.jetbrains",
+             "org.jetbrains.**->com.android.tools.r8.org.jetbrains",
              "--map",
-             "org.intellij->com.android.tools.r8.org.intellij",
+             "org.intellij.**->com.android.tools.r8.org.intellij",
              "--map",
-             "org.checkerframework->com.android.tools.r8.org.checkerframework",
+             "org.checkerframework.**->com.android.tools.r8.org.checkerframework",
              "--map",
-             "com.google.j2objc->com.android.tools.r8.com.google.j2objc",
+             "com.google.j2objc.**->com.android.tools.r8.com.google.j2objc",
              "--map",
-             "com.google.protobuf->com.android.tools.r8.com.google.protobuf",
+             "com.google.protobuf.**->com.android.tools.r8.com.google.protobuf",
              "--map",
-             "android.aapt->com.android.tools.r8.android.aaapt"
+             "android.aapt.**->com.android.tools.r8.android.aaapt"
       ))
   }
 }
diff --git a/src/main/java/com/android/tools/r8/graph/DexString.java b/src/main/java/com/android/tools/r8/graph/DexString.java
index 2ee60ff..24e13a6 100644
--- a/src/main/java/com/android/tools/r8/graph/DexString.java
+++ b/src/main/java/com/android/tools/r8/graph/DexString.java
@@ -486,10 +486,31 @@
     // Each string ends with a 0 terminating byte, hence the +/- 1.
     // Maintain the [[ at the beginning for array dimensions.
     assert prefix.startsWith("L") && rewrittenPrefix.startsWith("L");
+    if (prefix.equals(rewrittenPrefix)) {
+      return this;
+    }
+    // When concatinating two paths, we may end up adding a separator or removing a separator.
+    // 'L' -> 'Lfoo/bar', for a string 'Lbaz/Qux' the result should be 'Lfoo/bar' + '/' + 'baz/Qux'.
+    // 'Lfoo' -> 'L', for a string 'Lfoo/Qux' the result should be 'L' + Qux', thus we remove a '/'.
+    boolean insertSeparator =
+        prefix.size == 1 && !rewrittenPrefix.endsWith(factory.descriptorSeparator);
+    boolean removeSeparator =
+        rewrittenPrefix.size == 1 && !prefix.endsWith(factory.descriptorSeparator);
+    int sizeAdjustment = 0;
+    if (insertSeparator) {
+      sizeAdjustment += 1;
+    } else if (removeSeparator) {
+      sizeAdjustment -= 1;
+    }
     int arrayDim = getArrayDim();
-    int newSize = rewrittenPrefix.size + this.size - prefix.size;
+    // In the case that prefix is 'L' we also insert a separator '/' after the destination
+    int newSize = rewrittenPrefix.size + this.size - prefix.size + sizeAdjustment;
     byte[] newContent =
-        new byte[rewrittenPrefix.content.length + this.content.length - prefix.content.length];
+        new byte
+            [rewrittenPrefix.content.length
+                + this.content.length
+                - prefix.content.length
+                + sizeAdjustment];
     // Write array dim.
     for (int i = 0; i < arrayDim; i++) {
       newContent[i] = ARRAY_CHARACTER;
@@ -499,16 +520,16 @@
         rewrittenPrefix.content, 0, newContent, arrayDim, rewrittenPrefix.content.length - 1);
     // Account for target being an empty string as to not start the descriptor with '/'.
     int prefixIndex = prefix.content.length - 1;
-    if (rewrittenPrefix.content.length == 2 && !prefix.endsWith(factory.descriptorSeparator)) {
+    int rewrittenIndex = rewrittenPrefix.content.length - 1;
+    if (removeSeparator) {
       prefixIndex += 1;
+    } else if (insertSeparator) {
+      newContent[rewrittenIndex] = '/';
+      rewrittenIndex += 1;
     }
     // Write existing name - old prefix.
     System.arraycopy(
-        this.content,
-        prefixIndex,
-        newContent,
-        rewrittenPrefix.content.length - 1,
-        this.content.length - prefixIndex);
+        this.content, prefixIndex, newContent, rewrittenIndex, this.content.length - prefixIndex);
     return factory.createString(newSize, newContent);
   }
 
diff --git a/src/main/java/com/android/tools/r8/relocator/RelocatorCommand.java b/src/main/java/com/android/tools/r8/relocator/RelocatorCommand.java
index 7b7214a..1684616 100644
--- a/src/main/java/com/android/tools/r8/relocator/RelocatorCommand.java
+++ b/src/main/java/com/android/tools/r8/relocator/RelocatorCommand.java
@@ -24,7 +24,6 @@
 import com.android.tools.r8.utils.AbortException;
 import com.android.tools.r8.utils.AndroidApp;
 import com.android.tools.r8.utils.Box;
-import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.ExceptionDiagnostic;
 import com.android.tools.r8.utils.ExceptionUtils;
 import com.android.tools.r8.utils.FlagFile;
@@ -177,6 +176,8 @@
         ImmutableMap.builder();
     private final ImmutableMap.Builder<ClassReference, ClassReference> classMapping =
         ImmutableMap.builder();
+    private final ImmutableMap.Builder<PackageReference, PackageReference> subPackageMapping =
+        ImmutableMap.builder();
     private ClassFileConsumer consumer = null;
     private int threadCount = ThreadUtils.NOT_SPECIFIED;
     private boolean printVersion;
@@ -276,6 +277,11 @@
       return this;
     }
 
+    public Builder addSubPackageMapping(PackageReference source, PackageReference destination) {
+      subPackageMapping.put(source, destination);
+      return this;
+    }
+
     public Builder addClassMapping(ClassReference source, ClassReference destination) {
       classMapping.put(source, destination);
       return this;
@@ -317,7 +323,8 @@
             DexItemFactory factory = new DexItemFactory();
             result.set(
                 new RelocatorCommand(
-                    RelocatorMapping.create(packageMapping.build(), classMapping.build()),
+                    RelocatorMapping.create(
+                        packageMapping.build(), classMapping.build(), subPackageMapping.build()),
                     app.build(),
                     reporter,
                     factory,
@@ -424,14 +431,7 @@
             }
             String source = nextArg.substring(0, separator);
             String destination = nextArg.substring(separator + 2);
-            if (DescriptorUtils.isClassDescriptor(source)) {
-              builder.addClassMapping(
-                  Reference.classFromDescriptor(source),
-                  Reference.classFromDescriptor(destination));
-            } else {
-              builder.addPackageMapping(
-                  Reference.packageFromString(source), Reference.packageFromString(destination));
-            }
+            addMapping(source, destination, builder);
             break;
           default:
             builder.error(new StringDiagnostic("Unknown argument: " + arg, origin));
@@ -443,5 +443,20 @@
       builder.setOutputPath(outputPath);
       return builder;
     }
+
+    public static void addMapping(String source, String destination, Builder builder) {
+      if (source.endsWith(".**")) {
+        builder.addSubPackageMapping(
+            Reference.packageFromString(source.substring(0, source.length() - 3)),
+            Reference.packageFromString(destination));
+      } else if (source.endsWith(".*")) {
+        builder.addPackageMapping(
+            Reference.packageFromString(source.substring(0, source.length() - 2)),
+            Reference.packageFromString(destination));
+      } else {
+        builder.addClassMapping(
+            Reference.classFromTypeName(source), Reference.classFromTypeName(destination));
+      }
+    }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/relocator/RelocatorMapping.java b/src/main/java/com/android/tools/r8/relocator/RelocatorMapping.java
index 88f8f1e..5925d3f 100644
--- a/src/main/java/com/android/tools/r8/relocator/RelocatorMapping.java
+++ b/src/main/java/com/android/tools/r8/relocator/RelocatorMapping.java
@@ -31,21 +31,25 @@
 
   private final ImmutableMap<PackageReference, PackageReference> packageMappings;
   private final ImmutableMap<ClassReference, ClassReference> classMappings;
+  private final ImmutableMap<PackageReference, PackageReference> subPackageMappings;
 
   private final ConcurrentHashMap<String, PackageReference> stringToPackageReferenceCache =
       new ConcurrentHashMap<>();
 
   private RelocatorMapping(
       ImmutableMap<PackageReference, PackageReference> packageMappings,
-      ImmutableMap<ClassReference, ClassReference> classMappings) {
+      ImmutableMap<ClassReference, ClassReference> classMappings,
+      ImmutableMap<PackageReference, PackageReference> subPackageMappings) {
     this.packageMappings = packageMappings;
     this.classMappings = classMappings;
+    this.subPackageMappings = subPackageMappings;
   }
 
   public static RelocatorMapping create(
       ImmutableMap<PackageReference, PackageReference> packageMappings,
-      ImmutableMap<ClassReference, ClassReference> classMappings) {
-    return new RelocatorMapping(packageMappings, classMappings);
+      ImmutableMap<ClassReference, ClassReference> classMappings,
+      ImmutableMap<PackageReference, PackageReference> subPackageMappings) {
+    return new RelocatorMapping(packageMappings, classMappings, subPackageMappings);
   }
 
   public Map<PackageReference, PackageReference> getPackageMappings() {
@@ -99,11 +103,21 @@
       typeMappings.put(type, factory.createString(directClassMapping.getDescriptor()));
       return;
     }
-    // TODO(b/155618698): For now keep computing packages as prefixes but instead of matching
-    //  direct prefixes match parent packages. This will ensure that pattern foo/a will not match
-    //  package foo/aa.
-    computePackageMapping(
-        type, type.getPackageName(), factory, typeMappings, rewritePackageMappings);
+    PackageReference currentPackageReference = getPackageReference(type.getPackageName());
+    PackageReference packageReference = packageMappings.get(currentPackageReference);
+    if (packageReference != null) {
+      DexString sourceDescriptorPrefix = factory.createString("L" + type.getPackageDescriptor());
+      DexString targetDescriptor =
+          factory.createString("L" + packageReference.getPackageBinaryName());
+      DexString relocatedDescriptor =
+          type.descriptor.withNewPrefix(sourceDescriptorPrefix, targetDescriptor, factory);
+      typeMappings.put(type, relocatedDescriptor);
+      return;
+    }
+    if (!subPackageMappings.isEmpty() && !type.getPackageName().isEmpty()) {
+      computePackageMapping(
+          type, type.getPackageName(), factory, typeMappings, rewritePackageMappings);
+    }
   }
 
   /**
@@ -117,7 +131,7 @@
       Map<DexType, DexString> typeMappings,
       Map<String, String> rewritePackageMappings) {
     PackageReference currentPackageReference = getPackageReference(currentPackageName);
-    PackageReference packageReference = packageMappings.get(currentPackageReference);
+    PackageReference packageReference = subPackageMappings.get(currentPackageReference);
     if (packageReference != null) {
       DexString sourceDescriptorPrefix =
           factory.createString("L" + currentPackageReference.getPackageBinaryName());
diff --git a/src/test/bootstrap/com/android/tools/r8/bootstrap/SanityCheck.java b/src/test/bootstrap/com/android/tools/r8/bootstrap/SanityCheck.java
index 46a6379..f2eef8d 100644
--- a/src/test/bootstrap/com/android/tools/r8/bootstrap/SanityCheck.java
+++ b/src/test/bootstrap/com/android/tools/r8/bootstrap/SanityCheck.java
@@ -55,13 +55,7 @@
       ZipEntry entry = entries.nextElement();
       String name = entry.getName();
       if (ZipUtils.isClassFile(name) || name.endsWith(".kotlin_builtins")) {
-        // TODO(b/155618698): This should be "com/android/tools/r8, however, we cannot relacate
-        //  packages in that folder with the current handling of mappings.
-        assertThat(
-            name,
-            ToolHelper.isNewGradleSetup()
-                ? startsWith("com/android/")
-                : startsWith("com/android/tools/r8/"));
+        assertThat(name, startsWith("com/android/tools/r8/"));
       } else if (name.equals("META-INF/MANIFEST.MF")) {
         // Allow.
       } else if (name.equals("LICENSE")) {
diff --git a/src/test/java/com/android/tools/r8/RelocatorTestBuilder.java b/src/test/java/com/android/tools/r8/RelocatorTestBuilder.java
index 4d221c9..2fea242 100644
--- a/src/test/java/com/android/tools/r8/RelocatorTestBuilder.java
+++ b/src/test/java/com/android/tools/r8/RelocatorTestBuilder.java
@@ -135,17 +135,28 @@
   public RelocatorTestBuilder addPackageMapping(PackageReference from, PackageReference to) {
     if (isExternal) {
       commandLineArgs.add("--map");
-      commandLineArgs.add(from.getPackageName() + "->" + to.getPackageName());
+      commandLineArgs.add(from.getPackageName() + ".*->" + to.getPackageName());
     } else {
       commandBuilder.addPackageMapping(from, to);
     }
     return self();
   }
 
+  public RelocatorTestBuilder addPackageAndAllSubPackagesMapping(
+      PackageReference from, PackageReference to) {
+    if (isExternal) {
+      commandLineArgs.add("--map");
+      commandLineArgs.add(from.getPackageName() + ".**->" + to.getPackageName());
+    } else {
+      commandBuilder.addSubPackageMapping(from, to);
+    }
+    return self();
+  }
+
   public RelocatorTestBuilder addClassMapping(ClassReference from, ClassReference to) {
     if (isExternal) {
       commandLineArgs.add("--map");
-      commandLineArgs.add(from.getDescriptor() + "->" + to.getDescriptor());
+      commandLineArgs.add(from.getTypeName() + "->" + to.getTypeName());
     } else {
       commandBuilder.addClassMapping(from, to);
     }
@@ -155,4 +166,9 @@
   public RelocatorTestBuilder addPackageMapping(String from, String to) {
     return addPackageMapping(Reference.packageFromString(from), Reference.packageFromString(to));
   }
+
+  public RelocatorTestBuilder addPackageAndAllSubPackagesMapping(String from, String to) {
+    return addPackageAndAllSubPackagesMapping(
+        Reference.packageFromString(from), Reference.packageFromString(to));
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/relocator/RelocatorCommandTest.java b/src/test/java/com/android/tools/r8/relocator/RelocatorCommandTest.java
index 959b891..22ab8ad 100644
--- a/src/test/java/com/android/tools/r8/relocator/RelocatorCommandTest.java
+++ b/src/test/java/com/android/tools/r8/relocator/RelocatorCommandTest.java
@@ -182,7 +182,7 @@
                         "--output",
                         "output",
                         "--map",
-                        "invalid;package-name->bar",
+                        "invalid;package-name.*->bar",
                         "--input",
                         input1.toString()
                       },
diff --git a/src/test/java/com/android/tools/r8/relocator/RelocatorPathExpansionTest.java b/src/test/java/com/android/tools/r8/relocator/RelocatorPathExpansionTest.java
index 9216acb..b115979 100644
--- a/src/test/java/com/android/tools/r8/relocator/RelocatorPathExpansionTest.java
+++ b/src/test/java/com/android/tools/r8/relocator/RelocatorPathExpansionTest.java
@@ -5,10 +5,9 @@
 package com.android.tools.r8.relocator;
 
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
+import static org.hamcrest.CoreMatchers.startsWith;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.assertThrows;
 
-import com.android.tools.r8.RelocatorTestBuilder;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.references.ClassReference;
@@ -48,7 +47,7 @@
             Baz.dump(),
             BazImpl.dump(),
             transformer(Base.class).setClassDescriptor("Lfoo/Base;").transform())
-        .addPackageMapping("foo", "com.android.tools.r8.foo")
+        .addPackageAndAllSubPackagesMapping("foo", "com.android.tools.r8.foo")
         .addClassMapping(
             destination, Reference.classFromDescriptor("Lcom/android/tools/r8/foo/bar/Baz;"))
         .run()
@@ -87,7 +86,7 @@
     String destinationPackage = "com.android.tools.r8.foo2.bar";
     testForRelocator(external)
         .addProgramClassFileData(Baz.dump(), BazImpl.dump())
-        .addPackageMapping("foo.bar", destinationPackage)
+        .addPackageAndAllSubPackagesMapping("foo.bar", destinationPackage)
         .addClassMapping(Reference.classFromDescriptor("Lfoo/bar/Baz;"), destination)
         .run()
         .inspect(
@@ -99,12 +98,18 @@
   }
 
   @Test
-  public void rewriteAllSubPackages() {
-    RelocatorTestBuilder relocatorTestBuilder =
-        testForRelocator(external).addProgramClassFileData(Baz.dump(), BazImpl.dump());
-    // The language as input is very simple and ** to match sub-packages would be a feature request.
-    assertThrows(
-        IllegalArgumentException.class,
-        () -> relocatorTestBuilder.addPackageMapping("foo/**", "com.android.tools.r8.foo"));
+  public void rewriteAllSubPackages() throws Exception {
+    testForRelocator(external)
+        .addProgramClassFileData(Baz.dump(), BazImpl.dump())
+        .addPackageAndAllSubPackagesMapping("foo", "com.android.tools.r8.foo")
+        .run()
+        .inspect(
+            inspector ->
+                inspector
+                    .allClasses()
+                    .forEach(
+                        clazz ->
+                            assertThat(
+                                clazz.getFinalName(), startsWith("com.android.tools.r8.foo"))));
   }
 }
diff --git a/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java b/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java
index 9a0a971..ab21610 100644
--- a/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java
+++ b/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java
@@ -6,10 +6,8 @@
 
 import static com.android.tools.r8.ToolHelper.CHECKED_IN_R8_17_WITH_DEPS;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
-import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import com.android.tools.r8.RelocatorTestBuilder;
@@ -65,24 +63,28 @@
   }
 
   @Test
+  public void testRelocatorIdentityPackage() throws Exception {
+    testForRelocator(external)
+        .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
+        .addPackageMapping("com.android.tools.r8", "com.android.tools.r8")
+        .run()
+        .inspectAllClassesRelocated(CHECKED_IN_R8_17_WITH_DEPS, "", "");
+  }
+
+  @Test
   public void testRelocatorEmptyToSomething() throws Exception {
     String originalPrefix = "";
     String newPrefix = "foo.bar.baz";
-    RelocatorTestCompileResult compileResult =
-        testForRelocator(external)
-            .addPackageMapping(
-                Reference.packageFromString(originalPrefix), Reference.packageFromString(newPrefix))
-            .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-            .run();
-    // TODO(b/155618698): Extend relocator with a richer language such that java.lang.Object is not
-    //   relocated.
-    RuntimeException compilationError =
-        assertThrows(
-            RuntimeException.class,
-            () ->
-                compileResult.inspectAllClassesRelocated(
-                    CHECKED_IN_R8_17_WITH_DEPS, originalPrefix, newPrefix + "."));
-    assertThat(compilationError.getMessage(), containsString("must extend class java.lang.Object"));
+    testForRelocator(external)
+        .addClassMapping(
+            Reference.classFromClass(Object.class), Reference.classFromClass(Object.class))
+        .addClassMapping(
+            Reference.classFromClass(String.class), Reference.classFromClass(String.class))
+        .addPackageAndAllSubPackagesMapping(
+            Reference.packageFromString(originalPrefix), Reference.packageFromString(newPrefix))
+        .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
+        .run()
+        .inspectAllClassesRelocated(CHECKED_IN_R8_17_WITH_DEPS, originalPrefix, newPrefix + ".");
   }
 
   @Test
@@ -90,7 +92,7 @@
     String originalPrefix = "com.android.tools.r8";
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping(
+        .addPackageAndAllSubPackagesMapping(
             Reference.packageFromString(originalPrefix), Reference.packageFromString(""))
         .run()
         .inspectAllSignaturesNotContainingString(originalPrefix);
@@ -101,7 +103,7 @@
     PackageReference pkg = Reference.packageFromString("com.android.tools.r8");
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping(pkg, pkg)
+        .addPackageAndAllSubPackagesMapping(pkg, pkg)
         .run()
         .inspect(
             inspector -> {
@@ -153,9 +155,9 @@
     String newPrefix = "foo.bar.baz";
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping(
+        .addPackageAndAllSubPackagesMapping(
             Reference.packageFromString(originalPrefix), Reference.packageFromString(newPrefix))
-        .addPackageMapping("some.package.that.does.not.exist", "foo")
+        .addPackageAndAllSubPackagesMapping("some.package.that.does.not.exist", "foo")
         .run()
         .inspectAllClassesRelocated(CHECKED_IN_R8_17_WITH_DEPS, originalPrefix, newPrefix);
   }
@@ -164,8 +166,8 @@
   public void testOrderingOfPrefixes() throws Exception {
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping("com.android", "foo.bar.baz")
-        .addPackageMapping("com.android.tools.r8", "qux")
+        .addPackageAndAllSubPackagesMapping("com.android", "foo.bar.baz")
+        .addPackageAndAllSubPackagesMapping("com.android.tools.r8", "qux")
         .run()
         // Because we see "com.android.tools.r8" before seeing "com.android" we always choose qux.
         .inspectAllClassesRelocated(
@@ -175,11 +177,10 @@
 
   @Test
   public void testNoReEntry() throws Exception {
-    // TODO(b/154909222): Check if this is the behavior we would like.
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping("com.android", "foo.bar.baz")
-        .addPackageMapping("foo.bar.baz", "qux")
+        .addPackageAndAllSubPackagesMapping("com.android", "foo.bar.baz")
+        .addPackageAndAllSubPackagesMapping("foo.bar.baz", "qux")
         .run()
         .inspect(
             inspector -> {
@@ -212,7 +213,7 @@
         if (seenPackages.add(mappedPackageName)) {
           String relocatedPackageName = "number" + packageNameCounter++;
           packageMappings.add(new Pair<>(mappedPackageName, relocatedPackageName));
-          testBuilder.addPackageMapping(mappedPackageName, relocatedPackageName);
+          testBuilder.addPackageAndAllSubPackagesMapping(mappedPackageName, relocatedPackageName);
         }
       }
     }
@@ -228,7 +229,7 @@
     String originalPrefix = "com.android.tools.r";
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping(originalPrefix, "i_cannot_w")
+        .addPackageAndAllSubPackagesMapping(originalPrefix, "i_cannot_w")
         .run()
         .inspectAllClassesRelocated(CHECKED_IN_R8_17_WITH_DEPS, originalPrefix, originalPrefix);
   }
@@ -240,7 +241,7 @@
     RelocatorTestCompileResult result =
         testForRelocator(external)
             .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-            .addPackageMapping(originalPrefix, newPrefix)
+            .addPackageAndAllSubPackagesMapping(originalPrefix, newPrefix)
             .run();
     // Check that all classes has been remapped.
     result.inspectAllClassesRelocated(CHECKED_IN_R8_17_WITH_DEPS, originalPrefix, newPrefix);
@@ -278,7 +279,7 @@
     // Assert that all classes are the same, have the same methods and nest info.
     testForRelocator(external)
         .addProgramFiles(CHECKED_IN_R8_17_WITH_DEPS)
-        .addPackageMapping(originalPrefix, newPrefix)
+        .addPackageAndAllSubPackagesMapping(originalPrefix, newPrefix)
         .run()
         .inspect(
             relocatedInspector -> {