add addition and removal root to header

src/full_node.py

@@ -3,7 +3,6 @@ import concurrent
 import logging
 import time
 from asyncio import Event
-from secrets import token_bytes
 from typing import AsyncGenerator, List, Optional, Tuple, Dict
 
 from chiabip158 import PyBIP158
@@ -18,6 +17,7 @@ from src.consensus.weight_verifier import verify_weight
 from src.protocols.wallet_protocol import FullProofForHash, ProofHash
 from src.store import FullNodeStore
 from src.protocols import farmer_protocol, full_node_protocol, timelord_protocol
+from src.util.MerkleSet import MerkleSet
 from src.util.bundle_tools import best_solution_program
 from src.mempool_manager import MempoolManager
 from src.server.outbound_message import Delivery, Message, NodeType, OutboundMessage
@@ -25,7 +25,7 @@ from src.server.server import ChiaServer
 from src.types.body import Body
 from src.types.challenge import Challenge
 from src.types.full_block import FullBlock
-from src.types.hashable.Coin import Coin
+from src.types.hashable.Coin import Coin, hash_coin_list
 from src.types.hashable.BLSSignature import BLSSignature
 from src.util.hash import std_hash
 from src.types.hashable.SpendBundle import SpendBundle
@@ -1271,8 +1271,39 @@ class FullNode:
         iterations_needed: uint64 = calculate_iterations(
             request.proof_of_space, difficulty, vdf_ips, constants["MIN_BLOCK_TIME"],
         )
-        additions_root = token_bytes(32)  # TODO(straya)
-        removal_root = token_bytes(32)  # TODO(straya)
+
+        removal_merkle_set = MerkleSet()
+        addition_merkle_set = MerkleSet()
+
+        additions = []
+        removals = []
+
+        if spend_bundle:
+            additions = spend_bundle.additions()
+            removals = spend_bundle.removals()
+
+        additions.append(request.coinbase)
+        additions.append(fees_coin)
+
+        # Create removal Merkle set
+        for coin in removals:
+            removal_merkle_set.add_already_hashed(coin.name())
+
+        # Create addition Merkle set
+        puzzlehash_coins_map: Dict[bytes32, List[Coin]] = {}
+        for coin in additions:
+            if coin.puzzle_hash in puzzlehash_coins_map:
+                puzzlehash_coins_map[coin.puzzle_hash].append(coin)
+            else:
+                puzzlehash_coins_map[coin.puzzle_hash] = [coin]
+
+        # Addition Merkle set contains puzzlehash and hash of all coins with that puzzlehash
+        for puzzle, coins in puzzlehash_coins_map.items():
+            addition_merkle_set.add_already_hashed(puzzle)
+            addition_merkle_set.add_already_hashed(hash_coin_list(coins))
+
+        additions_root = addition_merkle_set.get_root()
+        removal_root = removal_merkle_set.get_root()
 
         block_header_data: HeaderData = HeaderData(
             uint32(target_tip.height + 1),

src/types/hashable/Coin.py

@@ -1,9 +1,11 @@
 import io
 from dataclasses import dataclass
+from typing import List
 
 from clvm.casts import int_to_bytes, int_from_bytes
 
 from src.types.sized_bytes import bytes32
+from src.util.hash import std_hash
 from src.util.ints import uint64
 from src.util.streamable import streamable, Streamable
 
@@ -22,6 +24,10 @@ class Coin(Streamable):
     def name(self) -> bytes32:
         return self.get_hash()
 
+    @property
+    def name_str(self) -> str:
+        return self.name().hex()
+
     @classmethod
     def from_bytes(cls, blob):
         parent_coin_info = blob[:32]
@@ -35,3 +41,13 @@ class Coin(Streamable):
         f.write(self.puzzle_hash)
         f.write(int_to_bytes(self.amount))
         return f.getvalue()
+
+
+def hash_coin_list(coin_list: List[Coin]) -> bytes32:
+    coin_list.sort(key=lambda x: x.name_str, reverse=True)
+    buffer = bytearray()
+
+    for coin in coin_list:
+        buffer.extend(coin.name())
+
+    return std_hash(buffer)

src/util/MerkleSet.py

@@ -73,7 +73,7 @@ def get_bit(mybytes, pos):
     return (mybytes[pos // 8] >> (7 - (pos % 8))) & 1
 
 
-class ReferenceMerkleSet:
+class MerkleSet:
     def __init__(self, root=None):
         self.root = root
         if root is None:
@@ -336,7 +336,7 @@ def deserialize_proof(proof):
         r, pos = _deserialize(proof, 0, [])
         if pos != len(proof):
             raise SetError()
-        return ReferenceMerkleSet(r)
+        return MerkleSet(r)
     except IndexError:
         raise SetError()
 

tests/block_tools.py

@@ -20,12 +20,13 @@ from src.types.challenge import Challenge
 from src.types.classgroup import ClassgroupElement
 from src.types.full_block import FullBlock, additions_for_npc
 from src.types.hashable.BLSSignature import BLSSignature
-from src.types.hashable.Coin import Coin
+from src.types.hashable.Coin import Coin, hash_coin_list
 from src.types.hashable.Program import Program
 from src.types.header import Header, HeaderData
 from src.types.proof_of_space import ProofOfSpace
 from src.types.proof_of_time import ProofOfTime
 from src.types.sized_bytes import bytes32
+from src.util.MerkleSet import MerkleSet
 from src.util.errors import NoProofsOfSpaceFound
 from src.util.ints import uint8, uint32, uint64
 from src.util.hash import std_hash
@@ -447,12 +448,16 @@ class BlockTools:
 
         # Create filter
         byte_array_tx: List[bytes32] = []
+        tx_additions: List[Coin] = []
+        tx_removals: List[bytes32] = []
         if transactions:
             error, npc_list, _ = get_name_puzzle_conditions(transactions)
             additions: List[Coin] = additions_for_npc(npc_list)
             for coin in additions:
+                tx_additions.append(coin)
                 byte_array_tx.append(bytearray(coin.puzzle_hash))
             for npc in npc_list:
+                tx_removals.append(npc.coin_name)
                 byte_array_tx.append(bytearray(npc.coin_name))
 
         byte_array_tx.append(bytearray(coinbase_coin.puzzle_hash))
@@ -461,6 +466,32 @@ class BlockTools:
         bip158: PyBIP158 = PyBIP158(byte_array_tx)
         encoded = bytes(bip158.GetEncoded())
 
+        removal_merkle_set = MerkleSet()
+        addition_merkle_set = MerkleSet()
+
+        tx_additions.append(coinbase_coin)
+        tx_additions.append(fees_coin)
+
+        # Create removal Merkle set
+        for coin_name in tx_removals:
+            removal_merkle_set.add_already_hashed(coin_name)
+
+        # Create addition Merkle set
+        puzzlehash_coin_map: Dict[bytes32, List[Coin]] = {}
+        for coin in tx_additions:
+            if coin.puzzle_hash in puzzlehash_coin_map:
+                puzzlehash_coin_map[coin.puzzle_hash].append(coin)
+            else:
+                puzzlehash_coin_map[coin.puzzle_hash] = [coin]
+
+        # Addition Merkle set contains puzzlehash and hash of all coins with that puzzlehash
+        for puzzle, coins in puzzlehash_coin_map.items():
+            addition_merkle_set.add_already_hashed(puzzle)
+            addition_merkle_set.add_already_hashed(hash_coin_list(coins))
+
+        additions_root = addition_merkle_set.get_root()
+        removal_root = removal_merkle_set.get_root()
+
         header_data: HeaderData = HeaderData(
             height,
             prev_header_hash,
@@ -470,8 +501,8 @@ class BlockTools:
             body.get_hash(),
             uint64(prev_weight + difficulty),
             uint64(prev_iters + number_iters),
-            bytes([0] * 32),
-            bytes([0] * 32),
+            additions_root,
+            removal_root,
         )
 
         header_hash_sig: PrependSignature = plot_sk.sign_prepend(header_data.get_hash())

tests/test_merkle_set.py

@@ -2,7 +2,7 @@ import asyncio
 
 import pytest
 
-from src.util.MerkleSet import ReferenceMerkleSet, confirm_included_already_hashed
+from src.util.MerkleSet import MerkleSet, confirm_included_already_hashed
 from tests.setup_nodes import test_constants, bt
 from tests.wallet_tools import WalletTool
 
@@ -27,7 +27,7 @@ class TestMerkleSet:
             reward_puzzlehash=wallet_tool.get_new_puzzlehash(),
         )
 
-        merkle_set = ReferenceMerkleSet()
+        merkle_set = MerkleSet()
         for block in blocks:
             merkle_set.add_already_hashed(block.body.coinbase.name())