use crate::{CatLayer, DriverError, Layer, Puzzle, Spend, SpendContext};
use chik_consensus::make_aggsig_final_message::u64_to_bytes;
use chik_protocol::{Bytes, Bytes32, Coin};
use chik_puzzle_types::{
cat::{CatArgs, CatSolution},
CoinProof, LineageProof, Memos,
};
use chik_sdk_types::{run_puzzle, Condition, Conditions};
use chik_sha2::Sha256;
use klvm_traits::FromKlvm;
use klvm_utils::{tree_hash, TreeHash};
use klvmr::{op_utils::u64_from_bytes, Allocator, NodePtr};
use crate::{StreamLayer, StreamPuzzleSolution};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StreamingPuzzleInfo {
pub recipient: Bytes32,
pub clawback_ph: Option<Bytes32>,
pub end_time: u64,
pub last_payment_time: u64,
}
impl StreamingPuzzleInfo {
pub fn new(
recipient: Bytes32,
clawback_ph: Option<Bytes32>,
end_time: u64,
last_payment_time: u64,
) -> Self {
Self {
recipient,
clawback_ph,
end_time,
last_payment_time,
}
}
pub fn amount_to_be_paid(&self, my_coin_amount: u64, payment_time: u64) -> u64 {
my_coin_amount * (payment_time - self.last_payment_time)
/ (self.end_time - self.last_payment_time)
}
pub fn get_hint(recipient: Bytes32) -> Bytes32 {
let mut s = Sha256::new();
s.update(b"s");
s.update(recipient.as_slice());
s.finalize().into()
}
pub fn get_launch_hints(&self) -> Vec<Bytes> {
let hint: Bytes = self.recipient.into();
let clawback_ph: Bytes = if let Some(clawback_ph) = self.clawback_ph {
clawback_ph.into()
} else {
Bytes::new(vec![])
};
let second_memo = u64_to_bytes(self.last_payment_time);
let third_memo = u64_to_bytes(self.end_time);
vec![hint, clawback_ph, second_memo.into(), third_memo.into()]
}
#[must_use]
pub fn with_last_payment_time(self, last_payment_time: u64) -> Self {
Self {
last_payment_time,
..self
}
}
pub fn parse(allocator: &Allocator, puzzle: Puzzle) -> Result<Option<Self>, DriverError> {
let Some(layer) = StreamLayer::parse_puzzle(allocator, puzzle)? else {
return Ok(None);
};
Ok(Some(Self::from_layer(layer)))
}
pub fn into_layer(self) -> StreamLayer {
StreamLayer::new(
self.recipient,
self.clawback_ph,
self.end_time,
self.last_payment_time,
)
}
pub fn from_layer(layer: StreamLayer) -> Self {
Self {
recipient: layer.recipient,
clawback_ph: layer.clawback_ph,
end_time: layer.end_time,
last_payment_time: layer.last_payment_time,
}
}
pub fn inner_puzzle_hash(&self) -> TreeHash {
self.into_layer().puzzle_hash()
}
pub fn from_memos(memos: &[Bytes]) -> Result<Option<Self>, DriverError> {
if memos.len() < 4 || memos.len() > 5 {
return Ok(None);
}
let (recipient, clawback_ph, last_payment_time, end_time): (
Bytes32,
Option<Bytes32>,
u64,
u64,
) = if memos.len() == 4 {
let Ok(recipient_b64): Result<Bytes32, _> = memos[0].clone().try_into() else {
return Ok(None);
};
let clawback_ph_b64: Option<Bytes32> = if memos[1].is_empty() {
None
} else {
let b32: Result<Bytes32, _> = memos[1].clone().try_into();
if let Ok(b32) = b32 {
Some(b32)
} else {
return Ok(None);
}
};
(
recipient_b64,
clawback_ph_b64,
u64_from_bytes(&memos[2]),
u64_from_bytes(&memos[3]),
)
} else {
let Ok(recipient_b64): Result<Bytes32, _> = memos[1].clone().try_into() else {
return Ok(None);
};
let clawback_ph_b64: Option<Bytes32> = if memos[2].is_empty() {
None
} else {
let b32: Result<Bytes32, _> = memos[2].clone().try_into();
if let Ok(b32) = b32 {
Some(b32)
} else {
return Ok(None);
}
};
(
recipient_b64,
clawback_ph_b64,
u64_from_bytes(&memos[3]),
u64_from_bytes(&memos[4]),
)
};
Ok(Some(Self::new(
recipient,
clawback_ph,
end_time,
last_payment_time,
)))
}
}
#[derive(Debug, Clone)]
#[must_use]
pub struct StreamedCat {
pub coin: Coin,
pub asset_id: Bytes32,
pub proof: LineageProof,
pub info: StreamingPuzzleInfo,
}
impl StreamedCat {
pub fn new(
coin: Coin,
asset_id: Bytes32,
proof: LineageProof,
info: StreamingPuzzleInfo,
) -> Self {
Self {
coin,
asset_id,
proof,
info,
}
}
pub fn layers(&self) -> CatLayer<StreamLayer> {
CatLayer::<StreamLayer>::new(self.asset_id, self.info.into_layer())
}
pub fn construct_puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
self.layers().construct_puzzle(ctx)
}
pub fn construct_solution(
&self,
ctx: &mut SpendContext,
payment_time: u64,
clawback: bool,
) -> Result<NodePtr, DriverError> {
self.layers().construct_solution(
ctx,
CatSolution {
inner_puzzle_solution: StreamPuzzleSolution {
my_amount: self.coin.amount,
payment_time,
to_pay: self.info.amount_to_be_paid(self.coin.amount, payment_time),
clawback,
},
lineage_proof: Some(self.proof),
prev_coin_id: self.coin.coin_id(),
this_coin_info: self.coin,
next_coin_proof: CoinProof {
parent_coin_info: self.coin.parent_coin_info,
inner_puzzle_hash: self.info.inner_puzzle_hash().into(),
amount: self.coin.amount,
},
prev_subtotal: 0,
extra_delta: 0,
},
)
}
pub fn spend(
&self,
ctx: &mut SpendContext,
payment_time: u64,
clawback: bool,
) -> Result<(), DriverError> {
let puzzle = self.construct_puzzle(ctx)?;
let solution = self.construct_solution(ctx, payment_time, clawback)?;
ctx.spend(self.coin, Spend::new(puzzle, solution))
}
pub fn from_parent_spend(
allocator: &mut Allocator,
parent_coin: Coin,
parent_puzzle: Puzzle,
parent_solution: NodePtr,
) -> Result<(Option<Self>, bool, u64), DriverError> {
let Some(layers) = CatLayer::<StreamLayer>::parse_puzzle(allocator, parent_puzzle)? else {
let parent_puzzle_ptr = parent_puzzle.ptr();
let output = run_puzzle(allocator, parent_puzzle_ptr, parent_solution)?;
let conds: Conditions<NodePtr> = Conditions::from_klvm(allocator, output)?;
let Some(parent_layer) = CatLayer::<NodePtr>::parse_puzzle(allocator, parent_puzzle)?
else {
return Ok((None, false, 0));
};
let mut found_stream_layer: Option<Self> = None;
for cond in conds {
let Condition::CreateCoin(cc) = cond else {
continue;
};
let Memos::Some(memos) = cc.memos else {
continue;
};
let memos = Vec::<Bytes>::from_klvm(allocator, memos)?;
let Some(candidate_info) = StreamingPuzzleInfo::from_memos(&memos)? else {
continue;
};
let candidate_inner_puzzle_hash = candidate_info.inner_puzzle_hash();
let candidate_puzzle_hash =
CatArgs::curry_tree_hash(parent_layer.asset_id, candidate_inner_puzzle_hash);
if cc.puzzle_hash != candidate_puzzle_hash.into() {
continue;
}
found_stream_layer = Some(Self::new(
Coin::new(
parent_coin.coin_id(),
candidate_puzzle_hash.into(),
cc.amount,
),
parent_layer.asset_id,
LineageProof {
parent_parent_coin_info: parent_coin.parent_coin_info,
parent_inner_puzzle_hash: tree_hash(allocator, parent_layer.inner_puzzle)
.into(),
parent_amount: parent_coin.amount,
},
candidate_info,
));
}
return Ok((found_stream_layer, false, 0));
};
let proof = LineageProof {
parent_parent_coin_info: parent_coin.parent_coin_info,
parent_inner_puzzle_hash: layers.inner_puzzle.puzzle_hash().into(),
parent_amount: parent_coin.amount,
};
let parent_solution =
CatSolution::<StreamPuzzleSolution>::from_klvm(allocator, parent_solution)?;
if parent_solution.inner_puzzle_solution.clawback {
return Ok((None, true, parent_solution.inner_puzzle_solution.to_pay));
}
let new_amount = parent_coin.amount - parent_solution.inner_puzzle_solution.to_pay;
let new_inner_layer = StreamLayer::new(
layers.inner_puzzle.recipient,
layers.inner_puzzle.clawback_ph,
layers.inner_puzzle.end_time,
parent_solution.inner_puzzle_solution.payment_time,
);
let new_puzzle_hash =
CatArgs::curry_tree_hash(layers.asset_id, new_inner_layer.puzzle_hash());
Ok((
Some(Self::new(
Coin::new(parent_coin.coin_id(), new_puzzle_hash.into(), new_amount),
layers.asset_id,
proof,
StreamingPuzzleInfo::from_layer(layers.inner_puzzle)
.with_last_payment_time(parent_solution.inner_puzzle_solution.payment_time),
)),
false,
0,
))
}
}
#[cfg(test)]
mod tests {
use chik_protocol::Bytes;
use chik_sdk_test::{BlsPair, Simulator};
use klvm_utils::tree_hash;
use klvmr::serde::node_from_bytes;
use crate::{Cat, StandardLayer, STREAM_PUZZLE, STREAM_PUZZLE_HASH};
use super::*;
#[test]
fn test_puzzle_hash() {
let mut allocator = Allocator::new();
let ptr = node_from_bytes(&mut allocator, &STREAM_PUZZLE).unwrap();
assert_eq!(tree_hash(&allocator, ptr), STREAM_PUZZLE_HASH);
}
#[test]
fn test_streamed_cat() -> anyhow::Result<()> {
let mut ctx = SpendContext::new();
let mut sim = Simulator::new();
let claim_intervals = [1000, 2000, 500, 1000, 10];
let clawback_offset = 1234;
let total_claim_time = claim_intervals.iter().sum::<u64>() + clawback_offset;
let user_key = BlsPair::new(0);
let user_p2 = StandardLayer::new(user_key.pk);
let user_puzzle_hash: Bytes32 = user_key.puzzle_hash;
let payment_cat_amount = 1000;
let minter_key = BlsPair::new(1);
let minter_coin = sim.new_coin(minter_key.puzzle_hash, payment_cat_amount);
let minter_p2 = StandardLayer::new(minter_key.pk);
let clawback_puzzle_ptr = ctx.alloc(&1)?;
let clawback_ph = ctx.tree_hash(clawback_puzzle_ptr);
let streaming_inner_puzzle = StreamLayer::new(
user_puzzle_hash,
Some(clawback_ph.into()),
total_claim_time + 1000,
1000,
);
let streaming_inner_puzzle_hash: Bytes32 = streaming_inner_puzzle.puzzle_hash().into();
let (issue_cat, cats) = Cat::issue_with_coin(
&mut ctx,
minter_coin.coin_id(),
payment_cat_amount,
Conditions::new().create_coin(
streaming_inner_puzzle_hash,
payment_cat_amount,
Memos::None,
),
)?;
minter_p2.spend(&mut ctx, minter_coin, issue_cat)?;
let initial_vesting_cat = cats[0];
sim.spend_coins(ctx.take(), &[minter_key.sk.clone()])?;
sim.set_next_timestamp(1000 + claim_intervals[0])?;
let mut streamed_cat = StreamedCat::new(
initial_vesting_cat.coin,
initial_vesting_cat.info.asset_id,
initial_vesting_cat.lineage_proof.unwrap(),
StreamingPuzzleInfo::new(
user_puzzle_hash,
Some(clawback_ph.into()),
total_claim_time + 1000,
1000,
),
);
let mut claim_time = sim.next_timestamp();
for (i, _interval) in claim_intervals.iter().enumerate() {
if i < claim_intervals.len() - 1 {
sim.pass_time(claim_intervals[i + 1]);
}
let user_coin = sim.new_coin(user_puzzle_hash, 0);
let message_to_send: Bytes = Bytes::new(u64_to_bytes(claim_time));
let coin_id_ptr = ctx.alloc(&streamed_cat.coin.coin_id())?;
user_p2.spend(
&mut ctx,
user_coin,
Conditions::new().send_message(23, message_to_send, vec![coin_id_ptr]),
)?;
streamed_cat.spend(&mut ctx, claim_time, false)?;
let spends = ctx.take();
let streamed_cat_spend = spends.last().unwrap().clone();
sim.spend_coins(spends, &[user_key.sk.clone()])?;
if i < claim_intervals.len() - 1 {
claim_time += claim_intervals[i + 1];
}
let parent_puzzle = ctx.alloc(&streamed_cat_spend.puzzle_reveal)?;
let parent_puzzle = Puzzle::from_klvm(&ctx, parent_puzzle)?;
let parent_solution = ctx.alloc(&streamed_cat_spend.solution)?;
let (Some(new_streamed_cat), clawback, _) = StreamedCat::from_parent_spend(
&mut ctx,
streamed_cat.coin,
parent_puzzle,
parent_solution,
)?
else {
panic!("Failed to parse new streamed cat");
};
assert!(!clawback);
streamed_cat = new_streamed_cat;
}
assert!(streamed_cat.coin.amount > 0);
let clawback_msg_coin = sim.new_coin(clawback_ph.into(), 0);
let claim_time = sim.next_timestamp() + 1;
let message_to_send: Bytes = Bytes::new(u64_to_bytes(claim_time));
let coin_id_ptr = ctx.alloc(&streamed_cat.coin.coin_id())?;
let solution =
ctx.alloc(&Conditions::new().send_message(23, message_to_send, vec![coin_id_ptr]))?;
ctx.spend(clawback_msg_coin, Spend::new(clawback_puzzle_ptr, solution))?;
streamed_cat.spend(&mut ctx, claim_time, true)?;
let spends = ctx.take();
let streamed_cat_spend = spends.last().unwrap().clone();
sim.spend_coins(spends, &[user_key.sk.clone()])?;
let parent_puzzle = ctx.alloc(&streamed_cat_spend.puzzle_reveal)?;
let parent_puzzle = Puzzle::from_klvm(&ctx, parent_puzzle)?;
let parent_solution = ctx.alloc(&streamed_cat_spend.solution)?;
let (new_streamed_cat, clawback, _paid_amount_if_clawback) =
StreamedCat::from_parent_spend(
&mut ctx,
streamed_cat.coin,
parent_puzzle,
parent_solution,
)?;
assert!(clawback);
assert!(new_streamed_cat.is_none());
Ok(())
}
}