fix issue 201 by a rewriting rule to inline unnest functions

git-svn-id: https://asterixdb.googlecode.com/svn/branches/asterix_stabilization_printerfix@958 eaa15691-b419-025a-1212-ee371bd00084
diff --git a/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/base/RuleCollections.java b/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/base/RuleCollections.java
index 0bfa5c4..a327703 100644
--- a/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/base/RuleCollections.java
+++ b/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/base/RuleCollections.java
@@ -28,6 +28,7 @@
 import edu.uci.ics.asterix.optimizer.rules.FuzzyEqRule;
 import edu.uci.ics.asterix.optimizer.rules.FuzzyJoinRule;
 import edu.uci.ics.asterix.optimizer.rules.IfElseToSwitchCaseFunctionRule;
+import edu.uci.ics.asterix.optimizer.rules.InlineUnnestFunctionRule;
 import edu.uci.ics.asterix.optimizer.rules.IntroduceDynamicTypeCastRule;
 import edu.uci.ics.asterix.optimizer.rules.IntroduceEnforcedTypeRule;
 import edu.uci.ics.asterix.optimizer.rules.IntroduceSecondaryIndexInsertDeleteRule;
@@ -94,6 +95,7 @@
 
     public final static List<IAlgebraicRewriteRule> buildTypeInferenceRuleCollection() {
         List<IAlgebraicRewriteRule> typeInfer = new LinkedList<IAlgebraicRewriteRule>();
+        typeInfer.add(new InlineUnnestFunctionRule());
         typeInfer.add(new InferTypesRule());
         return typeInfer;
     }
diff --git a/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/rules/InlineUnnestFunctionRule.java b/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/rules/InlineUnnestFunctionRule.java
new file mode 100644
index 0000000..2398c04
--- /dev/null
+++ b/asterix-algebra/src/main/java/edu/uci/ics/asterix/optimizer/rules/InlineUnnestFunctionRule.java
@@ -0,0 +1,145 @@
+package edu.uci.ics.asterix.optimizer.rules;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.commons.lang3.mutable.Mutable;
+import org.apache.commons.lang3.mutable.MutableObject;
+
+import edu.uci.ics.asterix.om.functions.AsterixBuiltinFunctions;
+import edu.uci.ics.hyracks.algebricks.common.exceptions.AlgebricksException;
+import edu.uci.ics.hyracks.algebricks.common.utils.Pair;
+import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalExpression;
+import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalOperator;
+import edu.uci.ics.hyracks.algebricks.core.algebra.base.IOptimizationContext;
+import edu.uci.ics.hyracks.algebricks.core.algebra.base.LogicalExpressionTag;
+import edu.uci.ics.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
+import edu.uci.ics.hyracks.algebricks.core.algebra.base.LogicalVariable;
+import edu.uci.ics.hyracks.algebricks.core.algebra.expressions.AbstractFunctionCallExpression;
+import edu.uci.ics.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
+import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator;
+import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
+import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.UnnestOperator;
+import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities;
+import edu.uci.ics.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;
+
+public class InlineUnnestFunctionRule implements IAlgebraicRewriteRule {
+
+    @Override
+    public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context) throws AlgebricksException {
+        return false;
+    }
+
+    @Override
+    public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
+            throws AlgebricksException {
+        AbstractLogicalOperator op1 = (AbstractLogicalOperator) opRef.getValue();
+        if (context.checkIfInDontApplySet(this, op1))
+            return false;
+        context.addToDontApplySet(this, op1);
+        if (op1.getOperatorTag() != LogicalOperatorTag.UNNEST)
+            return false;
+        UnnestOperator unnestOperator = (UnnestOperator) op1;
+        AbstractFunctionCallExpression expr = (AbstractFunctionCallExpression) unnestOperator.getExpressionRef()
+                .getValue();
+        if (expr.getFunctionIdentifier() != AsterixBuiltinFunctions.SCAN_COLLECTION)
+            return false;
+
+        // inline all variables from an unnesting function call
+        AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) expr;
+        List<Mutable<ILogicalExpression>> args = funcExpr.getArguments();
+        for (int i = 0; i < args.size(); i++) {
+            ILogicalExpression argExpr = args.get(i).getValue();
+            if (argExpr.getExpressionTag() == LogicalExpressionTag.VARIABLE) {
+                VariableReferenceExpression varExpr = (VariableReferenceExpression) argExpr;
+                inlineVariable(varExpr.getVariableReference(), unnestOperator);
+            }
+        }
+        return true;
+    }
+
+    private void inlineVariable(LogicalVariable usedVar, UnnestOperator unnestOp) throws AlgebricksException {
+        AbstractFunctionCallExpression expr = (AbstractFunctionCallExpression) unnestOp.getExpressionRef().getValue();
+        List<Pair<AbstractFunctionCallExpression, Integer>> parentAndIndexList = new ArrayList<Pair<AbstractFunctionCallExpression, Integer>>();
+        getParentFunctionExpression(usedVar, expr, parentAndIndexList);
+        ILogicalExpression usedVarOrginExpr = findUsedVarOrigin(usedVar, unnestOp, (AbstractLogicalOperator) unnestOp
+                .getInputs().get(0).getValue());
+        if (usedVarOrginExpr != null) {
+            for (Pair<AbstractFunctionCallExpression, Integer> parentAndIndex : parentAndIndexList) {
+                //we only rewrite the top scan-collection function
+                if (parentAndIndex.first.getFunctionIdentifier() == AsterixBuiltinFunctions.SCAN_COLLECTION
+                        && parentAndIndex.first == expr) {
+                    unnestOp.getExpressionRef().setValue(usedVarOrginExpr);
+                }
+            }
+        }
+    }
+
+    private void getParentFunctionExpression(LogicalVariable usedVar, ILogicalExpression expr,
+            List<Pair<AbstractFunctionCallExpression, Integer>> parentAndIndexList) {
+        AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) expr;
+        List<Mutable<ILogicalExpression>> args = funcExpr.getArguments();
+        for (int i = 0; i < args.size(); i++) {
+            ILogicalExpression argExpr = args.get(i).getValue();
+            if (argExpr.getExpressionTag() == LogicalExpressionTag.VARIABLE) {
+                VariableReferenceExpression varExpr = (VariableReferenceExpression) argExpr;
+                if (varExpr.getVariableReference().equals(usedVar))
+                    parentAndIndexList.add(new Pair<AbstractFunctionCallExpression, Integer>(funcExpr, i));
+            }
+            if (argExpr.getExpressionTag() == LogicalExpressionTag.FUNCTION_CALL) {
+                getParentFunctionExpression(usedVar, argExpr, parentAndIndexList);
+            }
+        }
+    }
+
+    public ILogicalExpression findUsedVarOrigin(LogicalVariable usedVar, AbstractLogicalOperator parentOp,
+            AbstractLogicalOperator currentOp) throws AlgebricksException {
+        ILogicalExpression ret = null;
+        if (currentOp.getOperatorTag() == LogicalOperatorTag.ASSIGN) {
+            List<LogicalVariable> producedVars = new ArrayList<LogicalVariable>();
+            VariableUtilities.getProducedVariables(currentOp, producedVars);
+            if (producedVars.contains(usedVar)) {
+                AssignOperator assignOp = (AssignOperator) currentOp;
+                int index = assignOp.getVariables().indexOf(usedVar);
+                ILogicalExpression returnedExpr = assignOp.getExpressions().get(index).getValue();
+                if (returnedExpr.getExpressionTag() == LogicalExpressionTag.FUNCTION_CALL) {
+                    AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) returnedExpr;
+                    if (AsterixBuiltinFunctions.isBuiltinUnnestingFunction(funcExpr.getFunctionIdentifier())) {
+                        // we only inline for unnest functions
+                        removeUnecessaryAssign(parentOp, currentOp, assignOp, index);
+                        ret = returnedExpr;
+                    }
+                } else if (returnedExpr.getExpressionTag() == LogicalExpressionTag.VARIABLE) {
+                    //recusively inline
+                    VariableReferenceExpression varExpr = (VariableReferenceExpression) returnedExpr;
+                    LogicalVariable var = varExpr.getVariableReference();
+                    ILogicalExpression finalExpr = findUsedVarOrigin(var, currentOp,
+                            (AbstractLogicalOperator) currentOp.getInputs().get(0).getValue());
+                    if (finalExpr != null) {
+                        removeUnecessaryAssign(parentOp, currentOp, assignOp, index);
+                        ret = finalExpr;
+                    }
+                }
+            }
+        } else {
+            for (Mutable<ILogicalOperator> child : currentOp.getInputs()) {
+                ILogicalExpression expr = findUsedVarOrigin(usedVar, currentOp,
+                        (AbstractLogicalOperator) child.getValue());
+                if (expr != null) {
+                    ret = expr;
+                }
+            }
+        }
+        return ret;
+    }
+
+    private void removeUnecessaryAssign(AbstractLogicalOperator parentOp, AbstractLogicalOperator currentOp,
+            AssignOperator assignOp, int index) {
+        assignOp.getVariables().remove(index);
+        assignOp.getExpressions().remove(index);
+        if (assignOp.getVariables().size() == 0) {
+            int opIndex = parentOp.getInputs().indexOf(new MutableObject<ILogicalOperator>(currentOp));
+            parentOp.getInputs().get(opIndex).setValue(assignOp.getInputs().get(0).getValue());
+        }
+    }
+}
diff --git a/asterix-app/src/test/resources/runtimets/queries/user-defined-functions/query-issue201.aql b/asterix-app/src/test/resources/runtimets/queries/user-defined-functions/query-issue201.aql
new file mode 100644
index 0000000..0297873
--- /dev/null
+++ b/asterix-app/src/test/resources/runtimets/queries/user-defined-functions/query-issue201.aql
@@ -0,0 +1,5 @@
+write output to nc1:"rttest/user-defined-functions_query-issue201.adm";
+
+let $x:=range(1,100)
+for $i in $x
+return $i
\ No newline at end of file
diff --git a/asterix-app/src/test/resources/runtimets/results/user-defined-functions/query-issue201.adm b/asterix-app/src/test/resources/runtimets/results/user-defined-functions/query-issue201.adm
new file mode 100644
index 0000000..190423f
--- /dev/null
+++ b/asterix-app/src/test/resources/runtimets/results/user-defined-functions/query-issue201.adm
@@ -0,0 +1,100 @@
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+20
+21
+22
+23
+24
+25
+26
+27
+28
+29
+30
+31
+32
+33
+34
+35
+36
+37
+38
+39
+40
+41
+42
+43
+44
+45
+46
+47
+48
+49
+50
+51
+52
+53
+54
+55
+56
+57
+58
+59
+60
+61
+62
+63
+64
+65
+66
+67
+68
+69
+70
+71
+72
+73
+74
+75
+76
+77
+78
+79
+80
+81
+82
+83
+84
+85
+86
+87
+88
+89
+90
+91
+92
+93
+94
+95
+96
+97
+98
+99
+100
diff --git a/asterix-app/src/test/resources/runtimets/testsuite.xml b/asterix-app/src/test/resources/runtimets/testsuite.xml
index b80c499..f32f0b8 100644
--- a/asterix-app/src/test/resources/runtimets/testsuite.xml
+++ b/asterix-app/src/test/resources/runtimets/testsuite.xml
@@ -3577,6 +3577,11 @@
   </test-group>
   <test-group name="user-defined-functions">
     <test-case FilePath="user-defined-functions">
+      <compilation-unit name="query-issue201">
+        <output-file compare="Text">query-issue201.adm</output-file>
+      </compilation-unit>
+    </test-case>
+    <test-case FilePath="user-defined-functions">
       <compilation-unit name="udf01">
         <output-file compare="Text">udf01.adm</output-file>
       </compilation-unit>