0003-Subscripting-for-array-v12.patch

From 1c852532355fddcc3dbeaca73a1ba2266831776a Mon Sep 17 00:00:00 2001
From: Dmitrii Dolgov <dmitrii.dolgov@zalando.de>
Date: Fri, 20 Jul 2018 22:22:56 +0200
Subject: [PATCH 3/5] Subscripting for array

---
 src/backend/utils/adt/arrayfuncs.c   | 290 +++++++++++++++++++++++++++++++++++
 src/include/catalog/pg_proc.dat      |   7 +
 src/test/regress/expected/arrays.out |  12 +-
 src/test/regress/sql/arrays.sql      |   4 +-
 4 files changed, 305 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.dat b/src/include/catalog/pg_proc.dat
index 40d54ed..b9d9a54 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -10176,6 +10176,13 @@
   proargnames => '{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}',
   prosrc => 'pg_control_init' },
 
+{ oid => '4004',
+  descr => 'Array subscripting logic',
+  proname => 'array_subscript_handler',
+  prorettype => 'internal',
+  proargtypes => 'internal',
+  prosrc => 'array_subscript_handler' },
+
 # collation management functions
 { oid => '3445', descr => 'import collations from operating system',
   proname => 'pg_import_system_collations', procost => '100',
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