/*
 * 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.matching.Captures;
import com.facebook.presto.matching.Pattern;
import com.facebook.presto.metadata.Metadata;
import com.facebook.presto.sql.planner.iterative.Rule;
import com.facebook.presto.sql.planner.iterative.RuleSet;
import com.facebook.presto.sql.planner.optimizations.TableLayoutRewriter;
import com.facebook.presto.sql.planner.plan.FilterNode;
import com.facebook.presto.sql.planner.plan.PlanNode;
import com.facebook.presto.sql.planner.plan.TableScanNode;
import com.facebook.presto.sql.planner.plan.ValuesNode;
import com.facebook.presto.sql.tree.BooleanLiteral;
import com.google.common.collect.ImmutableSet;

import java.util.Optional;
import java.util.Set;

import static com.facebook.presto.sql.planner.plan.Patterns.filter;
import static com.facebook.presto.sql.planner.plan.Patterns.tableScan;
import static java.util.Objects.requireNonNull;

public class PickTableLayout
        implements RuleSet
{
    private final ImmutableSet<Rule<?>> rules;

    public PickTableLayout(Metadata metadata)
    {
        rules = ImmutableSet.of(
                new PickTableLayoutForPredicate(metadata),
                new PickTableLayoutWithoutPredicate(metadata));
    }

    @Override
    public Set<Rule<?>> rules()
    {
        return rules;
    }

    private static final class PickTableLayoutForPredicate
            implements Rule<FilterNode>
    {
        private final Metadata metadata;

        private PickTableLayoutForPredicate(Metadata metadata)
        {
            this.metadata = requireNonNull(metadata, "metadata is null");
        }

        private static final Pattern<FilterNode> PATTERN = filter();

        @Override
        public Pattern<FilterNode> getPattern()
        {
            return PATTERN;
        }

        @Override
        public Optional<PlanNode> apply(FilterNode filterNode, Captures captures, Context context)
        {
            PlanNode source = context.getLookup().resolve(filterNode.getSource());

            if (!((source instanceof TableScanNode) && shouldRewriteTableLayout((TableScanNode) source))) {
                return Optional.empty();
            }

            TableLayoutRewriter tableLayoutRewriter = new TableLayoutRewriter(metadata, context.getSession(), context.getSymbolAllocator(), context.getIdAllocator());
            PlanNode rewrittenTableScan = tableLayoutRewriter.planTableScan((TableScanNode) source, filterNode.getPredicate());

            if (rewrittenTableScan instanceof TableScanNode || rewrittenTableScan instanceof ValuesNode || (((FilterNode) rewrittenTableScan).getPredicate() != filterNode.getPredicate())) {
                return Optional.of(rewrittenTableScan);
            }

            return Optional.empty();
        }

        private boolean shouldRewriteTableLayout(TableScanNode source)
        {
            return !source.getLayout().isPresent() || source.getOriginalConstraint() == BooleanLiteral.TRUE_LITERAL;
        }
    }

    private static final class PickTableLayoutWithoutPredicate
            implements Rule<TableScanNode>
    {
        private final Metadata metadata;

        private PickTableLayoutWithoutPredicate(Metadata metadata)
        {
            this.metadata = requireNonNull(metadata, "metadata is null");
        }

        private static final Pattern<TableScanNode> PATTERN = tableScan();

        @Override
        public Pattern<TableScanNode> getPattern()
        {
            return PATTERN;
        }

        @Override
        public Optional<PlanNode> apply(TableScanNode tableScanNode, Captures captures, Context context)
        {
            if (tableScanNode.getLayout().isPresent()) {
                return Optional.empty();
            }

            TableLayoutRewriter tableLayoutRewriter = new TableLayoutRewriter(metadata, context.getSession(), context.getSymbolAllocator(), context.getIdAllocator());
            return Optional.of(tableLayoutRewriter.planTableScan(tableScanNode, BooleanLiteral.TRUE_LITERAL));
        }
    }
}

@@ -612,16 +612,61 @@
 .rule;%0A%0A
+import com.facebook.presto.matching.Capture;%0A
 import c
@@ -1352,24 +1352,87 @@
 .util.Set;%0A%0A
+import static com.facebook.presto.matching.Capture.newCapture;%0A
 import stati
@@ -1483,24 +1483,92 @@
 rns.filter;%0A
+import static com.facebook.presto.sql.planner.plan.Patterns.source;%0A
 import stati
@@ -2415,123 +2415,160 @@
 nal 
-Pattern%3CFilterNode%3E PATTERN = filter();%0A%0A        @Override%0A        public Pattern%3CFilterNode%3E getPattern()%0A
+Capture%3CTableScanNode%3E TABLE_SCAN = newCapture();%0A%0A        private static final Pattern%3CFilterNode%3E PATTERN = filter().with(source().matching(%0A 
         
 %7B%0A  
@@ -2555,34 +2555,32 @@
 ching(%0A         
-%7B%0A
             retu
@@ -2579,245 +2579,351 @@
     
-return PATTERN;%0A        %7D%0A%0A        @Override%0A        public Optional%3CPlanNode%3E apply(FilterNode filterNode, Captures captures, Context context)%0A        %7B%0A            PlanNode source = context.getLookup().resolve(filterNode.getSource());%0A
+   tableScan().matching(PickTableLayoutForPredicate::shouldRewriteTableLayout)%0A                                .capturedAs(TABLE_SCAN)));%0A%0A        private static boolean shouldRewriteTableLayout(TableScanNode source)%0A        %7B%0A            return !source.getLayout().isPresent() %7C%7C source.getOriginalConstraint() == BooleanLiteral.TRUE_LITERAL;
 %0A   
@@ -2931,114 +2931,87 @@
     
+%7D%0A%0A
     
-if (!((source instanceof TableScanNode) && shouldRewriteTableLayout((TableScanNode) source))) %7B%0A
+    @Override%0A        public Pattern%3CFilterNode%3E getPattern()%0A  
       
+%7B%0A
     
@@ -3020,45 +3020,162 @@
     
+  
 return 
-Optional.empty();%0A    
+PATTERN;%0A        %7D%0A%0A        @Override%0A        public Optional%3CPlanNode%3E apply(FilterNode filterNode, Captures captures, Context context)%0A
         
 %7D%0A%0A 
@@ -3162,34 +3162,33 @@
 ontext)%0A        
-%7D%0A
+%7B
 %0A            Tab
@@ -3419,30 +3419,32 @@
 can(
-(TableScanNode) source
+captures.get(TABLE_SCAN)
 , fi
@@ -3781,217 +3781,8 @@
    %7D
-%0A%0A        private boolean shouldRewriteTableLayout(TableScanNode source)%0A        %7B%0A            return !source.getLayout().isPresent() %7C%7C source.getOriginalConstraint() == BooleanLiteral.TRUE_LITERAL;%0A        %7D
 %0A   

Use more precise pattern in PickTableLayout