Puzzle hash optimizations (#17995)

* use the NULL_HASHTREE constant instead of re-computing the hash on the fly

* optimize puzzle hash computations by using shatree_atom() and shatree_pair() directly, rather than constructing a Program to call get_tree_hash() on.

* cache synthetic public keys and private keys

* rename NULL_TREEHASH -> NIL_TREEHASH

chia/_tests/clvm/test_curry_and_treehash.py

@@ -1,8 +1,19 @@
 from __future__ import annotations
 
+from typing import List
+
+import pytest
+
 from chia.types.blockchain_format.program import Program
+from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.wallet.puzzles import p2_delegated_puzzle_or_hidden_puzzle  # import (puzzle_for_pk, puzzle_hash_for_pk, MOD)
-from chia.wallet.util.curry_and_treehash import calculate_hash_of_quoted_mod_hash, curry_and_treehash
+from chia.wallet.util.curry_and_treehash import (
+    calculate_hash_of_quoted_mod_hash,
+    curry_and_treehash,
+    shatree_atom,
+    shatree_atom_list,
+    shatree_int,
+)
 
 
 def test_curry_and_treehash() -> None:
@@ -21,3 +32,26 @@ def test_curry_and_treehash() -> None:
         hashed_args = [Program.to(_).get_tree_hash() for _ in args]
         puzzle_hash_via_f = curry_and_treehash(quoted_mod_hash, *hashed_args)
         assert puzzle_hash_via_curry == puzzle_hash_via_f
+
+
+@pytest.mark.parametrize(
+    "value", [[], [bytes32([3] * 32)], [bytes32([0] * 32), bytes32([1] * 32)], [bytes([1]), bytes([1, 2, 3])]]
+)
+def test_shatree_atom_list(value: List[bytes]) -> None:
+    h1 = shatree_atom_list(value)
+    h2 = Program.to(value).get_tree_hash()
+    assert h1 == h2
+
+
+@pytest.mark.parametrize("value", [0, -1, 1, 0x7F, 0x80, 100000000, -10000000])
+def test_shatree_int(value: int) -> None:
+    h1 = shatree_int(value)
+    h2 = Program.to(value).get_tree_hash()
+    assert h1 == h2
+
+
+@pytest.mark.parametrize("value", [bytes([1] * 1), bytes([]), bytes([5] * 1000)])
+def test_shatree_atom(value: bytes) -> None:
+    h1 = shatree_atom(value)
+    h2 = Program.to(value).get_tree_hash()
+    assert h1 == h2

chia/_tests/wallet/did_wallet/test_did.py

@@ -1313,7 +1313,7 @@ async def test_did_resync(self, self_hostname, two_wallet_nodes, trusted) -> Non
                 wallet_node_1.wallet_state_manager,
                 wallet,
                 uint64(101),
-                [bytes(ph)],
+                [bytes32(ph)],
                 uint64(1),
                 {"Twitter": "Test", "GitHub": "测试"},
                 fee=fee,

chia/pools/pool_puzzles.py

@@ -205,7 +205,7 @@ def create_travel_spend(
         log.debug(
             f"create_travel_spend: waitingroom: target PoolState bytes:\n{bytes(target).hex()}\n"
             f"{target}"
-            f"hash:{Program.to(bytes(target)).get_tree_hash()}"
+            f"hash:{shatree_atom(bytes(target))}"
         )
         # key_value_list is:
         # "p" -> poolstate as bytes

chia/rpc/wallet_rpc_api.py

@@ -110,6 +110,7 @@
 from chia.wallet.util.address_type import AddressType, is_valid_address
 from chia.wallet.util.compute_hints import compute_spend_hints_and_additions
 from chia.wallet.util.compute_memos import compute_memos
+from chia.wallet.util.curry_and_treehash import NIL_TREEHASH
 from chia.wallet.util.query_filter import HashFilter, TransactionTypeFilter
 from chia.wallet.util.transaction_type import CLAWBACK_INCOMING_TRANSACTION_TYPES, TransactionType
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG, CoinSelectionConfig, CoinSelectionConfigLoader, TXConfig
@@ -2332,7 +2333,7 @@ async def did_find_lost_did(self, request: Dict[str, Any]) -> EndpointResult:
             )
             full_puzzle = create_singleton_puzzle(did_puzzle, launcher_id)
             did_puzzle_empty_recovery = DID_INNERPUZ_MOD.curry(
-                our_inner_puzzle, Program.to([]).get_tree_hash(), uint64(0), singleton_struct, metadata
+                our_inner_puzzle, NIL_TREEHASH, uint64(0), singleton_struct, metadata
             )
             # Check if we have the DID wallet
             did_wallet: Optional[DIDWallet] = None
@@ -2422,7 +2423,7 @@ async def did_find_lost_did(self, request: Dict[str, Any]) -> EndpointResult:
                     inner_solution: Program = full_solution.rest().rest().first()
                     recovery_list: List[bytes32] = []
                     backup_required: int = num_verification.as_int()
-                    if recovery_list_hash != Program.to([]).get_tree_hash():
+                    if recovery_list_hash != NIL_TREEHASH:
                         try:
                             for did in inner_solution.rest().rest().rest().rest().rest().as_python():
                                 recovery_list.append(did[0])

chia/wallet/did_wallet/did_wallet.py

@@ -46,6 +46,7 @@
 from chia.wallet.transaction_record import TransactionRecord
 from chia.wallet.uncurried_puzzle import uncurry_puzzle
 from chia.wallet.util.compute_memos import compute_memos
+from chia.wallet.util.curry_and_treehash import NIL_TREEHASH, shatree_int, shatree_pair
 from chia.wallet.util.transaction_type import TransactionType
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG, CoinSelectionConfig, TXConfig
 from chia.wallet.util.wallet_sync_utils import fetch_coin_spend, fetch_coin_spend_for_coin_state
@@ -75,7 +76,7 @@ async def create_new_did_wallet(
         wallet_state_manager: Any,
         wallet: Wallet,
         amount: uint64,
-        backups_ids: List = [],
+        backups_ids: List[bytes32] = [],
         num_of_backup_ids_needed: uint64 = None,
         metadata: Dict[str, str] = {},
         name: Optional[str] = None,
@@ -227,10 +228,10 @@ async def create_new_did_wallet_from_coin_spend(
         inner_solution: Program = full_solution.rest().rest().first()
         recovery_list: List[bytes32] = []
         backup_required: int = num_verification.as_int()
-        if recovery_list_hash != Program.to([]).get_tree_hash():
+        if recovery_list_hash != NIL_TREEHASH:
             try:
                 for did in inner_solution.rest().rest().rest().rest().rest().as_python():
-                    recovery_list.append(did[0])
+                    recovery_list.append(bytes32(did[0]))
             except Exception:
                 self.log.warning(
                     f"DID {launch_coin.name().hex()} has a recovery list hash but missing a reveal,"
@@ -529,7 +530,9 @@ def puzzle_for_pk(self, pubkey: G1Element) -> Program:
     def puzzle_hash_for_pk(self, pubkey: G1Element) -> bytes32:
         if self.did_info.origin_coin is None:
             # TODO: this seem dumb. Why bother with this case? Is it ever used?
-            return puzzle_for_pk(pubkey).get_tree_hash()
+            # inner puzzle: (8 . 0)
+            innerpuz_hash = shatree_pair(shatree_int(8), NIL_TREEHASH)
+            return create_singleton_puzzle_hash(innerpuz_hash, bytes32([0] * 32))
         origin_coin_name = self.did_info.origin_coin.name()
         innerpuz_hash = did_wallet_puzzles.get_inner_puzhash_by_p2(
             p2_puzhash=puzzle_hash_for_pk(pubkey),

chia/wallet/did_wallet/did_wallet_puzzles.py

@@ -13,16 +13,26 @@
 from chia.wallet.puzzles.load_clvm import load_clvm_maybe_recompile
 from chia.wallet.singleton import (
     SINGLETON_LAUNCHER_PUZZLE_HASH,
+    SINGLETON_LAUNCHER_PUZZLE_HASH_TREE_HASH,
     SINGLETON_TOP_LAYER_MOD,
     SINGLETON_TOP_LAYER_MOD_HASH,
+    SINGLETON_TOP_LAYER_MOD_HASH_TREE_HASH,
     is_singleton,
 )
-from chia.wallet.util.curry_and_treehash import calculate_hash_of_quoted_mod_hash, curry_and_treehash
+from chia.wallet.util.curry_and_treehash import (
+    calculate_hash_of_quoted_mod_hash,
+    curry_and_treehash,
+    shatree_atom,
+    shatree_atom_list,
+    shatree_int,
+    shatree_pair,
+)
 
 DID_INNERPUZ_MOD = load_clvm_maybe_recompile(
     "did_innerpuz.clsp", package_or_requirement="chia.wallet.did_wallet.puzzles"
 )
 DID_INNERPUZ_MOD_HASH = DID_INNERPUZ_MOD.get_tree_hash()
+DID_INNERPUZ_MOD_HASH_QUOTED = calculate_hash_of_quoted_mod_hash(DID_INNERPUZ_MOD_HASH)
 INTERMEDIATE_LAUNCHER_MOD = load_clvm_maybe_recompile(
     "nft_intermediate_launcher.clsp", package_or_requirement="chia.wallet.nft_wallet.puzzles"
 )
@@ -74,17 +84,20 @@ def get_inner_puzhash_by_p2(
     :return: DID inner puzzle hash
     """
 
-    backup_ids_hash = Program(Program.to(recovery_list)).get_tree_hash()
-    singleton_struct = Program.to((SINGLETON_TOP_LAYER_MOD_HASH, (launcher_id, SINGLETON_LAUNCHER_PUZZLE_HASH)))
+    backup_ids_hash = shatree_atom_list(recovery_list)
 
-    quoted_mod_hash = calculate_hash_of_quoted_mod_hash(DID_INNERPUZ_MOD_HASH)
+    # singleton_struct = (MOD_HASH . (LAUNCHER_ID . LAUNCHER_PUZZLE_HASH))
+    singleton_struct = shatree_pair(
+        SINGLETON_TOP_LAYER_MOD_HASH_TREE_HASH,
+        shatree_pair(shatree_atom(launcher_id), SINGLETON_LAUNCHER_PUZZLE_HASH_TREE_HASH),
+    )
 
     return curry_and_treehash(
-        quoted_mod_hash,
+        DID_INNERPUZ_MOD_HASH_QUOTED,
         p2_puzhash,
-        Program.to(backup_ids_hash).get_tree_hash(),
-        Program.to(num_of_backup_ids_needed).get_tree_hash(),
-        Program.to(singleton_struct).get_tree_hash(),
+        shatree_atom(backup_ids_hash),
+        shatree_int(num_of_backup_ids_needed),
+        singleton_struct,
         metadata.get_tree_hash(),
     )
 

chia/wallet/puzzles/p2_delegated_puzzle_or_hidden_puzzle.py

@@ -59,14 +59,15 @@
 from __future__ import annotations
 
 import hashlib
+from functools import lru_cache
 from typing import Union
 
 from chia_rs import G1Element, PrivateKey
 from clvm.casts import int_from_bytes
 
 from chia.types.blockchain_format.program import Program
 from chia.types.blockchain_format.sized_bytes import bytes32
-from chia.wallet.util.curry_and_treehash import calculate_hash_of_quoted_mod_hash, curry_and_treehash
+from chia.wallet.util.curry_and_treehash import calculate_hash_of_quoted_mod_hash, curry_and_treehash, shatree_atom
 
 from .load_clvm import load_clvm_maybe_recompile
 from .p2_conditions import puzzle_for_conditions
@@ -91,13 +92,15 @@ def calculate_synthetic_offset(public_key: G1Element, hidden_puzzle_hash: bytes3
     return offset
 
 
+@lru_cache(maxsize=1000)
 def calculate_synthetic_public_key(public_key: G1Element, hidden_puzzle_hash: bytes32) -> G1Element:
     synthetic_offset: PrivateKey = PrivateKey.from_bytes(
         calculate_synthetic_offset(public_key, hidden_puzzle_hash).to_bytes(32, "big")
     )
     return public_key + synthetic_offset.get_g1()
 
 
+@lru_cache(maxsize=1000)
 def calculate_synthetic_secret_key(secret_key: PrivateKey, hidden_puzzle_hash: bytes32) -> PrivateKey:
     secret_exponent = int.from_bytes(bytes(secret_key), "big")
     public_key = secret_key.get_g1()
@@ -113,7 +116,7 @@ def puzzle_for_synthetic_public_key(synthetic_public_key: G1Element) -> Program:
 
 
 def puzzle_hash_for_synthetic_public_key(synthetic_public_key: G1Element) -> bytes32:
-    public_key_hash = Program.to(bytes(synthetic_public_key)).get_tree_hash()
+    public_key_hash = shatree_atom(bytes(synthetic_public_key))
     return curry_and_treehash(QUOTED_MOD_HASH, public_key_hash)
 
 

chia/wallet/singleton.py

@@ -8,13 +8,20 @@
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.coin_spend import CoinSpend, compute_additions
 from chia.wallet.puzzles.load_clvm import load_clvm_maybe_recompile
-from chia.wallet.util.curry_and_treehash import calculate_hash_of_quoted_mod_hash, curry_and_treehash
+from chia.wallet.util.curry_and_treehash import (
+    calculate_hash_of_quoted_mod_hash,
+    curry_and_treehash,
+    shatree_atom,
+    shatree_pair,
+)
 
 SINGLETON_TOP_LAYER_MOD = load_clvm_maybe_recompile("singleton_top_layer_v1_1.clsp")
 SINGLETON_TOP_LAYER_MOD_HASH = SINGLETON_TOP_LAYER_MOD.get_tree_hash()
+SINGLETON_TOP_LAYER_MOD_HASH_TREE_HASH = shatree_atom(SINGLETON_TOP_LAYER_MOD_HASH)
 SINGLETON_TOP_LAYER_MOD_HASH_QUOTED = calculate_hash_of_quoted_mod_hash(SINGLETON_TOP_LAYER_MOD_HASH)
 SINGLETON_LAUNCHER_PUZZLE = load_clvm_maybe_recompile("singleton_launcher.clsp")
 SINGLETON_LAUNCHER_PUZZLE_HASH = SINGLETON_LAUNCHER_PUZZLE.get_tree_hash()
+SINGLETON_LAUNCHER_PUZZLE_HASH_TREE_HASH = shatree_atom(SINGLETON_LAUNCHER_PUZZLE_HASH)
 
 
 def get_inner_puzzle_from_singleton(puzzle: Union[Program, SerializedProgram]) -> Optional[Program]:
@@ -66,9 +73,12 @@ def create_singleton_puzzle_hash(innerpuz_hash: bytes32, launcher_id: bytes32) -
     :return: Singleton full puzzle hash
     """
     # singleton_struct = (MOD_HASH . (LAUNCHER_ID . LAUNCHER_PUZZLE_HASH))
-    singleton_struct = Program.to((SINGLETON_TOP_LAYER_MOD_HASH, (launcher_id, SINGLETON_LAUNCHER_PUZZLE_HASH)))
+    singleton_struct = shatree_pair(
+        SINGLETON_TOP_LAYER_MOD_HASH_TREE_HASH,
+        shatree_pair(shatree_atom(launcher_id), SINGLETON_LAUNCHER_PUZZLE_HASH_TREE_HASH),
+    )
 
-    return curry_and_treehash(SINGLETON_TOP_LAYER_MOD_HASH_QUOTED, singleton_struct.get_tree_hash(), innerpuz_hash)
+    return curry_and_treehash(SINGLETON_TOP_LAYER_MOD_HASH_QUOTED, singleton_struct, innerpuz_hash)
 
 
 def create_singleton_puzzle(innerpuz: Union[Program, SerializedProgram], launcher_id: bytes32) -> Program:

chia/wallet/util/curry_and_treehash.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 from hashlib import sha256
-from typing import Callable, List
+from typing import Callable, List, Sequence
+
+from clvm.casts import int_to_bytes
 
 from chia.types.blockchain_format.sized_bytes import bytes32
 
@@ -35,7 +37,18 @@ def shatree_pair(left_hash: bytes32, right_hash: bytes32) -> bytes32:
 A_KW_TREEHASH = shatree_atom(A_KW)
 C_KW_TREEHASH = shatree_atom(C_KW)
 ONE_TREEHASH = shatree_atom(ONE)
-NULL_TREEHASH = shatree_atom(NULL)
+NIL_TREEHASH = shatree_atom(NULL)
+
+
+def shatree_atom_list(li: Sequence[bytes]) -> bytes32:
+    ret = NIL_TREEHASH
+    for item in reversed(li):
+        ret = shatree_pair(shatree_atom(item), ret)
+    return ret
+
+
+def shatree_int(val: int) -> bytes32:
+    return shatree_atom(int_to_bytes(val))
 
 
 # The environment `E = (F . R)` recursively expands out to
@@ -51,7 +64,7 @@ def curried_values_tree_hash(arguments: List[bytes32]) -> bytes32:
         C_KW_TREEHASH,
         shatree_pair(
             shatree_pair(Q_KW_TREEHASH, arguments[0]),
-            shatree_pair(curried_values_tree_hash(arguments[1:]), NULL_TREEHASH),
+            shatree_pair(curried_values_tree_hash(arguments[1:]), NIL_TREEHASH),
         ),
     )
 
@@ -69,7 +82,7 @@ def curry_and_treehash(hash_of_quoted_mod_hash: bytes32, *hashed_arguments: byte
     curried_values = curried_values_tree_hash(list(hashed_arguments))
     return shatree_pair(
         A_KW_TREEHASH,
-        shatree_pair(hash_of_quoted_mod_hash, shatree_pair(curried_values, NULL_TREEHASH)),
+        shatree_pair(hash_of_quoted_mod_hash, shatree_pair(curried_values, NIL_TREEHASH)),
     )
 
 

chia/wallet/wallet_state_manager.py

@@ -127,6 +127,7 @@
 from chia.wallet.util.address_type import AddressType
 from chia.wallet.util.compute_hints import compute_spend_hints_and_additions
 from chia.wallet.util.compute_memos import compute_memos
+from chia.wallet.util.curry_and_treehash import NIL_TREEHASH
 from chia.wallet.util.puzzle_decorator import PuzzleDecoratorManager
 from chia.wallet.util.query_filter import HashFilter
 from chia.wallet.util.transaction_type import CLAWBACK_INCOMING_TRANSACTION_TYPES, TransactionType
@@ -1254,7 +1255,7 @@ async def handle_did(
             full_puzzle = create_singleton_puzzle(did_puzzle, launch_id)
             did_puzzle_empty_recovery = DID_INNERPUZ_MOD.curry(
                 our_inner_puzzle,
-                Program.to([]).get_tree_hash(),
+                NIL_TREEHASH,
                 uint64(0),
                 parent_data.singleton_struct,
                 parent_data.metadata,