Introduce and use LineageProof.from_program (#17301)

* Introduce and use LineageProof.from_program.

* tweak from_program and add tests

---------

Co-authored-by: Amine Khaldi <amine.khaldi@reactos.org>

chia/wallet/lineage_proof.py

@@ -1,21 +1,49 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from enum import Enum
+from typing import Any, Dict, List, Optional
 
 from chia.types.blockchain_format.program import Program
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.ints import uint64
 from chia.util.streamable import Streamable, streamable
 
 
+class LineageProofField(Enum):
+    PARENT_NAME = 1
+    INNER_PUZZLE_HASH = 2
+    AMOUNT = 3
+
+
 @streamable
 @dataclass(frozen=True)
 class LineageProof(Streamable):
     parent_name: Optional[bytes32] = None
     inner_puzzle_hash: Optional[bytes32] = None
     amount: Optional[uint64] = None
 
+    @classmethod
+    def from_program(cls, program: Program, fields: List[LineageProofField]) -> LineageProof:
+        lineage_proof_info: Dict[str, Any] = {}
+        field_iter = iter(fields)
+        program_iter = program.as_iter()
+        for program_value in program_iter:
+            field = next(field_iter)
+            if field == LineageProofField.PARENT_NAME:
+                lineage_proof_info["parent_name"] = bytes32(program_value.as_atom())
+            elif field == LineageProofField.INNER_PUZZLE_HASH:
+                lineage_proof_info["inner_puzzle_hash"] = bytes32(program_value.as_atom())
+            elif field == LineageProofField.AMOUNT:
+                lineage_proof_info["amount"] = uint64(program_value.as_int())
+        try:
+            next(field_iter)
+            raise ValueError("Mismatch between program data and fields information")
+        except StopIteration:
+            pass
+
+        return LineageProof(**lineage_proof_info)
+
     def to_program(self) -> Program:
         final_list: List[Any] = []
         if self.parent_name is not None:

chia/wallet/vc_wallet/cr_cat_drivers.py

@@ -16,7 +16,7 @@
 from chia.util.ints import uint16, uint64
 from chia.util.streamable import Streamable, streamable
 from chia.wallet.cat_wallet.cat_utils import CAT_MOD, construct_cat_puzzle
-from chia.wallet.lineage_proof import LineageProof
+from chia.wallet.lineage_proof import LineageProof, LineageProofField
 from chia.wallet.payment import Payment
 from chia.wallet.puzzles.load_clvm import load_clvm_maybe_recompile
 from chia.wallet.puzzles.singleton_top_layer_v1_1 import SINGLETON_LAUNCHER_HASH, SINGLETON_MOD_HASH
@@ -314,10 +314,14 @@ def get_current_from_coin_spend(cls: Type[_T_CRCAT], spend: CoinSpend) -> CRCAT:
         uncurried_puzzle: UncurriedPuzzle = uncurry_puzzle(spend.puzzle_reveal.to_program())
         first_uncurried_cr_layer: UncurriedPuzzle = uncurry_puzzle(uncurried_puzzle.args.at("rrf"))
         second_uncurried_cr_layer: UncurriedPuzzle = uncurry_puzzle(first_uncurried_cr_layer.mod)
+        lineage_proof = LineageProof.from_program(
+            spend.solution.to_program().at("rf"),
+            [LineageProofField.PARENT_NAME, LineageProofField.INNER_PUZZLE_HASH, LineageProofField.AMOUNT],
+        )
         return CRCAT(
             spend.coin,
             bytes32(uncurried_puzzle.args.at("rf").as_atom()),
-            spend.solution.to_program().at("rf"),
+            lineage_proof,
             [bytes32(ap.as_atom()) for ap in second_uncurried_cr_layer.args.at("rf").as_iter()],
             second_uncurried_cr_layer.args.at("rrf"),
             first_uncurried_cr_layer.args.at("rf").get_tree_hash(),

tests/wallet/test_util.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any, Dict, List, Tuple
+
 import pytest
 
 from chia.consensus.default_constants import DEFAULT_CONSTANTS
@@ -9,6 +11,7 @@
 from chia.types.coin_spend import make_spend
 from chia.util.errors import ValidationError
 from chia.util.ints import uint64
+from chia.wallet.lineage_proof import LineageProof, LineageProofField
 from chia.wallet.util.compute_hints import HintedCoin, compute_spend_hints_and_additions
 from chia.wallet.util.merkle_utils import list_to_binary_tree
 from chia.wallet.util.tx_config import (
@@ -103,3 +106,76 @@ def test_list_to_binary_tree() -> None:
     assert list_to_binary_tree([1, 2, 3, 4, 5]) == (((1, 2), 3), (4, 5))
     with pytest.raises(ValueError):
         list_to_binary_tree([])
+
+
+@pytest.mark.parametrize(
+    "serializations",
+    [
+        (tuple(), Program.to(None), []),
+        ((bytes32([0] * 32),), Program.to([bytes32([0] * 32)]), [LineageProofField.PARENT_NAME]),
+        (
+            (bytes32([0] * 32), bytes32([0] * 32)),
+            Program.to([bytes32([0] * 32), bytes32([0] * 32)]),
+            [LineageProofField.PARENT_NAME, LineageProofField.INNER_PUZZLE_HASH],
+        ),
+        (
+            (bytes32([0] * 32), bytes32([0] * 32), uint64(0)),
+            Program.to([bytes32([0] * 32), bytes32([0] * 32), uint64(0)]),
+            [LineageProofField.PARENT_NAME, LineageProofField.INNER_PUZZLE_HASH, LineageProofField.AMOUNT],
+        ),
+    ],
+)
+def test_lineage_proof_varargs(serializations: Tuple[Tuple[Any, ...], Program, List[LineageProofField]]) -> None:
+    var_args, expected_program, lp_fields = serializations
+    assert LineageProof(*var_args).to_program() == expected_program
+    assert LineageProof(*var_args) == LineageProof.from_program(expected_program, lp_fields)
+
+
+@pytest.mark.parametrize(
+    "serializations",
+    [
+        ({}, Program.to(None), []),
+        ({"parent_name": bytes32([0] * 32)}, Program.to([bytes32([0] * 32)]), [LineageProofField.PARENT_NAME]),
+        (
+            {"parent_name": bytes32([0] * 32), "inner_puzzle_hash": bytes32([0] * 32)},
+            Program.to([bytes32([0] * 32), bytes32([0] * 32)]),
+            [LineageProofField.PARENT_NAME, LineageProofField.INNER_PUZZLE_HASH],
+        ),
+        (
+            {"parent_name": bytes32([0] * 32), "inner_puzzle_hash": bytes32([0] * 32), "amount": uint64(0)},
+            Program.to([bytes32([0] * 32), bytes32([0] * 32), uint64(0)]),
+            [LineageProofField.PARENT_NAME, LineageProofField.INNER_PUZZLE_HASH, LineageProofField.AMOUNT],
+        ),
+        (
+            {"parent_name": bytes32([0] * 32), "amount": uint64(0)},
+            Program.to([bytes32([0] * 32), uint64(0)]),
+            [LineageProofField.PARENT_NAME, LineageProofField.AMOUNT],
+        ),
+        (
+            {"inner_puzzle_hash": bytes32([0] * 32), "amount": uint64(0)},
+            Program.to([bytes32([0] * 32), uint64(0)]),
+            [LineageProofField.INNER_PUZZLE_HASH, LineageProofField.AMOUNT],
+        ),
+        ({"amount": uint64(0)}, Program.to([uint64(0)]), [LineageProofField.AMOUNT]),
+        (
+            {"inner_puzzle_hash": bytes32([0] * 32)},
+            Program.to([bytes32([0] * 32)]),
+            [LineageProofField.INNER_PUZZLE_HASH],
+        ),
+    ],
+)
+def test_lineage_proof_kwargs(serializations: Tuple[Dict[str, Any], Program, List[LineageProofField]]) -> None:
+    kwargs, expected_program, lp_fields = serializations
+    assert LineageProof(**kwargs).to_program() == expected_program
+    assert LineageProof(**kwargs) == LineageProof.from_program(expected_program, lp_fields)
+
+
+def test_lineage_proof_errors() -> None:
+    with pytest.raises(ValueError, match="Mismatch"):
+        LineageProof.from_program(Program.to([]), [LineageProofField.PARENT_NAME])
+    with pytest.raises(StopIteration):
+        LineageProof.from_program(Program.to([bytes32([0] * 32)]), [])
+    with pytest.raises(ValueError):
+        LineageProof.from_program(Program.to([bytes32([1] * 32)]), [LineageProofField.AMOUNT])
+    with pytest.raises(ValueError):
+        LineageProof.from_program(Program.to([uint64(0)]), [LineageProofField.PARENT_NAME])