Add support for casting objects and tuples to std::json

`<json>Object` now always returns a `std::json` value containing
the output representation of the object. That is, the result is
the same as the output of `SELECT Object` in JSON mode, including
the type shape.

Now, casting a non-named tuple to `std::json` returns a JSON array,
and casting a named tuple to `std::json` returns a JSON object.

Casting non-scalar JSON values to scalars is now an error, as well as
casting a non-matching JSON scalar value.

Issue: #251.

edb/lang/edgeql/compiler/expr.py

@@ -44,6 +44,7 @@
 from . import setgen
 from . import schemactx
 from . import typegen
+from . import viewgen
 
 from . import func  # NOQA
 
@@ -318,13 +319,20 @@ def compile_Coalesce(
 @dispatch.compile.register(qlast.TypeCast)
 def compile_TypeCast(
         expr: qlast.Base, *, ctx: context.ContextLevel) -> irast.Base:
-    maintype = expr.type.maintype
+    target_typeref = typegen.ql_typeref_to_ir_typeref(expr.type, ctx=ctx)
 
     if (isinstance(expr.expr, qlast.EmptyCollection) and
-            maintype.name == 'array'):
+            target_typeref.maintype == 'array'):
         ir_expr = irast.Array()
     else:
-        ir_expr = dispatch.compile(expr.expr, ctx=ctx)
+        with ctx.new() as subctx:
+            # We use "exposed" mode in case this is a type of a cast
+            # that wants view shapes, e.g. a std::json cast.  We do
+            # this wholesale to support tuple and array casts without
+            # having to analyze the target type (which is cumbersome
+            # in QL AST).
+            subctx.expr_exposed = True
+            ir_expr = dispatch.compile(expr.expr, ctx=subctx)
 
     return setgen.ensure_set(
         _cast_expr(expr.type, ir_expr, ctx=ctx,
@@ -344,57 +352,115 @@ def _cast_expr(
         # if the expr is an empty set (or a coalesce of empty sets).
         orig_type = None
 
+    new_type = typegen.ql_typeref_to_type(ql_type, ctx=ctx)
+    new_typeref = typegen.ql_typeref_to_ir_typeref(ql_type, ctx=ctx)
+    json_t = ctx.schema.get('std::json')
+
     if isinstance(orig_type, s_types.Tuple):
-        # For tuple-to-tuple casts we generate a new tuple
-        # to simplify things on sqlgen side.
-        new_type = typegen.ql_typeref_to_type(ql_type, ctx=ctx)
-        if not isinstance(new_type, s_types.Tuple):
-            raise errors.EdgeQLError(
-                f'cannot cast tuple to {new_type.name}',
-                context=source_context)
+        if new_type.issubclass(json_t):
+            # Casting to std::json involves casting each tuple
+            # element and also keeping the cast around the whole tuple.
+            # This is to trigger the downstream logic of casting
+            # objects (in elements of the tuple).
+            elements = []
+            for i, n in enumerate(orig_type.element_types):
+                val = setgen.generated_set(
+                    irast.TupleIndirection(
+                        expr=ir_expr,
+                        name=n
+                    ),
+                    ctx=ctx
+                )
+                val.path_id = irutils.tuple_indirection_path_id(
+                    ir_expr.path_id, n, orig_type.element_types[n])
 
-        if len(orig_type.element_types) != len(new_type.element_types):
-            raise errors.EdgeQLError(
-                f'cannot cast to {new_type.name}: '
-                f'number of elements is not the same',
-                context=source_context)
+                val_type = irutils.infer_type(val, ctx.schema)
+                # Element cast
+                val = _cast_expr(ql_type, val, ctx=ctx,
+                                 source_context=source_context)
 
-        new_names = list(new_type.element_types)
+                elements.append(irast.TupleElement(name=n, val=val))
 
-        elements = []
-        for i, n in enumerate(orig_type.element_types):
-            val = setgen.generated_set(
-                irast.TupleIndirection(
-                    expr=ir_expr,
-                    name=n
-                ),
-                ctx=ctx
-            )
-            val.path_id = irutils.tuple_indirection_path_id(
-                ir_expr.path_id, n, orig_type.element_types[n])
+            new_tuple = setgen.ensure_set(
+                irast.Tuple(named=orig_type.named, elements=elements), ctx=ctx)
 
-            val_type = irutils.infer_type(val, ctx.schema)
-            new_el_name = new_names[i]
-            if val_type != new_type.element_types[new_el_name]:
-                # Element cast
-                val = _cast_expr(ql_type.subtypes[i], val, ctx=ctx,
-                                 source_context=source_context)
+            return setgen.ensure_set(
+                irast.TypeCast(expr=new_tuple, type=new_typeref), ctx=ctx)
 
-            elements.append(irast.TupleElement(name=new_el_name, val=val))
+        else:
+            # For tuple-to-tuple casts we generate a new tuple
+            # to simplify things on sqlgen side.
+            if not isinstance(new_type, s_types.Tuple):
+                raise errors.EdgeQLError(
+                    f'cannot cast tuple to {new_type.name}',
+                    context=source_context)
+
+            if len(orig_type.element_types) != len(new_type.element_types):
+                raise errors.EdgeQLError(
+                    f'cannot cast to {new_type.name}: '
+                    f'number of elements is not the same',
+                    context=source_context)
+
+            new_names = list(new_type.element_types)
+
+            elements = []
+            for i, n in enumerate(orig_type.element_types):
+                val = setgen.generated_set(
+                    irast.TupleIndirection(
+                        expr=ir_expr,
+                        name=n
+                    ),
+                    ctx=ctx
+                )
+                val.path_id = irutils.tuple_indirection_path_id(
+                    ir_expr.path_id, n, orig_type.element_types[n])
+
+                val_type = irutils.infer_type(val, ctx.schema)
+                new_el_name = new_names[i]
+                if val_type != new_type.element_types[new_el_name]:
+                    # Element cast
+                    val = _cast_expr(ql_type.subtypes[i], val, ctx=ctx,
+                                     source_context=source_context)
 
-        return irast.Tuple(named=new_type.named, elements=elements)
+                elements.append(irast.TupleElement(name=new_el_name, val=val))
+
+            return irast.Tuple(named=new_type.named, elements=elements)
 
     elif isinstance(ir_expr, irast.EmptySet):
         # For the common case of casting an empty set, we simply
         # generate a new EmptySet node of the requested type.
-        scls = typegen.ql_typeref_to_type(ql_type, ctx=ctx)
-        return irutils.new_empty_set(ctx.schema, scls=scls,
+        return irutils.new_empty_set(ctx.schema, scls=new_type,
                                      alias=ir_expr.path_id.target.name.name)
 
+    elif (isinstance(ir_expr, irast.Set) and
+            isinstance(ir_expr.expr, irast.Array)):
+        if new_type.issubclass(json_t):
+            el_type = ql_type
+        elif not isinstance(new_type, s_types.Array):
+            raise errors.EdgeQLError(
+                f'cannot cast array to {new_type.name}',
+                context=source_context)
+        else:
+            el_type = ql_type.subtypes[0]
+
+        casted_els = []
+        for el in ir_expr.expr.elements:
+            el = _cast_expr(el_type, el, ctx=ctx,
+                            source_context=source_context)
+            casted_els.append(el)
+
+        ir_expr.expr = irast.Array(elements=casted_els)
+        return setgen.ensure_set(
+            irast.TypeCast(expr=ir_expr, type=new_typeref), ctx=ctx)
+
     else:
-        typ = typegen.ql_typeref_to_ir_typeref(ql_type, ctx=ctx)
+        if new_type.issubclass(json_t) and ir_expr.path_id.is_objtype_path():
+            # JSON casts of objects are special: we want the full shape
+            # and not just an identity.
+            viewgen.compile_view_shapes(ir_expr, ctx=ctx)
+
         return setgen.ensure_set(
-            irast.TypeCast(expr=ir_expr, type=typ), ctx=ctx)
+            irast.TypeCast(expr=ir_expr, type=new_typeref), ctx=ctx)
 
 
 @dispatch.compile.register(qlast.TypeFilter)
@@ -451,7 +517,7 @@ def compile_Indirection(
             raise ValueError('unexpected indirection node: '
                              '{!r}'.format(indirection_el))
 
-    return node
+    return setgen.ensure_set(node, ctx=ctx)
 
 
 def try_fold_arithmetic_binop(

edb/lang/edgeql/compiler/func.py

@@ -76,7 +76,8 @@ def compile_FunctionCall(
 
         fixup_param_scope(funcobj, args, kwargs, ctx=fctx)
 
-        node = irast.FunctionCall(func=funcobj, args=args, kwargs=kwargs)
+        node = irast.FunctionCall(func=funcobj, args=args, kwargs=kwargs,
+                                  context=expr.context)
 
         if funcobj.initial_value is not None:
             rtype = irutils.infer_type(node, fctx.schema)

edb/lang/edgeql/compiler/setgen.py

@@ -563,6 +563,7 @@ def new_expression_set(
         path_id=path_id,
         scls=result_type,
         expr=ir_expr,
+        context=ir_expr.context,
         ctx=ctx
     )
 

edb/lang/edgeql/compiler/viewgen.py

@@ -685,3 +685,13 @@ def _compile_view_shapes_in_tuple(
         ctx: context.ContextLevel) -> None:
     for element in expr.elements:
         compile_view_shapes(element.val, ctx=ctx)
+
+
+@compile_view_shapes.register(irast.Array)
+def _compile_view_shapes_in_array(
+        expr: irast.Array, *,
+        rptr: typing.Optional[irast.Pointer]=None,
+        parent_view_type: typing.Optional[s_types.ViewType]=None,
+        ctx: context.ContextLevel) -> None:
+    for element in expr.elements:
+        compile_view_shapes(element, ctx=ctx)

edb/lang/ir/utils.py

@@ -17,6 +17,9 @@
 #
 
 
+import json
+import typing
+
 from edb.lang.common import ast
 
 from edb.lang.schema import objtypes as s_objtypes
@@ -25,6 +28,7 @@
 from edb.lang.schema import pointers as s_pointers
 from edb.lang.schema import schema as s_schema
 from edb.lang.schema import sources as s_sources  # NOQA
+from edb.lang.schema import types as s_types  # NOQA
 
 from . import ast as irast
 from .inference import amend_empty_set_type  # NOQA
@@ -260,3 +264,17 @@ def type_indirection_path_id(path_id, target_type, *, optional: bool,
         s_pointers.PointerDirection.Outbound,
         target_type
     )
+
+
+def get_source_context_as_json(expr: irast.Base) -> typing.Optional[str]:
+    if expr.context:
+        details = json.dumps({
+            'line': expr.context.start.line,
+            'column': expr.context.start.column,
+            'name': expr.context.name,
+        })
+
+    else:
+        details = None
+
+    return details

edb/server/pgsql/ast.py

@@ -509,6 +509,12 @@ def _is_nullable(self, kwargs: typing.Dict[str, object],
         return nullable
 
 
+class NamedFuncArg(Base):
+
+    name: str
+    val: Base
+
+
 class Indices(Base):
     """Array subscript or slice bounds."""
 

edb/server/pgsql/codegen.py

@@ -570,6 +570,10 @@ def visit_FuncCall(self, node):
         if node.with_ordinality:
             self.write(' WITH ORDINALITY')
 
+    def visit_NamedFuncArg(self, node):
+        self.write(common.quote_ident(node.name), ' => ')
+        self.visit(node.val)
+
     def visit_SubLink(self, node):
         if node.type == pgast.SubLinkType.EXISTS:
             self.write('EXISTS')

edb/server/pgsql/compiler/expr.py

@@ -19,7 +19,6 @@
 
 """Compilation handlers for non-statement expressions."""
 
-import json
 import typing
 
 from edb.lang.common import ast
@@ -121,7 +120,7 @@ def compile_Parameter(
     result = pgast.ParamRef(number=index)
     return typecomp.cast(
         result, source_type=expr.type, target_type=expr.type,
-        force=True, env=ctx.env)
+        ir_expr=expr, force=True, env=ctx.env)
 
 
 @dispatch.compile.register(irast.Constant)
@@ -130,7 +129,7 @@ def compile_Constant(
     result = pgast.Constant(val=expr.value)
     result = typecomp.cast(
         result, source_type=expr.type, target_type=expr.type,
-        force=True, env=ctx.env)
+        ir_expr=expr, force=True, env=ctx.env)
     return result
 
 
@@ -148,13 +147,14 @@ def compile_TypeCast(
 
         return typecomp.cast(
             pg_expr, source_type=target_type,
-            target_type=target_type, force=True, env=ctx.env)
+            target_type=target_type, ir_expr=expr.expr,
+            force=True, env=ctx.env)
 
     else:
         source_type = _infer_type(expr.expr, ctx=ctx)
         return typecomp.cast(
             pg_expr, source_type=source_type, target_type=target_type,
-            env=ctx.env)
+            ir_expr=expr.expr, env=ctx.env)
 
 
 @dispatch.compile.register(irast.IndexIndirection)
@@ -170,7 +170,7 @@ def compile_IndexIndirection(
     arg_type = _infer_type(expr.expr, ctx=ctx)
     # line, column and filename are captured here to be used with the
     # error message
-    exc_details = get_exc_details(expr.index)
+    srcctx = pgast.Constant(val=irutils.get_source_context_as_json(expr.index))
 
     with ctx.new() as subctx:
         subctx.expr_exposed = False
@@ -198,7 +198,7 @@ def compile_IndexIndirection(
 
             return pgast.FuncCall(
                 name=('edgedb', '_json_index'),
-                args=[subj, index, exc_details]
+                args=[subj, index, srcctx]
             )
 
         is_string = b.name == 'std::str'
@@ -231,12 +231,12 @@ def compile_IndexIndirection(
     if is_string:
         result = pgast.FuncCall(
             name=('edgedb', '_string_index'),
-            args=[subj, index, exc_details]
+            args=[subj, index, srcctx]
         )
     else:
         result = pgast.FuncCall(
             name=('edgedb', '_array_index'),
-            args=[subj, index, exc_details]
+            args=[subj, index, srcctx]
         )
 
     return result
@@ -670,17 +670,3 @@ def _infer_type(
         expr: irast.Base, *,
         ctx: context.CompilerContextLevel) -> s_obj.Object:
     return irutils.infer_type(expr, schema=ctx.env.schema)
-
-
-def get_exc_details(expr: irast.Base) -> pgast.Base:
-    if expr.context:
-        details = pgast.Constant(val=json.dumps({
-            'line': expr.context.start.line,
-            'column': expr.context.start.column,
-            'name': expr.context.name,
-        }).encode('utf-8'))
-
-    else:
-        details = pgast.Constant(val=None)
-
-    return details