package com.facebook.presto.sql.planner.optimizations;

import com.facebook.presto.sql.analyzer.Symbol;
import com.facebook.presto.sql.analyzer.Type;
import com.facebook.presto.sql.planner.plan.AggregationNode;
import com.facebook.presto.sql.planner.plan.FilterNode;
import com.facebook.presto.sql.planner.plan.JoinNode;
import com.facebook.presto.sql.planner.plan.LimitNode;
import com.facebook.presto.sql.planner.plan.OutputNode;
import com.facebook.presto.sql.planner.plan.PlanNode;
import com.facebook.presto.sql.planner.plan.PlanVisitor;
import com.facebook.presto.sql.planner.plan.ProjectNode;
import com.facebook.presto.sql.planner.plan.SortNode;
import com.facebook.presto.sql.planner.plan.TableScanNode;
import com.facebook.presto.sql.planner.plan.TopNNode;
import com.facebook.presto.sql.tree.Expression;
import com.facebook.presto.sql.tree.QualifiedNameReference;

import java.util.Map;

/**
 * Removes pure identity projections (e.g., Project $0 := $0, $1 := $1, ...)
 */
public class PruneRedundantProjections
        extends PlanOptimizer
{
    @Override
    public PlanNode optimize(PlanNode plan, Map<Symbol, Type> types)
    {
        return plan.accept(new Visitor(), null);
    }

    private static class Visitor
            extends PlanVisitor<Void, PlanNode>
    {
        @Override
        public PlanNode visitProject(ProjectNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);

            if (node.getOutputSymbols().size() != source.getOutputSymbols().size()) {
                // Can't get rid of this projection. It constrains the output tuple from the underlying operator
                return new ProjectNode(source, node.getOutputMap());
            }

            boolean canElide = true;
            for (Map.Entry<Symbol, Expression> entry : node.getOutputMap().entrySet()) {
                Expression expression = entry.getValue();
                Symbol symbol = entry.getKey();
                if (!(expression instanceof QualifiedNameReference && ((QualifiedNameReference) expression).getName().equals(symbol.toQualifiedName()))) {
                    canElide = false;
                    break;
                }
            }

            if (canElide) {
                return source;
            }

            return new ProjectNode(source, node.getOutputMap());
        }

        @Override
        public PlanNode visitAggregation(AggregationNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new AggregationNode(source, node.getGroupBy(), node.getAggregations(), node.getFunctions());
        }

        @Override
        public PlanNode visitTableScan(TableScanNode node, Void context)
        {
            return node;
        }

        @Override
        public PlanNode visitFilter(FilterNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new FilterNode(source, node.getPredicate(), node.getOutputSymbols());
        }

        @Override
        public PlanNode visitOutput(OutputNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new OutputNode(source, node.getColumnNames(), node.getAssignments());
        }

        @Override
        public PlanNode visitLimit(LimitNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new LimitNode(source, node.getCount());
        }

        @Override
        public PlanNode visitTopN(TopNNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new TopNNode(source, node.getCount(), node.getOrderBy(), node.getOrderings());
        }

        @Override
        public PlanNode visitSort(SortNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new SortNode(source, node.getOrderBy(), node.getOrderings());
        }

        @Override
        public PlanNode visitJoin(JoinNode node, Void context)
        {
            PlanNode left = node.getLeft().accept(this, context);
            PlanNode right = node.getRight().accept(this, context);

            return new JoinNode(left, right, node.getCriteria());
        }

        @Override
        protected PlanNode visitPlan(PlanNode node, Void context)
        {
            throw new UnsupportedOperationException("not yet implemented: " + getClass().getName());
        }
    }
}

@@ -471,32 +471,152 @@
 .plan.PlanNode;%0A
+import com.facebook.presto.sql.planner.plan.PlanNodeRewriter;%0Aimport com.facebook.presto.sql.planner.plan.PlanRewriter;%0A
 import com.faceb
@@ -1276,39 +1276,53 @@
 urn 
-plan.accept
+PlanRewriter.rewriteWith
 (new 
-Visito
+Rewrite
 r(), 
-null
+plan
 );%0A 
@@ -1352,22 +1352,23 @@
 c class 
-Visito
+Rewrite
 r%0A      
@@ -1389,30 +1389,25 @@
 Plan
-Visito
+NodeRewrite
 r%3CVoid
-, PlanNode
 %3E%0A  
@@ -1452,21 +1452,23 @@
 lanNode 
-vis
+rewr
 it
+e
 Project(
@@ -1489,32 +1489,65 @@
 de, Void context
+, PlanRewriter%3CVoid%3E planRewriter
 )%0A        %7B%0A    
@@ -1564,32 +1564,53 @@
 anNode source = 
+planRewriter.rewrite(
 node.getSource()
@@ -1605,36 +1605,24 @@
 .getSource()
-.accept(this
 , context);%0A
@@ -2542,2238 +2542,8 @@
    %7D
-%0A%0A        @Override%0A        public PlanNode visitAggregation(AggregationNode node, Void context)%0A        %7B%0A            PlanNode source = node.getSource().accept(this, context);%0A            return new AggregationNode(source, node.getGroupBy(), node.getAggregations(), node.getFunctions());%0A        %7D%0A%0A        @Override%0A        public PlanNode visitTableScan(TableScanNode node, Void context)%0A        %7B%0A            return node;%0A        %7D%0A%0A        @Override%0A        public PlanNode visitFilter(FilterNode node, Void context)%0A        %7B%0A            PlanNode source = node.getSource().accept(this, context);%0A            return new FilterNode(source, node.getPredicate(), node.getOutputSymbols());%0A        %7D%0A%0A        @Override%0A        public PlanNode visitOutput(OutputNode node, Void context)%0A        %7B%0A            PlanNode source = node.getSource().accept(this, context);%0A            return new OutputNode(source, node.getColumnNames(), node.getAssignments());%0A        %7D%0A%0A        @Override%0A        public PlanNode visitLimit(LimitNode node, Void context)%0A        %7B%0A            PlanNode source = node.getSource().accept(this, context);%0A            return new LimitNode(source, node.getCount());%0A        %7D%0A%0A        @Override%0A        public PlanNode visitTopN(TopNNode node, Void context)%0A        %7B%0A            PlanNode source = node.getSource().accept(this, context);%0A            return new TopNNode(source, node.getCount(), node.getOrderBy(), node.getOrderings());%0A        %7D%0A%0A        @Override%0A        public PlanNode visitSort(SortNode node, Void context)%0A        %7B%0A            PlanNode source = node.getSource().accept(this, context);%0A            return new SortNode(source, node.getOrderBy(), node.getOrderings());%0A        %7D%0A%0A        @Override%0A        public PlanNode visitJoin(JoinNode node, Void context)%0A        %7B%0A            PlanNode left = node.getLeft().accept(this, context);%0A            PlanNode right = node.getRight().accept(this, context);%0A%0A            return new JoinNode(left, right, node.getCriteria());%0A        %7D%0A%0A        @Override%0A        protected PlanNode visitPlan(PlanNode node, Void context)%0A        %7B%0A            throw new UnsupportedOperationException(%22not yet implemented: %22 + getClass().getName());%0A        %7D
 %0A   

Implement PruneRedundantProjections as plan rewriter