1use std::sync::Arc;
2
3use dyn_clone::DynClone;
4
5use simplicityhl::CompiledProgram;
6use simplicityhl::WitnessValues;
7use simplicityhl::elements::pset::PartiallySignedTransaction;
8use simplicityhl::elements::{Address, Script, Transaction, TxOut, script, taproot};
9use simplicityhl::simplicity::bitcoin::{XOnlyPublicKey, secp256k1};
10use simplicityhl::simplicity::jet::Elements;
11use simplicityhl::simplicity::jet::elements::{ElementsEnv, ElementsUtxo};
12use simplicityhl::simplicity::{BitMachine, RedeemNode, Value};
13use simplicityhl::tracker::{DefaultTracker, TrackerLogLevel};
14
15use super::arguments::ArgumentsTrait;
16use super::error::ProgramError;
17
18use crate::provider::SimplicityNetwork;
19use crate::utils::hash_script;
20
21pub trait ProgramTrait: DynClone {
22 fn get_env(
23 &self,
24 pst: &PartiallySignedTransaction,
25 input_index: usize,
26 network: &SimplicityNetwork,
27 ) -> Result<ElementsEnv<Arc<Transaction>>, ProgramError>;
28
29 fn execute(
30 &self,
31 pst: &PartiallySignedTransaction,
32 witness: &WitnessValues,
33 input_index: usize,
34 network: &SimplicityNetwork,
35 ) -> Result<(Arc<RedeemNode<Elements>>, Value), ProgramError>;
36
37 fn finalize(
38 &self,
39 pst: &PartiallySignedTransaction,
40 witness: &WitnessValues,
41 input_index: usize,
42 network: &SimplicityNetwork,
43 ) -> Result<Vec<Vec<u8>>, ProgramError>;
44}
45
46#[derive(Clone)]
47pub struct Program {
48 source: &'static str,
49 pub_key: XOnlyPublicKey,
50 arguments: Box<dyn ArgumentsTrait>,
51}
52
53dyn_clone::clone_trait_object!(ProgramTrait);
54
55impl ProgramTrait for Program {
56 fn get_env(
57 &self,
58 pst: &PartiallySignedTransaction,
59 input_index: usize,
60 network: &SimplicityNetwork,
61 ) -> Result<ElementsEnv<Arc<Transaction>>, ProgramError> {
62 let genesis_hash = network.genesis_block_hash();
63 let cmr = self.load()?.commit().cmr();
64 let utxos: Vec<TxOut> = pst.inputs().iter().filter_map(|x| x.witness_utxo.clone()).collect();
65
66 if utxos.len() <= input_index {
67 return Err(ProgramError::UtxoIndexOutOfBounds {
68 input_index,
69 utxo_count: utxos.len(),
70 });
71 }
72
73 let target_utxo = &utxos[input_index];
74 let script_pubkey = self.get_tr_address(&network)?.script_pubkey();
75
76 if target_utxo.script_pubkey != script_pubkey {
77 return Err(ProgramError::ScriptPubkeyMismatch {
78 expected_hash: script_pubkey.script_hash().to_string(),
79 actual_hash: target_utxo.script_pubkey.script_hash().to_string(),
80 });
81 }
82
83 Ok(ElementsEnv::new(
84 Arc::new(pst.extract_tx()?),
85 utxos
86 .iter()
87 .map(|utxo| ElementsUtxo {
88 script_pubkey: utxo.script_pubkey.clone(),
89 asset: utxo.asset,
90 value: utxo.value,
91 })
92 .collect(),
93 u32::try_from(input_index)?,
94 cmr,
95 self.control_block()?,
96 None,
97 genesis_hash,
98 ))
99 }
100
101 fn execute(
102 &self,
103 pst: &PartiallySignedTransaction,
104 witness: &WitnessValues,
105 input_index: usize,
106 network: &SimplicityNetwork,
107 ) -> Result<(Arc<RedeemNode<Elements>>, Value), ProgramError> {
108 let satisfied = self
109 .load()?
110 .satisfy(witness.clone())
111 .map_err(ProgramError::WitnessSatisfaction)?;
112
113 let mut tracker = DefaultTracker::new(satisfied.debug_symbols()).with_log_level(TrackerLogLevel::Debug);
114
115 let env = self.get_env(pst, input_index, network)?;
116
117 let pruned = satisfied.redeem().prune_with_tracker(&env, &mut tracker)?;
118 let mut mac = BitMachine::for_program(&pruned)?;
119
120 let result = mac.exec(&pruned, &env)?;
121
122 Ok((pruned, result))
123 }
124
125 fn finalize(
126 &self,
127 pst: &PartiallySignedTransaction,
128 witness: &WitnessValues,
129 input_index: usize,
130 network: &SimplicityNetwork,
131 ) -> Result<Vec<Vec<u8>>, ProgramError> {
132 let pruned = self.execute(&pst, witness, input_index, network)?.0;
133
134 let (simplicity_program_bytes, simplicity_witness_bytes) = pruned.to_vec_with_witness();
135 let cmr = pruned.cmr();
136
137 Ok(vec![
138 simplicity_witness_bytes,
139 simplicity_program_bytes,
140 cmr.as_ref().to_vec(),
141 self.control_block()?.serialize(),
142 ])
143 }
144}
145
146impl Program {
147 pub fn new(source: &'static str, pub_key: XOnlyPublicKey, arguments: Box<dyn ArgumentsTrait>) -> Self {
148 Self {
149 source: source,
150 pub_key: pub_key,
151 arguments: arguments,
152 }
153 }
154
155 pub fn get_tr_address(&self, network: &SimplicityNetwork) -> Result<Address, ProgramError> {
156 let spend_info = self.taproot_spending_info()?;
157
158 Ok(Address::p2tr(
159 secp256k1::SECP256K1,
160 spend_info.internal_key(),
161 spend_info.merkle_root(),
162 None,
163 network.address_params(),
164 ))
165 }
166
167 pub fn get_script_pubkey(&self, network: &SimplicityNetwork) -> Result<Script, ProgramError> {
168 Ok(self.get_tr_address(network)?.script_pubkey())
169 }
170
171 pub fn get_script_hash(&self, network: &SimplicityNetwork) -> Result<[u8; 32], ProgramError> {
172 Ok(hash_script(&self.get_script_pubkey(network)?))
173 }
174
175 fn load(&self) -> Result<CompiledProgram, ProgramError> {
176 let compiled = CompiledProgram::new(self.source, self.arguments.build_arguments(), true)
177 .map_err(ProgramError::Compilation)?;
178 Ok(compiled)
179 }
180
181 fn script_version(&self) -> Result<(Script, taproot::LeafVersion), ProgramError> {
182 let cmr = self.load()?.commit().cmr();
183 let script = script::Script::from(cmr.as_ref().to_vec());
184
185 Ok((script, simplicityhl::simplicity::leaf_version()))
186 }
187
188 fn taproot_spending_info(&self) -> Result<taproot::TaprootSpendInfo, ProgramError> {
189 let builder = taproot::TaprootBuilder::new();
190 let (script, version) = self.script_version()?;
191
192 let builder = builder
193 .add_leaf_with_ver(0, script, version)
194 .expect("tap tree should be valid");
195
196 Ok(builder
197 .finalize(secp256k1::SECP256K1, self.pub_key)
198 .expect("tap tree should be valid"))
199 }
200
201 fn control_block(&self) -> Result<taproot::ControlBlock, ProgramError> {
202 let info = self.taproot_spending_info()?;
203 let script_ver = self.script_version()?;
204
205 Ok(info.control_block(&script_ver).expect("control block should exist"))
206 }
207}