From 2f75348b67faa5c7cea5c289b33e0fdd2d9dc47c Mon Sep 17 00:00:00 2001
From: kbkpbot <kbkpbot@sina.com>
Date: Thu, 31 Oct 2024 03:55:42 +0800
Subject: [PATCH] bitfield: add `shift_left/1` and `shift_right/1` methods to
 `BitField` (#22700)

---
 vlib/bitfield/bitfield.v      | 26 ++++++++++++++++++++++++++
 vlib/bitfield/bitfield_test.v | 14 ++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/vlib/bitfield/bitfield.v b/vlib/bitfield/bitfield.v
index 82b88d691c..ee69362d07 100644
--- a/vlib/bitfield/bitfield.v
+++ b/vlib/bitfield/bitfield.v
@@ -600,6 +600,32 @@ pub fn (instance BitField) rotate(offset int) BitField {
 	return output
 }
 
+// shift_left shift-left the bits by `count` positions.
+pub fn (instance BitField) shift_left(count int) BitField {
+	size := instance.size
+	if count <= 0 {
+		return instance
+	} else if count >= size {
+		// return zeroes
+		return new(size)
+	}
+	zeroes := new(count)
+	return join(instance.slice(count, size), zeroes)
+}
+
+// shift_right shift-right the bits by `count` positions.
+pub fn (instance BitField) shift_right(count int) BitField {
+	size := instance.size
+	if count <= 0 {
+		return instance
+	} else if count >= size {
+		// return zeroes
+		return new(size)
+	}
+	zeroes := new(count)
+	return join(zeroes, instance.slice(0, size - count))
+}
+
 // Internal functions
 // clear_tail clears the extra bits that are not part of the bitfield, but yet are allocated
 @[inline]
diff --git a/vlib/bitfield/bitfield_test.v b/vlib/bitfield/bitfield_test.v
index e1f0af36d1..0f8f18f91f 100644
--- a/vlib/bitfield/bitfield_test.v
+++ b/vlib/bitfield/bitfield_test.v
@@ -358,3 +358,17 @@ fn test_bf_printing() {
 	println(input)
 	assert true
 }
+
+fn test_bf_shift() {
+	str := '0001001101111111'
+	bf := bitfield.from_str(str)
+	bf_left := bf.shift_left(4)
+	assert bf_left.str() == '0011011111110000'
+	bf_right := bf.shift_right(4)
+	assert bf_right.str() == '0000000100110111'
+
+	bf_large_left := bf.shift_left(100)
+	bf_large_right := bf.shift_right(100)
+	assert bf_large_left.str() == '0000000000000000'
+	assert bf_large_right.str() == '0000000000000000'
+}