From ed9521ba4b26d59b55515a903fc3b816507e1b6d Mon Sep 17 00:00:00 2001
From: Ahmet Inan <inan@aicodix.de>
Date: Sat, 2 Mar 2024 12:00:43 +0100
Subject: [PATCH] reordered forked metrics to benefit from a stable sort

nth_element does not honor the established order, so previously good
candidates can still be thrown out when the scores are the same.
Using a stable_sort keeps the order but is more costly.
We don't have a stable_partial_sort unfortunately and the one
provided as example, implemented using partial_sort and two aditional
comparisions to make it stable, is slower than a stable_sort.
---
 polar_list_decoder.hh | 18 ++++++++++++------
 polar_parity_aided.hh | 18 ++++++++++++------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/polar_list_decoder.hh b/polar_list_decoder.hh
index 83ce286..56f41cf 100644
--- a/polar_list_decoder.hh
+++ b/polar_list_decoder.hh
@@ -55,24 +55,30 @@ struct PolarListNode<TYPE, 0>
 		TYPE sft = soft[1];
 		PATH fork[2*TYPE::SIZE];
 		for (int k = 0; k < TYPE::SIZE; ++k)
-			fork[k] = fork[k+TYPE::SIZE] = metric[k];
+			fork[2*k] = fork[2*k+1] = metric[k];
 		for (int k = 0; k < TYPE::SIZE; ++k)
 			if (sft.v[k] < 0)
-				fork[k] -= sft.v[k];
+				fork[2*k] -= sft.v[k];
 			else
-				fork[k+TYPE::SIZE] += sft.v[k];
+				fork[2*k+1] += sft.v[k];
 		int perm[2*TYPE::SIZE];
 		for (int k = 0; k < 2*TYPE::SIZE; ++k)
 			perm[k] = k;
-		std::nth_element(perm, perm+TYPE::SIZE, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; });
+#if 1
+		std::stable_sort(perm, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; });
+#else
+		std::partial_sort(perm, perm+TYPE::SIZE-1, perm+2*TYPE::SIZE, [fork](int a, int b){
+			return fork[a] < fork[b] ? true : fork[a] > fork[b] ? false : a < b;
+		});
+#endif
 		for (int k = 0; k < TYPE::SIZE; ++k)
 			metric[k] = fork[perm[k]];
 		MAP map;
 		for (int k = 0; k < TYPE::SIZE; ++k)
-			map.v[k] = perm[k] % TYPE::SIZE;
+			map.v[k] = perm[k] >> 1;
 		TYPE hrd;
 		for (int k = 0; k < TYPE::SIZE; ++k)
-			hrd.v[k] = perm[k] < TYPE::SIZE ? 1 : -1;
+			hrd.v[k] = 1 - 2 * (perm[k] & 1);
 		message[*count] = hrd;
 		maps[*count] = map;
 		++*count;
diff --git a/polar_parity_aided.hh b/polar_parity_aided.hh
index f913d95..4068075 100644
--- a/polar_parity_aided.hh
+++ b/polar_parity_aided.hh
@@ -103,24 +103,30 @@ struct PolarParityNode<TYPE, 0>
 		TYPE sft = soft[1];
 		PATH fork[2*TYPE::SIZE];
 		for (int k = 0; k < TYPE::SIZE; ++k)
-			fork[k] = fork[k+TYPE::SIZE] = metric[k];
+			fork[2*k] = fork[2*k+1] = metric[k];
 		for (int k = 0; k < TYPE::SIZE; ++k)
 			if (sft.v[k] < 0)
-				fork[k] -= sft.v[k];
+				fork[2*k] -= sft.v[k];
 			else
-				fork[k+TYPE::SIZE] += sft.v[k];
+				fork[2*k+1] += sft.v[k];
 		int perm[2*TYPE::SIZE];
 		for (int k = 0; k < 2*TYPE::SIZE; ++k)
 			perm[k] = k;
-		std::nth_element(perm, perm+TYPE::SIZE, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; });
+#if 1
+		std::stable_sort(perm, perm+2*TYPE::SIZE, [fork](int a, int b){ return fork[a] < fork[b]; });
+#else
+		std::partial_sort(perm, perm+TYPE::SIZE-1, perm+2*TYPE::SIZE, [fork](int a, int b){
+			return fork[a] < fork[b] ? true : fork[a] > fork[b] ? false : a < b;
+		});
+#endif
 		for (int k = 0; k < TYPE::SIZE; ++k)
 			metric[k] = fork[perm[k]];
 		MAP map;
 		for (int k = 0; k < TYPE::SIZE; ++k)
-			map.v[k] = perm[k] % TYPE::SIZE;
+			map.v[k] = perm[k] >> 1;
 		TYPE hrd;
 		for (int k = 0; k < TYPE::SIZE; ++k)
-			hrd.v[k] = perm[k] < TYPE::SIZE ? 1 : -1;
+			hrd.v[k] = 1 - 2 * (perm[k] & 1);
 		if (*count) {
 			message[*index] = hrd;
 			*parity = PH::qmul(vshuf(*parity, map), hrd);