/query/src/main/java/org/infinispan/query/dsl/embedded/impl/BaseEmbeddedQuery.java

package org.infinispan.query.dsl.embedded.impl;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;

import org.infinispan.AdvancedCache;
import org.infinispan.commons.util.CloseableIterator;
import org.infinispan.objectfilter.ObjectFilter;
import org.infinispan.query.dsl.QueryFactory;
import org.infinispan.query.dsl.impl.BaseQuery;

/**
 * Base class for embedded-mode query implementations. Subclasses need to implement {@link #getIterator()} and {@link
 * #getComparator()} methods and this class will take care of sorting (fully in-memory).
 *
 * @author anistor@redhat.com
 * @since 8.0
 */
abstract class BaseEmbeddedQuery extends BaseQuery {

   /**
    * Initial capacity of the collection used for collecting results when performing internal sorting.
    */
   private static final int INITIAL_CAPACITY = 1000;

   protected final AdvancedCache<?, ?> cache;

   /**
    * The cached results, lazily evaluated.
    */
   private List<Object> results;

   /**
    * The total number of results matching the query, ignoring pagination. This is lazily evaluated.
    */
   private int resultSize;

   protected BaseEmbeddedQuery(QueryFactory queryFactory, AdvancedCache<?, ?> cache, String jpaQuery, Map<String, Object> namedParameters,
                               String[] projection, long startOffset, int maxResults) {
      super(queryFactory, jpaQuery, namedParameters, projection, startOffset, maxResults);
      this.cache = cache;
   }

   @Override
   public void resetQuery() {
      results = null;
   }

   @Override
   public <T> List<T> list() {
      if (results == null) {
         results = listInternal();
      }
      return (List<T>) results;
   }

   private List<Object> listInternal() {
      List<Object> results;

      CloseableIterator<ObjectFilter.FilterResult> iterator = getIterator();
      if (!iterator.hasNext()) {
         return Collections.emptyList();
      }

      Comparator<Comparable[]> comparator = getComparator();
      if (comparator == null) {
         // collect unsorted results and get the requested page if any was specified
         try {
            if (iterator.hasNext()) {
               results = new ArrayList<>(INITIAL_CAPACITY);
               while (iterator.hasNext()) {
                  ObjectFilter.FilterResult entry = iterator.next();
                  resultSize++;
                  if (resultSize > startOffset && (maxResults == -1 || results.size() < maxResults)) {
                     results.add(projection != null ? entry.getProjection() : entry.getInstance());
                  }
               }
            } else {
               results = Collections.emptyList();
            }
         } finally {
            try {
               iterator.close();
            } catch (Exception e) {
               // exception ignored
            }
         }
      } else {
         // collect and sort results, in reverse order for now
         PriorityQueue<ObjectFilter.FilterResult> filterResults = new PriorityQueue<>(INITIAL_CAPACITY, new ReverseFilterResultComparator(comparator));
         try {
            while (iterator.hasNext()) {
               ObjectFilter.FilterResult entry = iterator.next();
               resultSize++;
               filterResults.add(entry);
               if (maxResults != -1 && filterResults.size() > startOffset + maxResults) {
                  // remove the head, which is actually the highest result
                  filterResults.remove();
               }
            }
         } finally {
            try {
               iterator.close();
            } catch (Exception e) {
               // exception ignored
            }
         }

         // collect and reverse
         if (filterResults.size() > startOffset) {
            Object[] res = new Object[filterResults.size() - startOffset];
            int i = filterResults.size();
            while (i-- > startOffset) {
               ObjectFilter.FilterResult r = filterResults.remove();
               res[i - startOffset] = projection != null ? r.getProjection() : r.getInstance();
            }
            results = Arrays.asList(res);
         } else {
            results = Collections.emptyList();
         }
      }

      return results;
   }

   /**
    * Create a comparator to be used for ordering the results returned by {@link #getIterator()}.
    *
    * @return the comparator or {@code null} if no sorting needs to be applied
    */
   protected abstract Comparator<Comparable[]> getComparator();

   /**
    * Create an iterator over the results of the query, in no particular order. Ordering will be provided if {@link
    * #getComparator()} returns a non-null {@link Comparator}.
    */
   protected abstract CloseableIterator<ObjectFilter.FilterResult> getIterator();

   @Override
   public int getResultSize() {
      list();
      return resultSize;
   }

   @Override
   public String toString() {
      return "BaseEmbeddedQuery{" +
            "jpaQuery=" + jpaQuery +
            ", namedParameters=" + namedParameters +
            ", projection=" + Arrays.toString(projection) +
            ", startOffset=" + startOffset +
            ", maxResults=" + maxResults +
            '}';
   }

   /**
    * Compares two {@link ObjectFilter.FilterResult} objects based on a given {@link Comparator} and reverses the
    * result.
    */
   private static class ReverseFilterResultComparator implements Comparator<ObjectFilter.FilterResult> {

      private final Comparator<Comparable[]> comparator;

      private ReverseFilterResultComparator(Comparator<Comparable[]> comparator) {
         this.comparator = comparator;
      }

      @Override
      public int compare(ObjectFilter.FilterResult o1, ObjectFilter.FilterResult o2) {
         return -comparator.compare(o1.getSortProjection(), o2.getSortProjection());
      }
   }
}