From 866a3e3ab723d37413b97edc42b7b098ba6f49a5 Mon Sep 17 00:00:00 2001
From: Liam Miller-Cushon <cushon@google.com>
Date: Sun, 11 Aug 2024 11:51:19 -0700
Subject: [PATCH] Implement `getPermittedSubclasses`

Startblock:
  * unknown commit is submitted
PiperOrigin-RevId: 661841593
---
 .../turbine/processing/TurbineElement.java    | 18 ++++++++++
 .../processing/ProcessingIntegrationTest.java | 34 +++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/java/com/google/turbine/processing/TurbineElement.java b/java/com/google/turbine/processing/TurbineElement.java
index 38745366..aa70f1cb 100644
--- a/java/com/google/turbine/processing/TurbineElement.java
+++ b/java/com/google/turbine/processing/TurbineElement.java
@@ -437,6 +437,24 @@ public Element getEnclosingElement() {
       return enclosing.get();
     }
 
+    private final Supplier<ImmutableList<TypeMirror>> permits =
+        memoize(
+            new Supplier<>() {
+              @Override
+              public ImmutableList<TypeMirror> get() {
+                ImmutableList.Builder<TypeMirror> result = ImmutableList.builder();
+                for (ClassSymbol p : infoNonNull().permits()) {
+                  result.add(factory.asTypeMirror(ClassTy.asNonParametricClassTy(p)));
+                }
+                return result.build();
+              }
+            });
+
+    @Override
+    public List<? extends TypeMirror> getPermittedSubclasses() {
+      return permits.get();
+    }
+
     private final Supplier<ImmutableList<Element>> enclosed =
         memoize(
             new Supplier<ImmutableList<Element>>() {
diff --git a/javatests/com/google/turbine/processing/ProcessingIntegrationTest.java b/javatests/com/google/turbine/processing/ProcessingIntegrationTest.java
index c2ca38ed..4f24788e 100644
--- a/javatests/com/google/turbine/processing/ProcessingIntegrationTest.java
+++ b/javatests/com/google/turbine/processing/ProcessingIntegrationTest.java
@@ -71,6 +71,7 @@
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
 import javax.lang.model.type.ExecutableType;
+import javax.lang.model.util.ElementFilter;
 import javax.tools.Diagnostic;
 import javax.tools.FileObject;
 import javax.tools.JavaFileObject;
@@ -899,6 +900,39 @@ public void modifiers() {
         .containsExactly("I [abstract, sealed]", "J [abstract, non-sealed]");
   }
 
+  @SupportedAnnotationTypes("*")
+  public static class PermitsProcessor extends AbstractProcessor {
+    @Override
+    public SourceVersion getSupportedSourceVersion() {
+      return SourceVersion.latestSupported();
+    }
+
+    @Override
+    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
+      for (TypeElement e : ElementFilter.typesIn(roundEnv.getRootElements())) {
+        processingEnv
+            .getMessager()
+            .printMessage(
+                Diagnostic.Kind.ERROR, String.format("%s %s", e, e.getPermittedSubclasses()), e);
+      }
+      return false;
+    }
+  }
+
+  @Test
+  public void permits() {
+    ImmutableList<Tree.CompUnit> units =
+        parseUnit(
+            "=== I.java ===", //
+            "interface I permits J, K {}",
+            "interface J {}",
+            "interface K {}");
+    TurbineError e1 = runProcessors(units, new PermitsProcessor());
+    TurbineError e = e1;
+    assertThat(e.diagnostics().stream().map(d -> d.message()))
+        .containsExactly("I [J, K]", "J []", "K []");
+  }
+
   private TurbineError runProcessors(ImmutableList<Tree.CompUnit> units, Processor... processors) {
     return assertThrows(
         TurbineError.class,