Allow construct_cat_puzzle to take either inner puzzle or hash and document the use cases (#17112)

Allow construct_cat_puzzle to take inner_puzzle (Program) or hash (bytes32) and document the use cases.

chia/wallet/cat_wallet/cat_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Iterator, List, Optional
+from typing import Iterator, List, Optional, Union
 
 from chia_rs import G2Element
 
@@ -65,14 +65,19 @@ def get_innerpuzzle_from_puzzle(puzzle: Program) -> Program:
 
 
 def construct_cat_puzzle(
-    mod_code: Program, limitations_program_hash: bytes32, inner_puzzle: Program, mod_code_hash: Optional[bytes32] = None
+    mod_code: Program,
+    limitations_program_hash: bytes32,
+    inner_puzzle_or_hash: Union[Program, bytes32],
+    mod_code_hash: Optional[bytes32] = None,
 ) -> Program:
     """
-    Given an inner puzzle hash and tail hash calculate a puzzle program for a specific cc.
+    Given an inner puzzle and a tail hash, calculate a puzzle program for a specific cc.
+    We can also receive an inner puzzle hash instead, which wouldn't calculate a valid
+    puzzle, but that can be useful if calling `.get_tree_hash_precalc()` on it."
     """
     if mod_code_hash is None:
         mod_code_hash = mod_code.get_tree_hash()
-    return mod_code.curry(mod_code_hash, limitations_program_hash, inner_puzzle)
+    return mod_code.curry(mod_code_hash, limitations_program_hash, inner_puzzle_or_hash)
 
 
 def subtotals_for_deltas(deltas: List[int]) -> List[int]:

chia/wallet/cat_wallet/cat_wallet.py

@@ -937,10 +937,6 @@ class CATWallet:
 
     async def match_hinted_coin(self, coin: Coin, hint: bytes32) -> bool:
         return (
-            construct_cat_puzzle(
-                CAT_MOD,
-                self.cat_info.limitations_program_hash,
-                hint,  # type: ignore[arg-type]
-            ).get_tree_hash_precalc(hint)
+            construct_cat_puzzle(CAT_MOD, self.cat_info.limitations_program_hash, hint).get_tree_hash_precalc(hint)
             == coin.puzzle_hash
         )

chia/wallet/vc_wallet/cr_cat_drivers.py

@@ -207,11 +207,9 @@ class CRCAT:
             proofs_checker,
             payment.puzzle_hash,  # type: ignore
         ).get_tree_hash_precalc(payment.puzzle_hash)
-        new_cat_puzhash: bytes32 = construct_cat_puzzle(
-            CAT_MOD,
-            tail_hash,
-            new_cr_layer_hash,  # type: ignore
-        ).get_tree_hash_precalc(new_cr_layer_hash)
+        new_cat_puzhash = construct_cat_puzzle(CAT_MOD, tail_hash, new_cr_layer_hash).get_tree_hash_precalc(
+            new_cr_layer_hash
+        )
 
         eve_innerpuz: Program = Program.to(
             (

chia/wallet/vc_wallet/cr_cat_wallet.py

@@ -891,7 +891,7 @@ class CRCATWallet(CATWallet):
             construct_cat_puzzle(
                 Program.to(CAT_MOD_HASH),
                 self.info.limitations_program_hash,
-                hint_inner_hash,  # type: ignore
+                hint_inner_hash,
                 mod_code_hash=CAT_MOD_HASH_HASH,
             ).get_tree_hash_precalc(hint, CAT_MOD_HASH, CAT_MOD_HASH_HASH, hint_inner_hash)
             == coin.puzzle_hash
@@ -907,7 +907,7 @@ class CRCATWallet(CATWallet):
             construct_cat_puzzle(
                 Program.to(CAT_MOD_HASH),
                 self.info.limitations_program_hash,
-                pending_approval_inner_hash,  # type: ignore
+                pending_approval_inner_hash,
                 mod_code_hash=CAT_MOD_HASH_HASH,
             ).get_tree_hash_precalc(CAT_MOD_HASH, CAT_MOD_HASH_HASH, pending_approval_inner_hash)
             == coin.puzzle_hash

tests/wallet/cat_wallet/test_cat_wallet.py

@@ -945,9 +945,7 @@ class TestCATWallet:
         pubkey_unhardened = _derive_path_unhardened(intermediate_sk_un, [100000000]).get_g1()
         inner_puzhash = puzzle_hash_for_pk(pubkey_unhardened)
         puzzlehash_unhardened = construct_cat_puzzle(
-            CAT_MOD,
-            Program.to(None).get_tree_hash(),
-            inner_puzhash,  # type: ignore[arg-type]
+            CAT_MOD, Program.to(None).get_tree_hash(), inner_puzhash
         ).get_tree_hash_precalc(inner_puzhash)
         change_derivation = DerivationRecord(
             uint32(0),