java/src/main/com/google/appengine/api/search/query/QueryTreeWalker.java

   1 package com.google.appengine.api.search.query;
   2
   3 import com.google.appengine.api.search.SearchQueryException;
   4 import com.google.common.collect.ImmutableSet;
   5
   6 import org.antlr.runtime.tree.Tree;
   7
   8 /**
   9  * The walking of the query tree. This class takes care of visiting
  10  * a tree resulting from parsing a query. As it traverses the tree
  11  * it calls appropriate methods of the visitor, set at the construction
  12  * time. The class uses a depth-first search, visiting all children
  13  * of a node, before visiting the node. The visit is done by calling
  14  * an appropriate method of the visitor. Typical code should match
  15  * the following pattern:
  16  * <pre>
  17  * class MyVisitor implements QueryTreeVisitor {
  18  *   ...
  19  * }
  20  * class MyContext extends QueryTreeContext&lt;MyContext&gt; {
  21  *   ...
  22  *   &#64;Override
  23  *   protected MyContext newChildContext() {
  24  *     return new MyContext();
  25  *   }
  26  * }
  27  *
  28  * MyContext context = new MyContext();
  29  * QueryTreeWalker&lt;MyContext&gt; walker = new QueryTreeWalker&lt;MyContext&gt;(new MyVisitor());
  30  * Tree root = parser.query(queryStr);
  31  * walker.walk(root, context);
  32  * // retrieve whatever information you need from context
  33  * </pre>
  34  *
  35  * @param <T> the context used by the visitor
  36  */
  37 public class QueryTreeWalker<T extends QueryTreeContext<T>> {
  38
  39   private final QueryTreeVisitor<T> visitor;
  40
  41   /**
  42    * Creates a new query walker that calls the given {@code visitor}.
  43    *
  44    * @param visitor the visitor to be called by this walker
  45    */
  46   public QueryTreeWalker(QueryTreeVisitor<T> visitor) {
  47     this.visitor = visitor;
  48   }
  49
  50   /**
  51    * @param tree the tree to be walked
  52    * @param context the context in which the tree is walked
  53    */
  54   public void walk(Tree tree, T context) throws QueryTreeException {
  55     tree = flatten(tree, null);
  56     validate(tree);
  57     walkInternal(tree, context);
  58   }
  59
  60   /**
  61    * Walks the tree and updates the context. This is a depth-first search,
  62    * always exploring children of the {@code tree} in the order of their
  63    * index 0 ... n - 1, where n is the number of children. Once all children
  64    * of the {@code node} are visited, the appropriate visitor's method
  65    * is called
  66    *
  67    * @param tree the tree to be walked
  68    * @param context the context of the visit
  69    */
  70   private void walkInternal(Tree tree, T context) {
  71     switch (tree.getType()) {
  72       case QueryLexer.CONJUNCTION:
  73         walkChildren(tree, context);
  74         visitor.visitConjunction(tree, context);
  75         postBooleanExpression(context);
  76         break;
  77       case QueryLexer.DISJUNCTION:
  78         walkChildren(tree, context);
  79         visitor.visitDisjunction(tree, context);
  80         postBooleanExpression(context);
  81         break;
  82       case QueryLexer.SEQUENCE:
  83         walkChildren(tree, context);
  84         visitor.visitSequence(tree, context);
  85         postBooleanExpression(context);
  86         break;
  87       case QueryLexer.NEGATION:
  88         walkChildren(tree, context);
  89         visitor.visitNegation(tree, context);
  90         postBooleanExpression(context);
  91         break;
  92       case QueryLexer.HAS:
  93         walkChildren(tree, context);
  94         visitor.visitContains(tree, context);
  95         postBooleanExpression(context);
  96         break;
  97       case QueryLexer.EQ:
  98         walkChildren(tree, context);
  99         visitor.visitEqual(tree, context);
 100         postBooleanExpression(context);
 101         break;
 102       case QueryLexer.NE:
 103         throw new SearchQueryException("!= comparison operator is not available");
 104       case QueryLexer.LESSTHAN:
 105         walkChildren(tree, context);
 106         visitor.visitLessThan(tree, context);
 107         postBooleanExpression(context);
 108         break;
 109       case QueryLexer.LE:
 110         walkChildren(tree, context);
 111         visitor.visitLessOrEqual(tree, context);
 112         postBooleanExpression(context);
 113         break;
 114       case QueryLexer.GT:
 115         walkChildren(tree, context);
 116         visitor.visitGreaterThan(tree, context);
 117         postBooleanExpression(context);
 118         break;
 119       case QueryLexer.GE:
 120         walkChildren(tree, context);
 121         visitor.visitGreaterOrEqual(tree, context);
 122         postBooleanExpression(context);
 123         break;
 124       case QueryLexer.FUZZY:
 125         walkInternal(tree.getChild(0), context);
 126         visitor.visitFuzzy(tree, context);
 127         break;
 128       case QueryLexer.LITERAL:
 129         walkInternal(tree.getChild(0), context);
 130         visitor.visitLiteral(tree, context);
 131         break;
 132       case QueryLexer.VALUE:
 133         visitor.visitValue(tree, context);
 134         break;
 135       case QueryLexer.FUNCTION:
 136         walkChildren(tree.getChild(1), context);
 137         visitor.visitFunction(tree, context);
 138         break;
 139       case QueryLexer.GLOBAL:
 140         visitor.visitGlobal(tree, context);
 141         break;
 142       default:
 143         visitor.visitOther(tree, context);
 144     }
 145   }
 146
 147   protected void postBooleanExpression(T context) {
 148     context.setReturnType(QueryTreeContext.Type.BOOL);
 149     context.setKind(QueryTreeContext.Kind.EXPRESSION);
 150   }
 151
 152   private void walkChildren(Tree parent, T context) {
 153     for (int i = 0; i < parent.getChildCount(); ++i) {
 154       walkInternal(parent.getChild(i), context.addChild());
 155     }
 156   }
 157
 158   private static final ImmutableSet<String> QUERY_FUNCTION_NAMES =
 159       ImmutableSet.of("distance", "geopoint");
 160
 161   /**
 162    * Basic stateless query tree validation.
 163    *
 164    * @param tree parsed query tree to validate
 165    * @throws QueryTreeException if the tree is invalid
 166    */
 167   private static void validate(Tree tree) throws QueryTreeException {
 168     for (int i = 0; i < tree.getChildCount(); ++i) {
 169       validate(tree.getChild(i));
 170     }
 171     switch (tree.getType()) {
 172       case QueryLexer.FUNCTION:
 173         Tree name = tree.getChild(0);
 174         if (!QUERY_FUNCTION_NAMES.contains(name.getText())) {
 175           throw new QueryTreeException("unknown function '" + name.getText() + "'",
 176               name.getCharPositionInLine());
 177         }
 178         break;
 179       default:
 180         break;
 181     }
 182   }
 183
 184   public static Tree simplify(Tree tree) {
 185     for (int i = 0; i < tree.getChildCount(); ++i) {
 186       Tree child = tree.getChild(i);
 187       Tree optimized = simplify(child);
 188       if (child != optimized) {
 189         tree.setChild(i, optimized);
 190       }
 191     }
 192     switch (tree.getType()) {
 193       case QueryLexer.CONJUNCTION:
 194       case QueryLexer.DISJUNCTION:
 195       case QueryLexer.SEQUENCE:
 196         if (tree.getChildCount() == 1) {
 197           return tree.getChild(0);
 198         }
 199         break;
 200     }
 201     return tree;
 202   }
 203
 204   /**
 205    * Flattens the tree by pushing down the field name. For example, if
 206    * the tree looks like this:
 207    * <pre>
 208    *                EQ
 209    *             /      \
 210    *        VALUE         EQ
 211    *        /    \      /    \
 212    *      TEXT  field GLOBAL (N)
 213    * </pre>
 214    * Then we will output tree that looks like this:
 215    * <pre>
 216    *                EQ
 217    *             /      \
 218    *        VALUE       (N)
 219    *        /    \
 220    *      TEXT  field
 221    * </pre>
 222    * Here <code>(N)</code> is an arbitrary node. We also drop EQ if it
 223    * is in front of conjunction or disjunction. We do not drop it for
 224    * other comparators, as we want parsing to fail for foo &lt; (1 2).
 225    */
 226   private static Tree flatten(Tree tree, Tree restriction) throws QueryTreeException {
 227     if (tree.getType() == QueryLexer.VALUE) {
 228       return tree;
 229     }
 230     if (tree.getType() == QueryLexer.HAS || tree.getType() == QueryLexer.EQ) {
 231       Tree lhs = tree.getChild(0);
 232       if (lhs.getType() == QueryLexer.VALUE) {
 233         String myField = lhs.getChild(1).getText();
 234         if (restriction == null) {
 235           restriction = lhs;
 236         } else {
 237           String otherField = restriction.getChild(1).getText();
 238           if (!myField.equals(otherField)) {
 239             throw new QueryTreeException(
 240                 String.format("Restriction on %s and %s", otherField, myField),
 241                 lhs.getChild(1).getCharPositionInLine());
 242           }
 243         }
 244       }
 245       Tree rhs = tree.getChild(1);
 246       Tree flattened = flatten(rhs, restriction);
 247       if (flattened.getType() == QueryLexer.HAS
 248           || flattened.getType() == QueryLexer.EQ
 249           || flattened.getType() == QueryLexer.CONJUNCTION
 250           || flattened.getType() == QueryLexer.DISJUNCTION
 251           || flattened.getType() == QueryLexer.SEQUENCE) {
 252         return flattened;
 253       }
 254       if (flattened != rhs) {
 255         tree.setChild(1, flattened);
 256       }
 257       if (restriction != lhs) {
 258         tree.setChild(0, restriction);
 259       }
 260       return tree;
 261     }
 262     for (int i = 0; i < tree.getChildCount(); ++i) {
 263       Tree original = tree.getChild(i);
 264       Tree flattened = flatten(tree.getChild(i), restriction);
 265       if (original != flattened) {
 266         tree.setChild(i, flattened);
 267       }
 268     }
 269     return tree;
 270   }
 271 }