[PL/pgSQL] %TYPE and array declaration

diff --git a/src/pl/plpgsql/src/gram.y b/src/pl/plpgsql/src/gram.y
index 92b54dd..efdc833 100644
--- a/src/pl/plpgsql/src/gram.y
+++ b/src/pl/plpgsql/src/gram.y
@@ -76,6 +76,8 @@ static	PLpgSQL_expr	*read_sql_expression2(int until, int until2,
 											  int *endtoken);
 static	PLpgSQL_expr	*read_sql_stmt(const char *sqlstart);
 static	PLpgSQL_type	*read_datatype(int tok);
+static	int				 read_array_dims(int tok);
+static	int				 read_single_array_dim(int tok);
 static	PLpgSQL_stmt	*make_execsql_stmt(int firsttoken, int location);
 static	PLpgSQL_stmt_fetch *read_fetch_direction(void);
 static	void			 complete_direction(PLpgSQL_stmt_fetch *fetch,
@@ -2503,7 +2505,13 @@ read_datatype(int tok)
 			if (tok_is_keyword(tok, &yylval,
 							   K_TYPE, "type"))
 			{
-				result = plpgsql_parse_wordtype(dtname);
+				/*
+				 * Now try to determine if there are any array indicators
+				 * [] or [ICONST] follow type definition, and then use
+				 * this information during type construction.
+				 */
+				const int arrndim = read_array_dims(tok);
+				result = plpgsql_parse_wordtype(dtname, arrndim);
 				if (result)
 					return result;
 			}
@@ -2527,7 +2535,8 @@ read_datatype(int tok)
 			if (tok_is_keyword(tok, &yylval,
 							   K_TYPE, "type"))
 			{
-				result = plpgsql_parse_cwordtype(dtnames);
+				const int arrndim = read_array_dims(tok);
+				result = plpgsql_parse_cwordtype(dtnames, arrndim);
 				if (result)
 					return result;
 			}
@@ -2582,6 +2591,67 @@ read_datatype(int tok)
 	return result;
 }
 
+/*
+ * Parse all items of array declaration: ('[' ICONST? ']')*
+ */
+static int
+read_array_dims(int tok) {
+	int n = 0;
+	int k;
+
+	while (true) {
+		tok = yylex();
+		k = read_single_array_dim(tok);
+		switch (k) {
+			case 1: /* ok */
+				n += 1;
+				break;
+
+			case 0: /* probably wrong syntax, but don't panic here */
+				plpgsql_push_back_token(tok);
+				return 0;
+
+			case -1: /* end of array def */
+				plpgsql_push_back_token(tok);
+				return n;
+
+			default:
+				Assert(false);
+		}
+	}
+}
+
+/* Parse single item of array declaration: '[' ']' or '[' ICONST ']' */
+static int
+read_single_array_dim(int tok) {
+	int tok2;
+	if (tok == '[') {
+		tok = yylex();
+		if (tok == ']') {
+			/* variant [] */
+			return 1;
+		}
+		else {
+			/* variant [ICONST] */
+			if (tok == ICONST) {
+				/* valid only if ICONST > 0 */
+				if (yylval.ival <= 0)
+					return 0;
+
+				tok2 = yylex();
+				if (tok2 == ']')
+					return 1;
+				else
+					plpgsql_push_back_token(tok2);
+			}
+			return 0;
+		}
+	}
+	else
+		return -1;
+}
+
+
 static PLpgSQL_stmt *
 make_execsql_stmt(int firsttoken, int location)
 {
diff --git a/src/pl/plpgsql/src/pl_comp.c b/src/pl/plpgsql/src/pl_comp.c
index d22fa68..826eabe 100644
--- a/src/pl/plpgsql/src/pl_comp.c
+++ b/src/pl/plpgsql/src/pl_comp.c
@@ -1588,6 +1588,35 @@ plpgsql_parse_tripword(char *word1, char *word2, char *word3,
 
 
 /* ----------
+ * build_arraydatatype	Try to construct array type from a type
+ *
+ * Function is used internally by plpgsql_parse_wordtype and
+ * plpgsql_parse_cwordtype to handle syntax word%TYPE[].
+ * 
+ * Returns datatype struct, or NULL if no match found for word.
+ * ----------
+ */
+static PLpgSQL_type *
+build_arraydatatype(Oid typoid, int32 typmod, Oid collation) {
+	Oid elem_typoid;
+
+	if (type_is_array(typoid))
+		/* copy array type */
+		return plpgsql_build_datatype(typoid, typmod, collation);
+	else
+	{
+		/* get array type for given scalar type */
+		elem_typoid = get_array_type(typoid);
+		if (elem_typoid != InvalidOid) 
+			return plpgsql_build_datatype(elem_typoid, typmod, collation);
+		else
+			/* typoid can't be element type of an array */
+			return NULL;
+	}
+}
+
+
+/* ----------
  * plpgsql_parse_wordtype	The scanner found word%TYPE. word can be
  *				a variable name or a basetype.
  *
@@ -1595,7 +1624,7 @@ plpgsql_parse_tripword(char *word1, char *word2, char *word3,
  * ----------
  */
 PLpgSQL_type *
-plpgsql_parse_wordtype(char *ident)
+plpgsql_parse_wordtype(char *ident, int arrndim)
 {
 	PLpgSQL_type *dtype;
 	PLpgSQL_nsitem *nse;
@@ -1613,7 +1642,12 @@ plpgsql_parse_wordtype(char *ident)
 		switch (nse->itemtype)
 		{
 			case PLPGSQL_NSTYPE_VAR:
-				return ((PLpgSQL_var *) (plpgsql_Datums[nse->itemno]))->datatype;
+				dtype = ((PLpgSQL_var *) (plpgsql_Datums[nse->itemno]))->datatype;
+				if (arrndim <= 0)
+					return dtype;
+				else
+					return build_arraydatatype(dtype->typoid,
+								-1, plpgsql_curr_compile->fn_input_collation);
 
 				/* XXX perhaps allow REC/ROW here? */
 
@@ -1638,7 +1672,11 @@ plpgsql_parse_wordtype(char *ident)
 			return NULL;
 		}
 
-		dtype = build_datatype(typeTup, -1,
+		if (arrndim <= 0)
+			dtype = build_datatype(typeTup, -1,
+							   plpgsql_curr_compile->fn_input_collation);
+		else
+			dtype = build_arraydatatype(HeapTupleGetOid(typeTup), -1,
 							   plpgsql_curr_compile->fn_input_collation);
 
 		ReleaseSysCache(typeTup);
@@ -1658,7 +1696,7 @@ plpgsql_parse_wordtype(char *ident)
  * ----------
  */
 PLpgSQL_type *
-plpgsql_parse_cwordtype(List *idents)
+plpgsql_parse_cwordtype(List *idents, int arrndim)
 {
 	PLpgSQL_type *dtype = NULL;
 	PLpgSQL_nsitem *nse;
@@ -1690,6 +1728,9 @@ plpgsql_parse_cwordtype(List *idents)
 		if (nse != NULL && nse->itemtype == PLPGSQL_NSTYPE_VAR)
 		{
 			dtype = ((PLpgSQL_var *) (plpgsql_Datums[nse->itemno]))->datatype;
+			if (arrndim > 0)
+				dtype = build_arraydatatype(dtype->typoid, -1,
+							plpgsql_curr_compile->fn_input_collation);
 			goto done;
 		}
 
@@ -1749,9 +1790,14 @@ plpgsql_parse_cwordtype(List *idents)
 	 * return it
 	 */
 	MemoryContextSwitchTo(oldCxt);
-	dtype = build_datatype(typetup,
-						   attrStruct->atttypmod,
-						   attrStruct->attcollation);
+	if (arrndim <= 0)
+		dtype = build_datatype(typetup,
+							   attrStruct->atttypmod,
+							   attrStruct->attcollation);
+	else
+		dtype = build_arraydatatype(HeapTupleGetOid(typetup),
+									 attrStruct->atttypmod,
+									 attrStruct->attcollation);
 	MemoryContextSwitchTo(compile_tmp_cxt);
 
 done:
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index c543f1c..f2ea428 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -871,8 +871,8 @@ extern bool plpgsql_parse_dblword(char *word1, char *word2,
 					  PLwdatum *wdatum, PLcword *cword);
 extern bool plpgsql_parse_tripword(char *word1, char *word2, char *word3,
 					   PLwdatum *wdatum, PLcword *cword);
-extern PLpgSQL_type *plpgsql_parse_wordtype(char *ident);
-extern PLpgSQL_type *plpgsql_parse_cwordtype(List *idents);
+extern PLpgSQL_type *plpgsql_parse_wordtype(char *ident, int arrndim);
+extern PLpgSQL_type *plpgsql_parse_cwordtype(List *idents, int arrndim);
 extern PLpgSQL_type *plpgsql_parse_wordrowtype(char *ident);
 extern PLpgSQL_type *plpgsql_parse_cwordrowtype(List *idents);
 extern PLpgSQL_type *plpgsql_build_datatype(Oid typeOid, int32 typmod,