0003-Subscripting-for-array-v9.patch

From 665b879383f44049edbe7ee7ded94d9db6dc65ed Mon Sep 17 00:00:00 2001
From: Dmitrii Dolgov <dmitrii.dolgov@zalando.de>
Date: Thu, 22 Feb 2018 17:56:39 +0100
Subject: [PATCH 3/5] Subscripting for array

---
 src/backend/utils/adt/arrayfuncs.c   | 290 +++++++++++++++++++++++++++++++++++
 src/include/catalog/pg_proc.h        |   3 +
 src/test/regress/expected/arrays.out |  12 +-
 src/test/regress/sql/arrays.sql      |   4 +-
 4 files changed, 301 insertions(+), 8 deletions(-)

diff --git a/src/backend/utils/adt/arrayfuncs.c b/src/backend/utils/adt/arrayfuncs.c
index 0cbdbe5..a4dbcad 100644
--- a/src/backend/utils/adt/arrayfuncs.c
+++ b/src/backend/utils/adt/arrayfuncs.c
@@ -25,13 +25,20 @@
 #include "catalog/pg_type.h"
 #include "funcapi.h"
 #include "libpq/pqformat.h"
+#include "nodes/makefuncs.h"
+#include "nodes/nodeFuncs.h"
+#include "executor/execExpr.h"
 #include "utils/array.h"
 #include "utils/arrayaccess.h"
 #include "utils/builtins.h"
 #include "utils/datum.h"
+#include "utils/fmgroids.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
+#include "utils/syscache.h"
 #include "utils/typcache.h"
+#include "parser/parse_node.h"
+#include "parser/parse_coerce.h"
 
 
 /*
@@ -158,7 +165,14 @@ static int width_bucket_array_variable(Datum operand,
 							ArrayType *thresholds,
 							Oid collation,
 							TypeCacheEntry *typentry);
+static SubscriptingRef *array_subscript_prepare(bool isAssignment, SubscriptingRef *sbsref);
+static SubscriptingRef *array_subscript_validate(bool isAssignment, SubscriptingRef *sbsref,
+												 ParseState *pstate);
 
+static Datum array_subscript_fetch(Datum containerSource,
+								   SubscriptingRefState *sbstate);
+static Datum array_subscript_assign(Datum containerSource,
+									SubscriptingRefState *sbstate);
 
 /*
  * array_in :
@@ -6558,3 +6572,279 @@ width_bucket_array_variable(Datum operand,
 
 	return left;
 }
+
+/*
+ * Perform an actual data extraction or modification for the array
+ * subscripting. As a result the extracted Datum or the modified containers
+ * value will be returned.
+ */
+Datum
+array_subscript_assign(Datum containerSource, SubscriptingRefState *sbstate)
+{
+	bool						is_slice = (sbstate->numlower != 0);
+	IntArray					u_index, l_index;
+	bool						eisnull = sbstate->resnull;
+	int							i = 0;
+
+	if (sbstate->refelemlength == 0)
+	{
+		/* do one-time catalog lookups for type info */
+		get_typlenbyvalalign(sbstate->refelemtype,
+							 &sbstate->refelemlength,
+							 &sbstate->refelembyval,
+							 &sbstate->refelemalign);
+	}
+
+	for(i = 0; i < sbstate->numupper; i++)
+		u_index.indx[i] = DatumGetInt32(sbstate->upperindex[i]);
+
+	if (is_slice)
+	{
+		for(i = 0; i < sbstate->numlower; i++)
+			l_index.indx[i] = DatumGetInt32(sbstate->lowerindex[i]);
+	}
+
+	/*
+	 * For assignment to varlena arrays, we handle a NULL original array
+	 * by substituting an empty (zero-dimensional) array; insertion of the
+	 * new element will result in a singleton array value.  It does not
+	 * matter whether the new element is NULL.
+	 */
+	if (eisnull)
+	{
+		containerSource = PointerGetDatum(construct_empty_array(sbstate->refelemtype));
+		sbstate->resnull = false;
+		eisnull = false;
+	}
+
+	if (!is_slice)
+		return array_set_element(containerSource, sbstate->numupper,
+								 u_index.indx,
+								 sbstate->replacevalue,
+								 sbstate->replacenull,
+								 sbstate->refattrlength,
+								 sbstate->refelemlength,
+								 sbstate->refelembyval,
+								 sbstate->refelemalign);
+	else
+		return array_set_slice(containerSource, sbstate->numupper,
+							   u_index.indx, l_index.indx,
+							   sbstate->upperprovided,
+							   sbstate->lowerprovided,
+							   sbstate->replacevalue,
+							   sbstate->replacenull,
+							   sbstate->refattrlength,
+							   sbstate->refelemlength,
+							   sbstate->refelembyval,
+							   sbstate->refelemalign);
+}
+
+Datum
+array_subscript_fetch(Datum containerSource, SubscriptingRefState *sbstate)
+{
+	bool							is_slice = (sbstate->numlower != 0);
+	IntArray						u_index, l_index;
+	int								i = 0;
+
+	if (sbstate->refelemlength == 0)
+	{
+		/* do one-time catalog lookups for type info */
+		get_typlenbyvalalign(sbstate->refelemtype,
+							 &sbstate->refelemlength,
+							 &sbstate->refelembyval,
+							 &sbstate->refelemalign);
+	}
+
+	for(i = 0; i < sbstate->numupper; i++)
+		u_index.indx[i] = DatumGetInt32(sbstate->upperindex[i]);
+
+	if (is_slice)
+	{
+		for(i = 0; i < sbstate->numlower; i++)
+			l_index.indx[i] = DatumGetInt32(sbstate->lowerindex[i]);
+	}
+
+	if (!is_slice)
+		return array_get_element(containerSource, sbstate->numupper,
+								 u_index.indx,
+								 sbstate->refattrlength,
+								 sbstate->refelemlength,
+								 sbstate->refelembyval,
+								 sbstate->refelemalign,
+								 &sbstate->resnull);
+	else
+		return array_get_slice(containerSource, sbstate->numupper,
+							   u_index.indx, l_index.indx,
+							   sbstate->upperprovided,
+							   sbstate->lowerprovided,
+							   sbstate->refattrlength,
+							   sbstate->refelemlength,
+							   sbstate->refelembyval,
+							   sbstate->refelemalign);
+}
+
+/*
+ * Handle array-type subscripting logic.
+ */
+Datum
+array_subscript_handler(PG_FUNCTION_ARGS)
+{
+	SubscriptRoutines *sbsroutines = (SubscriptRoutines *)
+									 palloc(sizeof(SubscriptRoutines));
+
+	sbsroutines->prepare = array_subscript_prepare;
+	sbsroutines->validate = array_subscript_validate;
+	sbsroutines->fetch = array_subscript_fetch;
+	sbsroutines->assign = array_subscript_assign;
+
+	PG_RETURN_POINTER(sbsroutines);
+}
+
+SubscriptingRef *
+array_subscript_prepare(bool isAssignment, SubscriptingRef *sbsref)
+{
+	Oid					array_type = sbsref->refcontainertype;
+	HeapTuple			type_tuple_container;
+	Form_pg_type		type_struct_container;
+	bool				is_slice = sbsref->reflowerindexpr != NIL;
+
+	/* Get the type tuple for the container */
+	type_tuple_container = SearchSysCache1(TYPEOID, ObjectIdGetDatum(array_type));
+	if (!HeapTupleIsValid(type_tuple_container))
+		elog(ERROR, "cache lookup failed for type %u", array_type);
+	type_struct_container = (Form_pg_type) GETSTRUCT(type_tuple_container);
+
+	/* needn't check typisdefined since this will fail anyway */
+	sbsref->refelemtype = type_struct_container->typelem;
+
+	/* Identify type that RHS must provide */
+	if (isAssignment)
+		sbsref->refassgntype = is_slice ? sbsref->refcontainertype : sbsref->refelemtype;
+
+	ReleaseSysCache(type_tuple_container);
+
+	return sbsref;
+}
+
+SubscriptingRef *
+array_subscript_validate(bool isAssignment, SubscriptingRef *sbsref,
+					  ParseState *pstate)
+{
+	bool				is_slice = sbsref->reflowerindexpr != NIL;
+	Oid					typeneeded = InvalidOid,
+						typesource = InvalidOid;
+	Node				*new_from;
+	Node				*subexpr;
+	List				*upperIndexpr = NIL;
+	List				*lowerIndexpr = NIL;
+	ListCell			*u, *l, *s;
+
+	foreach(u, sbsref->refupperindexpr)
+	{
+		subexpr = (Node *) lfirst(u);
+
+		if (subexpr == NULL)
+		{
+			upperIndexpr = lappend(upperIndexpr, subexpr);
+			continue;
+		}
+
+		subexpr = coerce_to_target_type(pstate,
+										subexpr, exprType(subexpr),
+										INT4OID, -1,
+										COERCION_ASSIGNMENT,
+										COERCE_IMPLICIT_CAST,
+										-1);
+		if (subexpr == NULL)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("array subscript must have type integer"),
+					 parser_errposition(pstate, exprLocation(subexpr))));
+
+		upperIndexpr = lappend(upperIndexpr, subexpr);
+	}
+
+	sbsref->refupperindexpr = upperIndexpr;
+
+	forboth(l, sbsref->reflowerindexpr, s, sbsref->refindexprslice)
+	{
+		A_Indices *ai = (A_Indices *) lfirst(s);
+		subexpr = (Node *) lfirst(l);
+
+		if (subexpr == NULL && !ai->is_slice)
+		{
+			/* Make a constant 1 */
+			subexpr = (Node *) makeConst(INT4OID,
+										 -1,
+										 InvalidOid,
+										 sizeof(int32),
+										 Int32GetDatum(1),
+										 false,
+										 true);		/* pass by value */
+		}
+
+		if (subexpr == NULL)
+		{
+			lowerIndexpr = lappend(lowerIndexpr, subexpr);
+			continue;
+		}
+
+		subexpr = coerce_to_target_type(pstate,
+										subexpr, exprType(subexpr),
+										INT4OID, -1,
+										COERCION_ASSIGNMENT,
+										COERCE_IMPLICIT_CAST,
+										-1);
+		if (subexpr == NULL)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("array subscript must have type integer"),
+					 parser_errposition(pstate, exprLocation(subexpr))));
+
+		lowerIndexpr = lappend(lowerIndexpr, subexpr);
+	}
+
+	sbsref->reflowerindexpr = lowerIndexpr;
+
+	if (isAssignment)
+	{
+		SubscriptingRef *assignRef = (SubscriptingRef *) sbsref;
+		Node *assignExpr = (Node *) assignRef->refassgnexpr;
+
+		typesource = exprType(assignExpr);
+		typeneeded = is_slice ? sbsref->refcontainertype : sbsref->refelemtype;
+		new_from = coerce_to_target_type(pstate,
+										assignExpr, typesource,
+										typeneeded, sbsref->reftypmod,
+										COERCION_ASSIGNMENT,
+										COERCE_IMPLICIT_CAST,
+										-1);
+		if (new_from == NULL)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("array assignment requires type %s"
+							" but expression is of type %s",
+							format_type_be(sbsref->refelemtype),
+							format_type_be(typesource)),
+				 errhint("You will need to rewrite or cast the expression."),
+					 parser_errposition(pstate, exprLocation(assignExpr))));
+		assignRef->refassgnexpr = (Expr *) new_from;
+	}
+
+	sbsref->refnestedfunc = F_ARRAY_SUBSCRIPT_HANDLER;
+
+	/* Verify subscript list lengths are within limit */
+	if (list_length(sbsref->refupperindexpr) > MAXDIM)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("number of array dimensions (%d) exceeds the maximum allowed (%d)",
+						list_length(sbsref->refupperindexpr), MAXDIM)));
+
+	if (list_length(sbsref->reflowerindexpr) > MAXDIM)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("number of array dimensions (%d) exceeds the maximum allowed (%d)",
+						list_length(sbsref->reflowerindexpr), MAXDIM)));
+
+	return sbsref;
+}
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index 2a53213..2cef67a 100644
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -5544,6 +5544,9 @@ DESCR("pg_controldata recovery state information as a function");
 DATA(insert OID = 3444 ( pg_control_init PGNSP PGUID 12 1 0 0 0 f f f f t f v s 0 0 2249 "" "{23,23,23,23,23,23,23,23,23,16,16,23}" "{o,o,o,o,o,o,o,o,o,o,o,o}" "{max_data_alignment,database_block_size,blocks_per_segment,wal_block_size,bytes_per_wal_segment,max_identifier_length,max_index_columns,max_toast_chunk_size,large_object_chunk_size,float4_pass_by_value,float8_pass_by_value,data_page_checksum_version}" _null_ _null_ pg_control_init _null_ _null_ _null_ ));
 DESCR("pg_controldata init state information as a function");
 
+DATA(insert OID = 4004 (  array_subscript_handler PGNSP PGUID 12 1 0 0 0 f f f t f i s 1 0 "2281" "2281" _null_ _null_ _null_ _null_ _null_ array_subscript_handler _null_ _null_ _null_ ));
+DESCR("Array subscripting logic");
+
 /* collation management functions */
 DATA(insert OID = 3445 ( pg_import_system_collations PGNSP PGUID 12 100 0 0 0 f f f f t f v r 1 0 23 "4089" _null_ _null_ _null_ _null_ _null_ pg_import_system_collations _null_ _null_ _null_ ));
 DESCR("import collations from operating system");
diff --git a/src/test/regress/expected/arrays.out b/src/test/regress/expected/arrays.out
index c730563..f54b73c 100644
--- a/src/test/regress/expected/arrays.out
+++ b/src/test/regress/expected/arrays.out
@@ -190,9 +190,9 @@ select ('[0:2][0:2]={{1,2,3},{4,5,6},{7,8,9}}'::int[])[1:2][2];
 --
 -- check subscription corner cases
 --
--- More subscripts than MAXDIMS(6)
-SELECT ('{}'::int[])[1][2][3][4][5][6][7];
-ERROR:  number of array dimensions (7) exceeds the maximum allowed (6)
+-- More subscripts than MAXDIMS(12)
+SELECT ('{}'::int[])[1][2][3][4][5][6][7][8][9][10][11][12][13];
+ERROR:  number of array dimensions (13) exceeds the maximum allowed (6)
 -- NULL index yields NULL when selecting
 SELECT ('{{{1},{2},{3}},{{4},{5},{6}}}'::int[])[1][NULL][1];
  int4 
@@ -216,15 +216,15 @@ SELECT ('{{{1},{2},{3}},{{4},{5},{6}}}'::int[])[1][1:NULL][1];
 UPDATE arrtest
   SET c[NULL] = '{"can''t assign"}'
   WHERE array_dims(c) is not null;
-ERROR:  array subscript in assignment must not be null
+ERROR:  subscript in assignment must not be null
 UPDATE arrtest
   SET c[NULL:1] = '{"can''t assign"}'
   WHERE array_dims(c) is not null;
-ERROR:  array subscript in assignment must not be null
+ERROR:  subscript in assignment must not be null
 UPDATE arrtest
   SET c[1:NULL] = '{"can''t assign"}'
   WHERE array_dims(c) is not null;
-ERROR:  array subscript in assignment must not be null
+ERROR:  subscript in assignment must not be null
 -- test slices with empty lower and/or upper index
 CREATE TEMP TABLE arrtest_s (
   a       int2[],
diff --git a/src/test/regress/sql/arrays.sql b/src/test/regress/sql/arrays.sql
index 25dd4e2..fb7a319 100644
--- a/src/test/regress/sql/arrays.sql
+++ b/src/test/regress/sql/arrays.sql
@@ -106,8 +106,8 @@ select ('[0:2][0:2]={{1,2,3},{4,5,6},{7,8,9}}'::int[])[1:2][2];
 --
 -- check subscription corner cases
 --
--- More subscripts than MAXDIMS(6)
-SELECT ('{}'::int[])[1][2][3][4][5][6][7];
+-- More subscripts than MAXDIMS(12)
+SELECT ('{}'::int[])[1][2][3][4][5][6][7][8][9][10][11][12][13];
 -- NULL index yields NULL when selecting
 SELECT ('{{{1},{2},{3}},{{4},{5},{6}}}'::int[])[1][NULL][1];
 SELECT ('{{{1},{2},{3}},{{4},{5},{6}}}'::int[])[1][NULL:1][1];
-- 
2.7.4