<commit_before>package com.facebook.presto.sql.tree;

import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;

import static com.facebook.presto.util.IterableUtils.sameElements;

public final class TreeRewriter<C>
{
    private final NodeRewriter<C> nodeRewriter;
    private final AstVisitor<Node, TreeRewriter.Context<C>> visitor;

    public static <C, T extends Node> T rewriteWith(NodeRewriter<C> rewriter, T node)
    {
        return new TreeRewriter<C>(rewriter).rewrite(node, null);
    }

    public TreeRewriter(NodeRewriter<C> nodeRewriter)
    {
        this.nodeRewriter = nodeRewriter;
        this.visitor = new RewritingVisitor();
    }

    public <T extends Node> T rewrite(T node, C context)
    {
        return (T) visitor.process(node, new Context<>(context, false));
    }

    /**
     * Invoke the default rewrite logic explicitly. Specifically, it skips the invocation of the node rewriter for the provided node.
     */
    public <T extends Node> T defaultRewrite(T node, C context)
    {
        return (T) visitor.process(node, new Context<>(context, true));
    }

    public static <C, T extends Node> Function<Node, T> rewriteFunction(final NodeRewriter<C> rewriter)
    {
        return new Function<Node, T>()
        {
            @Override
            public T apply(Node node)
            {
                return (T) rewriteWith(rewriter, node);
            }
        };
    }

    private class RewritingVisitor
            extends AstVisitor<Node, TreeRewriter.Context<C>>
    {
        @Override
        public Node visitNode(Node node, Context<C> context)
        {
            if (!context.isDefaultRewrite()) {
                Node result = nodeRewriter.rewriteNode(node, context.get(), TreeRewriter.this);
                if (result != null) {
                    return result;
                }
            }

            throw new UnsupportedOperationException("not yet implemented: " + getClass().getSimpleName() + " for " + node.getClass().getName());
        }

        @Override
        public Node visitQuery(Query node, Context<C> context)
        {
            if (!context.isDefaultRewrite()) {
                Node result = nodeRewriter.rewriteQuery(node, context.get(), TreeRewriter.this);
                if (result != null) {
                    return result;
                }
            }

            Select select = rewrite(node.getSelect(), context.get());

            ImmutableList.Builder<Relation> from = ImmutableList.builder();
            for (Relation relation : node.getFrom()) {
                from.add(rewrite(relation, context.get()));
            }

            Optional<Expression> where = null;
            if (node.getWhere().isPresent()) {
                where = Optional.fromNullable(rewrite(node.getWhere().get(), context.get()));
            }

            ImmutableList.Builder<Expression> groupBy = ImmutableList.builder();
            for (Expression expression : node.getGroupBy()) {
                groupBy.add(rewrite(expression, context.get()));
            }

            Optional<Expression> having = null;
            if (node.getHaving().isPresent()) {
                having = Optional.fromNullable(rewrite(node.getHaving().get(), context.get()));
            }

            ImmutableList.Builder<SortItem> orderBy = ImmutableList.builder();
            for (SortItem sortItem : node.getOrderBy()) {
                orderBy.add(rewrite(sortItem, context.get()));
            }

            if (select != node.getSelect() ||
                    !sameElements(node.getFrom(), from.build()) ||
                    where != node.getWhere() ||
                    !sameElements(node.getGroupBy(), groupBy.build()) ||
                    having.orNull() != node.getHaving().orNull() ||
                    !sameElements(orderBy.build(), node.getOrderBy())) {
                return new Query(select, from.build(), where, groupBy.build(), having, orderBy.build(), node.getLimit());
            }

            return node;
        }

        @Override
        public Node visitSelect(Select node, Context<C> context)
        {
            if (!context.isDefaultRewrite()) {
                Node result = nodeRewriter.rewriteSelect(node, context.get(), TreeRewriter.this);<commit_msg>Fix potential NPE and improve clarity<commit_after>package com.facebook.presto.sql.tree;

import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;

import static com.facebook.presto.util.IterableUtils.sameElements;

public final class TreeRewriter<C>
{
    private final NodeRewriter<C> nodeRewriter;
    private final AstVisitor<Node, TreeRewriter.Context<C>> visitor;

    public static <C, T extends Node> T rewriteWith(NodeRewriter<C> rewriter, T node)
    {
        return new TreeRewriter<C>(rewriter).rewrite(node, null);
    }

    public TreeRewriter(NodeRewriter<C> nodeRewriter)
    {
        this.nodeRewriter = nodeRewriter;
        this.visitor = new RewritingVisitor();
    }

    public <T extends Node> T rewrite(T node, C context)
    {
        return (T) visitor.process(node, new Context<>(context, false));
    }

    /**
     * Invoke the default rewrite logic explicitly. Specifically, it skips the invocation of the node rewriter for the provided node.
     */
    public <T extends Node> T defaultRewrite(T node, C context)
    {
        return (T) visitor.process(node, new Context<>(context, true));
    }

    public static <C, T extends Node> Function<Node, T> rewriteFunction(final NodeRewriter<C> rewriter)
    {
        return new Function<Node, T>()
        {
            @Override
            public T apply(Node node)
            {
                return (T) rewriteWith(rewriter, node);
            }
        };
    }

    private class RewritingVisitor
            extends AstVisitor<Node, TreeRewriter.Context<C>>
    {
        @Override
        public Node visitNode(Node node, Context<C> context)
        {
            if (!context.isDefaultRewrite()) {
                Node result = nodeRewriter.rewriteNode(node, context.get(), TreeRewriter.this);
                if (result != null) {
                    return result;
                }
            }

            throw new UnsupportedOperationException("not yet implemented: " + getClass().getSimpleName() + " for " + node.getClass().getName());
        }

        @Override
        public Node visitQuery(Query node, Context<C> context)
        {
            if (!context.isDefaultRewrite()) {
                Node result = nodeRewriter.rewriteQuery(node, context.get(), TreeRewriter.this);
                if (result != null) {
                    return result;
                }
            }

            Select select = rewrite(node.getSelect(), context.get());

            ImmutableList.Builder<Relation> from = ImmutableList.builder();
            for (Relation relation : node.getFrom()) {
                from.add(rewrite(relation, context.get()));
            }

            Expression where = null;
            if (node.getWhere().isPresent()) {
                where = rewrite(node.getWhere().get(), context.get());
            }

            ImmutableList.Builder<Expression> groupBy = ImmutableList.builder();
            for (Expression expression : node.getGroupBy()) {
                groupBy.add(rewrite(expression, context.get()));
            }

            Expression having = null;
            if (node.getHaving().isPresent()) {
                having = rewrite(node.getHaving().get(), context.get());
            }

            ImmutableList.Builder<SortItem> orderBy = ImmutableList.builder();
            for (SortItem sortItem : node.getOrderBy()) {
                orderBy.add(rewrite(sortItem, context.get()));
            }

            if (select != node.getSelect() ||
                    !sameElements(node.getFrom(), from.build()) ||
                    where != node.getWhere().orNull() ||
                    !sameElements(node.getGroupBy(), groupBy.build()) ||
                    having != node.getHaving().orNull() ||
                    !sameElements(orderBy.build(), node.getOrderBy())) {

                return new Query(select,
                        from.build(),
                        Optional.fromNullable(where),
                        groupBy.build(),
                        Optional.fromNullable(having),
                        orderBy.build(),
                        node.getLimit());
            }

            return node;
        }

        @Override
        public Node visitSelect(Select node, Context<C> context)
        {
            if (!context.isDefaultRewrite()) {
                Node result = nodeRewriter.rewriteSelect(node, context.get(), TreeRewriter.this);
{'fraction_non_alphanumeric': 0.11549327607639281, 'fraction_numerical': 0.0, 'mean_word_length': 1.9787950185122853, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2372080', 'n_tokens_mistral': 2408, 'n_tokens_neox': 2253, 'n_words': 518}