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 java.util.Map;

/**
 * Merges successive LIMIT operators into a single LIMIT that's the minimum of the entire chain
 */
public class CoalesceLimits
        extends PlanOptimizer
{
    @Override
    public PlanNode optimize(PlanNode plan, Map<Symbol, Type> types)
    {
        Visitor visitor = new Visitor();
        return plan.accept(visitor, null);
    }

    private static class Visitor
        extends PlanVisitor<Void, PlanNode>
    {
        @Override
        protected PlanNode visitPlan(PlanNode node, Void context)
        {
            throw new UnsupportedOperationException("not yet implemented");
        }

        @Override
        public PlanNode visitLimit(LimitNode node, Void context)
        {
            PlanNode child = node.getSource().accept(this, context);
            if (child instanceof LimitNode) {
                return new LimitNode(((LimitNode) child).getSource(), Math.min(node.getCount(), ((LimitNode) child).getCount()));
            }

            return node;
        }

        @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 visitProject(ProjectNode node, Void context)
        {
            PlanNode source = node.getSource().accept(this, context);
            return new ProjectNode(source, node.getOutputMap());
        }

        @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 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());
        }
    }
}

@@ -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
@@ -1169,81 +1169,60 @@
     
-Visitor visitor = new Visitor();%0A        return plan.accept(visitor, null
+return PlanRewriter.rewriteWith(new Rewriter(), plan
 );%0A 
@@ -1252,22 +1252,23 @@
 c class 
-Visito
+Rewrite
 r%0A      
@@ -1285,219 +1285,33 @@
 Plan
-Visitor%3CVoid, PlanNode%3E%0A    %7B%0A        @Override%0A        protected PlanNode visitPlan(PlanNode node, Void context)%0A 
+NodeRewriter%3CVoid%3E%0A
     
-   
 %7B%0A
-            throw new UnsupportedOperationException(%22not yet implemented%22);%0A        %7D%0A%0A
     
@@ -1348,21 +1348,23 @@
 lanNode 
-vis
+rewr
 it
+e
 Limit(Li
@@ -1393,1004 +1393,179 @@
 text
-)%0A        %7B%0A            PlanNode child = node.getSource().accept(this, context);%0A            if (child instanceof LimitNode) %7B%0A                return new LimitNode(((LimitNode) child).getSource(), Math.min(node.getCount(), ((LimitNode) child).getCount()));%0A            %7D%0A%0A            return node;%0A        %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());
+, PlanRewriter%3CVoid%3E planRewriter)%0A        %7B%0A            PlanNode source = planRewriter.rewrite(node.getSource(), context);%0A%0A            if (source instanceof LimitNode) %7B
 %0A   
@@ -1565,35 +1565,32 @@
 ) %7B%0A        
-%7D%0A%0A
         
 @Override%0A  
@@ -1581,763 +1581,156 @@
     
-@Override%0A        public PlanNode visitProject(ProjectNode node, Void context)%0A        %7B%0A            Plan
+LimitNode limitNode = (Limit
 Node
+)
  source
- = node.getSource().accept(this, context);%0A            return new ProjectNode(source, node.getOutputMap());%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 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                return new LimitNode(limitNode.getSource(), Math.min(node.getCount(), limitNode.getCount()
 ));%0A
@@ -1741,130 +1741,43 @@
     
-%7D%0A%0A
     
+%7D%0A
     
-@Override%0A        public PlanNode visitSort(SortNode node, Void context)%0A        %7B%0A            PlanNode 
+        else if (
 source 
+!
 = no
@@ -1794,32 +1794,14 @@
 ce()
-.accept(this, context);%0A
+) %7B%0A  
     
@@ -1800,32 +1800,34 @@
 %7B%0A              
+  
 return new SortN
@@ -1821,19 +1821,20 @@
 urn new 
-Sor
+Limi
 tNode(so
@@ -1851,342 +1851,56 @@
 .get
-OrderBy(), 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())
+Count());%0A            %7D%0A%0A            return node
 ;%0A  

Implement CoalesceLimits as a plan rewriter