Refactor and clean up IR lowering code. Add common IR lowering module. (kensho-technologies#324)

* Refactor and clean up IR lowering code. Add common IR lowering operations module.

* Delint.

* Improve expression comment around null values.

* Add defensive strip_non_null_from_type() call.

Author: obi1kenobi

graphql_compiler/compiler/expressions.py

@@ -205,6 +205,8 @@ def to_match(self):
 
     def to_gremlin(self):
         """Return a unicode object with the Gremlin representation of this expression."""
+        self.validate()
+
         # We can't directly pass a Date or a DateTime object, so we have to pass it as a string
         # and then parse it inline. For date format parameter meanings, see:
         # http://docs.oracle.com/javase/7/docs/api/java/text/SimpleDateFormat.html
@@ -559,8 +561,9 @@ def to_match(self):
         return template % template_data
 
     def to_gremlin(self):
-        """Must never be called."""
-        raise NotImplementedError()
+        """Not implemented, should not be used."""
+        raise AssertionError(u'FoldedContextField are not used during the query emission process '
+                             u'in Gremlin, so this is a bug. This function should not be called.')
 
     def __eq__(self, other):
         """Return True if the given object is equal to this one, and False otherwise."""
@@ -618,7 +621,7 @@ def to_match(self):
         return template % template_data
 
     def to_gremlin(self):
-        """Must never be called."""
+        """Not supported yet."""
         raise NotImplementedError()
 
 
@@ -799,9 +802,9 @@ def to_match(self):
         intersects_operator_format = '(%(operator)s(%(left)s, %(right)s).asList().size() > 0)'
         # pylint: enable=unused-variable
 
-        # Null literals use 'is/is not' as (in)equality operators, while other values use '=/<>'.
-        if any((isinstance(self.left, Literal) and self.left.value is None,
-                isinstance(self.right, Literal) and self.right.value is None)):
+        # Null literals use the OrientDB 'IS/IS NOT' (in)equality operators,
+        # while other values use the OrientDB '=/<>' operators.
+        if self.left == NullLiteral or self.right == NullLiteral:
             translation_table = {
                 u'=': (u'IS', regular_operator_format),
                 u'!=': (u'IS NOT', regular_operator_format),
@@ -947,6 +950,7 @@ def visitor_fn(expression):
     def to_gremlin(self):
         """Return a unicode object with the Gremlin representation of this expression."""
         self.validate()
+
         return u'({predicate} ? {if_true} : {if_false})'.format(
             predicate=self.predicate.to_gremlin(),
             if_true=self.if_true.to_gremlin(),

graphql_compiler/compiler/ir_lowering_common/__init__.py

@@ -0,0 +1 @@
+# Copyright 2019-present Kensho Technologies, LLC.

graphql_compiler/compiler/ir_lowering_common.py renamed to graphql_compiler/compiler/ir_lowering_common/common.py

@@ -2,13 +2,13 @@
 """Language-independent IR lowering and optimization functions."""
 import six
 
-from .blocks import (
+from ..blocks import (
     ConstructResult, EndOptional, Filter, Fold, MarkLocation, Recurse, Traverse, Unfold
 )
-from .expressions import (
+from ..expressions import (
     BinaryComposition, ContextField, ContextFieldExistence, FalseLiteral, NullLiteral, TrueLiteral
 )
-from .helpers import validate_safe_string
+from ..helpers import validate_safe_string
 
 
 def merge_consecutive_filter_clauses(ir_blocks):

graphql_compiler/compiler/ir_lowering_common/location_renaming.py

@@ -0,0 +1,67 @@
+# Copyright 2019-present Kensho Technologies, LLC.
+"""Utilities for rewriting IR to replace one set of locations with another."""
+import six
+
+from ..helpers import FoldScopeLocation, Location
+
+
+def make_revisit_location_translations(query_metadata_table):
+    """Return a dict mapping location revisits to the location being revisited, for rewriting."""
+    location_translations = dict()
+
+    for location, _ in query_metadata_table.registered_locations:
+        location_being_revisited = query_metadata_table.get_revisit_origin(location)
+        if location_being_revisited != location:
+            location_translations[location] = location_being_revisited
+
+    return location_translations
+
+
+def translate_potential_location(location_translations, potential_location):
+    """If the input is a BaseLocation object, translate it, otherwise return it as-is."""
+    if isinstance(potential_location, Location):
+        old_location_at_vertex = potential_location.at_vertex()
+        field = potential_location.field
+
+        new_location = location_translations.get(old_location_at_vertex, None)
+        if new_location is None:
+            # No translation needed.
+            return potential_location
+        else:
+            # If necessary, add the field component to the new location before returning it.
+            if field is None:
+                return new_location
+            else:
+                return new_location.navigate_to_field(field)
+    elif isinstance(potential_location, FoldScopeLocation):
+        old_base_location = potential_location.base_location
+        new_base_location = location_translations.get(old_base_location, old_base_location)
+        fold_path = potential_location.fold_path
+        fold_field = potential_location.field
+        return FoldScopeLocation(new_base_location, fold_path, field=fold_field)
+    else:
+        return potential_location
+
+
+def make_location_rewriter_visitor_fn(location_translations):
+    """Return a visitor function that is able to replace locations with equivalent locations."""
+    def visitor_fn(expression):
+        """Expression visitor function used to rewrite expressions with updated Location data."""
+        # All CompilerEntity objects store their exact constructor input args/kwargs.
+        # To minimize the chances that we forget to update a location somewhere in an expression,
+        # we rewrite all locations that we find as arguments to expression constructors.
+        # pylint: disable=protected-access
+        new_args = [
+            translate_potential_location(location_translations, arg)
+            for arg in expression._print_args
+        ]
+        new_kwargs = {
+            kwarg_name: translate_potential_location(location_translations, kwarg_value)
+            for kwarg_name, kwarg_value in six.iteritems(expression._print_kwargs)
+        }
+        # pylint: enable=protected-access
+
+        expression_cls = type(expression)
+        return expression_cls(*new_args, **new_kwargs)
+
+    return visitor_fn

graphql_compiler/compiler/ir_lowering_gremlin/__init__.py

@@ -1,9 +1,11 @@
 # Copyright 2018-present Kensho Technologies, LLC.
 from .ir_lowering import (lower_coerce_type_block_type_data, lower_coerce_type_blocks,
-                          lower_folded_outputs, rewrite_filters_in_optional_blocks)
+                          lower_folded_outputs_and_context_fields,
+                          rewrite_filters_in_optional_blocks)
 from ..ir_sanity_checks import sanity_check_ir_blocks_from_frontend
-from ..ir_lowering_common import (lower_context_field_existence, merge_consecutive_filter_clauses,
-                                  optimize_boolean_expression_comparisons)
+from ..ir_lowering_common.common import (lower_context_field_existence,
+                                         merge_consecutive_filter_clauses,
+                                         optimize_boolean_expression_comparisons)
 
 
 ##############
@@ -48,6 +50,6 @@ def lower_ir(ir_blocks, query_metadata_table, type_equivalence_hints=None):
     ir_blocks = lower_coerce_type_blocks(ir_blocks)
     ir_blocks = rewrite_filters_in_optional_blocks(ir_blocks)
     ir_blocks = merge_consecutive_filter_clauses(ir_blocks)
-    ir_blocks = lower_folded_outputs(ir_blocks)
+    ir_blocks = lower_folded_outputs_and_context_fields(ir_blocks)
 
     return ir_blocks

graphql_compiler/compiler/ir_lowering_gremlin/ir_lowering.py

@@ -1,27 +1,30 @@
 # Copyright 2017-present Kensho Technologies, LLC.
 """Perform optimizations and lowering of the IR that allows the compiler to emit Gremlin queries.
 
-The compiler IR allows blocks and expressions that cannot be directly compiled to Gremlin or MATCH.
-For example, ContextFieldExistence is an Expression that returns True iff its given vertex exists,
-but the produced Gremlin and MATCH outputs for this purpose are entirely different and not easy
-to generate directly from this Expression object. An output-language-aware IR lowering step allows
-us to convert this Expression into other Expressions, using data already present in the IR,
-to simplify the final code generation step.
+The compiler IR allows blocks and expressions that cannot be directly compiled to the underlying
+database query languages. For example, ContextFieldExistence is an Expression that returns
+True iff its given vertex exists, but the produced Gremlin and MATCH outputs for this purpose
+are entirely different and not easy to generate directly from this Expression object.
+An output-language-aware IR lowering step allows us to convert this Expression into
+other Expressions, using data already present in the IR, to simplify the final code generation step.
 """
-from graphql import GraphQLList
+from graphql import GraphQLInt, GraphQLList
 from graphql.type import GraphQLInterfaceType, GraphQLObjectType, GraphQLUnionType
 import six
 
 from ...exceptions import GraphQLCompilationError
 from ...schema import GraphQLDate, GraphQLDateTime
-from ..blocks import Backtrack, CoerceType, ConstructResult, Filter, MarkLocation, Traverse
+from ..blocks import Backtrack, CoerceType, Filter, GlobalOperationsStart, MarkLocation, Traverse
 from ..compiler_entities import Expression
-from ..expressions import BinaryComposition, FoldedContextField, Literal, LocalField, NullLiteral
+from ..expressions import (
+    BinaryComposition, FoldedContextField, Literal, LocalField, NullLiteral,
+    make_type_replacement_visitor
+)
 from ..helpers import (
     STANDARD_DATE_FORMAT, STANDARD_DATETIME_FORMAT, FoldScopeLocation,
     get_only_element_from_collection, strip_non_null_from_type, validate_safe_string
 )
-from ..ir_lowering_common import extract_folds_from_ir_blocks
+from ..ir_lowering_common.common import extract_folds_from_ir_blocks
 
 
 ##################################
@@ -148,15 +151,19 @@ def validate(self):
                     u'Allowed types are {}.'
                     .format(type(block), self.folded_ir_blocks, allowed_block_types))
 
-        if not isinstance(self.field_type, GraphQLList):
-            raise ValueError(u'Invalid value of "field_type", expected a list type but got: '
-                             u'{}'.format(self.field_type))
-
-        inner_type = strip_non_null_from_type(self.field_type.of_type)
-        if isinstance(inner_type, GraphQLList):
-            raise GraphQLCompilationError(
-                u'Outputting list-valued fields in a @fold context is currently '
-                u'not supported: {} {}'.format(self.fold_scope_location, self.field_type.of_type))
+        bare_field_type = strip_non_null_from_type(self.field_type)
+        if isinstance(bare_field_type, GraphQLList):
+            inner_type = strip_non_null_from_type(bare_field_type.of_type)
+            if isinstance(inner_type, GraphQLList):
+                raise GraphQLCompilationError(
+                    u'Outputting list-valued fields in a @fold context is currently not supported: '
+                    u'{} {}'.format(self.fold_scope_location, bare_field_type.of_type))
+        elif GraphQLInt.is_same_type(bare_field_type):
+            # This needs to be implemented for @fold _x_count support.
+            raise NotImplementedError()
+        else:
+            raise ValueError(u'Invalid value of "field_type", expected a (possibly non-null) '
+                             u'list or int type but got: {}'.format(self.field_type))
 
     def to_match(self):
         """Must never be called."""
@@ -316,40 +323,42 @@ def folded_context_visitor(expression):
     return new_folded_ir_blocks
 
 
-def lower_folded_outputs(ir_blocks):
-    """Lower standard folded output fields into GremlinFoldedContextField objects."""
+def lower_folded_outputs_and_context_fields(ir_blocks):
+    """Lower standard folded output / context fields into GremlinFoldedContextField objects."""
     folds, remaining_ir_blocks = extract_folds_from_ir_blocks(ir_blocks)
 
     if not remaining_ir_blocks:
         raise AssertionError(u'Expected at least one non-folded block to remain: {} {} '
                              u'{}'.format(folds, remaining_ir_blocks, ir_blocks))
-    output_block = remaining_ir_blocks[-1]
-    if not isinstance(output_block, ConstructResult):
-        raise AssertionError(u'Expected the last non-folded block to be ConstructResult, '
-                             u'but instead was: {} {} '
-                             u'{}'.format(type(output_block), output_block, ir_blocks))
 
     # Turn folded Filter blocks into GremlinFoldedFilter blocks.
     converted_folds = {
         base_fold_location.get_location_name()[0]: _convert_folded_blocks(folded_ir_blocks)
         for base_fold_location, folded_ir_blocks in six.iteritems(folds)
     }
 
-    new_output_fields = dict()
-    for output_name, output_expression in six.iteritems(output_block.fields):
-        new_output_expression = output_expression
+    def rewriter_fn(folded_context_field):
+        """Rewrite FoldedContextField objects into GremlinFoldedContextField ones."""
+        # Get the matching folded IR blocks and put them in the new context field.
+        base_fold_location_name = folded_context_field.fold_scope_location.get_location_name()[0]
+        folded_ir_blocks = converted_folds[base_fold_location_name]
+        return GremlinFoldedContextField(
+            folded_context_field.fold_scope_location, folded_ir_blocks,
+            folded_context_field.field_type)
+
+    visitor_fn = make_type_replacement_visitor(FoldedContextField, rewriter_fn)
+
+    # Start by just appending blocks to the output list.
+    new_ir_blocks = []
+    block_collection = new_ir_blocks
+    for block in remaining_ir_blocks:
+        block_collection.append(block)
 
-        # Turn FoldedContextField expressions into GremlinFoldedContextField ones.
-        if isinstance(output_expression, FoldedContextField):
-            # Get the matching folded IR blocks and put them in the new context field.
-            base_fold_location_name = output_expression.fold_scope_location.get_location_name()[0]
-            folded_ir_blocks = converted_folds[base_fold_location_name]
-            new_output_expression = GremlinFoldedContextField(
-                output_expression.fold_scope_location, folded_ir_blocks,
-                output_expression.field_type)
+        if isinstance(block, GlobalOperationsStart):
+            # Once we see the GlobalOperationsStart, start accumulating the blocks for rewriting.
+            block_collection = []
 
-        new_output_fields[output_name] = new_output_expression
+    for block in block_collection:
+        new_ir_blocks.append(block.visit_and_update_expressions(visitor_fn))
 
-    new_ir_blocks = remaining_ir_blocks[:-1]
-    new_ir_blocks.append(ConstructResult(new_output_fields))
     return new_ir_blocks

graphql_compiler/compiler/ir_lowering_match/__init__.py

@@ -2,10 +2,12 @@
 import six
 
 from ..blocks import Filter
-from ..ir_lowering_common import (extract_optional_location_root_info,
-                                  extract_simple_optional_location_info,
-                                  lower_context_field_existence, merge_consecutive_filter_clauses,
-                                  optimize_boolean_expression_comparisons, remove_end_optionals)
+from ..ir_lowering_common.common import (extract_optional_location_root_info,
+                                         extract_simple_optional_location_info,
+                                         lower_context_field_existence,
+                                         merge_consecutive_filter_clauses,
+                                         optimize_boolean_expression_comparisons,
+                                         remove_end_optionals)
 from .ir_lowering import (lower_backtrack_blocks,
                           lower_folded_coerce_types_into_filter_blocks,
                           lower_has_substring_binary_compositions,
@@ -100,7 +102,7 @@ def lower_ir(ir_blocks, query_metadata_table, type_equivalence_hints=None):
 
     match_query = lower_comparisons_to_between(match_query)
 
-    match_query = lower_backtrack_blocks(match_query, location_types)
+    match_query = lower_backtrack_blocks(match_query, query_metadata_table)
     match_query = truncate_repeated_single_step_traversals(match_query)
     match_query = orientdb_class_with_while.workaround_type_coercions_in_recursions(match_query)