На собеседованиях и литкоде любят вращать бинарные деревья. Но что насчёт трансформации обычного дерева в другое? Как решить эту задачу, и какие могут быть подходы? Рассмотрим на опыте трансляции одного синтаксического в другое, чтобы разобраться.

Как мы к этому пришли

Причины трансформировать одно дерево в другое придумать не очень сложно:

• Конвертация между разными форматами, например XML в HTML.

• Алгоритмические задачи, включая пресловутое вращение бинарных деревьев.

• Преобразования внутри компиляторов. Не обязательно дерево в дерево, а, например, синтаксическое дерево в граф потока управления. Зачем это может быть нужно, я уже писал.

Далее буду использовать более точный термин "трансляция" вместо трансформации.

Мы в компании PVS-Studio занимаемся разработкой статических анализаторов кода, так что как раз пересекаемся с компиляторами по инструментарию. В числе прочего, мы также работаем с абстрактным синтаксическим деревом (AST). Но все примеры будут упрощёнными, так что глубоко разбираться в теме не надо.

Сама задача родилась хитрым образом:

1. Нам хочется поддерживать больше языков, пока умеем только в C#, C/C++ и Java. А хочется, в том числе, JavaScript.

2. Нам, Java-команде, писать новый анализатор хочется на (внезапно) Java.

3. Самую полноценную поддержку js, а также поддержку ts на будущее даёт компилятор TypeScript. Но написан он на другом языке, так что AST приходится передавать по gRPC. Чтобы это было максимально быстро, используется protobuf.

И как же мы вообще пришли к трансляции?

• Представление protobuf, которое мы получаем, анализировать невозможно, так как структура неизменяема и объекты полностью теряют иерархию.

• Хорошей идеей кажется иметь свою обобщённую модель. Тогда мы не зависим от разбирающего код фронтенда, а также упрощаем поддержку новых языков.

Так и родилась наша задача: транслировать дерево ts-компилятора, пережёванное protobuf, в своё представление.

Что транслируем

Так что конкретно мы делаем? Опустим (пока) проблемы с protobuf — нам надо из одного дерева сделать другое. Они оба синтаксические и оба для одного языка, так что разница между ними небольшая. Большинство вершин будут отображаться изоморфно 1 к 1, имея только косметические отличия. Допустим, в одной модели может использоваться именование do...while, а в другой — repeat...until.

Т.е. в данном случае у нас преобразование оставляет структуру "как есть", мы её просто дублируем, но могут быть ситуации сложнее. В том же JavaScript есть интерполяция строк, там называемая шаблонными строками.

Если идти дальше, то почему-то в дереве TypeScript выражения внутри интерполяции и текст совмещаются в одну вершину, хотя в других языках интерполяция состоит из плоского списка.

Чтобы получить дерево справа, нам нужно из свойств вершины TemplateSpan сделать два (или меньше) узла, проведя разбиение. Так что, помимо прямого преобразования узла в узел, мы получаем преобразование из свойства узла в узел.

Кстати, обратная операция называлась бы агрегацией — объединение нескольких узлов в один — но в нашей задаче ей места не нашлось.

Пока всё довольно прямолинейно, но что, если исходная структура совсем не совпадает с целевой? В контексте AST мы можем захотеть нормализовать структуру, чтобы в дальнейшем нам было проще с ней работать. Например, всегда добавлять тело в стрелочные функции:

let a = () -> 1; // source
let a = () -> { return 1 }; // normalized

В данном случае это просто произвольное правило, вынуждающее нас добавлять узлы в дерево, которых там до этого не было.

Итого мы имеем:

1. Преобразование из узла в узел: прямое дублирование, максимум меняется идентификатор.

2. Преобразование из свойства в узел: создание дополнительных элементов дерева из свойств исходных узлов.

3. Набор произвольных правил для нормализации, смысл которых, как и эмоции от них, можно обобщить картинкой ниже.

Все эти ситуации мы рассматривали изолированно, но синтаксическое дерево строится целиком на весь файл с исходным кодом, так что на практике всё будет не просто вперемешку, а ещё и в огромном масштабе. Если только в одном объявлении и инициализации let a = 2 + 5 минимум 4 вершины, страшно представить, сколько их будет в файле на тысячи строк кода.

Надеюсь, к этому моменту у вас уже сформировалось впечатление касательно того, что мы пытаемся сделать, потому что пора переходить к практике.

Как залезть на дерево

Обход

Прежде чем определиться с тем, как делать такие преобразования, надо понять, как мы будем перемещаться по исходному дереву.

Первым в голову может прийти реализация итератора по дереву. Я как-то даже писал статью про то, как он может помочь при обходе AST. И в целом, если ваша цель — транслировать дерево с небольшим количеством типов узлов, то это наверняка правильный выбор, но не в нашем случае.

Я быстро посчитал количество разных типов вершин на рисунках выше, и оно там в районе десяти — а это буквально три примера. В нашем AST для JavaScript на данный момент немногим менее 100 типов, а в исходной модели их ещё больше. А каждый тип — дополнительный if при решении, как узел обрабатывать. Так что если нам хочется писать код, который можно поддерживать без слёз, то лучше обратиться к паттерну "посетитель".

На его реализации останавливаться не буду — можно ознакомиться с ней по ссылке выше или на другом любом ресурсе — вместо этого просто скажу, что за счёт двойной диспетчеризации количество условий при работе с разными типами узлов почти сводится к нулю. Также на нём довольно просто сделать рекурсивный спуск. Пример под спойлером.

interface Node {
    void accept(Visitor v);
}

interface Visitor {
    void visitBinary(Binary n);
    void visitLiteral(Literal n);
}

class Binary implements Node {
    private final Node left;
    private final Node right;

    public Binary(Node left, Node right) {
        this.left = left;
        this.right = right;
    }

    @Override public void accept(Visitor v) {
        v.visitBinary(this);
    }
}

class Literal implements Node {
    private final int value;

    public Literal(int value) {
        this.value = value;
    }

    @Override public void accept(Visitor v) {
        v.visitLiteral(this);
    }
}

class Scanner implements Visitor {
    public void scan(Node n) {
        n.accept(this); 
    }

    @Override public Binary visitBinary(Binary n) {
        scan(n.getLeft());
        scan(n.getRight());
    }

    @Override public void visitLiteral(Literal n) {
        
    }
}

А как же protobuf?

Внимательный читатель заметит, что модель для обхода нам пережевал protobuf, так о какой двойной диспетчеризации вообще может идти речь? Ни о какой. Почти.

Protobuf действительно причинил боль заставил пораскинуть мозгами. Но по итогу получилось эмулировать посетителя через кодогенерацию обходчика на основе .protoset файла, в котором описаны типы всех узлов. visit-ы у такого посетителя вполне обычные:

@Override
public void visitBinaryExpression(BinaryExpression binaryExpression) {
    scan(binaryExpression.getToken());
    scan(binaryExpression.getLeft());
    scan(binaryExpression.getRight());
}

Достигается это за счёт невинной реализации scan:

public final void scan(Object node) {
    switch (node) {
      case BinaryExpression binaryExpression ->
          visitBinaryExpression(binaryExpression);
      case UnaryExpression unaryExpression -> 
          visitUnaryExpression(unaryExpression);
      case LiteralExpression literalExpression -> 
          visitLiteralExpression(literalExpression);
      // 121 LOC more
    }
}

Элегантным решение назвать сложно, но свою задачу оно выполняет. Так что сойдёмся на следующем:

1. Посетитель у нас и правда немного необычный.

2. Принцип работы с ним прежний, так как возможность писать логику в visit-ы без написанных руками if-ов у нас сохранилась.

Трансляция деревьев

Со сценариями трансляции и с обходом разобрались. Пора наконец переходить к самой трансляции.

Узел-узел

Разберёмся с самым простым. Если мы "строим" дерево, то нам его нужно где-то хранить. Самый легкий способ — возвращать из каждого visit-а текущий узел, а также передавать в аргументах родителя. Если мы наследуемся от сканера выше, перед этим добавив ему обобщённые возвращаемое значение и второй аргумент в visit-ы, то мы получим что-то вроде этого:

interface Visitor<R, P> {
    R visitBinary(Binary n, P parent);
    R visitLiteral(Literal n, P parent);
}

class Builder extends Scanner<JsNode, JsNode> {
    @Override
    public JsNode visitBinary(Binary n, JsNode parent) {
        var left = (JsExpression) scan(n.getLeft());
        var right = (JsExpression) scan(n.getRight());
        return new JsBinary(left, right);
    }

    @Override 
    public JsNode visitLiteral(Literal n, JsNode parent) {
        return new JsLiteral(n.getValue());
    }
}

Таким крайне нехитрым образом мы можем дублировать структуру дерева. Чтобы не запутаться, узлы целевого дерева я помечаю префиксом Js в начале имени класса.

Свойство-узел

Этот случай ненамного сложнее, но тут нам как раз понадобится переданный родитель. Вместо копирования в лоб, мы добавляем немного логики. Возвращаясь к тому же примеру с интерполяцией:

@Override
public JsNode visitTemplateSpan(TemplateSpan templateSpan, JsNode parent) {
    var interpolation = (Interpolation) parent;
    var expr = new InterpolationExpression();
    interpolation.add(expr);
    interpolation.setExpression(
        scan(templateSpan.getExpression())
    );

    // tail is optional
    if (templateSpan.hasTemplateTail()) {
        interpolation.add(new InterpolationText(
            templateSpan.getTemplateTail().getValue()
        ));
    }

    return null;
}

Так как мы пролезли поперёк TemplateSpan, то возвращаем null — мы уже присоединили потомков через родителя, переданного сюда в аргументе.

Нормализация

Свои эмоции в прошлый раз я не просто так выразил такой картинкой — тут случай каждый раз уникальный. Иногда везёт, и нормализацию получается сделать бескровно, иногда приходится заморачиваться. В целом, за правило большого пальца мы взяли себе следующее: нормализацию не делать, если это меняет поток управления или это слишком сложно. При этом понятие "слишком сложно" остаётся на совести разработчика и его менеджера :)

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

Сам visit выглядит примерно так:

@Override
public JsNode visitArrowFunction(
    ArrowFunction arrowFunction,
    JsNode parent
) {
    var func = new JsArrowFunction();
    // link inside
    scan(arrowFunction.getParameterDeclarationList(), func);

    // OK
    if (arrowFunction.hasBlock()) {
        func.setBlock(
            scan(arrowFunction.getBlock())
        );
    } else if (arrowFunction.hasExpression()) {
        var normalized = normalizeArrow(
            arrowFunction
        );
        func.setBlock(normalized);
    }

    return func;
}

Трансформация спрятана в normalizeArrow:

public JsBlock normalizeArrow(ArrowFunction f) {
    var block = new JsCodeBlock();
    block.setFlag(Flag.SYNTHETIC);

    var ret = new JsReturn();
    ret.setFlag(Flag.SYNTHETIC);
    block.addStatement(ret);
    
    var expr = scan(f.getExpression());
    ret.setExpression(expr);

    return block;
}

Помимо самого создания синтетических узлов, мы также помечаем, что они "виртуальные", так как в исходном дереве их не было совсем. Это полезно при статическом анализе — при эвристиках в диагностических правилах мы сможем учитывать, что это не разработчик написал вещи в духе (var) -> { return console.log(var) }, а это просто артефакт наших преобразований. В остальном с преобразованиями на этом всё.

Слон в комнате

К этому моменту мы разобрали все основные принципы, связанные с решением поставленной выше задачи:

• обход ведём рекурсивно, сохраняя текущее состояние дерева через параметры и возвращаемые значения visit-ов;

• имеем две типовые трансформации: узел-узел и свойство-узел;

• имеем одну нетиповую: нормализация дерева;

• по итогу трансляции мы получаем похожее дерево, но новое, отличающееся в деталях.

Осталось только решить саму задачу. Напомню, что в целевом дереве в районе 100 типов узлов, в исходном — 263 типа. Такая разница во многом обусловлена спецификой protobuf. Это значит, что придётся писать руками visit-ы и scan-ы сотни и тысячи раз. Как вы там, готовы? Я — нет.

Решение проблемы, на самом деле, уже появлялось в тексте. Точно так же, как исходный обходчик можно было сгенерировать на основе .protoset файла, так и типовые правила трансляции можно делать при помощи кодогенерации — на основе какого-либо конфигурационного правила. Но об этом, как и о том, как в это втиснуть нетиповую нормализацию, мы поговорим как-нибудь в другой раз.

Послесловие

Вот с таким кейсом мне довелось столкнуться в своей работе. Хотя я и пробежал по верхам, получилось затронуть несколько интересных моментов: типовые преобразования, выбор между итератором и посетителем, неортодоксальная реализация посетителя, призма статического анализа (синтетические узлы, подход к нормализации).

В любом случае, буду рад, если вы подчерпнули для себя что-то новое по работе с деревьями. Чтобы не пропускать статьи по этой теме или новым анализаторам в целом, вы можете подписаться на Твиттер PVS-Studio и наш ежемесячный дайджест статей. Сам инструмент можно бесплатно попробовать здесь.

А если вам хочется следить за моими публикациями, то можете подписаться на личный блог.

Если хотите поделиться этой статьей с англоязычной аудиторией, то прошу использовать ссылку на перевод: Konstantin Volohovsky. How to copy a tree, but not word for word.