std.math: Add O(log N) implementation of log2(x) for comptime_int

Since Zig provides @clz and not @ffs (find-first-set), log2 for comptime
integers needs to be computed algorithmically. To avoid hitting the
backward branch quota, this updates log2(x) to use a simple O(log N)
algorithm.

1 files changed, 19 insertions, 5 deletions

diff --git a/lib/std/math/log2.zig b/lib/std/math/log2.zig
index c83b170208..4158b290f0 100644
--- a/lib/std/math/log2.zig
+++ b/lib/std/math/log2.zig
@@ -17,12 +17,20 @@ pub fn log2(x: anytype) @TypeOf(x) {
         },
         .Float => return @log2(x),
         .ComptimeInt => comptime {
-            var result = 0;
             var x_shifted = x;
-            while (b: {
-                x_shifted >>= 1;
-                break :b x_shifted != 0;
-            }) : (result += 1) {}
+            // First, calculate floorPowerOfTwo(x)
+            var shift_amt = 1;
+            while (x_shifted >> (shift_amt << 1) != 0) shift_amt <<= 1;
+
+            // Answer is in the range [shift_amt, 2 * shift_amt - 1]
+            // We can find it in O(log(N)) using binary search.
+            var result = 0;
+            while (shift_amt != 0) : (shift_amt >>= 1) {
+                if (x_shifted >> shift_amt != 0) {
+                    x_shifted >>= shift_amt;
+                    result += shift_amt;
+                }
+            }
             return result;
         },
         .Int => |IntType| switch (IntType.signedness) {
@@ -36,4 +44,10 @@ pub fn log2(x: anytype) @TypeOf(x) {
 test "log2" {
     try expect(log2(@as(f32, 0.2)) == @log2(0.2));
     try expect(log2(@as(f64, 0.2)) == @log2(0.2));
+    comptime {
+        try expect(log2(1) == 0);
+        try expect(log2(15) == 3);
+        try expect(log2(16) == 4);
+        try expect(log2(1 << 4073) == 4073);
+    }
 }