improve v2-plot support in plot-sync

chia/_tests/plot_sync/test_plot_sync.py

@@ -10,7 +10,7 @@
 from typing import Any, Callable, Optional
 
 import pytest
-from chia_rs import G1Element
+from chia_rs import G1Element, PlotParam
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint8, uint32, uint64
 
@@ -109,6 +109,10 @@ def remove_duplicates(self, list_paths: list[Path]) -> None:
         self.duplicates_delta.removals += [str(x) for x in list_paths]
 
 
+def equal_param(a: PlotParam, b: PlotParam) -> bool:
+    return a.size_v1 == b.size_v1 and a.strength_v2 == b.strength_v2
+
+
 @dataclass
 class Environment:
     root_path: Path
@@ -195,8 +199,7 @@ async def plot_sync_callback(self, peer_id: bytes32, delta: Optional[Delta]) ->
             plot = harvester.plot_manager.plots.get(Path(path), None)
             assert plot is not None
             assert plot.prover.get_filename() == delta.valid.additions[path].filename
-            # TODO: todo_v2_plots support v2 plots
-            assert plot.prover.get_param().size_v1 == delta.valid.additions[path].size
+            assert equal_param(plot.prover.get_param(), delta.valid.additions[path].param())
             assert plot.prover.get_id() == delta.valid.additions[path].plot_id
             assert plot.prover.get_compression_level() == delta.valid.additions[path].compression_level
             assert plot.pool_public_key == delta.valid.additions[path].pool_public_key
@@ -257,8 +260,7 @@ async def run_sync_test(self) -> None:
             for path, plot_info in plot_manager.plots.items():
                 assert str(path) in receiver.plots()
                 assert plot_info.prover.get_filename() == receiver.plots()[str(path)].filename
-                # TODO: todo_v2_plots support v2 plots
-                assert plot_info.prover.get_param().size_v1 == receiver.plots()[str(path)].size
+                assert equal_param(plot_info.prover.get_param(), receiver.plots()[str(path)].param())
                 assert plot_info.prover.get_id() == receiver.plots()[str(path)].plot_id
                 assert plot_info.prover.get_compression_level() == receiver.plots()[str(path)].compression_level
                 assert plot_info.pool_public_key == receiver.plots()[str(path)].pool_public_key

chia/_tests/plot_sync/test_receiver.py

@@ -7,7 +7,7 @@
 from typing import Any, Callable, Union
 
 import pytest
-from chia_rs import G1Element, PlotParam
+from chia_rs import G1Element
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint8, uint32, uint64
 
@@ -189,7 +189,7 @@ def plot_sync_setup(seeded_random: random.Random) -> tuple[Receiver, list[SyncSt
     # TODO: todo_v2_plots support v2 plots
     receiver._total_effective_plot_size = int(
         sum(
-            UI_ACTUAL_SPACE_CONSTANT_FACTOR * int(_expected_plot_size(PlotParam.make_v1(plot.size), DEFAULT_CONSTANTS))
+            UI_ACTUAL_SPACE_CONSTANT_FACTOR * int(_expected_plot_size(plot.param(), DEFAULT_CONSTANTS))
             for plot in receiver.plots().values()
         )
     )
@@ -272,9 +272,8 @@ async def test_to_dict(counts_only: bool, seeded_random: random.Random) -> None:
     assert get_list_or_len(plot_sync_dict_1["duplicates"], not counts_only) == 0
     assert plot_sync_dict_1["total_plot_size"] == sum(plot.file_size for plot in receiver.plots().values())
     assert plot_sync_dict_1["total_effective_plot_size"] == int(
-        # TODO: todo_v2_plots support v2 plots
         sum(
-            UI_ACTUAL_SPACE_CONSTANT_FACTOR * int(_expected_plot_size(PlotParam.make_v1(plot.size), DEFAULT_CONSTANTS))
+            UI_ACTUAL_SPACE_CONSTANT_FACTOR * int(_expected_plot_size(plot.param(), DEFAULT_CONSTANTS))
             for plot in receiver.plots().values()
         )
     )
@@ -322,10 +321,9 @@ async def test_to_dict(counts_only: bool, seeded_random: random.Random) -> None:
     assert get_list_or_len(sync_steps[State.duplicates].args[0], counts_only) == plot_sync_dict_3["duplicates"]
 
     assert plot_sync_dict_3["total_plot_size"] == sum(plot.file_size for plot in receiver.plots().values())
-    # TODO: todo_v2_plots support v2 plots
     assert plot_sync_dict_3["total_effective_plot_size"] == int(
         sum(
-            UI_ACTUAL_SPACE_CONSTANT_FACTOR * int(_expected_plot_size(PlotParam.make_v1(plot.size), DEFAULT_CONSTANTS))
+            UI_ACTUAL_SPACE_CONSTANT_FACTOR * int(_expected_plot_size(plot.param(), DEFAULT_CONSTANTS))
             for plot in receiver.plots().values()
         )
     )

chia/harvester/harvester.py

@@ -11,7 +11,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
 
 from chia_rs import ConsensusConstants
-from chia_rs.sized_ints import uint32
+from chia_rs.sized_ints import uint8, uint32
 from typing_extensions import Literal
 
 from chia.plot_sync.sender import Sender
@@ -195,9 +195,10 @@ def get_plots(self) -> tuple[list[dict[str, Any]], list[str], list[str]]:
                 prover = plot_info.prover
                 param = prover.get_param()
                 if param.size_v1 is not None:
-                    k = param.size_v1
+                    k = uint8(param.size_v1)
                 else:
-                    k = self.constants.PLOT_SIZE_V2
+                    assert param.strength_v2 is not None
+                    k = uint8(0x80 | param.strength_v2)
                 response_plots.append(
                     {
                         "filename": str(path),

chia/plot_sync/receiver.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Callable, Optional, Union
 
-from chia_rs import ConsensusConstants, PlotParam
+from chia_rs import ConsensusConstants
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import int16, uint32, uint64
 from typing_extensions import Protocol
@@ -354,10 +354,8 @@ async def _sync_done(self, data: PlotSyncDone) -> None:
         self._total_plot_size = sum(plot.file_size for plot in self._plots.values())
 
         self._total_effective_plot_size = int(
-            # TODO: todo_v2_plots support v2 plots
             sum(
-                UI_ACTUAL_SPACE_CONSTANT_FACTOR
-                * int(_expected_plot_size(PlotParam.make_v1(plot.size), self._constants))
+                UI_ACTUAL_SPACE_CONSTANT_FACTOR * _expected_plot_size(plot.param(), self._constants)
                 for plot in self._plots.values()
             )
         )

chia/plot_sync/sender.py

@@ -9,7 +9,7 @@
 from pathlib import Path
 from typing import Any, Generic, Optional, TypeVar
 
-from chia_rs.sized_ints import int16, uint32, uint64
+from chia_rs.sized_ints import int16, uint8, uint32, uint64
 from typing_extensions import Protocol
 
 from chia.plot_sync.exceptions import AlreadyStartedError, InvalidConnectionTypeError
@@ -37,13 +37,18 @@
 def _convert_plot_info_list(plot_infos: list[PlotInfo]) -> list[Plot]:
     converted: list[Plot] = []
     for plot_info in plot_infos:
-        # TODO: todo_v2_plots support v2 plots
-        k = plot_info.prover.get_param().size_v1
-        assert k is not None
+        param = plot_info.prover.get_param()
+        k: uint8
+        if param.size_v1 is not None:
+            k = param.size_v1
+        else:
+            assert param.strength_v2 is not None
+            k = uint8(0x80 | param.strength_v2)
+
         converted.append(
             Plot(
                 filename=plot_info.prover.get_filename(),
-                size=k,
+                size=uint8(k),
                 plot_id=plot_info.prover.get_id(),
                 pool_public_key=plot_info.pool_public_key,
                 pool_contract_puzzle_hash=plot_info.pool_contract_puzzle_hash,

chia/protocols/harvester_protocol.py

@@ -4,7 +4,7 @@
 from enum import IntEnum
 from typing import Optional
 
-from chia_rs import G1Element, G2Element, PartialProof, ProofOfSpace, RewardChainBlockUnfinished
+from chia_rs import G1Element, G2Element, PartialProof, PlotParam, ProofOfSpace, RewardChainBlockUnfinished
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import int16, uint8, uint32, uint64
 
@@ -142,6 +142,11 @@ class RespondSignatures(Streamable):
 @dataclass(frozen=True)
 class Plot(Streamable):
     filename: str
+    # for backwards compatibility with previous harvester (e.g. DrPlotter)
+    # this field is either k-size (for v1 plots) or strength (for v2 plots).
+    # the most significant bit is set for v2 plots
+    # TODO: after the phase-out, v1 harvester won't be relevant anymore and we
+    # can clean this up with a new protocol version
     size: uint8
     plot_id: bytes32
     pool_public_key: Optional[G1Element]
@@ -151,6 +156,12 @@ class Plot(Streamable):
     time_modified: uint64
     compression_level: Optional[uint8]
 
+    def param(self) -> PlotParam:
+        if (self.size & 0x80) != 0:
+            return PlotParam.make_v2(self.size & 0x7F)
+        else:
+            return PlotParam.make_v1(self.size)
+
 
 @streamable
 @dataclass(frozen=True)