Fix `fls_uint`

2 files changed, 49 insertions, 4 deletions

diff --git a/core/mem/tlsf/tlsf_internal.odin b/core/mem/tlsf/tlsf_internal.odin
index 1dc9c5da7..9c8f5c6d8 100644
--- a/core/mem/tlsf/tlsf_internal.odin
+++ b/core/mem/tlsf/tlsf_internal.odin
@@ -126,10 +126,14 @@ ffs :: proc "contextless" (word: u32) -> (bit: i32) {
 fls :: proc "contextless" (word: u32) -> (bit: i32) {
 	return i32(31 - intrinsics.count_leading_zeros(word))
 }
-@(require_results)
-fls_uint :: proc "contextless" (size: uint) -> (bit: i32) {
-	N :: size_of(uintptr)-1
-	return i32(N - intrinsics.count_leading_zeros(size))
+
+when size_of(uintptr) == 8 {
+	@(require_results)
+	fls_uint :: proc "contextless" (size: uint) -> (bit: i32) {
+		return i32(63 - intrinsics.count_leading_zeros(size))
+	}
+} else {
+	fls_uint :: fls_u32
 }
 
 @(require_results)
diff --git a/tests/core/mem/test_core_mem.odin b/tests/core/mem/test_core_mem.odin
new file mode 100644
index 000000000..d282ae1fd
--- /dev/null
+++ b/tests/core/mem/test_core_mem.odin
@@ -0,0 +1,41 @@
+package test_core_mem
+
+import "core:mem/tlsf"
+import "core:testing"
+
+@test
+test_tlsf_bitscan :: proc(t: ^testing.T) {
+	Vector :: struct {
+		op:  enum{ffs, fls, fls_uint},
+		v:   union{u32, uint},
+		exp: i32,
+	}
+	Tests := []Vector{
+		{.ffs,      u32 (0x0000_0000_0000_0000), -1},
+		{.ffs,      u32 (0x0000_0000_0000_0000), -1},
+		{.fls,      u32 (0x0000_0000_0000_0000), -1},
+		{.ffs,      u32 (0x0000_0000_0000_0001),  0},
+		{.fls,      u32 (0x0000_0000_0000_0001),  0},
+		{.ffs,      u32 (0x0000_0000_8000_0000), 31},
+		{.ffs,      u32 (0x0000_0000_8000_8000), 15},
+		{.fls,      u32 (0x0000_0000_8000_0008), 31},
+		{.fls,      u32 (0x0000_0000_7FFF_FFFF), 30},
+		{.fls_uint, uint(0x0000_0000_8000_0000), 31},
+		{.fls_uint, uint(0x0000_0001_0000_0000), 32},
+		{.fls_uint, uint(0xffff_ffff_ffff_ffff), 63},
+	}
+
+	for test in Tests {
+		switch test.op {
+		case .ffs:
+			res := tlsf.ffs(test.v.?)
+			testing.expectf(t, res == test.exp, "Expected tlsf.ffs(0x%08x) == %v, got %v", test.v, test.exp, res)
+		case .fls:
+			res := tlsf.fls(test.v.?)
+			testing.expectf(t, res == test.exp, "Expected tlsf.fls(0x%08x) == %v, got %v", test.v, test.exp, res)
+		case .fls_uint:
+			res := tlsf.fls_uint(test.v.?)
+			testing.expectf(t, res == test.exp, "Expected tlsf.fls_uint(0x%16x) == %v, got %v", test.v, test.exp, res)
+		}
+	}
+}
\ No newline at end of file