/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.facebook.presto.sql.planner.iterative.rule;

import com.facebook.presto.sql.planner.PlanNodeIdAllocator;
import com.facebook.presto.sql.planner.SymbolAllocator;
import com.facebook.presto.sql.planner.iterative.Lookup;
import com.facebook.presto.sql.planner.iterative.Rule;
import com.facebook.presto.sql.planner.plan.AggregationNode;
import com.facebook.presto.sql.planner.plan.DistinctLimitNode;
import com.facebook.presto.sql.planner.plan.LimitNode;
import com.facebook.presto.sql.planner.plan.PlanNode;

import java.util.Optional;

public class MergeLimitWithDistinct
    implements Rule
{
    @Override
    public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator)
    {
        if (!(node instanceof LimitNode)) {
            return Optional.empty();
        }

        LimitNode parent = (LimitNode) node;

        PlanNode input = lookup.resolve(parent.getSource());
        if (!(input instanceof AggregationNode)) {
            return Optional.empty();
        }

        AggregationNode child = (AggregationNode) input;

        if (isDistinct(child)) {
            return Optional.empty();
        }

        return Optional.of(
                new DistinctLimitNode(
                        parent.getId(),
                        child.getSource(),
                        parent.getCount(),
                        false,
                        child.getHashSymbol()));
    }

    /**
     * Whether this node corresponds to a DISTINCT operation in SQL
     */
    private boolean isDistinct(AggregationNode node)
    {
        return !node.getAggregations().isEmpty() ||
                node.getOutputSymbols().size() != node.getGroupingKeys().size() ||
                !node.getOutputSymbols().containsAll(node.getGroupingKeys());
    }
}

@@ -1102,16 +1102,72 @@
 ional;%0A%0A
+import static com.facebook.presto.util.Optionals.cast;%0A%0A
 public c
@@ -1194,16 +1194,20 @@
 istinct%0A
+    
     impl
@@ -1384,162 +1384,153 @@
     
-if (!(node instanceof LimitNode)) %7B%0A            return Optional.empty();%0A        %7D%0A%0A        LimitNode parent = (LimitNode) node;%0A%0A        PlanNode input =
+Optional%3CLimitNode%3E parent = cast(node, LimitNode.class);%0A%0A        Optional%3CAggregationNode%3E childDistinct = parent%0A                .flatMap(p -%3E
  loo
@@ -1542,21 +1542,16 @@
 esolve(p
-arent
 .getSour
@@ -1558,98 +1558,33 @@
 ce()
-);%0A        if (!(input instanceof AggregationNode)) %7B%0A            return Optional.empty();
+, AggregationNode.class))
 %0A   
@@ -1588,27 +1588,24 @@
         
-%7D%0A%0A
         
 Aggregat
@@ -1600,90 +1600,35 @@
     
-AggregationNode child = (AggregationNode) input;%0A%0A        if (isDistinct(child)) %7B
+.filter(this::isDistinct);%0A
 %0A   
@@ -1636,36 +1636,34 @@
     
-    
 return 
-Optional.empty();
+parent.flatMap(p -%3E
 %0A   
@@ -1671,19 +1671,16 @@
     
-%7D%0A%0A
         
 retu
@@ -1679,28 +1679,34 @@
     
-return Optional.of(%0A
+childDistinct.map(d -%3E%0A   
     
@@ -1717,16 +1717,21 @@
         
+     
 new Dist
@@ -1765,30 +1765,33 @@
             
-parent
+        p
 .getId(),%0A  
@@ -1808,28 +1808,32 @@
             
-chil
+        
 d.getSource(
@@ -1859,22 +1859,25 @@
         
-parent
+        p
 .getCoun
@@ -1877,24 +1877,32 @@
 getCount(),%0A
+        
             
@@ -1944,20 +1944,24 @@
         
-chil
+        
 d.getHas
@@ -1970,16 +1970,43 @@
 ymbol())
+%0A                )%0A        
 );%0A    %7D

Refactor MergeLimitWithDistinct to leverage Optional's methods