[KeepAnno] Introduce extraction options for supported back references

Bug: b/322114141
Change-Id: I95e686e324aa5a76b52208ed6da51cfe0008547e
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractor.java b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractor.java
index 815f6ba..c5d43c5 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractor.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractor.java
@@ -49,16 +49,17 @@
 /** Extract the PG keep rules that over-approximate a keep edge. */
 public class KeepRuleExtractor {
 
-  private final KeepRuleExtractorOptions options;
+  private final KeepRuleExtractorOptions extractorOptions;
   private final Consumer<String> ruleConsumer;
 
   public KeepRuleExtractor(Consumer<String> ruleConsumer) {
     this(ruleConsumer, KeepRuleExtractorOptions.getR8Options());
   }
 
-  public KeepRuleExtractor(Consumer<String> ruleConsumer, KeepRuleExtractorOptions options) {
+  public KeepRuleExtractor(
+      Consumer<String> ruleConsumer, KeepRuleExtractorOptions extractorOptions) {
     this.ruleConsumer = ruleConsumer;
-    this.options = options;
+    this.extractorOptions = extractorOptions;
   }
 
   public void extract(KeepDeclaration declaration) {
@@ -66,7 +67,7 @@
     PgRule.groupByKinds(rules);
     StringBuilder builder = new StringBuilder();
     for (PgRule rule : rules) {
-      rule.printRule(builder, options);
+      rule.printRule(builder, extractorOptions);
       builder.append("\n");
     }
     ruleConsumer.accept(builder.toString());
@@ -80,7 +81,7 @@
   }
 
   private List<PgRule> generateCheckRules(KeepCheck check) {
-    if (!options.hasCheckDiscardSupport()) {
+    if (!extractorOptions.hasCheckDiscardSupport()) {
       return Collections.emptyList();
     }
     KeepItemPattern itemPattern = check.getItemPattern();
@@ -114,7 +115,8 @@
             KeepOptions.keepAll(),
             memberPatterns,
             targetMembers,
-            TargetKeepKind.CHECK_DISCARD));
+            TargetKeepKind.CHECK_DISCARD,
+            extractorOptions));
     // If the check declaration is to ensure full removal we generate a soft-pin rule to disallow
     // moving/inlining the items.
     if (isRemovedPattern) {
@@ -130,7 +132,8 @@
                 allowShrinking,
                 Collections.singletonMap(memberSymbol, KeepMemberPattern.allMembers()),
                 Collections.singletonList(memberSymbol),
-                TargetKeepKind.CLASS_OR_MEMBERS));
+                TargetKeepKind.CLASS_OR_MEMBERS,
+                extractorOptions));
       } else {
         // A check removal on members just soft-pins the members.
         rules.add(
@@ -141,7 +144,8 @@
                 memberPatterns,
                 Collections.emptyList(),
                 targetMembers,
-                TargetKeepKind.JUST_MEMBERS));
+                TargetKeepKind.JUST_MEMBERS,
+                extractorOptions));
       }
     }
     return rules;
@@ -255,7 +259,7 @@
   }
 
   @SuppressWarnings("UnnecessaryParentheses")
-  private static List<PgRule> doSplit(KeepEdge edge) {
+  private List<PgRule> doSplit(KeepEdge edge) {
     List<PgRule> rules = new ArrayList<>();
     // Collection for all attribute constraints required for this edge.
     Set<KeepAttribute> allAttributeConstraints = new HashSet<>();
@@ -287,7 +291,8 @@
 
     // Generate at most one `-keepattributes` rule for the edge if needed.
     if (!allAttributeConstraints.isEmpty()) {
-      rules.add(new PgKeepAttributeRule(edge.getMetaInfo(), allAttributeConstraints));
+      rules.add(
+          new PgKeepAttributeRule(edge.getMetaInfo(), allAttributeConstraints, extractorOptions));
     }
 
     bindingUsers.forEach(
@@ -400,7 +405,7 @@
             callback.accept(newTargetHolder, memberPatterns, targetMembers, finalKeepKind));
   }
 
-  private static void createUnconditionalRules(
+  private void createUnconditionalRules(
       List<PgRule> rules,
       Holder holder,
       KeepEdgeMetaInfo metaInfo,
@@ -423,7 +428,8 @@
                     memberPatterns,
                     Collections.emptyList(),
                     targetMembers,
-                    targetKeepKind));
+                    targetKeepKind,
+                    extractorOptions));
           } else {
             rules.add(
                 new PgUnconditionalRule(
@@ -432,12 +438,13 @@
                     options,
                     memberPatterns,
                     targetMembers,
-                    targetKeepKind));
+                    targetKeepKind,
+                    extractorOptions));
           }
         });
   }
 
-  private static void createConditionalRules(
+  private void createConditionalRules(
       List<PgRule> rules,
       KeepEdgeMetaInfo metaInfo,
       Holder conditionHolder,
@@ -472,10 +479,11 @@
                     memberPatterns,
                     conditionMembers,
                     targetMembers,
-                    targetKeepKind)));
+                    targetKeepKind,
+                    extractorOptions)));
   }
 
-  private static void createDependentRules(
+  private void createDependentRules(
       List<PgRule> rules,
       Holder initialHolder,
       KeepEdgeMetaInfo metaInfo,
@@ -510,7 +518,8 @@
                       copyWithMethod,
                       conditionMembers,
                       Collections.singletonList(targetMember),
-                      targetKeepKind));
+                      targetKeepKind,
+                      extractorOptions));
               HashMap<KeepBindingSymbol, KeepMemberPattern> copyWithField =
                   new HashMap<>(memberPatterns);
               copyWithField.put(targetMember, copyFieldFromMember(memberPattern));
@@ -522,7 +531,8 @@
                       copyWithField,
                       conditionMembers,
                       Collections.singletonList(targetMember),
-                      targetKeepKind));
+                      targetKeepKind,
+                      extractorOptions));
             } else {
               nonAllMemberTargets.add(targetMember);
             }
@@ -538,7 +548,8 @@
                   memberPatterns,
                   conditionMembers,
                   nonAllMemberTargets,
-                  targetKeepKind));
+                  targetKeepKind,
+                  extractorOptions));
         });
   }
 
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractorOptions.java b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractorOptions.java
index c30b430..f155852 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractorOptions.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/KeepRuleExtractorOptions.java
@@ -6,12 +6,86 @@
 
 import com.android.tools.r8.keepanno.ast.KeepOptions.KeepOption;
 
-public class KeepRuleExtractorOptions {
+public abstract class KeepRuleExtractorOptions {
 
   private static final KeepRuleExtractorOptions PG_OPTIONS =
-      new KeepRuleExtractorOptions(false, false);
+      new KeepRuleExtractorOptions() {
+        @Override
+        public boolean hasCheckDiscardSupport() {
+          return false;
+        }
+
+        @Override
+        public boolean hasAllowAccessModificationOptionSupport() {
+          return false;
+        }
+
+        @Override
+        public boolean hasAllowAnnotationRemovalOptionSupport() {
+          return false;
+        }
+
+        @Override
+        public boolean hasFieldTypeBackReference() {
+          return false;
+        }
+
+        @Override
+        public boolean hasMethodReturnTypeBackReference() {
+          return false;
+        }
+
+        @Override
+        public boolean hasMethodParameterTypeBackReference() {
+          return false;
+        }
+
+        @Override
+        public boolean hasMethodParameterListBackReference() {
+          return false;
+        }
+      };
+
   private static final KeepRuleExtractorOptions R8_OPTIONS =
-      new KeepRuleExtractorOptions(true, true);
+      new KeepRuleExtractorOptions() {
+        @Override
+        public boolean hasCheckDiscardSupport() {
+          return true;
+        }
+
+        @Override
+        public boolean hasAllowAccessModificationOptionSupport() {
+          return true;
+        }
+
+        @Override
+        public boolean hasAllowAnnotationRemovalOptionSupport() {
+          // Allow annotation removal is currently a testing only option.
+          return false;
+        }
+
+        @Override
+        public boolean hasFieldTypeBackReference() {
+          return true;
+        }
+
+        @Override
+        public boolean hasMethodReturnTypeBackReference() {
+          return true;
+        }
+
+        @Override
+        public boolean hasMethodParameterTypeBackReference() {
+          return true;
+        }
+
+        @Override
+        public boolean hasMethodParameterListBackReference() {
+          // TODO(b/265892343): R8 does not support backrefs for `(...)`.
+          //  When resolving this the options need to be split in legacy R8 and current R8.
+          return false;
+        }
+      };
 
   public static KeepRuleExtractorOptions getPgOptions() {
     return PG_OPTIONS;
@@ -21,27 +95,21 @@
     return R8_OPTIONS;
   }
 
-  private final boolean allowCheckDiscard;
-  private final boolean allowAccessModificationOption;
-  private final boolean allowAnnotationRemovalOption = false;
+  private KeepRuleExtractorOptions() {}
 
-  private KeepRuleExtractorOptions(
-      boolean allowCheckDiscard, boolean allowAccessModificationOption) {
-    this.allowCheckDiscard = allowCheckDiscard;
-    this.allowAccessModificationOption = allowAccessModificationOption;
-  }
+  public abstract boolean hasCheckDiscardSupport();
 
-  public boolean hasCheckDiscardSupport() {
-    return allowCheckDiscard;
-  }
+  public abstract boolean hasAllowAccessModificationOptionSupport();
 
-  private boolean hasAllowAccessModificationOptionSupport() {
-    return allowAccessModificationOption;
-  }
+  public abstract boolean hasAllowAnnotationRemovalOptionSupport();
 
-  private boolean hasAllowAnnotationRemovalOptionSupport() {
-    return allowAnnotationRemovalOption;
-  }
+  public abstract boolean hasFieldTypeBackReference();
+
+  public abstract boolean hasMethodReturnTypeBackReference();
+
+  public abstract boolean hasMethodParameterTypeBackReference();
+
+  public abstract boolean hasMethodParameterListBackReference();
 
   public boolean isKeepOptionSupported(KeepOption keepOption) {
     switch (keepOption) {
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/PgRule.java b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/PgRule.java
index 75ea20c..328b6ca 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/PgRule.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/PgRule.java
@@ -94,16 +94,23 @@
 
   private final KeepEdgeMetaInfo metaInfo;
   private final KeepOptions options;
+  private final KeepRuleExtractorOptions extractorOptions;
 
-  private PgRule(KeepEdgeMetaInfo metaInfo, KeepOptions options) {
+  private PgRule(
+      KeepEdgeMetaInfo metaInfo, KeepOptions options, KeepRuleExtractorOptions extractorOptions) {
     this.metaInfo = metaInfo;
     this.options = options;
+    this.extractorOptions = extractorOptions;
   }
 
   public KeepEdgeMetaInfo getMetaInfo() {
     return metaInfo;
   }
 
+  public KeepRuleExtractorOptions getExtractorOptions() {
+    return extractorOptions;
+  }
+
   // Helper to print the class-name pattern in a class-item.
   public static BiConsumer<StringBuilder, KeepQualifiedClassNamePattern> classNamePrinter(
       KeepQualifiedClassNamePattern classNamePattern) {
@@ -203,8 +210,9 @@
         KeepOptions options,
         Map<KeepBindingSymbol, KeepMemberPattern> memberPatterns,
         List<KeepBindingSymbol> targetMembers,
-        TargetKeepKind targetKeepKind) {
-      super(metaInfo, options);
+        TargetKeepKind targetKeepKind,
+        KeepRuleExtractorOptions extractorOptions) {
+      super(metaInfo, options, extractorOptions);
       assert !targetKeepKind.equals(TargetKeepKind.JUST_MEMBERS);
       this.holderNamePattern = holder.getNamePattern();
       this.holderPattern = holder.getClassItemPattern();
@@ -234,7 +242,8 @@
     @Override
     void printTargetMember(StringBuilder builder, KeepBindingSymbol memberReference) {
       KeepMemberPattern memberPattern = memberPatterns.get(memberReference);
-      printMemberClause(memberPattern, RulePrinter.withoutBackReferences(builder));
+      printMemberClause(
+          memberPattern, RulePrinter.withoutBackReferences(builder), getExtractorOptions());
     }
   }
 
@@ -242,8 +251,11 @@
 
     private final Set<KeepAttribute> attributes;
 
-    public PgKeepAttributeRule(KeepEdgeMetaInfo metaInfo, Set<KeepAttribute> attributes) {
-      super(metaInfo, null);
+    public PgKeepAttributeRule(
+        KeepEdgeMetaInfo metaInfo,
+        Set<KeepAttribute> attributes,
+        KeepRuleExtractorOptions extractorOptions) {
+      super(metaInfo, null, extractorOptions);
       assert !attributes.isEmpty();
       this.attributes = attributes;
     }
@@ -308,8 +320,9 @@
         Map<KeepBindingSymbol, KeepMemberPattern> memberPatterns,
         List<KeepBindingSymbol> memberConditions,
         List<KeepBindingSymbol> memberTargets,
-        TargetKeepKind keepKind) {
-      super(metaInfo, options);
+        TargetKeepKind keepKind,
+        KeepRuleExtractorOptions extractorOptions) {
+      super(metaInfo, options, extractorOptions);
       this.classCondition = classCondition.getClassItemPattern();
       this.classTarget = classTarget.getClassItemPattern();
       this.memberPatterns = memberPatterns;
@@ -336,7 +349,8 @@
     @Override
     void printConditionMember(StringBuilder builder, KeepBindingSymbol member) {
       KeepMemberPattern memberPattern = memberPatterns.get(member);
-      printMemberClause(memberPattern, RulePrinter.withoutBackReferences(builder));
+      printMemberClause(
+          memberPattern, RulePrinter.withoutBackReferences(builder), getExtractorOptions());
     }
 
     @Override
@@ -360,7 +374,8 @@
     @Override
     void printTargetMember(StringBuilder builder, KeepBindingSymbol member) {
       KeepMemberPattern memberPattern = memberPatterns.get(member);
-      printMemberClause(memberPattern, RulePrinter.withoutBackReferences(builder));
+      printMemberClause(
+          memberPattern, RulePrinter.withoutBackReferences(builder), getExtractorOptions());
     }
 
     private void printClassName(StringBuilder builder, KeepQualifiedClassNamePattern clazzName) {
@@ -402,8 +417,9 @@
         Map<KeepBindingSymbol, KeepMemberPattern> memberPatterns,
         List<KeepBindingSymbol> memberConditions,
         List<KeepBindingSymbol> memberTargets,
-        TargetKeepKind keepKind) {
-      super(metaInfo, options);
+        TargetKeepKind keepKind,
+        KeepRuleExtractorOptions extractorOptions) {
+      super(metaInfo, options, extractorOptions);
       this.holderNamePattern = holder.getNamePattern();
       this.holderPattern = holder.getClassItemPattern();
       this.memberPatterns = memberPatterns;
@@ -458,7 +474,7 @@
       KeepMemberPattern memberPattern = memberPatterns.get(member);
       BackReferencePrinter printer =
           RulePrinter.withBackReferences(builder, this::getNextBackReferenceNumber);
-      printMemberClause(memberPattern, printer);
+      printMemberClause(memberPattern, printer, getExtractorOptions());
       membersBackReferencePatterns.put(member, printer.getBackReference());
     }
 
@@ -492,7 +508,8 @@
         }
       }
       KeepMemberPattern memberPattern = memberPatterns.get(member);
-      printMemberClause(memberPattern, RulePrinter.withoutBackReferences(builder));
+      printMemberClause(
+          memberPattern, RulePrinter.withoutBackReferences(builder), getExtractorOptions());
     }
   }
 }
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrinter.java b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrinter.java
index 518eeba..78513c6 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrinter.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrinter.java
@@ -22,6 +22,10 @@
     this.builder = builder;
   }
 
+  public RulePrinter allowBackReferencesIf(boolean isBackReferenceSupported) {
+    return this;
+  }
+
   public RulePrinter append(String str) {
     assert !str.contains("*");
     assert !str.contains("(...)");
@@ -102,7 +106,51 @@
 
     @Override
     public RulePrinter appendAnyParameters() {
-      // TODO(b/265892343): R8 does not yet support back reference to `...`.
+      return addBackRef("(...)");
+    }
+
+    @Override
+    public RulePrinter allowBackReferencesIf(boolean isBackReferenceSupported) {
+      return isBackReferenceSupported ? this : new SkipBackreferencePrinter(this);
+    }
+  }
+
+  private static class SkipBackreferencePrinter extends RulePrinter {
+    final BackReferencePrinter printer;
+
+    private SkipBackreferencePrinter(BackReferencePrinter printer) {
+      super(((RulePrinter) printer).builder);
+      this.printer = printer;
+    }
+
+    @Override
+    public RulePrinter appendWithoutBackReferenceAssert(String str) {
+      printer.appendWithoutBackReferenceAssert(str);
+      return this;
+    }
+
+    @Override
+    public RulePrinter appendStar() {
+      return appendWithoutBackReferenceAssert("*");
+    }
+
+    @Override
+    public RulePrinter appendDoubleStar() {
+      return appendWithoutBackReferenceAssert("**");
+    }
+
+    @Override
+    public RulePrinter appendTripleStar() {
+      return appendWithoutBackReferenceAssert("***");
+    }
+
+    @Override
+    public RulePrinter appendPercent() {
+      return appendWithoutBackReferenceAssert("%");
+    }
+
+    @Override
+    public RulePrinter appendAnyParameters() {
       return appendWithoutBackReferenceAssert("(...)");
     }
   }
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrintingUtils.java b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrintingUtils.java
index 636dbfc..07c7c9b 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrintingUtils.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/keeprules/RulePrintingUtils.java
@@ -115,7 +115,8 @@
     return builder;
   }
 
-  public static RulePrinter printMemberClause(KeepMemberPattern member, RulePrinter printer) {
+  public static RulePrinter printMemberClause(
+      KeepMemberPattern member, RulePrinter printer, KeepRuleExtractorOptions options) {
     if (member.isAllMembers()) {
       // Note: the rule language does not allow backref to a full member. A rule matching all
       // members via a binding must be split in two up front: one for methods and one for fields.
@@ -127,10 +128,10 @@
       printer.append(" ");
     }
     if (member.isMethod()) {
-      return printMethod(member.asMethod(), printer);
+      return printMethod(member.asMethod(), printer, options);
     }
     if (member.isField()) {
-      return printField(member.asField(), printer);
+      return printField(member.asField(), printer, options);
     }
     // The pattern is a restricted member pattern, e.g., it must apply to fields and methods
     // without any specifics not common to both. For now that is annotated-by and access patterns.
@@ -139,27 +140,37 @@
     return printer.appendWithoutBackReferenceAssert("*").append(";");
   }
 
-  private static RulePrinter printField(KeepFieldPattern fieldPattern, RulePrinter printer) {
+  private static RulePrinter printField(
+      KeepFieldPattern fieldPattern, RulePrinter printer, KeepRuleExtractorOptions options) {
     printFieldAccess(printer, fieldPattern.getAccessPattern());
-    printType(printer, fieldPattern.getTypePattern().asType());
+    printType(
+        printer.allowBackReferencesIf(options.hasFieldTypeBackReference()),
+        fieldPattern.getTypePattern().asType());
     printer.append(" ");
     printFieldName(printer, fieldPattern.getNamePattern());
     return printer.append(";");
   }
 
-  private static RulePrinter printMethod(KeepMethodPattern methodPattern, RulePrinter printer) {
+  private static RulePrinter printMethod(
+      KeepMethodPattern methodPattern, RulePrinter printer, KeepRuleExtractorOptions options) {
     printMethodAccess(printer, methodPattern.getAccessPattern());
-    printReturnType(printer, methodPattern.getReturnTypePattern());
+    printReturnType(
+        printer.allowBackReferencesIf(options.hasMethodReturnTypeBackReference()),
+        methodPattern.getReturnTypePattern());
     printer.append(" ");
     printMethodName(printer, methodPattern.getNamePattern());
-    printParameters(printer, methodPattern.getParametersPattern());
+    printParameters(printer, methodPattern.getParametersPattern(), options);
     return printer.append(";");
   }
 
   private static RulePrinter printParameters(
-      RulePrinter builder, KeepMethodParametersPattern parametersPattern) {
+      RulePrinter builder,
+      KeepMethodParametersPattern parametersPattern,
+      KeepRuleExtractorOptions options) {
     if (parametersPattern.isAny()) {
-      return builder.appendAnyParameters();
+      return builder
+          .allowBackReferencesIf(options.hasMethodParameterListBackReference())
+          .appendAnyParameters();
     }
     builder.append("(");
     List<KeepTypePattern> patterns = parametersPattern.asList();
@@ -167,7 +178,9 @@
       if (i > 0) {
         builder.append(", ");
       }
-      printType(builder, patterns.get(i));
+      printType(
+          builder.allowBackReferencesIf(options.hasMethodParameterTypeBackReference()),
+          patterns.get(i));
     }
     return builder.append(")");
   }
diff --git a/src/test/java/com/android/tools/r8/keepanno/KeepSameMethodTest.java b/src/test/java/com/android/tools/r8/keepanno/KeepSameMethodTest.java
index aebeeb2..b0c50d2 100644
--- a/src/test/java/com/android/tools/r8/keepanno/KeepSameMethodTest.java
+++ b/src/test/java/com/android/tools/r8/keepanno/KeepSameMethodTest.java
@@ -27,9 +27,6 @@
 
   static final String EXPECTED = StringUtils.lines("foo");
 
-  // TODO(b/265893433): The use of backreferences does not work in PG.
-  static final String UNEXPECTED_PG = StringUtils.lines("main");
-
   @Parameter public KeepAnnoParameters parameters;
 
   @Parameterized.Parameters(name = "{0}")
@@ -47,7 +44,7 @@
         // The "all members" target will create an unused "all fields" rule.
         .allowUnusedProguardConfigurationRules()
         .run(TestClass.class)
-        .assertSuccessWithOutput(parameters.isPG() ? UNEXPECTED_PG : EXPECTED)
+        .assertSuccessWithOutput(EXPECTED)
         .applyIf(parameters.isR8(), r -> r.inspect(this::checkOutput));
   }