From dbcb1e15648ef7b050b9b59b40d413038b4888ec Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Fri, 30 Apr 2021 10:03:54 -0700
Subject: [PATCH] loader: avoid UB when doing count trailing zeros

---
 .../include/vapours/util/util_bitutil.hpp     | 57 ++++++++++++++++---
 .../loader/source/ldr_capabilities.cpp        |  2 +-
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/libraries/libvapours/include/vapours/util/util_bitutil.hpp b/libraries/libvapours/include/vapours/util/util_bitutil.hpp
index 7b7b2ddde..fc6be79e7 100644
--- a/libraries/libvapours/include/vapours/util/util_bitutil.hpp
+++ b/libraries/libvapours/include/vapours/util/util_bitutil.hpp
@@ -175,20 +175,59 @@ namespace ams::util {
         } else {
             using U = typename std::make_unsigned<T>::type;
             const U u = static_cast<U>(x);
-            if constexpr (std::is_same<U, unsigned long long>::value) {
-                return __builtin_clzll(u);
-            } else if constexpr (std::is_same<U, unsigned long>::value) {
-                return __builtin_clzl(u);
-            }  else if constexpr(std::is_same<U, unsigned int>::value) {
-                return __builtin_clz(u);
+
+            if (u != 0) {
+                if constexpr (std::is_same<U, unsigned long long>::value) {
+                    return __builtin_clzll(u);
+                } else if constexpr (std::is_same<U, unsigned long>::value) {
+                    return __builtin_clzl(u);
+                }  else if constexpr(std::is_same<U, unsigned int>::value) {
+                    return __builtin_clz(u);
+                } else {
+                    static_assert(sizeof(U) < sizeof(unsigned int));
+                    constexpr size_t BitDiff = BITSIZEOF(unsigned int) - BITSIZEOF(U);
+                    return __builtin_clz(static_cast<unsigned int>(u)) - BitDiff;
+                }
             } else {
-                static_assert(sizeof(U) < sizeof(unsigned int));
-                constexpr size_t BitDiff = BITSIZEOF(unsigned int) - BITSIZEOF(U);
-                return __builtin_clz(static_cast<unsigned int>(u)) - BitDiff;
+                return BITSIZEOF(U);
             }
         }
     }
 
+    static_assert(CountLeadingZeros(~static_cast<u64>(0)) == 0);
+    static_assert(CountLeadingZeros(static_cast<u64>(0)) == BITSIZEOF(u64));
+
+    template<typename T> requires std::integral<T>
+    constexpr ALWAYS_INLINE int CountTrailingZeros(T x) {
+        if (std::is_constant_evaluated()) {
+            auto count = 0;
+            for (size_t i = 0; i < BITSIZEOF(T) && (x & 1) == 0; ++i) {
+                ++count;
+            }
+            return count;
+        } else {
+            using U = typename std::make_unsigned<T>::type;
+            const U u = static_cast<U>(x);
+            if (u != 0) {
+                if constexpr (std::is_same<U, unsigned long long>::value) {
+                    return __builtin_ctzll(u);
+                } else if constexpr (std::is_same<U, unsigned long>::value) {
+                    return __builtin_ctzl(u);
+                }  else if constexpr(std::is_same<U, unsigned int>::value) {
+                    return __builtin_ctz(u);
+                } else {
+                    static_assert(sizeof(U) < sizeof(unsigned int));
+                    return __builtin_ctz(static_cast<unsigned int>(u));
+                }
+            } else {
+                return BITSIZEOF(U);
+            }
+        }
+    }
+
+    static_assert(CountTrailingZeros(~static_cast<u64>(0)) == 0);
+    static_assert(CountTrailingZeros(static_cast<u64>(0)) == BITSIZEOF(u64));
+
     template<typename T> requires std::integral<T>
     constexpr ALWAYS_INLINE bool IsPowerOfTwo(T x) {
         return x > 0 && ResetLeastSignificantOneBit(x) == 0;
diff --git a/stratosphere/loader/source/ldr_capabilities.cpp b/stratosphere/loader/source/ldr_capabilities.cpp
index e24dcbc21..08b1516cc 100644
--- a/stratosphere/loader/source/ldr_capabilities.cpp
+++ b/stratosphere/loader/source/ldr_capabilities.cpp
@@ -44,7 +44,7 @@ namespace ams::ldr::caps {
             constexpr ALWAYS_INLINE typename name::Type Get##name() const { return this->Get<name>(); }
 
         constexpr ALWAYS_INLINE CapabilityId GetCapabilityId(util::BitPack32 cap) {
-            return static_cast<CapabilityId>(__builtin_ctz(~cap.value));
+            return static_cast<CapabilityId>(util::CountTrailingZeros<u32>(~cap.value));
         }
 
         constexpr inline util::BitPack32 EmptyCapability = {~u32{}};