intarray sort returns wrong result

From 262ca22b6efd1bd2dc2437bf2c2c83cd71096ccf Mon Sep 17 00:00:00 2001
From: Junwang Zhao <zhjwpku@gmail.com>
Date: Tue, 12 Nov 2024 00:22:00 +0000
Subject: [PATCH v1] int array support multi dim

Signed-off-by: Junwang Zhao <zhjwpku@gmail.com>
---
 contrib/intarray/_int.h      | 12 ++++++++++++
 contrib/intarray/_int_op.c   | 36 +++++++++++++++++++++++++++++++++---
 contrib/intarray/_int_tool.c | 28 ++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/contrib/intarray/_int.h b/contrib/intarray/_int.h
index 0352cbd368..2b2d6906ae 100644
--- a/contrib/intarray/_int.h
+++ b/contrib/intarray/_int.h
@@ -188,4 +188,16 @@ int			compDESC(const void *a, const void *b);
 				  (direction) ? compASC : compDESC ); \
 	} while(0)
 
+
+int			compArrayASC(const void *a, const void *b, void *arg);
+int			compArrayDESC(const void *a, const void *b, void *arg);
+
+/* sort, either ascending or descending */
+#define QSORT_ARRAY(a, direction, dim0, step) \
+	do { \
+		if (dim0 > 1) \
+			qsort_arg((void*) ARRPTR(a), dim0, sizeof(int32) * step, \
+				  (direction) ? compArrayASC : compArrayDESC, &step); \
+	} while(0)
+
 #endif							/* ___INT_H__ */
diff --git a/contrib/intarray/_int_op.c b/contrib/intarray/_int_op.c
index 5b164f6788..03f0deea84 100644
--- a/contrib/intarray/_int_op.c
+++ b/contrib/intarray/_int_op.c
@@ -199,11 +199,21 @@ sort(PG_FUNCTION_ARGS)
 	int32		dc = (dirstr) ? VARSIZE_ANY_EXHDR(dirstr) : 0;
 	char	   *d = (dirstr) ? VARDATA_ANY(dirstr) : NULL;
 	int			dir = -1;
+	int ndim,
+	nelems,
+	step;
 
 	CHECKARRVALID(a);
 	if (ARRNELEMS(a) < 2)
 		PG_RETURN_POINTER(a);
 
+	ndim = ARR_NDIM(a);
+	if (ndim < 1)
+		PG_RETURN_POINTER(a);
+
+	nelems = ARRNELEMS(a);
+	step = nelems / ARR_DIMS(a)[0];
+
 	if (dirstr == NULL || (dc == 3
 						   && (d[0] == 'A' || d[0] == 'a')
 						   && (d[1] == 'S' || d[1] == 's')
@@ -219,7 +229,7 @@ sort(PG_FUNCTION_ARGS)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 				 errmsg("second parameter must be \"ASC\" or \"DESC\"")));
-	QSORT(a, dir);
+	QSORT_ARRAY(a, dir, ARR_DIMS(a)[0], step);
 	PG_RETURN_POINTER(a);
 }
 
@@ -227,9 +237,19 @@ Datum
 sort_asc(PG_FUNCTION_ARGS)
 {
 	ArrayType  *a = PG_GETARG_ARRAYTYPE_P_COPY(0);
+	int ndim,
+	nelems,
+	step;
 
 	CHECKARRVALID(a);
-	QSORT(a, 1);
+
+	ndim = ARR_NDIM(a);
+	if (ndim < 1)
+		PG_RETURN_POINTER(a);
+
+	nelems = ARRNELEMS(a);
+	step = nelems / ARR_DIMS(a)[0];
+	QSORT_ARRAY(a, 1, ARR_DIMS(a)[0], step);
 	PG_RETURN_POINTER(a);
 }
 
@@ -237,9 +257,19 @@ Datum
 sort_desc(PG_FUNCTION_ARGS)
 {
 	ArrayType  *a = PG_GETARG_ARRAYTYPE_P_COPY(0);
+	int ndim,
+	nelems,
+	step;
 
 	CHECKARRVALID(a);
-	QSORT(a, 0);
+
+	ndim = ARR_NDIM(a);
+	if (ndim < 1)
+		PG_RETURN_POINTER(a);
+
+	nelems = ARRNELEMS(a);
+	step = nelems / ARR_DIMS(a)[0];
+	QSORT_ARRAY(a, 0, ARR_DIMS(a)[0], step);
 	PG_RETURN_POINTER(a);
 }
 
diff --git a/contrib/intarray/_int_tool.c b/contrib/intarray/_int_tool.c
index c85280c842..cae032d051 100644
--- a/contrib/intarray/_int_tool.c
+++ b/contrib/intarray/_int_tool.c
@@ -405,3 +405,31 @@ compDESC(const void *a, const void *b)
 {
 	return pg_cmp_s32(*(const int32 *) b, *(const int32 *) a);
 }
+
+int
+compArrayASC(const void *a, const void *b, void *arg)
+{
+	int cnt = *(int*)arg;
+	for (int i = 0; i < cnt; i++) {
+		int cmp = pg_cmp_s32(((const int32 *) a)[i], ((const int32 *) b)[i]);
+		if (cmp == 0) {
+			continue;
+		}
+		return cmp;
+	}
+	return 0;
+}
+
+int
+compArrayDESC(const void *a, const void *b, void *arg)
+{
+	int cnt = *(int*)arg;
+	for (int i = 0; i < cnt; i++) {
+		int cmp = pg_cmp_s32(((const int32 *) b)[i], ((const int32 *) a)[i]);
+		if (cmp == 0) {
+			continue;
+		}
+		return cmp;
+	}
+	return 0;
+}
-- 
2.39.5