Collapse RANK and NOT

1 files changed, 20 insertions, 11 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java b/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
index 8f49fa2bf48..e9806c2f614 100644
--- a/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
+++ b/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
@@ -63,13 +63,13 @@ public class QueryCanonicalizer {
         if ( ! (item instanceof CompositeItem)) return CanonicalizationResult.success();
         CompositeItem composite = (CompositeItem)item;
 
-        collapseLevels(composite);
-        
-        if (composite.getItemCount() == 0) {
+        if (composite.getItemCount() == 0) { // TODO: Remove
             parentIterator.remove();
             return CanonicalizationResult.success();
         }
-            
+
+        collapseLevels(composite);
+
         if (composite instanceof EquivItem) {
             removeDuplicates((EquivItem) composite);
         }
@@ -94,16 +94,25 @@ public class QueryCanonicalizer {
     }
     
     private static void collapseLevels(CompositeItem composite) {
-        if ( ! (composite instanceof AndItem || composite instanceof OrItem)) return;
-
-        for (ListIterator<Item> i = composite.getItemIterator(); i.hasNext(); ) {
-            Item child = i.next();
-            if (child.getClass() != composite.getClass()) continue;
-            i.remove();
-            moveChildren((CompositeItem)child, i);
+        if (composite instanceof RankItem || composite instanceof NotItem) {
+            collapseLevels(composite, composite.getItemIterator()); // collapse the first item only
+        }
+        else if (composite instanceof AndItem || composite instanceof OrItem) {
+            for (ListIterator<Item> i = composite.getItemIterator(); i.hasNext(); )
+                collapseLevels(composite, i);
         }
     }
     
+    /** Collapse the next item of this iterator into the given parent, if appropriate */
+    private static void collapseLevels(CompositeItem composite, ListIterator<Item> i) {
+        if ( ! i.hasNext()) return;
+        Item child = i.next();
+        if (child == null) return;
+        if (child.getClass() != composite.getClass()) return;
+        i.remove();
+        moveChildren((CompositeItem) child, i);
+    }
+    
     private static void moveChildren(CompositeItem from, ListIterator<Item> toIterator) {
         for (ListIterator<Item> i = from.getItemIterator(); i.hasNext(); )
             toIterator.add(i.next());