uncurry the puzzle once instead of 3 times (#12781)

chia/rpc/wallet_rpc_api.py

@@ -1659,7 +1659,7 @@ class WalletRpcApi:
             # Check if the metadata is updated
             full_puzzle: Program = Program.from_bytes(bytes(coin_spend.puzzle_reveal))
 
-            uncurried_nft: Optional[UncurriedNFT] = UncurriedNFT.uncurry(full_puzzle)
+            uncurried_nft: Optional[UncurriedNFT] = UncurriedNFT.uncurry(*full_puzzle.uncurry())
             if uncurried_nft is None:
                 return {"success": False, "error": "The coin is not a NFT."}
             metadata, p2_puzzle_hash = get_metadata_and_phs(uncurried_nft, coin_spend.solution)

chia/wallet/cat_wallet/cat_outer_puzzle.py

@@ -27,31 +27,29 @@ class CATOuterPuzzle:
     _get_inner_solution: Any
 
     def match(self, puzzle: Program) -> Optional[PuzzleInfo]:
-        matched, curried_args = match_cat_puzzle(puzzle)
-        if matched:
-            _, tail_hash, inner_puzzle = curried_args
-            constructor_dict = {
-                "type": "CAT",
-                "tail": "0x" + tail_hash.as_python().hex(),
-            }
-            next_constructor = self._match(inner_puzzle)
-            if next_constructor is not None:
-                constructor_dict["also"] = next_constructor.info
-            return PuzzleInfo(constructor_dict)
-        else:
+        args = match_cat_puzzle(*puzzle.uncurry())
+        if args is None:
             return None
+        _, tail_hash, inner_puzzle = args
+        constructor_dict = {
+            "type": "CAT",
+            "tail": "0x" + tail_hash.as_python().hex(),
+        }
+        next_constructor = self._match(inner_puzzle)
+        if next_constructor is not None:
+            constructor_dict["also"] = next_constructor.info
+        return PuzzleInfo(constructor_dict)
 
     def get_inner_puzzle(self, constructor: PuzzleInfo, puzzle_reveal: Program) -> Optional[Program]:
-        matched, curried_args = match_cat_puzzle(puzzle_reveal)
-        if matched:
-            _, _, inner_puzzle = curried_args
-            if constructor.also() is not None:
-                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(constructor.also(), inner_puzzle)
-                return deep_inner_puzzle
-            else:
-                return inner_puzzle
-        else:
+        args = match_cat_puzzle(*puzzle_reveal.uncurry())
+        if args is None:
             raise ValueError("This driver is not for the specified puzzle reveal")
+        _, _, inner_puzzle = args
+        if constructor.also() is not None:
+            deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(constructor.also(), inner_puzzle)
+            return deep_inner_puzzle
+        else:
+            return inner_puzzle
 
     def get_inner_solution(self, constructor: PuzzleInfo, solution: Program) -> Optional[Program]:
         my_inner_solution: Program = solution.first()
@@ -96,9 +94,9 @@ class CATOuterPuzzle:
             if constructor.also() is not None:
                 puzzle = self._construct(constructor.also(), puzzle)
                 solution = self._solve(constructor.also(), solver, inner_puzzle, inner_solution)
-            matched, curried_args = match_cat_puzzle(parent_spend.puzzle_reveal.to_program())
-            assert matched
-            _, _, parent_inner_puzzle = curried_args
+            args = match_cat_puzzle(*parent_spend.puzzle_reveal.to_program().uncurry())
+            assert args is not None
+            _, _, parent_inner_puzzle = args
             spendable_cats.append(
                 SpendableCAT(
                     coin,

chia/wallet/cat_wallet/cat_utils.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import List, Tuple, Iterator, Optional
+from typing import List, Iterator, Optional
 
 from blspy import G2Element
 
@@ -30,15 +30,15 @@ class SpendableCAT:
     limitations_program_reveal: Program = Program.to([])
 
 
-def match_cat_puzzle(puzzle: Program) -> Tuple[bool, Iterator[Program]]:
+def match_cat_puzzle(mod: Program, curried_args: Program) -> Optional[Iterator[Program]]:
     """
-    Given a puzzle test if it's a CAT and, if it is, return the curried arguments
+    Given the curried puzzle and args, test if it's a CAT and,
+    if it is, return the curried arguments
     """
-    mod, curried_args = puzzle.uncurry()
     if mod == CAT_MOD:
-        return True, curried_args.as_iter()
+        return curried_args.as_iter()
     else:
-        return False, iter(())
+        return None
 
 
 def get_innerpuzzle_from_puzzle(puzzle: Program) -> Program:

chia/wallet/cat_wallet/cat_wallet.py

@@ -340,9 +340,9 @@ class CATWallet:
     async def puzzle_solution_received(self, coin_spend: CoinSpend, parent_coin: Coin):
         coin_name = coin_spend.coin.name()
         puzzle: Program = Program.from_bytes(bytes(coin_spend.puzzle_reveal))
-        matched, curried_args = match_cat_puzzle(puzzle)
-        if matched:
-            mod_hash, genesis_coin_checker_hash, inner_puzzle = curried_args
+        args = match_cat_puzzle(*puzzle.uncurry())
+        if args is not None:
+            mod_hash, genesis_coin_checker_hash, inner_puzzle = args
             self.log.info(f"parent: {coin_name} inner_puzzle for parent is {inner_puzzle}")
 
             await self.add_lineage(
@@ -458,9 +458,9 @@ class CATWallet:
     async def sign(self, spend_bundle: SpendBundle) -> SpendBundle:
         sigs: List[G2Element] = []
         for spend in spend_bundle.coin_spends:
-            matched, puzzle_args = match_cat_puzzle(spend.puzzle_reveal.to_program())
-            if matched:
-                _, _, inner_puzzle = puzzle_args
+            args = match_cat_puzzle(*spend.puzzle_reveal.to_program().uncurry())
+            if args is not None:
+                _, _, inner_puzzle = args
                 puzzle_hash = inner_puzzle.get_tree_hash()
                 ret = await self.wallet_state_manager.get_keys(puzzle_hash)
                 if ret is None:

chia/wallet/did_wallet/did_wallet.py

@@ -1078,8 +1078,9 @@ class DIDWallet:
     async def sign(self, spend_bundle: SpendBundle) -> SpendBundle:
         sigs: List[G2Element] = []
         for spend in spend_bundle.coin_spends:
-            matched, puzzle_args = did_wallet_puzzles.match_did_puzzle(spend.puzzle_reveal.to_program())
-            if matched:
+
+            puzzle_args = did_wallet_puzzles.match_did_puzzle(*spend.puzzle_reveal.to_program().uncurry())
+            if puzzle_args is not None:
                 p2_puzzle, _, _, _, _ = puzzle_args
                 puzzle_hash = p2_puzzle.get_tree_hash()
                 pubkey, private = await self.wallet_state_manager.get_keys(puzzle_hash)

chia/wallet/did_wallet/did_wallet_puzzles.py

@@ -163,24 +163,22 @@ def create_spend_for_message(
     return coinsol
 
 
-def match_did_puzzle(puzzle: Program) -> Tuple[bool, Iterator[Program]]:
+def match_did_puzzle(mod: Program, curried_args: Program) -> Optional[Iterator[Program]]:
     """
         Given a puzzle test if it's a DID, if it is, return the curried arguments
     :param puzzle: Puzzle
     :return: Curried parameters
     """
     try:
-        mod, curried_args = puzzle.uncurry()
         if mod == SINGLETON_TOP_LAYER_MOD:
             mod, curried_args = curried_args.rest().first().uncurry()
             if mod == DID_INNERPUZ_MOD:
-                return True, curried_args.as_iter()
+                return curried_args.as_iter()
     except Exception:
         import traceback
 
         print(f"exception: {traceback.format_exc()}")
-        return False, iter(())
-    return False, iter(())
+    return None
 
 
 def check_is_did_puzzle(puzzle: Program) -> bool:

chia/wallet/nft_wallet/nft_puzzles.py

@@ -88,7 +88,7 @@ def get_nft_info_from_puzzle(nft_coin_info: NFTCoinInfo) -> NFTInfo:
     :param nft_coin_info NFTCoinInfo in local database
     :return: NFTInfo
     """
-    uncurried_nft: Optional[UncurriedNFT] = UncurriedNFT.uncurry(nft_coin_info.full_puzzle)
+    uncurried_nft: Optional[UncurriedNFT] = UncurriedNFT.uncurry(*nft_coin_info.full_puzzle.uncurry())
     assert uncurried_nft is not None
     data_uris: List[str] = []
 

chia/wallet/nft_wallet/nft_wallet.py

@@ -166,7 +166,7 @@ class NFTWallet:
         # At this point, the puzzle must be a NFT puzzle.
         # This method will be called only when the wallet state manager uncurried this coin as a NFT puzzle.
 
-        uncurried_nft = UncurriedNFT.uncurry(puzzle)
+        uncurried_nft = UncurriedNFT.uncurry(*puzzle.uncurry())
         assert uncurried_nft is not None
         self.log.info(
             f"found the info for NFT coin {coin_name} {uncurried_nft.inner_puzzle} {uncurried_nft.singleton_struct}"
@@ -414,7 +414,7 @@ class NFTWallet:
         for spend in spend_bundle.coin_spends:
             pks = {}
             if not puzzle_hashes:
-                uncurried_nft = UncurriedNFT.uncurry(spend.puzzle_reveal.to_program())
+                uncurried_nft = UncurriedNFT.uncurry(*spend.puzzle_reveal.to_program().uncurry())
                 if uncurried_nft is not None:
                     self.log.debug("Found a NFT state layer to sign")
                     puzzle_hashes.append(uncurried_nft.p2_puzzle.get_tree_hash())
@@ -450,7 +450,7 @@ class NFTWallet:
     async def update_metadata(
         self, nft_coin_info: NFTCoinInfo, key: str, uri: str, fee: uint64 = uint64(0)
     ) -> Optional[SpendBundle]:
-        uncurried_nft = UncurriedNFT.uncurry(nft_coin_info.full_puzzle)
+        uncurried_nft = UncurriedNFT.uncurry(*nft_coin_info.full_puzzle.uncurry())
         assert uncurried_nft is not None
         puzzle_hash = uncurried_nft.p2_puzzle.get_tree_hash()
 
@@ -695,7 +695,7 @@ class NFTWallet:
             puzzle_announcements_to_assert=puzzle_announcements_bytes,
         )
 
-        unft = UncurriedNFT.uncurry(nft_coin.full_puzzle)
+        unft = UncurriedNFT.uncurry(*nft_coin.full_puzzle.uncurry())
         assert unft is not None
         magic_condition = None
         if unft.supports_did:
@@ -898,7 +898,7 @@ class NFTWallet:
 
     async def set_nft_did(self, nft_coin_info: NFTCoinInfo, did_id: bytes, fee: uint64 = uint64(0)) -> SpendBundle:
         self.log.debug("Setting NFT DID with parameters: nft=%s did=%s", nft_coin_info, did_id)
-        unft = UncurriedNFT.uncurry(nft_coin_info.full_puzzle)
+        unft = UncurriedNFT.uncurry(*nft_coin_info.full_puzzle.uncurry())
         assert unft is not None
         nft_id = unft.singleton_launcher_id
         puzzle_hashes_to_sign = [unft.p2_puzzle.get_tree_hash()]

chia/wallet/nft_wallet/uncurry_nft.py

@@ -85,14 +85,14 @@ class UncurriedNFT:
     trade_price_percentage: Optional[uint16]
 
     @classmethod
-    def uncurry(cls: Type[_T_UncurriedNFT], puzzle: Program) -> Optional[_T_UncurriedNFT]:
+    def uncurry(cls: Type[_T_UncurriedNFT], mod: Program, curried_args: Program) -> Optional[_T_UncurriedNFT]:
         """
         Try to uncurry a NFT puzzle
         :param cls UncurriedNFT class
-        :param puzzle: Puzzle program
+        :param mod: uncurried Puzzle program
+        :param uncurried_args: uncurried arguments to program
         :return Uncurried NFT
         """
-        mod, curried_args = puzzle.uncurry()
         if mod != SINGLETON_TOP_LAYER_MOD:
             log.debug("Cannot uncurry NFT puzzle, failed on singleton top layer: Mod %s", mod)
             return None

chia/wallet/wallet_state_manager.py

@@ -636,21 +636,23 @@ class WalletStateManager:
 
         puzzle = Program.from_bytes(bytes(coin_spend.puzzle_reveal))
 
+        mod, curried_args = puzzle.uncurry()
+
         # Check if the coin is a CAT
-        cat_matched, cat_curried_args = match_cat_puzzle(puzzle)
-        if cat_matched:
+        cat_curried_args = match_cat_puzzle(mod, curried_args)
+        if cat_curried_args is not None:
             return await self.handle_cat(cat_curried_args, parent_coin_state, coin_state, coin_spend)
 
         # Check if the coin is a NFT
         #                                                        hint
         # First spend where 1 mojo coin -> Singleton launcher -> NFT -> NFT
-        uncurried_nft = UncurriedNFT.uncurry(puzzle)
+        uncurried_nft = UncurriedNFT.uncurry(mod, curried_args)
         if uncurried_nft is not None:
             return await self.handle_nft(coin_spend, uncurried_nft)
 
         # Check if the coin is a DID
-        did_matched, did_curried_args = match_did_puzzle(puzzle)
-        if did_matched:
+        did_curried_args = match_did_puzzle(mod, curried_args)
+        if did_curried_args is not None:
             return await self.handle_did(did_curried_args, parent_coin_state, coin_state, coin_spend)
 
         return None, None

tests/wallet/nft_wallet/test_nft_puzzles.py

@@ -67,7 +67,7 @@ def test_nft_transfer_puzzle_hashes():
     nft_info = match_puzzle(nft_puz)
     assert nft_info.also().also() is not None
 
-    unft = uncurry_nft.UncurriedNFT.uncurry(nft_puz)
+    unft = uncurry_nft.UncurriedNFT.uncurry(*nft_puz.uncurry())
     assert unft is not None
     assert unft.supports_did
 
@@ -183,7 +183,7 @@ def test_transfer_puzzle_builder() -> None:
         ownership_puzzle,
     )
     clvm_puzzle_hash = get_updated_nft_puzzle(clvm_nft_puzzle, solution.at("rrf"))
-    unft = uncurry_nft.UncurriedNFT.uncurry(puzzle)
+    unft = uncurry_nft.UncurriedNFT.uncurry(*puzzle.uncurry())
     assert unft is not None
     assert unft.nft_state_layer == clvm_nft_puzzle
     assert unft.inner_puzzle == ownership_puzzle

tools/run_block.py

@@ -106,12 +106,12 @@ def run_generator(block_generator: BlockGenerator, constants: ConsensusConstants
     for spend in coin_spends.as_iter():
 
         parent, puzzle, amount, solution = spend.as_iter()
-        matched, curried_args = match_cat_puzzle(puzzle)
+        args = match_cat_puzzle(*puzzle.uncurry())
 
-        if not matched:
+        if args is None:
             continue
 
-        _, asset_id, _ = curried_args
+        _, asset_id, _ = args
         memo = ""
 
         puzzle_result = puzzle.run(solution)