From ec3a5513066c0d37ccda6d4ebc122ff1a4d45148 Mon Sep 17 00:00:00 2001
From: David Rowley <dgrowley@gmail.com>
Date: Mon, 5 Jul 2021 20:03:27 +1200
Subject: [PATCH v2] Use lookup table for units in pg_size_pretty and
 pg_size_bytes

We've grown 2 versions of pg_size_pretty over the years, one for BIGINT
and one for NUMERIC.  Both should output the same, but keeping them in
sync would be harder than needed due to neither function sharing a source
of truth about which units to use.

Here we add a static array that defines the units that we recognize and
have both pg_size_pretty and pg_size_pretty_numeric use it.  This will
make adding any units in the future a very simple task.

The table contains all information required to allow us to also modify
pg_size_bytes to also use it there too, so adjust that too.

There are no behavioral changes here.
---
 src/backend/utils/adt/dbsize.c | 158 +++++++++++++++------------------
 1 file changed, 73 insertions(+), 85 deletions(-)

diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index 9c39e7d3b3..9a15625e89 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -34,6 +34,27 @@
 /* Divide by two and round towards positive infinity. */
 #define half_rounded(x)   (((x) + ((x) < 0 ? 0 : 1)) / 2)
 
+/* Units used in pg_size_pretty functions.  All units must be powers of 2 */
+struct size_pretty_unit
+{
+	const char *name;			/* bytes, kB, MB, GB etc */
+	uint32		limit;			/* upper limit, prior to half rounding after
+								 * converting to this unit. */
+	bool		round;			/* do half rounding for this unit */
+	uint8		unitbits;		/* (1 << unitbits) bytes to make 1 of this
+								 * unit */
+};
+
+/* When adding units here also update the error message in pg_size_bytes */
+static const struct size_pretty_unit size_pretty_units[] = {
+	{"bytes", 10 * 1024, false, 0},
+	{"kB", 20 * 1024 - 1, true, 10},
+	{"MB", 20 * 1024 - 1, true, 20},
+	{"GB", 20 * 1024 - 1, true, 30},
+	{"TB", 20 * 1024 - 1, true, 40},
+	{NULL, 0, false, 0}
+};
+
 /* Return physical size of directory contents, or 0 if dir doesn't exist */
 static int64
 db_dir_size(const char *path)
@@ -535,37 +556,28 @@ pg_size_pretty(PG_FUNCTION_ARGS)
 {
 	int64		size = PG_GETARG_INT64(0);
 	char		buf[64];
-	int64		limit = 10 * 1024;
-	int64		limit2 = limit * 2 - 1;
+	const struct size_pretty_unit *unit;
 
-	if (Abs(size) < limit)
-		snprintf(buf, sizeof(buf), INT64_FORMAT " bytes", size);
-	else
+	for (unit = size_pretty_units; unit->name != NULL; unit++)
 	{
-		size >>= 9;				/* keep one extra bit for rounding */
-		if (Abs(size) < limit2)
-			snprintf(buf, sizeof(buf), INT64_FORMAT " kB",
-					 half_rounded(size));
-		else
+		/* use this unit there are no more units or we're below the limit */
+		if (unit[1].name == NULL || Abs(size) < unit->limit)
 		{
-			size >>= 10;
-			if (Abs(size) < limit2)
-				snprintf(buf, sizeof(buf), INT64_FORMAT " MB",
-						 half_rounded(size));
-			else
-			{
-				size >>= 10;
-				if (Abs(size) < limit2)
-					snprintf(buf, sizeof(buf), INT64_FORMAT " GB",
-							 half_rounded(size));
-				else
-				{
-					size >>= 10;
-					snprintf(buf, sizeof(buf), INT64_FORMAT " TB",
-							 half_rounded(size));
-				}
-			}
+			if (unit->round)
+				size = half_rounded(size);
+
+			snprintf(buf, sizeof(buf), INT64_FORMAT " %s", size, unit->name);
+			break;
 		}
+
+		/*
+		 * We may need to shift 1 bit less than the difference between this
+		 * and the next unit if the next unit has half rounding.  Or we may
+		 * need to shift an extra bit if this unit has half rounding and the
+		 * next one does not.
+		 */
+		size >>= (unit[1].unitbits - unit->unitbits - (unit[1].round == true)
+				  + (unit->round == true));
 	}
 
 	PG_RETURN_TEXT_P(cstring_to_text(buf));
@@ -636,56 +648,34 @@ Datum
 pg_size_pretty_numeric(PG_FUNCTION_ARGS)
 {
 	Numeric		size = PG_GETARG_NUMERIC(0);
-	Numeric		limit,
-				limit2;
-	char	   *result;
-
-	limit = int64_to_numeric(10 * 1024);
-	limit2 = int64_to_numeric(10 * 1024 * 2 - 1);
+	char	   *result = NULL;
+	const struct size_pretty_unit *unit;
 
-	if (numeric_is_less(numeric_absolute(size), limit))
+	for (unit = size_pretty_units; unit->name != NULL; unit++)
 	{
-		result = psprintf("%s bytes", numeric_to_cstring(size));
-	}
-	else
-	{
-		/* keep one extra bit for rounding */
-		/* size >>= 9 */
-		size = numeric_shift_right(size, 9);
+		unsigned int shiftby;
 
-		if (numeric_is_less(numeric_absolute(size), limit2))
-		{
-			size = numeric_half_rounded(size);
-			result = psprintf("%s kB", numeric_to_cstring(size));
-		}
-		else
+		/* use this unit there are no more units or we're below the limit */
+		if (unit[1].name == NULL ||
+			numeric_is_less(numeric_absolute(size),
+							int64_to_numeric(unit->limit)))
 		{
-			/* size >>= 10 */
-			size = numeric_shift_right(size, 10);
-			if (numeric_is_less(numeric_absolute(size), limit2))
-			{
+			if (unit->round)
 				size = numeric_half_rounded(size);
-				result = psprintf("%s MB", numeric_to_cstring(size));
-			}
-			else
-			{
-				/* size >>= 10 */
-				size = numeric_shift_right(size, 10);
-
-				if (numeric_is_less(numeric_absolute(size), limit2))
-				{
-					size = numeric_half_rounded(size);
-					result = psprintf("%s GB", numeric_to_cstring(size));
-				}
-				else
-				{
-					/* size >>= 10 */
-					size = numeric_shift_right(size, 10);
-					size = numeric_half_rounded(size);
-					result = psprintf("%s TB", numeric_to_cstring(size));
-				}
-			}
+
+			result = psprintf("%s %s", numeric_to_cstring(size), unit->name);
+			break;
 		}
+
+		/*
+		 * We may need to shift 1 bit less than the difference between this
+		 * and the next unit if the next unit has half rounding.  Or we may
+		 * need to shift an extra bit if this unit has half rounding and the
+		 * next one does not.
+		 */
+		shiftby = (unit[1].unitbits - unit->unitbits - (unit[1].round == true)
+				   + (unit->round == true));
+		size = numeric_shift_right(size, shiftby);
 	}
 
 	PG_RETURN_TEXT_P(cstring_to_text(result));
@@ -786,6 +776,7 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 	/* Handle possible unit */
 	if (*strptr != '\0')
 	{
+		const struct size_pretty_unit *unit;
 		int64		multiplier = 0;
 
 		/* Trim any trailing whitespace */
@@ -797,21 +788,18 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 		endptr++;
 		*endptr = '\0';
 
-		/* Parse the unit case-insensitively */
-		if (pg_strcasecmp(strptr, "bytes") == 0)
-			multiplier = (int64) 1;
-		else if (pg_strcasecmp(strptr, "kb") == 0)
-			multiplier = (int64) 1024;
-		else if (pg_strcasecmp(strptr, "mb") == 0)
-			multiplier = ((int64) 1024) * 1024;
-
-		else if (pg_strcasecmp(strptr, "gb") == 0)
-			multiplier = ((int64) 1024) * 1024 * 1024;
-
-		else if (pg_strcasecmp(strptr, "tb") == 0)
-			multiplier = ((int64) 1024) * 1024 * 1024 * 1024;
+		for (unit = size_pretty_units; unit->name != NULL; unit++)
+		{
+			/* Parse the unit case-insensitively */
+			if (pg_strcasecmp(strptr, unit->name) == 0)
+			{
+				multiplier = ((int64) 1) << unit->unitbits;
+				break;
+			}
+		}
 
-		else
+		/* Verify we found a valid unit in the loop above */
+		if (unit->name == NULL)
 			ereport(ERROR,
 					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 					 errmsg("invalid size: \"%s\"", text_to_cstring(arg)),
-- 
2.30.2

