[KeepAnno] Refactor with consistent apply/match helpers

Change-Id: Ibb670ff22fbedbdeae91781dd911112e80f27740
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/asm/ClassNameParser.java b/src/keepanno/java/com/android/tools/r8/keepanno/asm/ClassNameParser.java
index b3fdddb..30f7fec 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/asm/ClassNameParser.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/asm/ClassNameParser.java
@@ -61,7 +61,7 @@
 
   KeepQualifiedClassNamePattern typeToClassType(
       KeepTypePattern typePattern, PropertyParsingContext parsingContext) {
-    return typePattern.match(
+    return typePattern.apply(
         KeepQualifiedClassNamePattern::any,
         primitiveTypePattern -> {
           throw parsingContext.error("Invalid use of primitive type where class type was expected");
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/asm/KeepEdgeWriter.java b/src/keepanno/java/com/android/tools/r8/keepanno/asm/KeepEdgeWriter.java
index 7f8bdde..a0a37e3 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/asm/KeepEdgeWriter.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/asm/KeepEdgeWriter.java
@@ -25,7 +25,6 @@
 import com.android.tools.r8.keepanno.ast.KeepAnnotationPattern;
 import com.android.tools.r8.keepanno.ast.KeepBindingReference;
 import com.android.tools.r8.keepanno.ast.KeepBindings;
-import com.android.tools.r8.keepanno.ast.KeepCheck;
 import com.android.tools.r8.keepanno.ast.KeepClassItemPattern;
 import com.android.tools.r8.keepanno.ast.KeepClassItemReference;
 import com.android.tools.r8.keepanno.ast.KeepConsequences;
@@ -150,25 +149,24 @@
     KeepEdgeMetaInfo metaInfo = decl.getMetaInfo();
     visitor.visit(ExtractedAnnotation.version, metaInfo.getVersion().toVersionString());
     visitor.visit(ExtractedAnnotation.context, metaInfo.getContextDescriptorString());
-    if (decl.isKeepEdge()) {
-      KeepEdge edge = decl.asKeepEdge();
-      withNewVisitor(
-          visitor.visitAnnotation(ExtractedAnnotation.edge, Edge.DESCRIPTOR),
-          v -> new KeepEdgeWriter().writeEdge(edge, v));
-      return;
-    }
-    assert decl.isKeepCheck();
-    KeepCheck check = decl.asKeepCheck();
-    switch (check.getKind()) {
-      case REMOVED:
-        visitor.visit(ExtractedAnnotation.checkRemoved, true);
-        break;
-      case OPTIMIZED_OUT:
-        visitor.visit(ExtractedAnnotation.checkOptimizedOut, true);
-        break;
-      default:
-        throw new KeepEdgeException("Unexpected keep check kind: " + check.getKind());
-    }
+    decl.match(
+        edge -> {
+          withNewVisitor(
+              visitor.visitAnnotation(ExtractedAnnotation.edge, Edge.DESCRIPTOR),
+              v -> new KeepEdgeWriter().writeEdge(edge, v));
+        },
+        check -> {
+          switch (check.getKind()) {
+            case REMOVED:
+              visitor.visit(ExtractedAnnotation.checkRemoved, true);
+              break;
+            case OPTIMIZED_OUT:
+              visitor.visit(ExtractedAnnotation.checkOptimizedOut, true);
+              break;
+            default:
+              throw new KeepEdgeException("Unexpected keep check kind: " + check.getKind());
+          }
+        });
   }
 
   private void writeEdge(KeepEdge edge, AnnotationVisitor visitor) {
@@ -476,22 +474,18 @@
             typePattern.match(
                 () -> {
                   // The empty type pattern matches any type.
-                  return null;
                 },
                 primitive -> {
                   if (primitive.isAny()) {
                     throw new Unimplemented("No support for any-primitive.");
                   }
                   v.visit(TypePattern.name, Type.getType(primitive.getDescriptor()).getClassName());
-                  return null;
                 },
                 array -> {
                   v.visit(TypePattern.name, Type.getType(array.getDescriptor()).getClassName());
-                  return null;
                 },
                 clazz -> {
                   writeClassNamePattern(clazz, TypePattern.classNamePattern, v);
-                  return null;
                 }));
   }
 
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/AstUtils.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/AstUtils.java
new file mode 100644
index 0000000..fb08f88
--- /dev/null
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/AstUtils.java
@@ -0,0 +1,27 @@
+// Copyright (c) 2024, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.keepanno.ast;
+
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+public abstract class AstUtils {
+  private AstUtils() {}
+
+  public static <T> Function<T, Void> toVoidFunction(Consumer<T> fn) {
+    return t -> {
+      fn.accept(t);
+      return null;
+    };
+  }
+
+  public static Supplier<Void> toVoidSupplier(Runnable fn) {
+    return () -> {
+      fn.run();
+      return null;
+    };
+  }
+}
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepArrayTypePattern.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepArrayTypePattern.java
index ff7a89b..6194fb3 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepArrayTypePattern.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepArrayTypePattern.java
@@ -43,7 +43,7 @@
       throw new KeepEdgeException("No descriptor exists for 'any' array");
     }
     return Strings.repeat("[", dimensions)
-        + baseType.match(
+        + baseType.apply(
             () -> {
               throw new KeepEdgeException("No descriptor exists for 'any primitive' array");
             },
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepDeclaration.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepDeclaration.java
index 4425317..e80bfbd 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepDeclaration.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepDeclaration.java
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.keepanno.ast;
 
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /** Base class for the declarations represented in the keep annoations library. */
@@ -10,13 +11,17 @@
 
   public abstract KeepEdgeMetaInfo getMetaInfo();
 
-  public final <T> T match(Function<KeepEdge, T> onEdge, Function<KeepCheck, T> onCheck) {
+  public final <T> T apply(Function<KeepEdge, T> onEdge, Function<KeepCheck, T> onCheck) {
     if (isKeepEdge()) {
       return onEdge.apply(asKeepEdge());
     }
     return onCheck.apply(asKeepCheck());
   }
 
+  public final void match(Consumer<KeepEdge> onEdge, Consumer<KeepCheck> onCheck) {
+    apply(AstUtils.toVoidFunction(onEdge), AstUtils.toVoidFunction(onCheck));
+  }
+
   public final boolean isKeepEdge() {
     return asKeepEdge() != null;
   }
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemPattern.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemPattern.java
index c718995..a455344 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemPattern.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemPattern.java
@@ -4,6 +4,7 @@
 package com.android.tools.r8.keepanno.ast;
 
 import java.util.Collection;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /**
@@ -46,7 +47,7 @@
 
   public abstract KeepItemReference toItemReference();
 
-  public <T> T match(
+  public final <T> T apply(
       Function<KeepClassItemPattern, T> onClass, Function<KeepMemberItemPattern, T> onMember) {
     if (isClassItemPattern()) {
       return onClass.apply(asClassItemPattern());
@@ -54,5 +55,10 @@
     assert isMemberItemPattern();
     return onMember.apply(asMemberItemPattern());
   }
+
+  public final void match(
+      Consumer<KeepClassItemPattern> onClass, Consumer<KeepMemberItemPattern> onMember) {
+    apply(AstUtils.toVoidFunction(onClass), AstUtils.toVoidFunction(onMember));
+  }
 }
 
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemReference.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemReference.java
index e1e0f8a..d17f6ee 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemReference.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepItemReference.java
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.keepanno.ast;
 
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /**
@@ -64,7 +65,7 @@
     return null;
   }
 
-  public <T> T match(
+  public <T> T apply(
       Function<KeepBindingReference, T> onBinding, Function<KeepItemPattern, T> onItem) {
     if (isBindingReference()) {
       return onBinding.apply(asBindingReference());
@@ -72,4 +73,8 @@
     assert isItemPattern();
     return onItem.apply(asItemPattern());
   }
+
+  public void match(Consumer<KeepBindingReference> onBinding, Consumer<KeepItemPattern> onItem) {
+    apply(AstUtils.toVoidFunction(onBinding), AstUtils.toVoidFunction(onItem));
+  }
 }
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepMemberPattern.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepMemberPattern.java
index 57e8499..dfce0f6 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepMemberPattern.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepMemberPattern.java
@@ -4,6 +4,8 @@
 package com.android.tools.r8.keepanno.ast;
 
 import java.util.Objects;
+import java.util.function.Consumer;
+import java.util.function.Function;
 
 public abstract class KeepMemberPattern {
 
@@ -119,4 +121,28 @@
   public abstract KeepMemberAccessPattern getAccessPattern();
 
   public abstract OptionalPattern<KeepQualifiedClassNamePattern> getAnnotatedByPattern();
+
+  public <T> T apply(
+      Function<KeepMemberPattern, T> onGeneralMember,
+      Function<KeepFieldPattern, T> onFieldMember,
+      Function<KeepMethodPattern, T> onMethodMember) {
+    if (isGeneralMember()) {
+      return onGeneralMember.apply(this);
+    }
+    if (isField()) {
+      return onFieldMember.apply(asField());
+    }
+    assert isMethod();
+    return onMethodMember.apply(asMethod());
+  }
+
+  public void match(
+      Consumer<KeepMemberPattern> onGeneralMember,
+      Consumer<KeepFieldPattern> onFieldMember,
+      Consumer<KeepMethodPattern> onMethodMember) {
+    apply(
+        AstUtils.toVoidFunction(onGeneralMember),
+        AstUtils.toVoidFunction(onFieldMember),
+        AstUtils.toVoidFunction(onMethodMember));
+  }
 }
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepTypePattern.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepTypePattern.java
index bc1904f..a0edf23 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepTypePattern.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/KeepTypePattern.java
@@ -6,6 +6,7 @@
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -51,12 +52,24 @@
     throw new KeepEdgeException("Invalid type descriptor: " + typeDescriptor);
   }
 
-  public abstract <T> T match(
+  public abstract <T> T apply(
       Supplier<T> onAny,
       Function<KeepPrimitiveTypePattern, T> onPrimitive,
       Function<KeepArrayTypePattern, T> onArray,
       Function<KeepQualifiedClassNamePattern, T> onClass);
 
+  public final void match(
+      Runnable onAny,
+      Consumer<KeepPrimitiveTypePattern> onPrimitive,
+      Consumer<KeepArrayTypePattern> onArray,
+      Consumer<KeepQualifiedClassNamePattern> onClass) {
+    apply(
+        AstUtils.toVoidSupplier(onAny),
+        AstUtils.toVoidFunction(onPrimitive),
+        AstUtils.toVoidFunction(onArray),
+        AstUtils.toVoidFunction(onClass));
+  }
+
   public boolean isAny() {
     return false;
   }
@@ -70,7 +83,7 @@
     }
 
     @Override
-    public <T> T match(
+    public <T> T apply(
         Supplier<T> onAny,
         Function<KeepPrimitiveTypePattern, T> onPrimitive,
         Function<KeepArrayTypePattern, T> onArray,
@@ -135,7 +148,7 @@
     }
 
     @Override
-    public <T> T match(
+    public <T> T apply(
         Supplier<T> onAny,
         Function<KeepPrimitiveTypePattern, T> onPrimitive,
         Function<KeepArrayTypePattern, T> onArray,
@@ -152,7 +165,7 @@
     }
 
     @Override
-    public <T> T match(
+    public <T> T apply(
         Supplier<T> onAny,
         Function<KeepPrimitiveTypePattern, T> onPrimitive,
         Function<KeepArrayTypePattern, T> onArray,
@@ -191,7 +204,7 @@
     }
 
     @Override
-    public <T> T match(
+    public <T> T apply(
         Supplier<T> onAny,
         Function<KeepPrimitiveTypePattern, T> onPrimitive,
         Function<KeepArrayTypePattern, T> onArray,
diff --git a/src/keepanno/java/com/android/tools/r8/keepanno/ast/OptionalPattern.java b/src/keepanno/java/com/android/tools/r8/keepanno/ast/OptionalPattern.java
index 6a3832e..dfce78b 100644
--- a/src/keepanno/java/com/android/tools/r8/keepanno/ast/OptionalPattern.java
+++ b/src/keepanno/java/com/android/tools/r8/keepanno/ast/OptionalPattern.java
@@ -4,7 +4,9 @@
 
 package com.android.tools.r8.keepanno.ast;
 
+import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 public abstract class OptionalPattern<T> {
 
@@ -39,6 +41,17 @@
     return defaultValue;
   }
 
+  public final <S> S apply(Supplier<S> onAbsent, Function<T, S> onPresent) {
+    if (isAbsent()) {
+      return onAbsent.get();
+    }
+    return onPresent.apply(get());
+  }
+
+  public final void match(Runnable onAbsent, Consumer<T> onPresent) {
+    apply(AstUtils.toVoidSupplier(onAbsent), AstUtils.toVoidFunction(onPresent));
+  }
+
   private static final class Absent extends OptionalPattern {
     private static final Absent INSTANCE = new Absent();
 
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 07c7c9b..b7c8fbb 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
@@ -217,7 +217,7 @@
   }
 
   private static RulePrinter printType(RulePrinter printer, KeepTypePattern typePattern) {
-    return typePattern.match(
+    return typePattern.apply(
         printer::appendTripleStar,
         primitivePattern -> printPrimitiveType(printer, primitivePattern),
         arrayTypePattern -> printArrayType(printer, arrayTypePattern),