Add support for annotations to TypeTable

rewrite-java/src/test/java/org/openrewrite/java/internal/parser/TypeTableTest.java

@@ -24,12 +24,17 @@
 import org.openrewrite.InMemoryExecutionContext;
 import org.openrewrite.java.JavaParser;
 import org.openrewrite.java.JavaParserExecutionContextView;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaType;
 import org.openrewrite.test.RewriteTest;
 
 import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.URL;
 import java.nio.file.*;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
 
 import static io.micrometer.core.instrument.util.DoubleFormat.decimalOrNan;
@@ -131,7 +136,7 @@ class Test {
                   @BeforeEach
                   void before() {
                   }
-
+              
                   @Test
                   void foo() {
                       Assertions.assertTrue(true);
@@ -142,6 +147,45 @@ void foo() {
         );
     }
 
+    @Test
+    void writeReadWithAnnotations() throws IOException, URISyntaxException {
+        URL resource = TypeTableTest.class.getClassLoader().getResource("jakarta.validation-api-3.1.1.jar");
+        try (TypeTable.Writer writer = TypeTable.newWriter(Files.newOutputStream(tsv))) {
+            writeJar(Path.of(resource.toURI()), writer);
+        }
+
+        TypeTable table = new TypeTable(ctx, tsv.toUri().toURL(), List.of("jakarta.validation-api"));
+        Path classesDir = table.load("jakarta.validation-api");
+        assertThat(classesDir).isNotNull();
+
+        // Test that we can use the annotations in code
+        rewriteRun(
+          spec -> spec.parser(JavaParser.fromJavaVersion()
+            .classpath(List.of(classesDir))),
+          java(
+            """
+              import jakarta.validation.constraints.AssertFalse;
+
+              class AnnotatedTest {
+                  @AssertFalse
+                  int i;
+              }
+              """,
+            spec -> spec.afterRecipe(cu -> {
+                // Assert that the JavaType.Class has the annotation
+                J.VariableDeclarations field = (J.VariableDeclarations) cu.getClasses().get(0).getBody().getStatements().get(0);
+                JavaType.Class assertFalseType = (JavaType.Class) field.getLeadingAnnotations().get(0).getType();
+                assertThat(assertFalseType).isNotNull();
+                assertThat(assertFalseType.getFullyQualifiedName()).isEqualTo("jakarta.validation.constraints.AssertFalse");
+                assertThat(assertFalseType.getAnnotations().size()).isEqualTo(5);
+                assertThat(assertFalseType.getMethods().stream().filter(m -> m.getName().equals("message"))).satisfiesExactly(
+                  m -> assertThat(m.getDefaultValue()).containsExactly("{jakarta.validation.constraints.AssertFalse.message}")
+                );
+            })
+          )
+        );
+    }
+
     private static long writeJar(Path classpath, TypeTable.Writer writer) throws IOException {
         String fileName = classpath.toFile().getName();
         if (fileName.endsWith(".jar")) {
@@ -167,4 +211,118 @@ private String humanReadableByteCount(double bytes) {
         String pre = "KMGTPE".charAt(exp - 1) + "i";
         return decimalOrNan(bytes / Math.pow(unit, exp)) + " " + pre + "B";
     }
+
+    @Test
+    void writeReadWithCustomAnnotations() throws IOException, URISyntaxException {
+        // Find the directory containing the compiled classes
+        URL classUrl = TypeTableTest.class.getResource("TypeTableTest.class");
+        Path classesDir = Paths.get(classUrl.toURI()).getParent().getParent().getParent().getParent().getParent().getParent();
+        Path dataDir = classesDir.resolve("org/openrewrite/java/internal/parser/data");
+
+        // Create a temporary jar file
+        Path jarFile = temp.resolve("custom-annotations-1.jar");
+
+        // Create jar file from the compiled classes
+        try (FileSystem zipfs = FileSystems.newFileSystem(
+          jarFile,
+          Map.of("create", "true"))) {
+
+            Files.walk(dataDir)
+              .filter(Files::isRegularFile)
+              .filter(p -> p.toString().endsWith(".class"))
+              .forEach(classFile -> {
+                  try {
+                      Path pathInZip = zipfs.getPath("/").resolve(
+                        classesDir.relativize(classFile).toString());
+                      Files.createDirectories(pathInZip.getParent());
+                      Files.copy(classFile, pathInZip);
+                  } catch (IOException e) {
+                      throw new RuntimeException(e);
+                  }
+              });
+        }
+
+        // Write the jar to the type table
+        try (TypeTable.Writer writer = TypeTable.newWriter(Files.newOutputStream(tsv))) {
+            writeJar(jarFile, writer);
+        }
+
+        // Load the type table
+        TypeTable table = new TypeTable(ctx, tsv.toUri().toURL(), List.of("custom-annotations"));
+        Path loadedClassesDir = table.load("custom-annotations");
+        assertThat(loadedClassesDir).isNotNull();
+
+        // Test that we can use the annotations in code
+        rewriteRun(
+          spec -> spec.parser(JavaParser.fromJavaVersion()
+            .classpath(List.of(loadedClassesDir))),
+          java(
+            """
+              import org.openrewrite.java.internal.parser.data.*;
+
+              @BasicAnnotation
+              @StringAnnotation
+              @NestedAnnotation
+              @ArrayAnnotation
+              @ClassRefAnnotation
+              @EnumAnnotation
+              @ConstantAnnotation
+              class TestClass {
+              }
+              """,
+            spec -> spec.afterRecipe(cu -> {
+                // Verify that all annotations are properly loaded with their attributes
+                J.ClassDeclaration clazz = cu.getClasses().get(0);
+                List<J.Annotation> annotations = clazz.getLeadingAnnotations();
+
+                // Verify BasicAnnotation
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.BasicAnnotation",
+                  "intValue", "42");
+
+                // Verify StringAnnotation
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.StringAnnotation",
+                  "value", "Default value");
+
+                // Verify ClassRefAnnotation
+                // TODO `JavaType.Method#defaultValue` is of type `List<String>`
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.ClassRefAnnotation",
+                  "value", "java.lang.String");
+
+                // Verify NestedAnnotation
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.NestedAnnotation",
+                  "nested", "@org.openrewrite.java.internal.parser.data.NestedLevel2");
+
+                // Verify ArrayAnnotation
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.ArrayAnnotation",
+                  "strings", "one", "two", "three");
+
+                // Verify EnumAnnotation
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.EnumAnnotation",
+                  "value", ".ONE");
+
+                // Verify ConstantAnnotation
+                verifyAnnotation(annotations, "org.openrewrite.java.internal.parser.data.ConstantAnnotation",
+                  "value", "This is a constant string");
+            })
+          )
+        );
+    }
+
+    private void verifyAnnotation(List<J.Annotation> annotations, String fqn, String attributeName, String... expectedValues) {
+        J.Annotation annotation = annotations.stream()
+          .filter(a -> {
+              JavaType.Class type = (JavaType.Class) a.getType();
+              return type != null && type.getFullyQualifiedName().equals(fqn);
+          })
+          .findFirst()
+          .orElseThrow(() -> new AssertionError("Annotation " + fqn + " not found"));
+
+        JavaType.Class annotationType = (JavaType.Class) annotation.getType();
+        assertThat(annotationType).isNotNull();
+        assertThat(annotationType.getMethods().stream()
+          .filter(m -> m.getName().equals(attributeName)))
+          .satisfiesExactly(
+            m -> assertThat(m.getDefaultValue()).containsExactly(expectedValues)
+          );
+    }
 }

rewrite-java/src/test/java/org/openrewrite/java/internal/parser/data/AnnotatedClass.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2025 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.internal.parser.data;
+
+import java.lang.annotation.ElementType;
+
+@BasicAnnotation(
+    intValue = 100,
+    longValue = 200L,
+    floatValue = 1.5f,
+    doubleValue = 2.5,
+    boolValue = false,
+    charValue = 'Y',
+    byteValue = 10,
+    shortValue = 20
+)
+@StringAnnotation(
+    value = "Custom value",
+//    escaped = "Testing \"quotes\" and \\ backslashes",
+    unicode = "More unicode: \u2764\uFE0F"
+//    multiLine = "Line 1\nLine 2\nLine 3"
+)
+@NestedAnnotation(
+    basic = @BasicAnnotation(intValue = 999),
+    string = @StringAnnotation(value = "Nested string"),
+    nested = @NestedLevel2(
+        value = "Custom nested",
+        deepNested = @NestedLevel3(value = "Deep nested custom")
+    )
+)
+@ArrayAnnotation(
+    strings = {"A", "B", "C"},
+    ints = {10, 20, 30},
+    classes = {String.class, Integer.class, Object.class},
+    annotations = {
+        @BasicAnnotation(intValue = 1),
+        @BasicAnnotation(intValue = 2)
+    }
+)
+@ClassRefAnnotation(
+    value = Integer.class,
+    primitive = boolean.class,
+    array = String[].class,
+    innerClass = AnnotatedClass.InnerClass.class
+)
+@EnumAnnotation(
+    value = TestEnum.THREE,
+    values = {TestEnum.TWO, TestEnum.THREE},
+    elementType = ElementType.FIELD
+)
+@ConstantAnnotation(
+    value = Constants.STRING_CONSTANT,
+    number = Constants.INT_CONSTANT,
+    decimal = Constants.DOUBLE_CONSTANT
+)
+public class AnnotatedClass {
+
+    @BasicAnnotation
+    private String field;
+
+    @StringAnnotation
+    public void method() {
+        // Method body
+    }
+
+    public static class InnerClass {
+        // Inner class for testing class references
+    }
+}

rewrite-java/src/test/java/org/openrewrite/java/internal/parser/data/ArrayAnnotation.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2025 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.internal.parser.data;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD, ElementType.FIELD})
+public @interface ArrayAnnotation {
+    String[] emptyArray() default {};
+    String[] singleElement() default {"single"};
+    String[] strings() default {"one", "two", "three"};
+    int[] ints() default {1, 2, 3};
+    Class<?>[] classes() default {String.class, Integer.class};
+    BasicAnnotation[] annotations() default {@BasicAnnotation, @BasicAnnotation(intValue = 100)};
+}

rewrite-java/src/test/java/org/openrewrite/java/internal/parser/data/BasicAnnotation.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2025 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.internal.parser.data;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD, ElementType.FIELD})
+public @interface BasicAnnotation {
+    int intValue() default 42;
+    long longValue() default Long.MAX_VALUE;
+    float floatValue() default 3.14f;
+    double doubleValue() default Math.PI;
+    boolean boolValue() default true;
+    char charValue() default 'X';
+    byte byteValue() default 8;
+    short shortValue() default 16;
+}

rewrite-java/src/test/java/org/openrewrite/java/internal/parser/data/ClassRefAnnotation.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2025 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.internal.parser.data;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD, ElementType.FIELD})
+public @interface ClassRefAnnotation {
+    Class<?> value() default String.class;
+    Class<?> primitive() default int.class;
+    Class<?> array() default byte[].class;
+    Class<?> innerClass() default Thread.State.class;
+}