From 09e1edbd171f60bf7a19079d299012fcabf7f07c Mon Sep 17 00:00:00 2001
From: Mike <45243121+tankf33der@users.noreply.github.com>
Date: Tue, 8 Jul 2025 17:22:45 +0300
Subject: [PATCH] math.unsigned: fix rsh() for uint256 too, add tests (#24865)

---
 vlib/math/unsigned/uint256.v      | 35 ++++++++++++++++++++++---------
 vlib/math/unsigned/uint256_test.v | 11 ++++++++++
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/vlib/math/unsigned/uint256.v b/vlib/math/unsigned/uint256.v
index 0faf23c059..b0c8ccb275 100644
--- a/vlib/math/unsigned/uint256.v
+++ b/vlib/math/unsigned/uint256.v
@@ -230,17 +230,32 @@ pub fn (u Uint256) quo_rem_64(v u64) (Uint256, u64) {
 }
 
 // rsh returns a new Uint256 that has been right bit shifted
-pub fn (u Uint256) rsh(n_ u32) Uint256 {
-	mut n := n_
-	if n > 128 {
-		return Uint256{u.hi.rsh(n - 128), uint128_zero}
+pub fn (u Uint256) rsh(n u32) Uint256 {
+	mut s := Uint256{}
+	if n == 0 {
+		s.lo = u.lo
+		s.hi = u.hi
+	} else if n >= 256 {
+		s.lo = uint128_zero
+		s.hi = uint128_zero
+	} else if n == 128 {
+		s.hi = uint128_zero
+		s.lo = u.hi
+	} else if n > 128 {
+		s.hi = uint128_zero
+		s.lo = u.hi.rsh(n - 128)
+	} else if n == 64 {
+		s.lo = Uint128{u.lo.hi, u.hi.lo}
+		s.hi = Uint128{u.hi.hi, 0}
+	} else if n > 64 {
+		shift := n - 64
+		s.lo = Uint128{u.lo.hi >> shift | u.hi.lo << (64 - shift), u.hi.lo >> shift | u.hi.hi << (64 - shift)}
+		s.hi = Uint128{u.hi.hi >> shift, 0}
+	} else {
+		s.lo = Uint128{u.lo.lo >> n | u.lo.hi << (64 - n), u.lo.hi >> n | u.hi.lo << (64 - n)}
+		s.hi = Uint128{u.hi.lo >> n | u.hi.hi << (64 - n), u.hi.hi >> n}
 	}
-
-	if n > 64 {
-		n -= 64
-		return Uint256{Uint128{u.lo.hi >> n | u.hi.lo << (64 - n), u.hi.lo >> n | u.hi.hi << (64 - n)}, Uint128{u.hi.hi >> n, 0}}
-	}
-	return Uint256{Uint128{u.lo.lo >> n | u.lo.hi << (64 - n), u.lo.hi >> n | u.hi.lo << (64 - n)}, Uint128{u.hi.lo >> n | u.hi.hi << (64 - n), u.hi.hi >> n}}
+	return s
 }
 
 // lsh returns a new Uint256 that has been left bit shifted
diff --git a/vlib/math/unsigned/uint256_test.v b/vlib/math/unsigned/uint256_test.v
index 13af5fb0ae..12987e9d6e 100644
--- a/vlib/math/unsigned/uint256_test.v
+++ b/vlib/math/unsigned/uint256_test.v
@@ -280,3 +280,14 @@ fn test_separators() {
 fn test_new() {
 	assert unsigned.uint256_new(unsigned.uint128_max, unsigned.uint128_max) == unsigned.uint256_max
 }
+
+fn test_rsh() {
+	a := unsigned.uint256_from_dec_str('115792089237316195423570985008687907853269984665640564039457584007913129639935')!
+	assert a.str() == a.rsh(0).str()
+	assert '57896044618658097711785492504343953926634992332820282019728792003956564819967' == a.rsh(1).str()
+	assert '6277101735386680763835789423207666416102355444464034512895' == a.rsh(64).str()
+	assert '91343852333181432387730302044767688728495783935' == a.rsh(100).str()
+	assert '340282366920938463463374607431768211455' == a.rsh(128).str()
+	assert '170141183460469231731687303715884105727' == a.rsh(129).str()
+	assert unsigned.uint256_zero == a.rsh(256)
+}