feat: implemented (EMA) difficulty adjustment

genesis.json

@@ -2,6 +2,8 @@
   "chain_id": "minichain-default",
   "timestamp": 1716880000000,
   "difficulty": 4,
+  "target_block_time": 10000,
+  "alpha": 0.1,
   "alloc": {
     "0000000000000000000000000000000000000001": {
       "balance": 1000000000

main.py

@@ -92,6 +92,7 @@ def mine_and_process_block(chain, mempool, miner_pk):
         receipt_root=calculate_receipt_root(receipts),
         receipts=receipts,
         miner=miner_pk,
+        difficulty=chain.current_difficulty,
     )
 
     mined_block = mine_block(block)

minichain/chain.py

@@ -70,6 +70,11 @@ def _create_genesis_block(self, genesis_path):
         timestamp = config.get("timestamp")
         difficulty = config.get("difficulty")
 
+        self.target_block_time = config.get("target_block_time", 10000)
+        self.alpha = config.get("alpha", 0.1)
+        self.current_difficulty = difficulty
+        self.avg_block_time = self.target_block_time
+
         genesis_block = Block(
             index=0,
             previous_hash="0",
@@ -128,6 +133,10 @@ def add_block(self, block):
                 logger.warning("Block %s rejected: %s", block.index, exc)
                 return False
 
+            if block.difficulty != self.current_difficulty:
+                logger.warning("Block %s rejected: Invalid difficulty. Expected %s, got %s", block.index, self.current_difficulty, block.difficulty)
+                return False
+
             # Validate transactions on a temporary state copy
             temp_state = self.state.copy()
             temp_state.chain_id = self.chain_id
@@ -161,6 +170,15 @@ def add_block(self, block):
                 logger.warning("Block %s rejected: Invalid state root. Expected %s, got %s", block.index, temp_state.state_root(), block.state_root)
                 return False
 
+            # Update EMA difficulty state
+            time_diff = block.timestamp - self.last_block.timestamp
+            self.avg_block_time = self.alpha * time_diff + (1 - self.alpha) * self.avg_block_time
+
+            if self.avg_block_time > self.target_block_time:
+                self.current_difficulty = max(1, self.current_difficulty - 1)
+            elif self.avg_block_time < self.target_block_time:
+                self.current_difficulty += 1
+
             # All transactions valid → commit state and append block
             self.state = temp_state
             self.chain.append(block)
@@ -198,11 +216,18 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
             temp_state.chain_id = self.chain_id
             temp_state.restore(self._genesis_state_snapshot)
 
+            temp_difficulty = new_chain_list[0].difficulty
+            temp_avg_block_time = self.target_block_time
+
             # Verify and apply blocks 1 to N
             for i in range(1, len(new_chain_list)):
                 prev_block = new_chain_list[i-1]
                 block = new_chain_list[i]
 
+                if block.difficulty != temp_difficulty:
+                    logger.warning("Reorg failed at block %s: Invalid difficulty. Expected %s, got %s", block.index, temp_difficulty, block.difficulty)
+                    return False, []
+
                 try:
                     validate_block_link_and_hash(prev_block, block)
                 except ValueError as exc:
@@ -230,12 +255,22 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
                     logger.warning("Reorg failed: Invalid state root at block %s", block.index)
                     return False, []
 
+                # Update EMA difficulty state for reorg validation
+                time_diff = block.timestamp - prev_block.timestamp
+                temp_avg_block_time = self.alpha * time_diff + (1 - self.alpha) * temp_avg_block_time
+                if temp_avg_block_time > self.target_block_time:
+                    temp_difficulty = max(1, temp_difficulty - 1)
+                elif temp_avg_block_time < self.target_block_time:
+                    temp_difficulty += 1
+
             # 4. Success! Compute orphaned transactions.
             old_txs = {tx.tx_id: tx for b in original_chain[1:] for tx in b.transactions}
             new_tx_ids = {tx.tx_id for b in new_chain_list[1:] for tx in b.transactions}
             orphans = [tx for tx_id, tx in old_txs.items() if tx_id not in new_tx_ids]
 
             self.chain = new_chain_list
             self.state = temp_state
+            self.current_difficulty = temp_difficulty
+            self.avg_block_time = temp_avg_block_time
             logger.info("Reorg successful! Switched to new chain tip: Block %s", self.last_block.index)
             return True, orphans

minichain/pow.py

@@ -13,14 +13,15 @@ def calculate_hash(block_dict):
 
 def mine_block(
     block,
-    difficulty=4,
+    difficulty=None,
     max_nonce=10_000_000,
     timeout_seconds=None,
     logger=None,
     progress_callback=None
 ):
     """Mines a block using Proof-of-Work without mutating input block until success."""
 
+    difficulty = difficulty if difficulty is not None else block.difficulty
     if not isinstance(difficulty, int) or difficulty <= 0:
         raise ValueError("Difficulty must be a positive integer.")
 

tests/test_difficulty.py

@@ -0,0 +1,63 @@
+import unittest
+from minichain import Blockchain, Block
+from minichain.pow import mine_block
+
+class TestEMADifficulty(unittest.TestCase):
+    def test_difficulty_adjustment(self):
+        chain = Blockchain()
+        chain.target_block_time = 1000
+        chain.alpha = 0.5
+        chain.avg_block_time = 1000
+        chain.current_difficulty = 3
+        chain.chain[0].difficulty = 3
+
+        # Fast mining: timestamps only 1ms apart
+        # avg = 0.5 * 1 + 0.5 * 1000 = 500.5 (which is < 1000) => difficulty increments to 4
+        ts = chain.last_block.timestamp + 1
+        block1 = Block(index=1, previous_hash=chain.last_block.hash, transactions=[], timestamp=ts, difficulty=chain.current_difficulty, state_root=chain.state.state_root())
+        mined_block1 = mine_block(block1)
+        self.assertTrue(chain.add_block(mined_block1))
+        self.assertEqual(chain.current_difficulty, 4)
+
+        # Slow mining: timestamp 5000ms apart
+        # avg = 0.5 * 5000 + 0.5 * 500.5 = 2750.25 (which is > 1000) => difficulty decrements to 3
+        ts = chain.last_block.timestamp + 5000
+        block2 = Block(index=2, previous_hash=chain.last_block.hash, transactions=[], timestamp=ts, difficulty=chain.current_difficulty, state_root=chain.state.state_root())
+        mined_block2 = mine_block(block2)
+        self.assertTrue(chain.add_block(mined_block2))
+        self.assertEqual(chain.current_difficulty, 3)
+
+    def test_reorg_difficulty_validation(self):
+        chain1 = Blockchain()
+        chain1.target_block_time = 1000
+        chain1.alpha = 0.5
+        chain1.avg_block_time = 1000
+        chain1.current_difficulty = 1
+        chain1.chain[0].difficulty = 1
+
+        chain2 = Blockchain()
+        chain2.target_block_time = 1000
+        chain2.alpha = 0.5
+        chain2.avg_block_time = 1000
+        chain2.current_difficulty = 1
+        chain2.chain[0].difficulty = 1
+
+        # Chain 2 mines a fast block, difficulty goes to 2
+        block1 = Block(1, chain2.last_block.hash, [], timestamp=chain2.last_block.timestamp + 1, difficulty=chain2.current_difficulty, state_root=chain2.state.state_root())
+        mine_block(block1)
+        chain2.add_block(block1)
+        self.assertEqual(chain2.current_difficulty, 2)
+
+        # Reorg chain1 to chain2
+        success, orphans = chain1.resolve_conflicts(chain2.chain)
+        self.assertTrue(success)
+        self.assertEqual(chain1.current_difficulty, 2)
+
+        # Forging a chain with wrong difficulty should be rejected
+        forged_chain = list(chain2.chain)
+        forged_block = Block(2, chain2.last_block.hash, [], timestamp=chain2.last_block.timestamp + 1000, difficulty=1, state_root=chain2.state.state_root())
+        mine_block(forged_block)
+        forged_chain.append(forged_block)
+
+        success, _ = chain1.resolve_conflicts(forged_chain)
+        self.assertFalse(success) # Rejected because difficulty should have been 2!

tests/test_persistence.py

@@ -5,6 +5,7 @@
 import shutil
 import sqlite3
 import tempfile
+import time
 import unittest
 
 from nacl.encoding import HexEncoder
@@ -49,12 +50,13 @@ def _chain_with_tx(self):
             index=1,
             previous_hash=bc.last_block.hash,
             transactions=[tx],
-            difficulty=1,
+            difficulty=bc.current_difficulty,
             state_root=temp_state.state_root(),
             receipt_root=calculate_receipt_root([receipt]),
             receipts=[receipt],
+            timestamp=int(time.time()),
         )
-        mine_block(block, difficulty=1)
+        mine_block(block)
         bc.add_block(block)
         return bc, alice_pk, bob_pk
 
@@ -243,12 +245,13 @@ def test_loaded_chain_can_add_new_block(self):
             index=len(restored.chain),
             previous_hash=restored.last_block.hash,
             transactions=[tx2],
-            difficulty=1,
+            difficulty=restored.current_difficulty,
             state_root=temp_state.state_root(),
             receipt_root=calculate_receipt_root([receipt2]),
             receipts=[receipt2],
+            timestamp=int(time.time()),
         )
-        mine_block(block2, difficulty=1)
+        mine_block(block2)
 
         self.assertTrue(restored.add_block(block2))
         self.assertEqual(len(restored.chain), len(bc.chain) + 1)

tests/test_reorg.py

@@ -21,7 +21,7 @@ def genesis_file(tmp_path):
     pk = sk.verify_key.encode(encoder=HexEncoder).decode()
     data = {
         "timestamp": int(time.time()),
-        "difficulty": 0,
+        "difficulty": 1,
         "alloc": {
             pk: {"balance": 1000}
         }