Define Math.pw2ceil for all integer types (#13127)

crystal-lang · Mar 2, 2023 · 6e9c356 · 6e9c356
1 parent 13a3d84
commit 6e9c356
Show file tree

Hide file tree

Showing 5 changed files with 104 additions and 57 deletions.
diff --git a/spec/std/big/big_int_spec.cr b/spec/std/big/big_int_spec.cr
@@ -666,4 +666,22 @@ describe "BigInt Math" do
   it "isqrt" do
     Math.isqrt(BigInt.new("1" + "0"*48)).should eq(BigInt.new("1" + "0"*24))
   end
+
+  it "pw2ceil" do
+    Math.pw2ceil("-100000000000000000000000000000000".to_big_i).should eq(1.to_big_i)
+    Math.pw2ceil(-1234567.to_big_i).should eq(1.to_big_i)
+    Math.pw2ceil(-1.to_big_i).should eq(1.to_big_i)
+    Math.pw2ceil(0.to_big_i).should eq(1.to_big_i)
+    Math.pw2ceil(1.to_big_i).should eq(1.to_big_i)
+    Math.pw2ceil(2.to_big_i).should eq(2.to_big_i)
+    Math.pw2ceil(3.to_big_i).should eq(4.to_big_i)
+    Math.pw2ceil(4.to_big_i).should eq(4.to_big_i)
+    Math.pw2ceil(5.to_big_i).should eq(8.to_big_i)
+    Math.pw2ceil(32.to_big_i).should eq(32.to_big_i)
+    Math.pw2ceil(33.to_big_i).should eq(64.to_big_i)
+    Math.pw2ceil(64.to_big_i).should eq(64.to_big_i)
+    Math.pw2ceil(2.to_big_i ** 12345 - 1).should eq(2.to_big_i ** 12345)
+    Math.pw2ceil(2.to_big_i ** 12345).should eq(2.to_big_i ** 12345)
+    Math.pw2ceil(2.to_big_i ** 12345 + 1).should eq(2.to_big_i ** 12346)
+  end
 end
diff --git a/spec/std/math_spec.cr b/spec/std/math_spec.cr
@@ -266,44 +266,48 @@ describe "Math" do
 
   # div rem
 
-  # pw2ceil
-
   describe ".pw2ceil" do
-    it "Int32" do
-      Math.pw2ceil(-1).should eq 1
-      Math.pw2ceil(33).should eq(64)
-      Math.pw2ceil(128).should eq(128)
-      Math.pw2ceil(0).should eq 1
-      Math.pw2ceil(1).should eq 1
-      Math.pw2ceil(2).should eq 2
-      Math.pw2ceil(3).should eq 4
-      Math.pw2ceil(4).should eq 4
-      Math.pw2ceil(5).should eq 8
-      # 1073741824 is the largest power of 2 that fits into Int32
-      Math.pw2ceil(1073741824).should eq 1073741824
-      Math.pw2ceil(1073741824 - 1).should eq 1073741824
-      expect_raises(OverflowError) do
-        Math.pw2ceil(1073741824 + 1)
+    {% for int in %w(Int8 Int16 Int32 Int64 Int128) %}
+      it {{ int }} do
+        Math.pw2ceil({{ int.id }}::MIN).should eq 1
+        Math.pw2ceil({{ int.id }}::MIN + 1).should eq 1
+        Math.pw2ceil({{ int.id }}.new(-11)).should eq 1
+        Math.pw2ceil({{ int.id }}.new(-1)).should eq 1
+        Math.pw2ceil({{ int.id }}.new(0)).should eq 1
+        Math.pw2ceil({{ int.id }}.new(1)).should eq 1
+        Math.pw2ceil({{ int.id }}.new(2)).should eq 2
+        Math.pw2ceil({{ int.id }}.new(3)).should eq 4
+        Math.pw2ceil({{ int.id }}.new(4)).should eq 4
+        Math.pw2ceil({{ int.id }}.new(5)).should eq 8
+        Math.pw2ceil({{ int.id }}.new(32)).should eq(32)
+        Math.pw2ceil({{ int.id }}.new(33)).should eq(64)
+        Math.pw2ceil({{ int.id }}.new(64)).should eq(64)
+
+        Math.pw2ceil({{ int.id }}::MAX // 2).should eq({{ int.id }}::MAX // 2 + 1)
+        Math.pw2ceil({{ int.id }}::MAX // 2 + 1).should eq({{ int.id }}::MAX // 2 + 1)
+        expect_raises(OverflowError) { Math.pw2ceil({{ int.id }}::MAX // 2 + 2) }
+        expect_raises(OverflowError) { Math.pw2ceil({{ int.id }}::MAX) }
       end
-    end
-
-    it "Int64" do
-      Math.pw2ceil(-1_i64).should eq 1
-      Math.pw2ceil(33_i64).should eq(64)
-      Math.pw2ceil(128_i64).should eq(128)
-      Math.pw2ceil(0_i64).should eq 1
-      Math.pw2ceil(1_i64).should eq 1
-      Math.pw2ceil(2_i64).should eq 2
-      Math.pw2ceil(3_i64).should eq 4
-      Math.pw2ceil(4_i64).should eq 4
-      Math.pw2ceil(5_i64).should eq 8
-      # 4611686018427387904 is the largest power of 2 that fits into Int64
-      Math.pw2ceil(4611686018427387904).should eq 4611686018427387904
-      Math.pw2ceil(4611686018427387904 - 1).should eq 4611686018427387904
-      expect_raises(OverflowError) do
-        Math.pw2ceil(4611686018427387904 + 1)
+    {% end %}
+
+    {% for uint in %w(UInt8 UInt16 UInt32 UInt64 UInt128) %}
+      it {{ uint }} do
+        Math.pw2ceil({{ uint.id }}.new(0)).should eq 1
+        Math.pw2ceil({{ uint.id }}.new(1)).should eq 1
+        Math.pw2ceil({{ uint.id }}.new(2)).should eq 2
+        Math.pw2ceil({{ uint.id }}.new(3)).should eq 4
+        Math.pw2ceil({{ uint.id }}.new(4)).should eq 4
+        Math.pw2ceil({{ uint.id }}.new(5)).should eq 8
+        Math.pw2ceil({{ uint.id }}.new(32)).should eq(32)
+        Math.pw2ceil({{ uint.id }}.new(33)).should eq(64)
+        Math.pw2ceil({{ uint.id }}.new(64)).should eq(64)
+
+        Math.pw2ceil({{ uint.id }}::MAX // 2).should eq({{ uint.id }}::MAX // 2 + 1)
+        Math.pw2ceil({{ uint.id }}::MAX // 2 + 1).should eq({{ uint.id }}::MAX // 2 + 1)
+        expect_raises(OverflowError) { Math.pw2ceil({{ uint.id }}::MAX // 2 + 2) }
+        expect_raises(OverflowError) { Math.pw2ceil({{ uint.id }}::MAX) }
       end
-    end
+    {% end %}
   end
 
   # ** (float and int)

diff --git a/src/big/big_int.cr b/src/big/big_int.cr
@@ -585,6 +585,14 @@ struct BigInt < Int
     LibGMP.scan1(self, 0)
   end
 
+  # :nodoc:
+  def next_power_of_two : self
+    one = BigInt.new(1)
+    return one if self <= 0
+
+    popcount == 1 ? self : one << bit_length
+  end
+
   def to_i : Int32
     to_i32
   end
@@ -847,6 +855,20 @@ module Math
   def isqrt(value : BigInt)
     BigInt.new { |mpz| LibGMP.sqrt(mpz, value) }
   end
+
+  # Computes the smallest nonnegative power of 2 that is greater than or equal
+  # to *v*.
+  #
+  # The returned value has the same type as the argument.
+  #
+  # ```
+  # Math.pw2ceil(33) # => 64
+  # Math.pw2ceil(64) # => 64
+  # Math.pw2ceil(-5) # => 1
+  # ```
+  def pw2ceil(v : BigInt) : BigInt
+    v.next_power_of_two
+  end
 end
 
 module Random

diff --git a/src/int.cr b/src/int.cr
@@ -448,6 +448,22 @@ struct Int
     end
   end
 
+  # :nodoc:
+  def next_power_of_two : self
+    one = self.class.new!(1)
+
+    bits = sizeof(self) * 8
+    shift = bits &- (self &- 1).leading_zeros_count
+    if self.is_a?(Int::Signed)
+      shift = 0 if shift >= bits &- 1
+    else
+      shift = 0 if shift == bits
+    end
+
+    result = one << shift
+    result >= self ? result : raise OverflowError.new
+  end
+
   # Returns the greatest common divisor of `self` and *other*. Signed
   # integers may raise `OverflowError` if either has value equal to `MIN` of
   # its type.

diff --git a/src/math/math.cr b/src/math/math.cr
@@ -722,31 +722,18 @@ module Math
     value1 <= value2 ? value1 : value2
   end
 
-  # Computes the next highest power of 2 of *v*.
+  # Computes the smallest nonnegative power of 2 that is greater than or equal
+  # to *v*.
+  #
+  # The returned value has the same type as the argument. Raises `OverflowError`
+  # if the result does not fit into the argument's type.
   #
   # ```
   # Math.pw2ceil(33) # => 64
+  # Math.pw2ceil(64) # => 64
+  # Math.pw2ceil(-5) # => 1
   # ```
-  def pw2ceil(v : Int32)
-    # Taken from http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
-    v -= 1
-    v |= v >> 1
-    v |= v >> 2
-    v |= v >> 4
-    v |= v >> 8
-    v |= v >> 16
-    v += v == -1 ? 2 : 1
-  end
-
-  def pw2ceil(v : Int64)
-    # Taken from http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
-    v -= 1
-    v |= v >> 1
-    v |= v >> 2
-    v |= v >> 4
-    v |= v >> 8
-    v |= v >> 16
-    v |= v >> 32
-    v += v == -1 ? 2 : 1
+  def pw2ceil(v : Int::Primitive)
+    v.next_power_of_two
   end
 end