Implement annotating classes, interfaces and enums with source lines …

…data (#6)

README.md

@@ -6,7 +6,7 @@ augments compiled classes with some additional data used by Datadog products.
 List of things that the plugin does:
 
 - Annotate every class with a `@SourcePath("...")` annotation that has the path to the class' source code
-- Annotate every public method with a `@MethodLines("...")` annotation that has method start and end lines (taking into account method modifiers and annotations)
+- Annotate every class, interface, enum and public method with a `@SourceLines("...")` annotation that has start and end lines (taking into account modifiers and annotations)
 
 ## Configuration
 
@@ -29,7 +29,7 @@ Add plugin-client JAR to the project's classpath:
     <dependency>
         <groupId>com.datadoghq</groupId>
         <artifactId>dd-javac-plugin-client</artifactId>
-        <version>0.2.0</version>
+        <version>0.2.2</version>
     </dependency>
 </dependencies>
 ```
@@ -49,7 +49,7 @@ Add plugin JAR to the compiler's annotation processor path and pass the plugin a
                     <annotationProcessorPath>
                         <groupId>com.datadoghq</groupId>
                         <artifactId>dd-javac-plugin</artifactId>
-                        <version>0.2.0</version>
+                        <version>0.2.2</version>
                     </annotationProcessorPath>
                 </annotationProcessorPaths>
                 <compilerArgs>
@@ -73,9 +73,9 @@ the plugin argument:
 
 ```groovy
 dependencies {
-    implementation 'com.datadoghq:dd-javac-plugin-client:0.2.0'
-    annotationProcessor 'com.datadoghq:dd-javac-plugin:0.2.0'
-    testAnnotationProcessor 'com.datadoghq:dd-javac-plugin:0.2.0'
+    implementation 'com.datadoghq:dd-javac-plugin-client:0.2.2'
+    annotationProcessor 'com.datadoghq:dd-javac-plugin:0.2.2'
+    testAnnotationProcessor 'com.datadoghq:dd-javac-plugin:0.2.2'
 }
 
 tasks.withType(JavaCompile).configureEach {
@@ -91,7 +91,7 @@ Below is an example for direct compiler invocation:
 
 ```shell
 javac \
-    -classpath dd-javac-plugin-client-0.2.1-all.jar \
+    -classpath dd-javac-plugin-client-0.2.2-all.jar \
     -Xplugin:DatadogCompilerPlugin \
     <PATH_TO_SOURCES>
 ```
@@ -103,13 +103,16 @@ To access the injected information, use `CompilerUtils` class from the `dd-javac
 ```java
 String sourcePath = CompilerUtils.getSourcePath(MyClass.class);
 
-int startLine = CompilerUtils.getStartLine(method);
-int endLine = CompilerUtils.getEndLine(method);
+int classStartLine = CompilerUtils.getStartLine(MyClass.class);
+int classEndLine = CompilerUtils.getEndLine(MyClass.class);
+
+int methodStartLine = CompilerUtils.getStartLine(method);
+int methodEndLine = CompilerUtils.getEndLine(method);
 ```
 
 ## Additional configuration
 
-Specify `disableMethodAnnotation` plugin argument if you want to disable annotating public methods.
+Specify `disableSourceLinesAnnotation` plugin argument if you want to disable annotating source lines.
 The argument can be specified in `javac` command line after the `-Xplugin` clause.
 
 ## Limitations

dd-javac-plugin-client/src/main/java/datadog/compiler/annotations/MethodLines.java → dd-javac-plugin-client/src/main/java/datadog/compiler/annotations/SourceLines.java

@@ -5,9 +5,9 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
-@Target({ElementType.METHOD, ElementType.CONSTRUCTOR})
+@Target({ElementType.METHOD, ElementType.CONSTRUCTOR, ElementType.TYPE})
 @Retention(RetentionPolicy.RUNTIME)
-public @interface MethodLines {
+public @interface SourceLines {
     int start();
     int end();
 }

dd-javac-plugin-client/src/main/java/datadog/compiler/utils/CompilerUtils.java

@@ -1,6 +1,6 @@
 package datadog.compiler.utils;
 
-import datadog.compiler.annotations.MethodLines;
+import datadog.compiler.annotations.SourceLines;
 import datadog.compiler.annotations.SourcePath;
 import java.lang.reflect.Executable;
 
@@ -41,7 +41,7 @@ public static String getSourcePath(Class<?> clazz) {
      * @return Start line of the method or {@link CompilerUtils#LINE_UNKNOWN} if the line cannot be determined
      */
     public static int getStartLine(Executable executable) {
-        MethodLines methodLines = executable.getAnnotation(MethodLines.class);
+        SourceLines methodLines = executable.getAnnotation(SourceLines.class);
         return methodLines != null ? methodLines.start() : LINE_UNKNOWN;
     }
 
@@ -52,7 +52,29 @@ public static int getStartLine(Executable executable) {
      * @return End line of the method or {@link CompilerUtils#LINE_UNKNOWN} if the line cannot be determined
      */
     public static int getEndLine(Executable executable) {
-        MethodLines methodLines = executable.getAnnotation(MethodLines.class);
+        SourceLines methodLines = executable.getAnnotation(SourceLines.class);
         return methodLines != null ? methodLines.end() : LINE_UNKNOWN;
     }
+
+    /**
+     * Returns start line of the provided class (class name, modifiers and annotations are taken into account).
+     *
+     * @param clazz Class
+     * @return Start line of the class or {@link CompilerUtils#LINE_UNKNOWN} if the line cannot be determined
+     */
+    public static int getStartLine(Class<?> clazz) {
+        SourceLines classLines = clazz.getAnnotation(SourceLines.class);
+        return classLines != null ? classLines.start() : LINE_UNKNOWN;
+    }
+
+    /**
+     * Returns end line of the provided class.
+     *
+     * @param clazz Class
+     * @return End line of the class or {@link CompilerUtils#LINE_UNKNOWN} if the line cannot be determined
+     */
+    public static int getEndLine(Class<?> clazz) {
+        SourceLines classLines = clazz.getAnnotation(SourceLines.class);
+        return classLines != null ? classLines.end() : LINE_UNKNOWN;
+    }
 }

dd-javac-plugin-client/src/test/java/datadog/compiler/utils/CompilerUtilsTest.java

@@ -1,15 +1,26 @@
 package datadog.compiler.utils;
 
+import datadog.compiler.annotations.SourceLines;
 import datadog.compiler.annotations.SourcePath;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import java.lang.reflect.Method;
 
 public class CompilerUtilsTest {
 
     private static final String TEST_CLASS_SOURCE_PATH = "/repo/src/package/TestClass.java";
+    private static final int TEST_CLASS_SOURCE_LINES_START = 19;
+    private static final int TEST_CLASS_SOURCE_LINES_END = 24;
+    private static final int TEST_METHOD_SOURCE_LINES_START = 21;
+    private static final int TEST_METHOD_SOURCE_LINES_END = 23;
 
     @SourcePath(TEST_CLASS_SOURCE_PATH)
+    @SourceLines(start = TEST_CLASS_SOURCE_LINES_START, end = TEST_CLASS_SOURCE_LINES_END)
     private static final class TestClass {
+        @SourceLines(start = TEST_METHOD_SOURCE_LINES_START, end = TEST_METHOD_SOURCE_LINES_END)
+        public static void testMethod() {
+            // no op
+        }
     }
 
     @Test
@@ -18,4 +29,20 @@ public void testSourcePathExtraction() {
         Assertions.assertEquals(TEST_CLASS_SOURCE_PATH, sourcePath);
     }
 
+    @Test
+    public void testMethodLinesExtraction() throws Exception {
+        Method method = TestClass.class.getDeclaredMethod("testMethod");
+        int startLine = CompilerUtils.getStartLine(method);
+        int endLine = CompilerUtils.getEndLine(method);
+        Assertions.assertEquals(TEST_METHOD_SOURCE_LINES_START, startLine);
+        Assertions.assertEquals(TEST_METHOD_SOURCE_LINES_END, endLine);
+    }
+
+    @Test
+    public void testClassLinesExtraction() {
+        int startLine = CompilerUtils.getStartLine(TestClass.class);
+        int endLine = CompilerUtils.getEndLine(TestClass.class);
+        Assertions.assertEquals(TEST_CLASS_SOURCE_LINES_START, startLine);
+        Assertions.assertEquals(TEST_CLASS_SOURCE_LINES_END, endLine);
+    }
 }

src/main/java/datadog/compiler/AnnotationsInjectingClassVisitor.java

@@ -18,52 +18,82 @@ public class AnnotationsInjectingClassVisitor extends TreeScanner<Void, Void> {
     private final TreeMaker maker;
     private final Names names;
     private final JCTree.JCAnnotation sourcePathAnnotation;
-    private final JCTree.JCExpression methodLinesAnnotationType;
-    private final boolean methodAnnotationDisabled;
+    private final JCTree.JCExpression sourceLinesAnnotationType;
+    private final boolean sourceLinesAnnotationDisabled;
     private final LineMap lineMap;
     private final EndPosTable endPositions;
 
     AnnotationsInjectingClassVisitor(TreeMaker maker,
                                      Names names,
                                      JCTree.JCAnnotation sourcePathAnnotation,
-                                     JCTree.JCExpression methodLinesAnnotationType,
-                                     boolean methodAnnotationDisabled,
+                                     JCTree.JCExpression sourceLinesAnnotationType,
+                                     boolean sourceLinesAnnotationDisabled,
                                      LineMap lineMap,
                                      EndPosTable endPositions) {
         this.maker = maker;
         this.names = names;
         this.sourcePathAnnotation = sourcePathAnnotation;
-        this.methodLinesAnnotationType = methodLinesAnnotationType;
-        this.methodAnnotationDisabled = methodAnnotationDisabled;
+        this.sourceLinesAnnotationType = sourceLinesAnnotationType;
+        this.sourceLinesAnnotationDisabled = sourceLinesAnnotationDisabled;
         this.lineMap = lineMap;
         this.endPositions = endPositions;
     }
 
     @Override
     public Void visitClass(ClassTree node, Void aVoid) {
+        boolean sourcePathDetected = false;
+        boolean sourceLinesDetected = false;
         JCTree.JCClassDecl classDeclaration = (JCTree.JCClassDecl) node;
+
         for (JCTree.JCAnnotation annotation : classDeclaration.mods.annotations) {
-            if (annotation.annotationType.toString().endsWith("SourcePath")) {
-                // The method is already annotated with @SourcePath.
+            if (annotation.annotationType.toString().endsWith("SourceLines")) {
+                // The class is already annotated with @SourceLines.
                 // This can happen, for instance, when code-generation tools are used
                 // that copy annotations from interface to class
-                return super.visitClass(node, aVoid);
+                sourceLinesDetected = true;
             }
+            if (annotation.annotationType.toString().endsWith("SourcePath")) {
+                // The class is already annotated with @SourcePath.
+                sourcePathDetected = true;
+            }
+        }
+
+        if (node.getSimpleName().length() == 0) {
+            // Anonymous
+            return super.visitClass(node, aVoid);
         }
 
-        if (node.getSimpleName().length() > 0) {
+        if (!sourcePathDetected) {
             classDeclaration.mods.annotations = classDeclaration.mods.annotations.prepend(sourcePathAnnotation);
         }
+
+        if (!sourceLinesAnnotationDisabled && !sourceLinesDetected) {
+            JCTree.JCModifiers modifiers = classDeclaration.getModifiers();
+
+            int startPosition = modifiers.getStartPosition();
+            if (startPosition == Position.NOPOS) {
+                startPosition = classDeclaration.getStartPosition();
+            }
+
+            int endPosition = classDeclaration.getEndPosition(endPositions);
+            if (endPosition == Position.NOPOS) {
+                return super.visitClass(node, aVoid);
+            }
+
+            JCTree.JCAnnotation sourceLinesAnnotation = sourceLinesAnnotation(startPosition, endPosition);
+            classDeclaration.mods.annotations = classDeclaration.mods.annotations.prepend(sourceLinesAnnotation);
+        }
+
         return super.visitClass(node, aVoid);
     }
 
     public Void visitMethod(MethodTree node, Void aVoid) {
-        if (!methodAnnotationDisabled && (node instanceof JCTree.JCMethodDecl)) {
+        if (!sourceLinesAnnotationDisabled && (node instanceof JCTree.JCMethodDecl)) {
             JCTree.JCMethodDecl methodDecl = (JCTree.JCMethodDecl) node;
 
             for (JCTree.JCAnnotation annotation : methodDecl.mods.annotations) {
-                if (annotation.annotationType.toString().endsWith("MethodLines")) {
-                    // The method is already annotated with @MethodLines.
+                if (annotation.annotationType.toString().endsWith("SourceLines")) {
+                    // The method is already annotated with @SourceLines.
                     // This can happen, for instance, when code-generation tools are used
                     // that copy annotations from interface methods to class methods
                     return super.visitMethod(node, aVoid);
@@ -87,26 +117,30 @@ public Void visitMethod(MethodTree node, Void aVoid) {
                     }
                 }
 
-                int startLine = (int) lineMap.getLineNumber(startPosition);
-                int endLine = (int) lineMap.getLineNumber(endPosition);
-
-                Name startName = names.fromString("start");
-                JCTree.JCIdent startIdent = maker.Ident(startName);
-                JCTree.JCLiteral startLiteral = maker.Literal(startLine);
-                JCTree.JCAssign startAssign = maker.Assign(startIdent, startLiteral);
-
-                Name endName = names.fromString("end");
-                JCTree.JCIdent endIdent = maker.Ident(endName);
-                JCTree.JCLiteral endLiteral = maker.Literal(endLine);
-                JCTree.JCAssign endAssign = maker.Assign(endIdent, endLiteral);
-
-                JCTree.JCAnnotation methodLinesAnnotation = annotation(maker, methodLinesAnnotationType, startAssign, endAssign);
-                methodDecl.mods.annotations = methodDecl.mods.annotations.prepend(methodLinesAnnotation);
+                JCTree.JCAnnotation sourceLinesAnnotation = sourceLinesAnnotation(startPosition, endPosition);
+                methodDecl.mods.annotations = methodDecl.mods.annotations.prepend(sourceLinesAnnotation);
             }
         }
         return super.visitMethod(node, aVoid);
     }
 
+    private JCTree.JCAnnotation sourceLinesAnnotation(int startPosition, int endPosition) {
+        int startLine = (int) lineMap.getLineNumber(startPosition);
+        int endLine = (int) lineMap.getLineNumber(endPosition);
+
+        Name startName = names.fromString("start");
+        JCTree.JCIdent startIdent = maker.Ident(startName);
+        JCTree.JCLiteral startLiteral = maker.Literal(startLine);
+        JCTree.JCAssign startAssign = maker.Assign(startIdent, startLiteral);
+
+        Name endName = names.fromString("end");
+        JCTree.JCIdent endIdent = maker.Ident(endName);
+        JCTree.JCLiteral endLiteral = maker.Literal(endLine);
+        JCTree.JCAssign endAssign = maker.Assign(endIdent, endLiteral);
+
+        return annotation(maker, sourceLinesAnnotationType, startAssign, endAssign);
+    }
+
     private static JCTree.JCAnnotation annotation(TreeMaker maker, JCTree type, JCTree.JCExpression... arguments) {
         return maker.Annotation(type, List.from(arguments));
     }

src/main/java/datadog/compiler/DatadogCompilerPlugin.java

@@ -10,7 +10,7 @@
 
 public class DatadogCompilerPlugin implements Plugin {
 
-    static final String DISABLE_METHOD_ANNOTATION = "disableMethodAnnotation";
+    static final String DISABLE_SOURCE_LINES_ANNOTATION = "disableSourceLinesAnnotation";
 
     static {
         CompilerModuleOpener.setup();
@@ -30,8 +30,8 @@ public void init(JavacTask task, String... strings) {
             Context context = basicJavacTask.getContext();
 
             Collection<String> arguments = Arrays.asList(strings);
-            boolean methodAnnotationDisabled = arguments.contains(DISABLE_METHOD_ANNOTATION);
-            task.addTaskListener(new DatadogTaskListener(basicJavacTask, methodAnnotationDisabled));
+            boolean sourceLinesAnnotationDisabled = arguments.contains(DISABLE_SOURCE_LINES_ANNOTATION);
+            task.addTaskListener(new DatadogTaskListener(basicJavacTask, sourceLinesAnnotationDisabled));
 
             Log.instance(context).printRawLines(Log.WriterKind.NOTICE, NAME + " initialized");
         }

src/main/java/datadog/compiler/DatadogTaskListener.java

@@ -20,11 +20,11 @@
 
 final class DatadogTaskListener implements TaskListener {
     private final BasicJavacTask basicJavacTask;
-    private final boolean methodAnnotationDisabled;
+    private final boolean sourceLinesAnnotationDisabled;
 
-    DatadogTaskListener(BasicJavacTask basicJavacTask, boolean methodAnnotationDisabled) {
+    DatadogTaskListener(BasicJavacTask basicJavacTask, boolean sourceLinesAnnotationDisabled) {
         this.basicJavacTask = basicJavacTask;
-        this.methodAnnotationDisabled = methodAnnotationDisabled;
+        this.sourceLinesAnnotationDisabled = sourceLinesAnnotationDisabled;
     }
 
     @Override
@@ -54,7 +54,7 @@ public void finished(TaskEvent e) {
             JCTree.JCExpression sourcePathAnnotationType = select(maker, names, "datadog", "compiler", "annotations", "SourcePath");
             JCTree.JCAnnotation sourcePathAnnotation = annotation(maker, sourcePathAnnotationType, maker.Literal(sourcePath.toString()));
 
-            JCTree.JCExpression methodLinesAnnotationType = select(maker, names, "datadog", "compiler", "annotations", "MethodLines");
+            JCTree.JCExpression sourceLinesAnnotationType = select(maker, names, "datadog", "compiler", "annotations", "SourceLines");
 
             CompilationUnitTree compilationUnit = e.getCompilationUnit();
             LineMap lineMap = compilationUnit.getLineMap();
@@ -67,7 +67,7 @@ public void finished(TaskEvent e) {
             }
 
             AnnotationsInjectingClassVisitor treeVisitor = new AnnotationsInjectingClassVisitor(
-                    maker, names, sourcePathAnnotation, methodLinesAnnotationType, methodAnnotationDisabled, lineMap, endPositions);
+                    maker, names, sourcePathAnnotation, sourceLinesAnnotationType, sourceLinesAnnotationDisabled, lineMap, endPositions);
             compilationUnit.accept(treeVisitor, null);
 
         } catch (Throwable t) {