Add `bit_size()` and `starts_with()` to `gleam/bit_array`

4 files changed, 124 insertions, 2 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 37e314f..fb04001 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+
+- The `bit_array` module gains the `bit_size` and `starts_with` functions.
+
 ## v0.41.0 - 2024-10-31
 
 Happy Samhain! 🎃
diff --git a/src/gleam/bit_array.gleam b/src/gleam/bit_array.gleam
index 2c70d14..e31709c 100644
--- a/src/gleam/bit_array.gleam
+++ b/src/gleam/bit_array.gleam
@@ -10,6 +10,13 @@ import gleam/string
 @external(javascript, "../gleam_stdlib.mjs", "bit_array_from_string")
 pub fn from_string(x: String) -> BitArray
 
+/// Returns an integer which is the number of bits in the bit array.
+///
+@external(erlang, "erlang", "bit_size")
+pub fn bit_size(x: BitArray) -> Int {
+  byte_size(x) * 8
+}
+
 /// Returns an integer which is the number of bytes in the bit array.
 ///
 @external(erlang, "erlang", "byte_size")
@@ -203,8 +210,6 @@ fn inspect_loop(input: BitArray, accumulator: String) -> String {
 /// // -> Eq
 /// ```
 ///
-/// Only supported on Erlang target for now.
-///
 @external(javascript, "../gleam_stdlib.mjs", "bit_array_compare")
 pub fn compare(a: BitArray, with b: BitArray) -> order.Order {
   case a, b {
@@ -235,3 +240,22 @@ pub fn compare(a: BitArray, with b: BitArray) -> order.Order {
 
 @external(erlang, "gleam_stdlib", "bit_array_to_int_and_size")
 fn bit_array_to_int_and_size(a: BitArray) -> #(Int, Int)
+
+/// Checks whether the first `BitArray` starts with the second one.
+///
+/// ## Examples
+///
+/// ```gleam
+/// starts_with(<<1, 2, 3, 4>>, <<1, 2>>)
+/// // -> True
+/// ```
+///
+@external(javascript, "../gleam_stdlib.mjs", "bit_array_starts_with")
+pub fn starts_with(bits: BitArray, prefix: BitArray) -> Bool {
+  let prefix_size = bit_size(prefix)
+
+  case bits {
+    <<pref:bits-size(prefix_size), _:bits>> if pref == prefix -> True
+    _ -> False
+  }
+}
diff --git a/src/gleam_stdlib.mjs b/src/gleam_stdlib.mjs
index 74df3d0..4d198cb 100644
--- a/src/gleam_stdlib.mjs
+++ b/src/gleam_stdlib.mjs
@@ -990,3 +990,17 @@ export function bit_array_compare(first, second) {
   }
   return new Lt(); // second has more items
 }
+
+export function bit_array_starts_with(bits, prefix) {
+  if (prefix.length > bits.length) {
+    return false;
+  }
+
+  for (let i = 0; i < prefix.length; i++) {
+    if (bits.buffer[i] !== prefix.buffer[i]) {
+      return false;
+    }
+  }
+
+  return true;
+}
diff --git a/test/gleam/bit_array_test.gleam b/test/gleam/bit_array_test.gleam
index a563c15..555c61a 100644
--- a/test/gleam/bit_array_test.gleam
+++ b/test/gleam/bit_array_test.gleam
@@ -4,6 +4,37 @@ import gleam/result
 import gleam/should
 import gleam/string
 
+pub fn bit_size_test() {
+  bit_array.bit_size(<<>>)
+  |> should.equal(0)
+
+  bit_array.bit_size(<<0>>)
+  |> should.equal(8)
+
+  bit_array.bit_size(<<-1:32>>)
+  |> should.equal(32)
+
+  bit_array.bit_size(<<0:-8>>)
+  |> should.equal(0)
+}
+
+// This test is target specific since it's using non byte-aligned BitArrays
+// and those are not supported on the JavaScript target.
+@target(erlang)
+pub fn bit_size_erlang_only_test() {
+  bit_array.bit_size(<<0:1>>)
+  |> should.equal(1)
+
+  bit_array.bit_size(<<7:3>>)
+  |> should.equal(3)
+
+  bit_array.bit_size(<<-1:190>>)
+  |> should.equal(190)
+
+  bit_array.bit_size(<<0:-1>>)
+  |> should.equal(0)
+}
+
 pub fn byte_size_test() {
   bit_array.byte_size(bit_array.from_string("hello"))
   |> should.equal(5)
@@ -379,3 +410,52 @@ pub fn compare_different_sizes_test() {
   bit_array.compare(<<0:2>>, <<0:1>>)
   |> should.equal(order.Gt)
 }
+
+pub fn starts_with_test() {
+  bit_array.starts_with(<<>>, <<>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<0>>, <<>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<>>, <<0>>)
+  |> should.be_false
+
+  bit_array.starts_with(<<0, 1, 2>>, <<0>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<0, 1, 2>>, <<0, 1>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<0, 1, 2>>, <<0, 1, 2>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<0, 1, 2>>, <<0, 1, 2, 3>>)
+  |> should.be_false
+
+  bit_array.starts_with(<<0, 1, 2>>, <<1>>)
+  |> should.be_false
+}
+
+// This test is target specific since it's using non byte-aligned BitArrays
+// and those are not supported on the JavaScript target.
+@target(erlang)
+pub fn starts_with_erlang_only_test() {
+  bit_array.starts_with(<<1:1>>, <<1:1>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<1:1>>, <<>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<1:1>>, <<1:2>>)
+  |> should.be_false
+
+  bit_array.starts_with(<<-1:127>>, <<-1:33>>)
+  |> should.be_true
+
+  bit_array.starts_with(<<-1:127>>, <<-1:128>>)
+  |> should.be_false
+
+  bit_array.starts_with(<<0:127>>, <<1:127>>)
+  |> should.be_false
+}