package org.basex.query.expr;

import static org.basex.query.QueryText.*;

import java.util.*;
import java.util.function.*;

import org.basex.query.*;
import org.basex.query.expr.path.*;
import org.basex.query.func.Function;
import org.basex.query.iter.*;
import org.basex.query.util.*;
import org.basex.query.util.list.*;
import org.basex.query.value.*;
import org.basex.query.value.item.*;
import org.basex.query.value.node.*;
import org.basex.query.value.seq.*;
import org.basex.query.value.type.*;
import org.basex.query.var.*;
import org.basex.util.*;
import org.basex.util.hash.*;

/**
 * Union expression.
 *
 * @author BaseX Team 2005-20, BSD License
 * @author Christian Gruen
 */
public final class Union extends Set {
  /**
   * Constructor.
   * @param info input info
   * @param exprs expressions
   */
  public Union(final InputInfo info, final Expr... exprs) {
    super(info, exprs);
  }

  @Override
  public Expr optimize(final CompileContext cc) throws QueryException {
    super.optimize(cc);

    final ExprList list = new ExprList(exprs.length);
    for(final Expr expr : exprs) {
      if(expr == Empty.VALUE) {
        // remove empty operands
        // example: * union ()  ->  *
        cc.info(OPTREMOVE_X_X, expr, (Supplier<?>) this::description);
      } else if(expr.seqType().type instanceof NodeType && !expr.has(Flag.CNS, Flag.NDT) &&
         ((Checks<Expr>) ex -> ex.equals(expr)).any(list)) {
        // remove duplicate
        // example: * union *  ->  *
        cc.info(OPTREMOVE_X_X, expr, (Supplier<?>) this::description);
      } else {
        list.add(expr);
      }
    }
    exprs = list.finish();

    switch(exprs.length) {
      case 0:  return Empty.VALUE;
      case 1:  return iterative ? exprs[0] : cc.function(Function._UTIL_DDO, info, exprs[0]);
      default: return merge(cc) ? cc.replaceWith(this, exprs[0]) : this;
    }
  }

  /**
   * Tries to merge steps with identical axes and nodes types to a single step.
   * @param cc compilation context
   * @return success flag
   * @throws QueryException query exception
   */
  private boolean merge(final CompileContext cc) throws QueryException {
    Expr root = null;
    Axis axis = null;
    Expr[] preds = null;
    final ArrayList<Test> tests = new ArrayList<>(2);

    for(final Expr expr : exprs) {
      if(!(expr instanceof Path)) return false;
      final Path path = (Path) expr;
      if(path.steps.length > 1) return false;
      if(root != null ? !Objects.equals(root, path.root) :
        path.root != null && path.root.has(Flag.CNS, Flag.NDT, Flag.POS)) return false;

      root = path.root;

      final Step step = (Step) path.steps[0];
      if(axis != null && axis != step.axis || step.has(Flag.CNS, Flag.NDT, Flag.POS)) return false;
      if(preds != null && !Arrays.equals(preds, step.exprs)) return false;

      axis = step.axis;
      preds = step.exprs;
      tests.add(step.test);
    }

    final Test test = Test.get(tests);
    if(test == null) return false;

    final Expr step = new StepBuilder(info).axis(axis).test(test).preds(preds).finish(cc, root);
    exprs[0] = Path.get(info, root, step).optimize(cc);
    return true;
  }

  @Override
  protected Value nodes(final QueryContext qc) throws QueryException {
    final ANodeBuilder nodes = new ANodeBuilder();
    for(final Expr expr : exprs) {
      final Iter iter = expr.iter(qc);
      for(Item item; (item = qc.next(iter)) != null;) nodes.add(toNode(item));
    }
    return nodes.value(this);
  }

  @Override
  protected NodeIter iterate(final QueryContext qc) throws QueryException {
    return new SetIter(qc, iters(qc)) {
      @Override
      public ANode next() throws QueryException {
        if(nodes == null) {
          final int il = iter.length;
          nodes = new ANode[il];
          for(int i = 0; i < il; i++) next(i);
        }

        int m = -1;
        final int il = nodes.length;
        for(int i = 0; i < il; i++) {
          if(nodes[i] == null) continue;
          final int d = m == -1 ? 1 : nodes[m].diff(nodes[i]);
          if(d == 0) {
            next(i--);
          } else if(d > 0) {
            m = i;
          }
        }
        if(m == -1) return null;

        final ANode node = nodes[m];
        next(m);
        return node;
      }
    };
  }

  @Override
  public Expr copy(final CompileContext cc, final IntObjMap<Var> vm) {
    final Union un = new Union(info, copyAll(cc, vm, exprs));
    un.iterative = iterative;
    return copyType(un);
  }

  @Override
  public boolean equals(final Object obj) {
    return this == obj || obj instanceof Union && super.equals(obj);
  }
}

@@ -2337,16 +2337,20 @@
 tanceof 
+Axis
 Path)) r
@@ -2374,16 +2374,20 @@
   final 
+Axis
 Path pat
@@ -2391,16 +2391,20 @@
 path = (
+Axis
 Path) ex
@@ -2653,15 +2653,8 @@
 ep =
- (Step)
  pat
@@ -2659,20 +2659,19 @@
 ath.step
-s%5B0%5D
+(0)
 ;%0A      

rewrite union to step