Плагин javac или добавляем Extension Methods в Java

от
Java    javac, compiler, extension methods

Screenshot from 2017-03-22 21-32-47.png
В Java 8 появилась возможность писать плагины к компилятору javac. С их помощью можно получать управление на нужном этапе компиляции и производить дополнительные проверки или изменения. Каждый плагин имеет название и может принимать аргументы для настройки своей работы.

Если сравнивать с процессорами аннотаций, то плагины к компилятору более гибкие и простые в использовании. Они не вызывают перекомпиляцию, если был сгенерирован какой-то класс, получить управление можно практически на любом этапе компиляции, вплоть до кодогенерации.


Настройка
Для реализации плагина нужно добавить в зависимости tools.jar, который находится в составе JDK.
  1. dependencies {
  2.     compile files("${System.env.JAVA_HOME}/lib/tools.jar")
  3. }

После этого создаём класс плагина com.example.javacplugin.helloworld.HelloWorldPlugin, который реализует интерфейс com.sun.source.util.Plugin. В методе getName задаём удобное имя, а метод init будет стартовой точкой входа нашего плагина.
  1. package com.example.javacplugin.helloworld;
  2.  
  3. import com.sun.source.util.JavacTask;
  4. import com.sun.source.util.Plugin;
  5.  
  6. public class HelloWorldPlugin implements Plugin {
  7.  
  8.     @Override
  9.     public String getName() {
  10.         return "Hello";
  11.     }
  12.  
  13.     @Override
  14.     public void init(JavacTask task, String... args) {
  15.         System.out.println("Hello, world!");
  16.     }
  17.  
  18. }

Далее, регистрируем плагин для загрузки ServiceLoader'ом, добавив в файл resources/META-INF/services/com.sun.source.util.Plugin строку: com.example.javacplugin.helloworld.HelloWorldPlugin.

Собираем проект, запускаем компилятор с параметрами -processorpath - путь к jar-файлу плагина и -Xplugin - имя плагина, которое мы задали в getName:

javac -processorpath plugin.jar -Xplugin:Hello Main.java

В результате получаем заветную строчку Hello, world!

Плагину можно передавать аргументы: -Xplugin:"Hello arg1 arg2", это может пригодиться, например, для передачи некоторого числа методов, превысив которые будет выдана ошибка компиляции.


Стадии компиляции
При компиляции java-исходника javac проходит следующие стадии:

1. Parse. Читаются исходные файлы, затем токены преобразуются в элементы абстрактного синтаксического дерева.
2. Enter. Строится таблица символов путём обхода дерева.
3. Process (annotations). Если в компиляции участвуют процессоры аннотаций, то происходит обработка аннотаций, затем, если сгенерировались новые классы, компилятор вновь переходит к стадии Parse.
4. Attribute. Синтаксические деревья помечаются атрибутами. На этом этапе происходит разрешение имён, проверка типов, и оптимизация свёртывания констант.
5. Flow. Выполняется анализ потока данных. Проверяется правильность присвоений, цепочек вызовов, а также достижимость кода.
6. Desugar. Происходит преобразование АСД с целью убрать синтаксический сахар. На этой стадии перечисления преобразуются в специальный класс, а вместо лямбд и ссылок на методы подставляются фабрики и прочие классы из пакета java.lang.invoke.
7. Generate. Генерируются class-файлы.

При написании плагина мы можем добавить обработчик перед началом и после стадий: Parse, Enter, Process, Flow, Generate.


Пример 1. Плагин подсчёта символов в исходнике
  1. public class CharactersCounterPlugin implements Plugin {
  2.  
  3.     @Override
  4.     public String getName() {
  5.         return "CharStat";
  6.     }
  7.  
  8.     @Override
  9.     public void init(JavacTask task, String... args) {
  10.         final CharactersCounterTask counterTask = new CharactersCounterTask(new ConcurrentHashMap<>());
  11.         task.addTaskListener(counterTask);
  12.     }
  13.  
  14. }
  15.  
  16. public class CharactersCounterTask implements TaskListener {
  17.  
  18.     private final Map<String, Integer> stat;
  19.  
  20.     public CharactersCounterTask(Map<String, Integer> stat) {
  21.         this.stat = stat;
  22.     }
  23.  
  24.     @Override
  25.     public void started(TaskEvent event) {
  26.         if (event.getKind() == TaskEvent.Kind.PARSE) {
  27.             final JavaFileObject source = event.getSourceFile();
  28.             final String name = source.getName();
  29.             try {
  30.                 stat.put(name, source.getCharContent(true).length());
  31.                 System.out.println(name + " processed");
  32.             } catch (IOException ex) {
  33.                 System.out.println(name + " fail");
  34.             }
  35.         }
  36.     }
  37.  
  38.     @Override
  39.     public void finished(TaskEvent event) {
  40.     }
  41. }

Плагин вызывается перед началом стадии Parse, считает количество символов в исходнике и добавляет эту информацию в глобальное хранилище, чтобы потом по окончанию компиляции вывести статистику. Хранилищем выступает обычный Map, а из объекта TaskEvent мы можем получить объект SourceFile, в котором уже будет путь к файлу и всё его содержимое. Остатся только положить данные в Map:
  1. String name = source.getName();
  2. stat.put(name, source.getCharContent(true).length());

Статистику покажем в самом конце, после того как будет выполнена последняя стадия Generate:
  1. public class CharactersCounterTask implements TaskListener {
  2.  
  3.     private Consumer<Map<String, Integer>> consumer;
  4.  
  5.     public void onComplete(Consumer<Map<String, Integer>> consumer) {
  6.         this.consumer = consumer;
  7.     }
  8.  
  9.     //...
  10.  
  11.     @Override
  12.     public void finished(TaskEvent event) {
  13.         if (event.getKind() == TaskEvent.Kind.GENERATE) {
  14.             if (consumer != null) {
  15.                 consumer.accept(stat);
  16.                 consumer = null;
  17.             }
  18.         }
  19.     }
  20. }
  21.  
  22. // CharactersCounterPlugin init
  23. counterTask.onComplete(classCharsCount -> {
  24.     System.out.println("Characters count statistics: ");
  25.     System.out.format("\t processed %d files%n", classCharsCount.size());
  26.     System.out.format("\t summary %d chars%n", classCharsCount.values().stream()
  27.             .mapToInt(Integer::intValue)
  28.             .sum());
  29. });

Выполнив компиляцию на каком-нибудь тестовом примере, получим подобный результат:
  1. $ ./gradlew :testcharstat:build
  2. ...
  3. :testcharstat:compileJava
  4. ../charstat/Main.java processed
  5. ../charstat/Test.java processed
  6. Characters count statistics:
  7.          processed 2 files
  8.          summary 239 chars
  9. :testcharstat:processResources UP-TO-DATE
  10. :testcharstat:classes
  11. ...


Пример 2. Плагин подсчёта количества методов
Плагин подсчитывает количество методов в каждом из исходников, а потом выводит общее число скомпилированных методов и названия классов в порядке убывания числа методов.
  1. public class MethodsCounterTask implements TaskListener {
  2.  
  3.     @Override
  4.     public void started(TaskEvent taskEvent) { }
  5.  
  6.     @Override
  7.     public void finished(TaskEvent event) {
  8.         if (event.getKind() == TaskEvent.Kind.PARSE) {
  9.             CompilationUnitTree compilationUnit = event.getCompilationUnit();
  10.             MutableInteger counter = new MutableInteger();
  11.             compilationUnit.accept(new MethodsVisitor(), counter);
  12.             stat.put(event.getSourceFile().getName(), counter.get());
  13.         }
  14.     }
  15. }
  16.  
  17. public class MethodsVisitor extends TreeScanner<Void, MutableInteger> {
  18.  
  19.     @Override
  20.     public Void visitMethod(MethodTree methodTree, MutableInteger input) {
  21.         input.increment();
  22.         return super.visitMethod(methodTree, input);
  23.     }
  24. }

Обработка происходит после завершения стадии Parse, когда у нас уже есть готовое АСД. Запускается MethodVisitor, который увеличивает значение счётчика каждый раз при посещении метода. Затем, как и в предыдущем примере, статистика помещается в Map.

На тестовом примере получаем следующий вывод:
  1. $ ./gradlew :testmethodscount:build
  2. ...
  3. :testmethodscount:compileJava
  4. Methods count statistics:
  5.          processed 2 files
  6.          summary 15 methods
  7. ../methodscount/Test.java: 10
  8. ../methodscount/Main.java: 5
  9. :testmethodscount:processResources UP-TO-DATE
  10. ,,,


Пример 3. Операторы доступа к Map как к массиву
Плагин добавляет возможность индексации Map по стороковому ключу, как будто это массив:
  1. Map<String, String> map = new HashMap<>();
  2. map["key"] = "ten"; // map.put("key", "ten")
  3. System.out.println(map["key"]); // map.get("key")

  1. public class MapAccessOperatorTask implements TaskListener {
  2.  
  3.     private final MapAccessReplacer replacer;
  4.  
  5.     public MapAccessOperatorTask(JavacTask task) {
  6.         this.replacer = new MapAccessReplacer(task);
  7.     }
  8.  
  9.     @Override
  10.     public void started(TaskEvent taskEvent) { }
  11.  
  12.     @Override
  13.     public void finished(TaskEvent event) {
  14.         if (event.getKind() == TaskEvent.Kind.PARSE) {
  15.             CompilationUnitTree compilationUnit = event.getCompilationUnit();
  16.             ((JCTree.JCCompilationUnit) compilationUnit).accept(replacer);
  17.         }
  18.     }
  19. }
  20.  
  21. public class MapAccessReplacer extends TreeTranslator {
  22.  
  23.     @Override
  24.     public void visitAssign(JCAssign tree) {
  25.         if (tree.getVariable().getKind() == Tree.Kind.ARRAY_ACCESS) {
  26.             JCArrayAccess arrayAccess = (JCArrayAccess) tree.getVariable();
  27.             if (arrayAccess.getIndex().getKind() == Tree.Kind.STRING_LITERAL) {
  28.                 JCExpression methodSelect = make.Select(arrayAccess.getExpression(), names.fromString("put"));
  29.                 result = make.Apply(List.nil(), methodSelect, List.of(arrayAccess.getIndex(), tree.getExpression()));
  30.                 return;
  31.             }
  32.         }
  33.         super.visitAssign(tree);
  34.     }
  35.  
  36.     @Override
  37.     public void visitIndexed(JCArrayAccess tree) {
  38.         if (tree.getIndex().getKind() == Tree.Kind.STRING_LITERAL) {
  39.             JCExpression methodSelect = make.Select(tree.getExpression(), names.fromString("get"));
  40.             result = make.Apply(List.nil(), methodSelect, List.of(tree.getIndex()));
  41.             return;
  42.         }
  43.         super.visitIndexed(tree);
  44.     }
  45. }

Здесь используется обход дерева с модификацией.
В первом случае, если в выражении присваивания JCAssign используется присваивание некоторого результата элементу массива, например x["key"] = ..., и ключом выступает строковый литерал, то мы преобразуем дерево так, чтобы на переменной x вызывался метод put с ключом key и значением, равным выражению, которое было после знака равно.
  1. // Было
  2. x["key"] = expr;
  3. // Стало
  4. x.put("key", expr)

Во втором случае, если массив индексируется строковым литералом, к примеру s = x["key"], заменяем это выражение на вызов метода get.
  1. // Было
  2. s = x["key"];
  3. // Стало
  4. s = x.get("key")

Кстати, плагин не ограничивает нас только классами Map, можно использовать любые объекты, у которых есть методы put или get.
  1. class Test {
  2.     public int get(String s) {
  3.         System.out.println("get " + s);
  4.         return 10;
  5.     }
  6. }
  7.  
  8. Test test = new Test();
  9. int x = test["5"];


Пример 4. Перегрузка операторов
Плагин позволяет вместо методов add, subtract, multiply и divide использовать операторы +, -, *, / соответственно.
  1. public final class OperatorOverloadingTask implements TaskListener {
  2.  
  3.     private final OperatorOverloadingVisitor visitor;
  4.  
  5.     public OperatorOverloadingTask(JavacTask task) {
  6.         this.visitor = new OperatorOverloadingVisitor(task);
  7.     }
  8.  
  9.     @Override
  10.     public void started(TaskEvent event) {
  11.         if (event.getKind() == TaskEvent.Kind.ANALYZE) {
  12.             CompilationUnitTree compilationUnit = event.getCompilationUnit();
  13.             ((JCTree.JCCompilationUnit) compilationUnit).accept(visitor);
  14.         }
  15.     }
  16.  
  17.     @Override
  18.     public void finished(TaskEvent event) { }
  19. }
  20.  
  21. public final class OperatorOverloadingVisitor extends TreeTranslator {
  22.  
  23.     public OperatorOverloadingVisitor(JavacTask task) {
  24.         operators = new EnumMap<>(Tree.Kind.class);
  25.         operators.put(Tree.Kind.PLUS, names.fromString("add"));
  26.         operators.put(Tree.Kind.MINUS, names.fromString("subtract"));
  27.         operators.put(Tree.Kind.MULTIPLY, names.fromString("multiply"));
  28.         operators.put(Tree.Kind.DIVIDE, names.fromString("divide"));
  29.     }
  30.  
  31.     @Override
  32.     public void visitBinary(JCTree.JCBinary bt) {
  33.         super.visitBinary(bt);
  34.         if (!operators.containsKey(bt.getKind())) return;
  35.  
  36.         JCExpression methodCall;
  37.         final JCExpression left = bt.getLeftOperand();
  38.         switch (left.getKind()) {
  39.             case IDENTIFIER:
  40.                 JCIdent lIdent = (JCIdent) left;
  41.                 methodCall = treeMaker.Ident(lIdent.getName());
  42.                 break;
  43.             case METHOD_INVOCATION:
  44.                 methodCall = left;
  45.                 break;
  46.             default:
  47.                 return;
  48.         }
  49.  
  50.         methodCall = treeMaker.Select(methodCall, operators.get(bt.getKind()));
  51.         final JCMethodInvocation app = treeMaker.Apply(
  52.                 List.nil(), methodCall, List.of(bt.getRightOperand())
  53.         );
  54.         result = app;
  55.     }
  56. }

Здесь для бинарной операции производится замена поддерживаемых операторов на вызов метода:
  1. // Было
  2. t = x + y * z()
  3. // Стало
  4. t = x.add(y.multiply(z()))

Отличной демонстрацией может послужить класс BigInteger:
  1. BigInteger x1 = BigInteger.TEN;
  2. BigInteger x2 = BigInteger.valueOf(120);
  3. BigInteger x3 = x1 + x2;
  4. // BigInteger x3 = x1.add(x2);
  5. BigInteger x4 = BigInteger.valueOf(2) * x3;
  6. // BigInteger x4 = BigInteger.valueOf(2).multiply(x3);
  7. BigInteger x5 = x4 - x2 / x1 + x3;
  8. // BigInteger x5 = x4.subtract(x2.divide(x1)).add(x3);

Однако стоит обратить внимание, что проверка типов не делается, так что на более серьёзном коде плагин может заменить не то, что нужно.


Пример 5. Поддержка методов-расширений
Наконец, последний пример - плагин, добавляющий поддержку extension-методов. С его помощью можно добавить метод в любой класс так, как будто он там действительно есть:
  1. "string".reverse();
  2. random.nextInt(0, 100);
  3. stream.sortBy(Person::name).forEachIndexed((i, person) -> ...);

Сперва обусловимся, что искать методы-расширения будем в том же классе, в котором использовать их. Если метод имеет модификаторы public static и хотя бы один аргумент, то он становится кандидатом в метод-расширение.

Модификация AST будет при всё том же обходе дерева. Для начала соберём всех кандидатов при посещении класса:
  1. @Override
  2. public Void visitClass(ClassTree ct, Void p) {
  3.     final TypeElement currentClass = (TypeElement) TreeInfo.symbolFor((JCTree) ct);
  4.     extensionMethods = Optional.ofNullable(currentClass)
  5.             .map(TypeElement::getEnclosedElements)
  6.             .map(enclosedElements -> enclosedElements
  7.                     .stream()
  8.                     .filter(e -> e.getModifiers().containsAll(PUBLIC_STATIC))
  9.                     .filter(e -> e.getKind() == ElementKind.METHOD)
  10.                     .map(ExecutableElement.class::cast)
  11.                     .filter(e -> !e.getParameters().isEmpty())
  12.                     .map(ExtensionMethod::new)
  13.                     .collect(Collectors.toCollection(ExtensionMethods::new))
  14.             )
  15.             .orElseGet(ExtensionMethods::new);
  16.  
  17.     return super.visitClass(ct, p);
  18. }

Берём все элементы, которые находятся в классе, затем:
  - фильтруем сначала те, которые имеют модификатор public static
  - затем, фильтруем сами методы
  - каждый оставшийся элемент, это гарантированно элемент-метод, преобразуем к типу ExecutableElement
  - если метод не принимает параметров - он нам не подходит,
  - оборачиваем элемент в собственную обёртку - класс ExtensionMethod
  - наконец, собираем все оставшиеся методы в список.
  1. public final class ExtensionMethods extends ArrayList<ExtensionMethod> {
  2.  
  3.     public Stream<String> names() {
  4.         return stream().map(ExtensionMethod::getName);
  5.     }
  6. }

Следующий этап, обработать вызов метода:
  1. @Override
  2. public Void visitMethodInvocation(MethodInvocationTree tree, Void p) {
  3.     super.visitMethodInvocation(tree, p);
  4.     if (!extensionMethods.isEmpty()) {
  5.         processMethodInvocation((JCMethodInvocation) tree);
  6.     }
  7.     return p;
  8. }
  9.  
  10. private void processMethodInvocation(JCMethodInvocation tree) {
  11.     if (tree == null) return;
  12.     if (tree.meth.getKind() != Tree.Kind.MEMBER_SELECT) return;
  13.  
  14.     final String methodName = ((JCFieldAccess) tree.meth).name.toString();
  15.     if (extensionMethods.names().noneMatch(m -> m.equals(methodName))) return;
  16.  
  17.     final JCExpression receiver = ((JCFieldAccess) tree.meth).getExpression();
  18.  
  19.     extensionMethods.stream()
  20.             .filter(m -> methodName.equals(m.getName()))
  21.             .filter(m -> m.getParametersCount() == tree.getArguments().length() + 1)
  22.             .filter(m -> checkTypes(m, receiver))
  23.             .findAny()
  24.             .ifPresent(me -> {
  25.                 tree.args = tree.args.prepend(receiver);
  26.  
  27.                 Symbol symbol = (Symbol) me.getElement();
  28.                 tree.meth = make.Ident(symbol);
  29.                 System.out.println(methodName + " processed");
  30.             });
  31. }
  32.  
  33. private boolean checkTypes(ExtensionMethod m, JCExpression receiver) {
  34.     attr.attribExpr(receiver, methodEnv);
  35.     TypeMirror type = receiver.type;
  36.     TypeMirror param = typeUtils.erasure(m.getElement().getParameters().get(0).asType());
  37.     if (type == null || param == null) return false;
  38.     return typeUtils.isAssignable(param, type);
  39. }

Если мы встретили вызов метода, проверяем, есть ли extension-метод с таким именем. Если есть, проверяем, соответствует ли он текущему вызову. Для этого у метода должны совпадать количество аргументов и их типы. Типы сравниваются по стёртому типу (type erasure) и присваиваемости (isAssignable), а значит не обязательно должны строго совпадать. Если такой метод найден, проводим замену:
  1. // Было
  2. "string".method(1, 2, x.ext())
  3. // Стало
  4. method("string", 1, 2, ext(x))

И теперь можно встраивать нужные методы. Например, добавить в Stream API sortBy и filterNot:
  1. public static IntStream filterNot(IntStream stream, IntPredicate predicate) {
  2.     return stream.filter(predicate.negate());
  3. }
  4.  
  5. public static <T, R extends Comparable<? super R>> Stream<T> sortBy(Stream<T> stream, Function<? super T, ? extends R> f) {
  6.     return stream.sorted((o1, o2) -> f.apply(o1).compareTo(f.apply(o2)));
  7. }
  8.  
  9. apps.stream()
  10.         .filterNot(Application::isInstalled)
  11.         .sortBy(Application::getDownloadingTime())

Вот так, с помощью нехитрых модификаций АСД, Java можно превратить в Kotlin... Но зачем?!

Исходный код: GitHub
  • +4
  • views 6440