Fix/diff unpack bit shift (#878)

* fix incorrect bit shift

* move inner function out and add test for diff pack

* fix test

* fix call arg check in test

* add assert

* fix test for py37

* refactor the diff pack into two functions
move the test to a unit test

* fix test

bittensor/utils/__init__.py

@@ -220,7 +220,7 @@ def run(self):
                 with self.check_block:
                     block_number = self.curr_block_num.value
                     block_bytes = bytes(self.curr_block)
-                    block_difficulty = int(self.curr_diff[0] >> 32 | self.curr_diff[1])
+                    block_difficulty = registration_diff_unpack(self.curr_diff)
 
                 self.newBlockEvent.clear()
                 # reset nonces to start from random point
@@ -270,6 +270,25 @@ def solve_for_nonce_block(solver: Solver, nonce_start: int, nonce_end: int, bloc
     return None, time.time() - start
 
 
+def registration_diff_unpack(packed_diff: multiprocessing.Array) -> int:
+    """Unpacks the packed two 32-bit integers into one 64-bit integer. Little endian."""
+    return int(packed_diff[0] << 32 | packed_diff[1])
+
+
+def registration_diff_pack(diff: int, packed_diff: multiprocessing.Array):
+    """Packs the difficulty into two 32-bit integers. Little endian."""
+    packed_diff[0] = diff >> 32
+    packed_diff[1] = diff & 0xFFFFFFFF # low 32 bits
+
+
+def update_curr_block(curr_diff: multiprocessing.Array, curr_block: multiprocessing.Array, curr_block_num: multiprocessing.Value, block_number: int, block_bytes: bytes, diff: int, lock: multiprocessing.Lock):
+    with lock:
+        curr_block_num.value = block_number
+        for i in range(64):
+            curr_block[i] = block_bytes[i]
+        registration_diff_pack(diff, curr_diff)
+
+
 def solve_for_difficulty_fast( subtensor, wallet, num_processes: Optional[int] = None, update_interval: Optional[int] = None ) -> Optional[POWSolution]:
     """
     Solves the POW for registration using multiprocessing.
@@ -306,15 +325,7 @@ def solve_for_difficulty_fast( subtensor, wallet, num_processes: Optional[int] =
     curr_block = multiprocessing.Array('h', 64, lock=True) # byte array
     curr_block_num = multiprocessing.Value('i', 0, lock=True) # int
     curr_diff = multiprocessing.Array('Q', [0, 0], lock=True) # [high, low]
-
-    def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: multiprocessing.Lock):
-        with lock:
-            curr_block_num.value = block_number
-            for i in range(64):
-                curr_block[i] = block_bytes[i]
-            curr_diff[0] = diff >> 32
-            curr_diff[1] = diff & 0xFFFFFFFF # low 32 bits
-
+
     status.start()
 
     # Establish communication queues
@@ -339,7 +350,7 @@ def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: mu
     block_bytes = block_hash.encode('utf-8')[2:]
     old_block_number = block_number
     # Set to current block
-    update_curr_block(block_number, block_bytes, difficulty, check_block)
+    update_curr_block(curr_diff, curr_block, curr_block_num, block_number, block_bytes, difficulty, check_block)
 
     # Set new block events for each solver to start
     for w in solvers:
@@ -373,7 +384,7 @@ def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: mu
             block_bytes = block_hash.encode('utf-8')[2:]
             difficulty = subtensor.difficulty
 
-            update_curr_block(block_number, block_bytes, difficulty, check_block)
+            update_curr_block(curr_diff, curr_block, curr_block_num, block_number, block_bytes, difficulty, check_block)
             # Set new block events for each solver
             for w in solvers:
                 w.newBlockEvent.set()

tests/integration_tests/test_subtensor.py

@@ -15,17 +15,17 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
 # DEALINGS IN THE SOFTWARE.
 
-from typing import DefaultDict
-from unittest import mock
+
+import multiprocessing
 from unittest.mock import patch
 import bittensor
 import pytest
-import psutil  
 import unittest
 import time
 import random
 from unittest.mock import MagicMock
 from bittensor.utils.balance import Balance
+from bittensor.utils import Solver, update_curr_block
 from substrateinterface import Keypair
 from bittensor._subtensor.subtensor_mock import mock_subtensor
 class TestSubtensor(unittest.TestCase):
@@ -411,11 +411,10 @@ def process_events(self):
             self.subtensor.substrate.submit_extrinsic = MagicMock(return_value = success())
 
             # should return True
-            assert self.subtensor.register(wallet=wallet,)
+            assert self.subtensor.register(wallet=wallet, num_processes=3, update_interval=5) == True
             # calls until True and once again before exiting subtensor class
             # This assertion is currently broken when difficulty is too low
-            #assert wallet.is_registered.call_count == workblocks_before_is_registered + 2      
-
+            #assert wallet.is_registered.call_count == workblocks_before_is_registered + 2   
 
     def test_registration_partly_failed( self ):
         class failed():
@@ -448,7 +447,7 @@ def process_events(self):
             self.subtensor.substrate.submit_extrinsic = MagicMock(side_effect = submit_extrinsic)
 
             # should return True
-            assert self.subtensor.register(wallet=wallet,) == True
+            assert self.subtensor.register(wallet=wallet, num_processes=3, update_interval=5) == True
 
     def test_registration_failed( self ):
         class failed():

tests/unit_tests/bittensor_tests/utils/test_utils.py

@@ -8,6 +8,7 @@
 import os 
 import random
 import torch
+import multiprocessing
 
 from sys import platform   
 from substrateinterface.base import Keypair
@@ -233,5 +234,18 @@ def test_is_valid_ed25519_pubkey():
     assert bittensor.utils.is_valid_ed25519_pubkey(good_pubkey)
     assert not bittensor.utils.is_valid_ed25519_pubkey(bad_pubkey)
 
+def test_registration_diff_pack_unpack():
+        fake_diff = pow(2, 31)# this is under 32 bits
+
+        mock_diff = multiprocessing.Array('Q', [0, 0], lock=True) # [high, low]
+
+        bittensor.utils.registration_diff_pack(fake_diff, mock_diff)
+        assert bittensor.utils.registration_diff_unpack(mock_diff) == fake_diff
+
+        fake_diff = pow(2, 32) * pow(2, 4) # this should be too large if the bit shift is wrong (32 + 4 bits)
+
+        bittensor.utils.registration_diff_pack(fake_diff, mock_diff)
+        assert bittensor.utils.registration_diff_unpack(mock_diff) == fake_diff
+
 if __name__ == "__main__":
     test_solve_for_difficulty_fast_registered_already()