Update common.hlsl

Author: CrabExtra

include/nbl/builtin/hlsl/bitonic_sort/common.hlsl

@@ -4,6 +4,7 @@
 #include <nbl/builtin/hlsl/cpp_compat.hlsl>
 #include <nbl/builtin/hlsl/concepts.hlsl>
 #include <nbl/builtin/hlsl/math/intutil.hlsl>
+#include <nbl/builtin/hlsl/pair.hlsl>
 
 namespace nbl
 {
@@ -14,24 +15,22 @@ namespace bitonic_sort
 
 template<typename KeyType, typename ValueType, typename Comparator>
 void compareExchangeWithPartner(
-    bool takeLarger,
-    NBL_REF_ARG(KeyType) loKey,
-    NBL_CONST_REF_ARG(KeyType) partnerLoKey,
-    NBL_REF_ARG(KeyType) hiKey,
-    NBL_CONST_REF_ARG(KeyType) partnerHiKey,
-    NBL_REF_ARG(ValueType) loVal,
-    NBL_CONST_REF_ARG(ValueType) partnerLoVal,
-    NBL_REF_ARG(ValueType) hiVal,
-    NBL_CONST_REF_ARG(ValueType) partnerHiVal,
-    NBL_CONST_REF_ARG(Comparator) comp)
+bool takeLarger,
+NBL_REF_ARG(KeyType) loKey,
+NBL_CONST_REF_ARG(KeyType) partnerLoKey,
+NBL_REF_ARG(KeyType) hiKey,
+NBL_CONST_REF_ARG(KeyType) partnerHiKey,
+NBL_REF_ARG(ValueType) loVal,
+NBL_CONST_REF_ARG(ValueType) partnerLoVal,
+NBL_REF_ARG(ValueType) hiVal,
+NBL_CONST_REF_ARG(ValueType) partnerHiVal,
+NBL_CONST_REF_ARG(Comparator) comp)
 {
-    // Process lo pair
     const bool loSelfSmaller = comp(loKey, partnerLoKey);
     const bool takePartnerLo = takeLarger ? loSelfSmaller : !loSelfSmaller;
     loKey = takePartnerLo ? partnerLoKey : loKey;
     loVal = takePartnerLo ? partnerLoVal : loVal;
 
-    // Process hi pair
     const bool hiSelfSmaller = comp(hiKey, partnerHiKey);
     const bool takePartnerHi = takeLarger ? hiSelfSmaller : !hiSelfSmaller;
     hiKey = takePartnerHi ? partnerHiKey : hiKey;
@@ -41,12 +40,12 @@ void compareExchangeWithPartner(
 
 template<typename KeyType, typename ValueType, typename Comparator>
 void compareSwap(
-    bool ascending,
-    NBL_REF_ARG(KeyType) loKey,
-    NBL_REF_ARG(KeyType) hiKey,
-    NBL_REF_ARG(ValueType) loVal,
-    NBL_REF_ARG(ValueType) hiVal,
-    NBL_CONST_REF_ARG(Comparator) comp)
+bool ascending,
+NBL_REF_ARG(KeyType) loKey,
+NBL_REF_ARG(KeyType) hiKey,
+NBL_REF_ARG(ValueType) loVal,
+NBL_REF_ARG(ValueType) hiVal,
+NBL_CONST_REF_ARG(Comparator) comp)
 {
     const bool shouldSwap = comp(hiKey, loKey);
 
@@ -63,10 +62,10 @@ void compareSwap(
 
 template<typename KeyType, typename ValueType>
 void swap(
-    NBL_REF_ARG(KeyType) loKey,
-    NBL_REF_ARG(KeyType) hiKey,
-    NBL_REF_ARG(ValueType) loVal,
-    NBL_REF_ARG(ValueType) hiVal)
+NBL_REF_ARG(KeyType) loKey,
+NBL_REF_ARG(KeyType) hiKey,
+NBL_REF_ARG(ValueType) loVal,
+NBL_REF_ARG(ValueType) hiVal)
 {
     KeyType tempKey = loKey;
     loKey = hiKey;
@@ -77,6 +76,55 @@ void swap(
     hiVal = tempVal;
 }
 
+
+
+template<typename KeyType, typename ValueType, typename Comparator>
+void compareExchangeWithPartner(
+bool takeLarger,
+NBL_REF_ARG(pair<KeyType, ValueType>) loPair,
+NBL_CONST_REF_ARG(pair<KeyType, ValueType>) partnerLoPair,
+NBL_REF_ARG(pair<KeyType, ValueType>) hiPair,
+NBL_CONST_REF_ARG(pair<KeyType, ValueType>) partnerHiPair,
+NBL_CONST_REF_ARG(Comparator) comp)
+{
+    const bool loSelfSmaller = comp(loPair.first, partnerLoPair.first);
+    const bool takePartnerLo = takeLarger ? loSelfSmaller : !loSelfSmaller;
+    loPair.first = takePartnerLo ? partnerLoPair.first : loPair.first;
+    loPair.second = takePartnerLo ? partnerLoPair.second : loPair.second;
+
+    const bool hiSelfSmaller = comp(hiPair.first, partnerHiPair.first);
+    const bool takePartnerHi = takeLarger ? hiSelfSmaller : !hiSelfSmaller;
+    hiPair.first = takePartnerHi ? partnerHiPair.first : hiPair.first;
+    hiPair.second = takePartnerHi ? partnerHiPair.second : hiPair.second;
+}
+
+template<typename KeyType, typename ValueType, typename Comparator>
+void compareSwap(
+bool ascending,
+NBL_REF_ARG(pair<KeyType, ValueType>) loPair,
+NBL_REF_ARG(pair<KeyType, ValueType>) hiPair,
+NBL_CONST_REF_ARG(Comparator) comp)
+{
+    const bool shouldSwap = comp(hiPair.first, loPair.first);
+    const bool doSwap = (shouldSwap == ascending);
+
+    KeyType tempKey = loPair.first;
+    ValueType tempVal = loPair.second;
+    loPair.first = doSwap ? hiPair.first : loPair.first;
+    loPair.second = doSwap ? hiPair.second : loPair.second;
+    hiPair.first = doSwap ? tempKey : hiPair.first;
+    hiPair.second = doSwap ? tempVal : hiPair.second;
+}
+
+template<typename KeyType, typename ValueType>
+void swap(
+NBL_REF_ARG(pair<KeyType, ValueType>) loPair,
+NBL_REF_ARG(pair<KeyType, ValueType>) hiPair)
+{
+    pair<KeyType, ValueType> temp = loPair;
+    loPair = hiPair;
+    hiPair = temp;
+}
 }
 }
 }