use chia_bls::Signature;
use chia_protocol::{Bytes32, Coin, CoinSpend};
use chia_puzzle_types::singleton::{LauncherSolution, SingletonArgs};
use chia_puzzle_types::{
LineageProof, Proof,
singleton::{SingletonSolution, SingletonStruct},
};
use chia_sdk_types::puzzles::{
RawActionLayerSolution, ReserveFinalizerSolution, RewardDistributorCommitmentSlotValue,
RewardDistributorEntrySlotValue, RewardDistributorRewardSlotValue, RewardDistributorSlotNonce,
SlotInfo,
};
use chia_sdk_types::{Condition, Conditions};
use clvm_traits::{FromClvm, clvm_list, match_tuple};
use clvm_utils::{ToTreeHash, tree_hash};
use clvmr::NodePtr;
use crate::{
ActionLayer, ActionLayerSolution, ActionSingleton, Cat, CatSpend, DriverError, Layer, Puzzle,
RewardDistributorAddEntryAction, RewardDistributorAddIncentivesAction,
RewardDistributorCommitIncentivesAction, RewardDistributorInitiatePayoutAction,
RewardDistributorNewEpochAction, RewardDistributorRemoveEntryAction,
RewardDistributorStakeAction, RewardDistributorSyncAction, RewardDistributorUnstakeAction,
RewardDistributorWithdrawIncentivesAction, SingletonAction, SingletonLayer, Slot, Spend,
SpendContext,
};
use super::{Reserve, RewardDistributorConstants, RewardDistributorInfo, RewardDistributorState};
#[derive(Debug, Clone)]
pub struct RewardDistributorPendingSpendInfo {
pub actions: Vec<Spend>,
pub spent_reward_slots: Vec<RewardDistributorRewardSlotValue>,
pub spent_commitment_slots: Vec<RewardDistributorCommitmentSlotValue>,
pub spent_entry_slots: Vec<RewardDistributorEntrySlotValue>,
pub created_reward_slots: Vec<RewardDistributorRewardSlotValue>,
pub created_commitment_slots: Vec<RewardDistributorCommitmentSlotValue>,
pub created_entry_slots: Vec<RewardDistributorEntrySlotValue>,
pub latest_state: (NodePtr, RewardDistributorState),
pub signature: Signature,
pub other_cats: Vec<CatSpend>,
}
impl RewardDistributorPendingSpendInfo {
pub fn new(latest_state: RewardDistributorState) -> Self {
Self {
actions: vec![],
created_reward_slots: vec![],
created_commitment_slots: vec![],
created_entry_slots: vec![],
spent_reward_slots: vec![],
spent_commitment_slots: vec![],
spent_entry_slots: vec![],
latest_state: (NodePtr::NIL, latest_state),
signature: Signature::default(),
other_cats: vec![],
}
}
pub fn add_delta(&mut self, delta: RewardDistributorPendingSpendInfo) {
self.actions.extend(delta.actions);
self.spent_reward_slots.extend(delta.spent_reward_slots);
self.spent_commitment_slots
.extend(delta.spent_commitment_slots);
self.spent_entry_slots.extend(delta.spent_entry_slots);
self.created_reward_slots.extend(delta.created_reward_slots);
self.created_commitment_slots
.extend(delta.created_commitment_slots);
self.created_entry_slots.extend(delta.created_entry_slots);
self.latest_state = delta.latest_state;
}
}
#[derive(Debug, Clone)]
#[must_use]
pub struct RewardDistributor {
pub coin: Coin,
pub proof: Proof,
pub info: RewardDistributorInfo,
pub reserve: Reserve,
pub pending_spend: RewardDistributorPendingSpendInfo,
}
impl RewardDistributor {
pub fn new(coin: Coin, proof: Proof, info: RewardDistributorInfo, reserve: Reserve) -> Self {
Self {
coin,
proof,
info,
reserve,
pending_spend: RewardDistributorPendingSpendInfo::new(info.state),
}
}
}
impl RewardDistributor {
#[allow(clippy::type_complexity)]
pub fn pending_info_delta_from_spend(
ctx: &mut SpendContext,
action_spend: Spend,
current_state_and_ephemeral: (NodePtr, RewardDistributorState),
constants: RewardDistributorConstants,
) -> Result<RewardDistributorPendingSpendInfo, DriverError> {
let mut spent_reward_slots = vec![];
let mut spent_commitment_slots = vec![];
let mut spent_entry_slots = vec![];
let mut created_reward_slots = vec![];
let mut created_commitment_slots = vec![];
let mut created_entry_slots = vec![];
let new_epoch_action = RewardDistributorNewEpochAction::from_constants(&constants);
let new_epoch_hash = new_epoch_action.tree_hash();
let commit_incentives_action =
RewardDistributorCommitIncentivesAction::from_constants(&constants);
let commit_incentives_hash = commit_incentives_action.tree_hash();
let add_entry_action = RewardDistributorAddEntryAction::from_constants(&constants);
let add_entry_hash = add_entry_action.tree_hash();
let remove_entry_action = RewardDistributorRemoveEntryAction::from_constants(&constants);
let remove_entry_hash = remove_entry_action.tree_hash();
let stake_action = RewardDistributorStakeAction::from_constants(&constants);
let stake_hash = stake_action.tree_hash();
let unstake_action = RewardDistributorUnstakeAction::from_constants(&constants);
let unstake_hash = unstake_action.tree_hash();
let withdraw_incentives_action =
RewardDistributorWithdrawIncentivesAction::from_constants(&constants);
let withdraw_incentives_hash = withdraw_incentives_action.tree_hash();
let initiate_payout_action =
RewardDistributorInitiatePayoutAction::from_constants(&constants);
let initiate_payout_hash = initiate_payout_action.tree_hash();
let add_incentives_action =
RewardDistributorAddIncentivesAction::from_constants(&constants);
let add_incentives_hash = add_incentives_action.tree_hash();
let sync_action = RewardDistributorSyncAction::from_constants(&constants);
let sync_hash = sync_action.tree_hash();
let actual_solution = ctx.alloc(&clvm_list!(
current_state_and_ephemeral,
action_spend.solution
))?;
let output = ctx.run(action_spend.puzzle, actual_solution)?;
let (new_state_and_ephemeral, _) =
ctx.extract::<match_tuple!((NodePtr, RewardDistributorState), NodePtr)>(output)?;
let raw_action_hash = ctx.tree_hash(action_spend.puzzle);
if raw_action_hash == new_epoch_hash {
created_reward_slots.push(RewardDistributorNewEpochAction::created_slot_value(
ctx,
action_spend.solution,
)?);
spent_reward_slots.push(RewardDistributorNewEpochAction::spent_slot_value(
ctx,
action_spend.solution,
)?);
} else if raw_action_hash == commit_incentives_hash {
let (comm, rews) = RewardDistributorCommitIncentivesAction::created_slot_values(
ctx,
constants.epoch_seconds,
action_spend.solution,
)?;
created_commitment_slots.push(comm);
created_reward_slots.extend(rews);
spent_reward_slots.push(RewardDistributorCommitIncentivesAction::spent_slot_value(
ctx,
action_spend.solution,
)?);
} else if raw_action_hash == add_entry_hash {
created_entry_slots.push(RewardDistributorAddEntryAction::created_slot_value(
ctx,
¤t_state_and_ephemeral.1,
action_spend.solution,
)?);
} else if raw_action_hash == stake_hash {
created_entry_slots.push(RewardDistributorStakeAction::created_slot_value(
ctx,
¤t_state_and_ephemeral.1,
action_spend.solution,
)?);
} else if raw_action_hash == remove_entry_hash {
spent_entry_slots.push(RewardDistributorRemoveEntryAction::spent_slot_value(
ctx,
action_spend.solution,
)?);
} else if raw_action_hash == unstake_hash {
spent_entry_slots.push(RewardDistributorUnstakeAction::spent_slot_value(
ctx,
action_spend.solution,
)?);
} else if raw_action_hash == withdraw_incentives_hash {
let (rew, cmt) = RewardDistributorWithdrawIncentivesAction::spent_slot_values(
ctx,
action_spend.solution,
)?;
spent_reward_slots.push(rew);
spent_commitment_slots.push(cmt);
created_reward_slots.push(
RewardDistributorWithdrawIncentivesAction::created_slot_value(
ctx,
constants.withdrawal_share_bps,
action_spend.solution,
)?,
);
} else if raw_action_hash == initiate_payout_hash {
created_entry_slots.push(RewardDistributorInitiatePayoutAction::created_slot_value(
ctx,
¤t_state_and_ephemeral.1,
action_spend.solution,
)?);
spent_entry_slots.push(RewardDistributorInitiatePayoutAction::spent_slot_value(
ctx,
action_spend.solution,
)?);
} else if raw_action_hash != add_incentives_hash && raw_action_hash != sync_hash {
return Err(DriverError::InvalidMerkleProof);
}
Ok(RewardDistributorPendingSpendInfo {
actions: vec![action_spend],
spent_reward_slots,
spent_commitment_slots,
spent_entry_slots,
created_reward_slots,
created_commitment_slots,
created_entry_slots,
latest_state: new_state_and_ephemeral,
signature: Signature::default(),
other_cats: vec![],
})
}
pub fn pending_info_from_spend(
ctx: &mut SpendContext,
inner_solution: NodePtr,
initial_state: RewardDistributorState,
constants: RewardDistributorConstants,
) -> Result<RewardDistributorPendingSpendInfo, DriverError> {
let mut pending_spend_info = RewardDistributorPendingSpendInfo::new(initial_state);
let inner_solution =
ActionLayer::<RewardDistributorState, NodePtr>::parse_solution(ctx, inner_solution)?;
for raw_action in &inner_solution.action_spends {
let delta = Self::pending_info_delta_from_spend(
ctx,
*raw_action,
pending_spend_info.latest_state,
constants,
)?;
pending_spend_info.add_delta(delta);
}
Ok(pending_spend_info)
}
pub fn from_spend(
ctx: &mut SpendContext,
spend: &CoinSpend,
reserve_lineage_proof: Option<LineageProof>,
constants: RewardDistributorConstants,
) -> Result<Option<Self>, DriverError> {
let coin = spend.coin;
let puzzle_ptr = ctx.alloc(&spend.puzzle_reveal)?;
let puzzle = Puzzle::parse(ctx, puzzle_ptr);
let solution_ptr = ctx.alloc(&spend.solution)?;
let Some(info) = RewardDistributorInfo::parse(ctx, puzzle, constants)? else {
return Ok(None);
};
let solution = ctx.extract::<SingletonSolution<NodePtr>>(solution_ptr)?;
let proof = solution.lineage_proof;
let pending_spend =
Self::pending_info_from_spend(ctx, solution.inner_solution, info.state, constants)?;
let inner_solution =
RawActionLayerSolution::<NodePtr, NodePtr, ReserveFinalizerSolution>::from_clvm(
ctx,
solution.inner_solution,
)?;
let reserve = Reserve::new(
inner_solution.finalizer_solution.reserve_parent_id,
reserve_lineage_proof.unwrap_or(LineageProof {
parent_parent_coin_info: Bytes32::default(),
parent_inner_puzzle_hash: Bytes32::default(),
parent_amount: 0,
}), constants.reserve_asset_id,
SingletonStruct::new(info.constants.launcher_id)
.tree_hash()
.into(),
0,
info.state.total_reserves,
);
Ok(Some(RewardDistributor {
coin,
proof,
info,
reserve,
pending_spend,
}))
}
pub fn child_lineage_proof(&self) -> LineageProof {
LineageProof {
parent_parent_coin_info: self.coin.parent_coin_info,
parent_inner_puzzle_hash: self.info.inner_puzzle_hash().into(),
parent_amount: self.coin.amount,
}
}
pub fn from_parent_spend(
ctx: &mut SpendContext,
parent_spend: &CoinSpend,
constants: RewardDistributorConstants,
) -> Result<Option<Self>, DriverError>
where
Self: Sized,
{
let Some(parent_registry) = Self::from_spend(ctx, parent_spend, None, constants)? else {
return Ok(None);
};
let new_info = parent_registry
.info
.with_state(parent_registry.pending_spend.latest_state.1);
Ok(Some(RewardDistributor {
coin: Coin::new(
parent_registry.coin.coin_id(),
new_info.puzzle_hash().into(),
1,
),
proof: Proof::Lineage(parent_registry.child_lineage_proof()),
info: new_info,
reserve: parent_registry.reserve.child(new_info.state.total_reserves),
pending_spend: RewardDistributorPendingSpendInfo::new(new_info.state),
}))
}
pub fn child(&self, child_state: RewardDistributorState) -> Self {
let new_info = self.info.with_state(child_state);
let new_coin = Coin::new(self.coin.coin_id(), new_info.puzzle_hash().into(), 1);
let new_reserve = self.reserve.child(child_state.total_reserves);
RewardDistributor {
coin: new_coin,
proof: Proof::Lineage(self.child_lineage_proof()),
info: new_info,
reserve: new_reserve,
pending_spend: RewardDistributorPendingSpendInfo::new(new_info.state),
}
}
#[allow(clippy::type_complexity)]
pub fn from_launcher_solution(
ctx: &mut SpendContext,
launcher_coin: Coin,
launcher_solution: NodePtr,
) -> Result<Option<(RewardDistributorConstants, RewardDistributorState, Coin)>, DriverError>
where
Self: Sized,
{
let Ok(launcher_solution) =
ctx.extract::<LauncherSolution<(u64, RewardDistributorConstants)>>(launcher_solution)
else {
return Ok(None);
};
let launcher_id = launcher_coin.coin_id();
let (first_epoch_start, constants) = launcher_solution.key_value_list;
if constants != constants.with_launcher_id(launcher_id) {
return Err(DriverError::Custom(
"Distributor constants invalid".to_string(),
));
}
let distributor_eve_coin =
Coin::new(launcher_id, launcher_solution.singleton_puzzle_hash, 1);
let initial_state = RewardDistributorState::initial(first_epoch_start);
Ok(Some((constants, initial_state, distributor_eve_coin)))
}
#[allow(clippy::type_complexity)]
pub fn from_eve_coin_spend(
ctx: &mut SpendContext,
constants: RewardDistributorConstants,
initial_state: RewardDistributorState,
eve_coin_spend: &CoinSpend,
reserve_parent_id: Bytes32,
reserve_lineage_proof: LineageProof,
) -> Result<Option<(RewardDistributor, Slot<RewardDistributorRewardSlotValue>)>, DriverError>
where
Self: Sized,
{
let eve_coin_puzzle_ptr = ctx.alloc(&eve_coin_spend.puzzle_reveal)?;
let eve_coin_puzzle = Puzzle::parse(ctx, eve_coin_puzzle_ptr);
let Some(eve_coin_puzzle) = SingletonLayer::<NodePtr>::parse_puzzle(ctx, eve_coin_puzzle)?
else {
return Err(DriverError::Custom("Eve coin not a singleton".to_string()));
};
let eve_coin_inner_puzzle_hash = tree_hash(ctx, eve_coin_puzzle.inner_puzzle);
let eve_coin_solution_ptr = ctx.alloc(&eve_coin_spend.solution)?;
let eve_coin_output = ctx.run(eve_coin_puzzle_ptr, eve_coin_solution_ptr)?;
let eve_coin_output = ctx.extract::<Conditions<NodePtr>>(eve_coin_output)?;
let Some(Condition::CreateCoin(odd_create_coin)) = eve_coin_output.into_iter().find(|c| {
if let Condition::CreateCoin(create_coin) = c {
create_coin.amount % 2 == 1
} else {
false
}
}) else {
return Err(DriverError::Custom(
"Eve coin did not create a coin".to_string(),
));
};
let new_coin = Coin::new(
eve_coin_spend.coin.coin_id(),
odd_create_coin.puzzle_hash,
odd_create_coin.amount,
);
let lineage_proof = LineageProof {
parent_parent_coin_info: eve_coin_spend.coin.parent_coin_info,
parent_inner_puzzle_hash: eve_coin_inner_puzzle_hash.into(),
parent_amount: eve_coin_spend.coin.amount,
};
let reserve = Reserve::new(
reserve_parent_id,
reserve_lineage_proof,
constants.reserve_asset_id,
SingletonStruct::new(constants.launcher_id)
.tree_hash()
.into(),
0,
0,
);
let new_distributor = RewardDistributor::new(
new_coin,
Proof::Lineage(lineage_proof),
RewardDistributorInfo::new(initial_state, constants),
reserve,
);
if SingletonArgs::curry_tree_hash(
constants.launcher_id,
new_distributor.info.inner_puzzle_hash(),
) != new_distributor.coin.puzzle_hash.into()
{
return Err(DriverError::Custom(
"Distributor singleton puzzle hash mismatch".to_string(),
));
}
let slot_value = RewardDistributorRewardSlotValue {
epoch_start: initial_state.round_time_info.epoch_end,
next_epoch_initialized: false,
rewards: 0,
};
let slot = Slot::new(
lineage_proof,
SlotInfo::from_value(
constants.launcher_id,
RewardDistributorSlotNonce::REWARD.to_u64(),
slot_value,
),
);
Ok(Some((new_distributor, slot)))
}
pub fn set_pending_signature(&mut self, signature: Signature) {
self.pending_spend.signature = signature;
}
pub fn set_pending_other_cats(&mut self, other_cats: Vec<CatSpend>) {
self.pending_spend.other_cats = other_cats;
}
}
impl ActionSingleton for RewardDistributor {
type State = RewardDistributorState;
type Constants = RewardDistributorConstants;
}
impl RewardDistributor {
pub fn finish_spend(
self,
ctx: &mut SpendContext,
other_cat_spends: Vec<CatSpend>,
) -> Result<(Self, Signature), DriverError> {
let layers = self.info.into_layers(ctx)?;
let puzzle = layers.construct_puzzle(ctx)?;
let action_puzzle_hashes = self
.pending_spend
.actions
.iter()
.map(|a| ctx.tree_hash(a.puzzle).into())
.collect::<Vec<Bytes32>>();
let finalizer_solution = ctx.alloc(&ReserveFinalizerSolution {
reserve_parent_id: self.reserve.coin.parent_coin_info,
})?;
let child = self.child(self.pending_spend.latest_state.1);
let solution = layers.construct_solution(
ctx,
SingletonSolution {
lineage_proof: self.proof,
amount: self.coin.amount,
inner_solution: ActionLayerSolution {
proofs: layers
.inner_puzzle
.get_proofs(
&RewardDistributorInfo::action_puzzle_hashes(&self.info.constants),
&action_puzzle_hashes,
)
.ok_or(DriverError::Custom(
"Couldn't build proofs for one or more actions".to_string(),
))?,
action_spends: self.pending_spend.actions,
finalizer_solution,
},
},
)?;
let my_spend = Spend::new(puzzle, solution);
ctx.spend(self.coin, my_spend)?;
let cat_spend = self.reserve.cat_spend_for_reserve_finalizer_controller(
ctx,
self.info.state,
self.info.inner_puzzle_hash().into(),
solution,
)?;
let mut cat_spends = other_cat_spends;
cat_spends.push(cat_spend);
cat_spends.extend(self.pending_spend.other_cats);
Cat::spend_all(ctx, &cat_spends)?;
Ok((child, self.pending_spend.signature))
}
pub fn new_action<A>(&self) -> A
where
A: SingletonAction<Self>,
{
A::from_constants(&self.info.constants)
}
pub fn created_slot_value_to_slot<SlotValue>(
&self,
slot_value: SlotValue,
nonce: RewardDistributorSlotNonce,
) -> Slot<SlotValue>
where
SlotValue: Copy + ToTreeHash,
{
Slot::new(
LineageProof {
parent_parent_coin_info: self.coin.parent_coin_info,
parent_inner_puzzle_hash: self.info.inner_puzzle_hash().into(),
parent_amount: self.coin.amount,
},
SlotInfo::from_value(self.info.constants.launcher_id, nonce.to_u64(), slot_value),
)
}
pub fn insert_action_spend(
&mut self,
ctx: &mut SpendContext,
action_spend: Spend,
) -> Result<(), DriverError> {
let delta = Self::pending_info_delta_from_spend(
ctx,
action_spend,
self.pending_spend.latest_state,
self.info.constants,
)?;
self.pending_spend.add_delta(delta);
Ok(())
}
pub fn actual_reward_slot_value(
&self,
slot: Slot<RewardDistributorRewardSlotValue>,
) -> Slot<RewardDistributorRewardSlotValue> {
let mut slot = slot;
for slot_value in &self.pending_spend.created_reward_slots {
if slot_value.epoch_start == slot.info.value.epoch_start {
slot = self
.created_slot_value_to_slot(*slot_value, RewardDistributorSlotNonce::REWARD);
}
}
slot
}
pub fn actual_entry_slot_value(
&self,
slot: Slot<RewardDistributorEntrySlotValue>,
) -> Slot<RewardDistributorEntrySlotValue> {
let mut slot = slot;
for slot_value in &self.pending_spend.created_entry_slots {
if slot_value.payout_puzzle_hash == slot.info.value.payout_puzzle_hash {
slot =
self.created_slot_value_to_slot(*slot_value, RewardDistributorSlotNonce::ENTRY);
}
}
slot
}
pub fn actual_commitment_slot_value(
&self,
slot: Slot<RewardDistributorCommitmentSlotValue>,
) -> Slot<RewardDistributorCommitmentSlotValue> {
let mut slot = slot;
for slot_value in &self.pending_spend.created_commitment_slots {
if slot_value.epoch_start == slot.info.value.epoch_start {
slot = self.created_slot_value_to_slot(
*slot_value,
RewardDistributorSlotNonce::COMMITMENT,
);
}
}
slot
}
}