use chia_bls::Signature;
use chia_protocol::{Bytes32, Coin, CoinSpend};
use chia_puzzle_types::{LineageProof, Proof, singleton::SingletonSolution};
use chia_sdk_types::puzzles::{CatalogSlotValue, SlotInfo};
use clvm_traits::{clvm_list, match_tuple};
use clvm_utils::ToTreeHash;
use clvmr::NodePtr;
use crate::{
ActionLayer, ActionLayerSolution, ActionSingleton, CatalogRefundAction, CatalogRegisterAction,
DelegatedStateAction, DriverError, Layer, Puzzle, SingletonAction, Spend, SpendContext,
};
use super::{CatalogRegistryConstants, CatalogRegistryInfo, CatalogRegistryState, Slot};
#[derive(Debug, Clone)]
pub struct CatalogPendingSpendInfo {
pub actions: Vec<Spend>,
pub created_slots: Vec<CatalogSlotValue>,
pub spent_slots: Vec<CatalogSlotValue>,
pub latest_state: (NodePtr, CatalogRegistryState),
pub signature: Signature,
}
impl CatalogPendingSpendInfo {
pub fn new(latest_state: CatalogRegistryState) -> Self {
Self {
actions: vec![],
created_slots: vec![],
spent_slots: vec![],
latest_state: (NodePtr::NIL, latest_state),
signature: Signature::default(),
}
}
}
#[derive(Debug, Clone)]
#[must_use]
pub struct CatalogRegistry {
pub coin: Coin,
pub proof: Proof,
pub info: CatalogRegistryInfo,
pub pending_spend: CatalogPendingSpendInfo,
}
impl CatalogRegistry {
pub fn new(coin: Coin, proof: Proof, info: CatalogRegistryInfo) -> Self {
Self {
coin,
proof,
info,
pending_spend: CatalogPendingSpendInfo::new(info.state),
}
}
}
impl CatalogRegistry {
#[allow(clippy::type_complexity)]
pub fn pending_info_delta_from_spend(
ctx: &mut SpendContext,
action_spend: Spend,
current_state_and_ephemeral: (NodePtr, CatalogRegistryState),
constants: CatalogRegistryConstants,
) -> Result<
(
(NodePtr, CatalogRegistryState),
Vec<CatalogSlotValue>, // created slot values
Vec<CatalogSlotValue>, // spent slot values
),
DriverError,
> {
let mut created_slots = vec![];
let mut spent_slots = vec![];
let register_action = CatalogRegisterAction::from_constants(&constants);
let register_hash = register_action.tree_hash();
let refund_action = CatalogRefundAction::from_constants(&constants);
let refund_hash = refund_action.tree_hash();
let delegated_state_action =
<DelegatedStateAction as SingletonAction<CatalogRegistry>>::from_constants(&constants);
let delegated_state_hash = delegated_state_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, CatalogRegistryState), NodePtr)>(output)?;
let raw_action_hash = ctx.tree_hash(action_spend.puzzle);
if raw_action_hash == register_hash {
spent_slots.extend(register_action.spent_slot_values(ctx, action_spend.solution)?);
created_slots.extend(register_action.created_slot_values(ctx, action_spend.solution)?);
} else if raw_action_hash == refund_hash {
if let (Some(spent_slot), Some(created_slot)) = (
refund_action.spent_slot_value(ctx, action_spend.solution)?,
refund_action.created_slot_value(ctx, action_spend.solution)?,
) {
spent_slots.push(spent_slot);
created_slots.push(created_slot);
}
} else if raw_action_hash != delegated_state_hash {
return Err(DriverError::InvalidMerkleProof);
}
Ok((new_state_and_ephemeral, created_slots, spent_slots))
}
pub fn pending_info_from_spend(
ctx: &mut SpendContext,
inner_solution: NodePtr,
initial_state: CatalogRegistryState,
constants: CatalogRegistryConstants,
) -> Result<CatalogPendingSpendInfo, DriverError> {
let mut created_slots = vec![];
let mut spent_slots = vec![];
let mut state_incl_ephemeral: (NodePtr, CatalogRegistryState) =
(NodePtr::NIL, initial_state);
let inner_solution =
ActionLayer::<CatalogRegistryState, NodePtr>::parse_solution(ctx, inner_solution)?;
for raw_action in &inner_solution.action_spends {
let res = Self::pending_info_delta_from_spend(
ctx,
*raw_action,
state_incl_ephemeral,
constants,
)?;
state_incl_ephemeral = res.0;
created_slots.extend(res.1);
spent_slots.extend(res.2);
}
Ok(CatalogPendingSpendInfo {
actions: inner_solution.action_spends,
created_slots,
spent_slots,
latest_state: state_incl_ephemeral,
signature: Signature::default(),
})
}
pub fn set_pending_signature(&mut self, signature: Signature) {
self.pending_spend.signature = signature;
}
pub fn from_spend(
ctx: &mut SpendContext,
spend: &CoinSpend,
constants: CatalogRegistryConstants,
) -> 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) = CatalogRegistryInfo::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)?;
Ok(Some(CatalogRegistry {
coin,
proof,
info,
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: CatalogRegistryConstants,
) -> Result<Option<Self>, DriverError>
where
Self: Sized,
{
let Some(parent_registry) = CatalogRegistry::from_spend(ctx, parent_spend, constants)?
else {
return Ok(None);
};
let proof = Proof::Lineage(parent_registry.child_lineage_proof());
let new_info = parent_registry
.info
.with_state(parent_registry.pending_spend.latest_state.1);
let new_coin = Coin::new(
parent_registry.coin.coin_id(),
new_info.puzzle_hash().into(),
1,
);
Ok(Some(CatalogRegistry {
coin: new_coin,
proof,
info: new_info,
pending_spend: CatalogPendingSpendInfo::new(new_info.state),
}))
}
pub fn child(&self, child_state: CatalogRegistryState) -> 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);
CatalogRegistry {
coin: new_coin,
proof: Proof::Lineage(self.child_lineage_proof()),
info: new_info,
pending_spend: CatalogPendingSpendInfo::new(new_info.state),
}
}
}
impl ActionSingleton for CatalogRegistry {
type State = CatalogRegistryState;
type Constants = CatalogRegistryConstants;
}
impl CatalogRegistry {
pub fn finish_spend(self, ctx: &mut SpendContext) -> Result<(Self, Signature), DriverError> {
let layers = self.info.into_layers();
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 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(
&CatalogRegistryInfo::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: NodePtr::NIL,
},
},
)?;
let my_spend = Spend::new(puzzle, solution);
ctx.spend(self.coin, my_spend)?;
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(
&self,
slot_value: CatalogSlotValue,
) -> Slot<CatalogSlotValue> {
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, 0, slot_value),
)
}
pub fn actual_neigbors(
&self,
new_tail_hash: Bytes32,
on_chain_left_slot: Slot<CatalogSlotValue>,
on_chain_right_slot: Slot<CatalogSlotValue>,
) -> (Slot<CatalogSlotValue>, Slot<CatalogSlotValue>) {
let mut left = on_chain_left_slot;
let mut right = on_chain_right_slot;
for slot_value in &self.pending_spend.created_slots {
if slot_value.asset_id < new_tail_hash
&& slot_value.asset_id >= left.info.value.asset_id
{
left = self.created_slot_value_to_slot(*slot_value);
}
if slot_value.asset_id > new_tail_hash
&& slot_value.asset_id <= right.info.value.asset_id
{
right = self.created_slot_value_to_slot(*slot_value);
}
}
(left, right)
}
pub fn actual_slot(&self, slot: Slot<CatalogSlotValue>) -> Slot<CatalogSlotValue> {
let mut slot = slot;
for slot_value in &self.pending_spend.created_slots {
if slot.info.value.asset_id == slot_value.asset_id {
slot = self.created_slot_value_to_slot(*slot_value);
}
}
slot
}
pub fn insert_action_spend(
&mut self,
ctx: &mut SpendContext,
action_spend: Spend,
) -> Result<(), DriverError> {
let res = Self::pending_info_delta_from_spend(
ctx,
action_spend,
self.pending_spend.latest_state,
self.info.constants,
)?;
self.pending_spend.latest_state = res.0;
self.pending_spend.created_slots.extend(res.1);
self.pending_spend.spent_slots.extend(res.2);
self.pending_spend.actions.push(action_spend);
Ok(())
}
}