1use std::iter;
2use std::sync::Arc;
3
4use dyn_clone::DynClone;
5
6use simplicityhl::CompiledProgram;
7use simplicityhl::elements::pset::PartiallySignedTransaction;
8use simplicityhl::elements::{Address, Script, Transaction, TxOut, 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, leaf_version};
13use simplicityhl::tracker::{DefaultTracker, TrackerLogLevel};
14use simplicityhl::{Parameters, WitnessTypes, WitnessValues};
15
16use super::arguments::ArgumentsTrait;
17use super::error::ProgramError;
18
19use crate::provider::SimplicityNetwork;
20use crate::utils::{hash_script, tap_data_hash, tr_unspendable_key};
21
22pub trait ProgramTrait: DynClone {
23 fn get_argument_types(&self) -> Result<Parameters, ProgramError>;
24
25 fn get_witness_types(&self) -> Result<WitnessTypes, ProgramError>;
26
27 fn get_env(
28 &self,
29 pst: &PartiallySignedTransaction,
30 input_index: usize,
31 network: &SimplicityNetwork,
32 ) -> Result<ElementsEnv<Arc<Transaction>>, ProgramError>;
33
34 fn execute(
35 &self,
36 pst: &PartiallySignedTransaction,
37 witness: &WitnessValues,
38 input_index: usize,
39 network: &SimplicityNetwork,
40 ) -> Result<(Arc<RedeemNode<Elements>>, Value), ProgramError>;
41
42 fn finalize(
43 &self,
44 pst: &PartiallySignedTransaction,
45 witness: &WitnessValues,
46 input_index: usize,
47 network: &SimplicityNetwork,
48 ) -> Result<Vec<Vec<u8>>, ProgramError>;
49}
50
51#[derive(Clone)]
52pub struct Program {
53 source: &'static str,
54 pub_key: XOnlyPublicKey,
55 arguments: Box<dyn ArgumentsTrait>,
56 storage: Vec<[u8; 32]>,
57}
58
59dyn_clone::clone_trait_object!(ProgramTrait);
60
61impl ProgramTrait for Program {
62 fn get_argument_types(&self) -> Result<Parameters, ProgramError> {
63 self.get_argument_types()
64 }
65
66 fn get_witness_types(&self) -> Result<WitnessTypes, ProgramError> {
67 self.get_witness_types()
68 }
69
70 fn get_env(
71 &self,
72 pst: &PartiallySignedTransaction,
73 input_index: usize,
74 network: &SimplicityNetwork,
75 ) -> Result<ElementsEnv<Arc<Transaction>>, ProgramError> {
76 let genesis_hash = network.genesis_block_hash();
77 let cmr = self.load()?.commit().cmr();
78 let utxos: Vec<TxOut> = pst.inputs().iter().filter_map(|x| x.witness_utxo.clone()).collect();
79
80 if utxos.len() <= input_index {
81 return Err(ProgramError::UtxoIndexOutOfBounds {
82 input_index,
83 utxo_count: utxos.len(),
84 });
85 }
86
87 let target_utxo = &utxos[input_index];
88 let script_pubkey = self.get_tr_address(network).script_pubkey();
89
90 if target_utxo.script_pubkey != script_pubkey {
91 return Err(ProgramError::ScriptPubkeyMismatch {
92 expected_hash: script_pubkey.script_hash().to_string(),
93 actual_hash: target_utxo.script_pubkey.script_hash().to_string(),
94 });
95 }
96
97 Ok(ElementsEnv::new(
98 Arc::new(pst.extract_tx()?),
99 utxos
100 .iter()
101 .map(|utxo| ElementsUtxo {
102 script_pubkey: utxo.script_pubkey.clone(),
103 asset: utxo.asset,
104 value: utxo.value,
105 })
106 .collect(),
107 u32::try_from(input_index)?,
108 cmr,
109 self.control_block()?,
110 None,
111 genesis_hash,
112 ))
113 }
114
115 fn execute(
116 &self,
117 pst: &PartiallySignedTransaction,
118 witness: &WitnessValues,
119 input_index: usize,
120 network: &SimplicityNetwork,
121 ) -> Result<(Arc<RedeemNode<Elements>>, Value), ProgramError> {
122 let satisfied = self
123 .load()?
124 .satisfy(witness.clone())
125 .map_err(ProgramError::WitnessSatisfaction)?;
126
127 let mut tracker = DefaultTracker::new(satisfied.debug_symbols()).with_log_level(TrackerLogLevel::Debug);
129
130 let env = self.get_env(pst, input_index, network)?;
131
132 let pruned = satisfied.redeem().prune_with_tracker(&env, &mut tracker)?;
133 let mut mac = BitMachine::for_program(&pruned)?;
134
135 let result = mac.exec(&pruned, &env)?;
136
137 Ok((pruned, result))
138 }
139
140 fn finalize(
141 &self,
142 pst: &PartiallySignedTransaction,
143 witness: &WitnessValues,
144 input_index: usize,
145 network: &SimplicityNetwork,
146 ) -> Result<Vec<Vec<u8>>, ProgramError> {
147 let pruned = self.execute(pst, witness, input_index, network)?.0;
148
149 let (simplicity_program_bytes, simplicity_witness_bytes) = pruned.to_vec_with_witness();
150 let cmr = pruned.cmr();
151
152 Ok(vec![
153 simplicity_witness_bytes,
154 simplicity_program_bytes,
155 cmr.as_ref().to_vec(),
156 self.control_block()?.serialize(),
157 ])
158 }
159}
160
161impl Program {
162 pub fn new(source: &'static str, arguments: Box<dyn ArgumentsTrait>) -> Self {
163 Self {
164 source,
165 pub_key: tr_unspendable_key(),
166 arguments,
167 storage: Vec::new(),
168 }
169 }
170
171 pub fn with_pub_key(mut self, pub_key: XOnlyPublicKey) -> Self {
172 self.pub_key = pub_key;
173
174 self
175 }
176
177 pub fn with_storage_capacity(mut self, capacity: usize) -> Self {
178 self.storage = vec![[0u8; 32]; capacity];
179
180 self
181 }
182
183 pub fn set_storage_at(&mut self, index: usize, new_value: [u8; 32]) {
184 let slot = self.storage.get_mut(index).expect("Index out of bounds");
185
186 *slot = new_value;
187 }
188
189 pub fn get_storage_len(&self) -> usize {
190 self.storage.len()
191 }
192
193 pub fn get_storage(&self) -> &[[u8; 32]] {
194 &self.storage
195 }
196
197 pub fn get_storage_at(&self, index: usize) -> [u8; 32] {
198 self.storage[index]
199 }
200
201 pub fn get_tr_address(&self, network: &SimplicityNetwork) -> Address {
202 let spend_info = self.taproot_spending_info().unwrap();
203
204 Address::p2tr(
205 secp256k1::SECP256K1,
206 spend_info.internal_key(),
207 spend_info.merkle_root(),
208 None,
209 network.address_params(),
210 )
211 }
212
213 pub fn get_script_pubkey(&self, network: &SimplicityNetwork) -> Script {
214 self.get_tr_address(network).script_pubkey()
215 }
216
217 pub fn get_script_hash(&self, network: &SimplicityNetwork) -> [u8; 32] {
218 hash_script(&self.get_script_pubkey(network))
219 }
220
221 pub fn get_argument_types(&self) -> Result<Parameters, ProgramError> {
222 let compiled = self.load()?;
223 let abi_meta = compiled.generate_abi_meta().map_err(ProgramError::ProgramGenAbiMeta)?;
224
225 Ok(abi_meta.param_types)
226 }
227
228 pub fn get_witness_types(&self) -> Result<WitnessTypes, ProgramError> {
229 let compiled = self.load()?;
230 let abi_meta = compiled.generate_abi_meta().map_err(ProgramError::ProgramGenAbiMeta)?;
231
232 Ok(abi_meta.witness_types)
233 }
234
235 fn load(&self) -> Result<CompiledProgram, ProgramError> {
236 let compiled = CompiledProgram::new(self.source, self.arguments.build_arguments(), true)
237 .map_err(ProgramError::Compilation)?;
238 Ok(compiled)
239 }
240
241 fn script_version(&self) -> Result<(Script, taproot::LeafVersion), ProgramError> {
242 let cmr = self.load()?.commit().cmr();
243 let script = Script::from(cmr.as_ref().to_vec());
244
245 Ok((script, leaf_version()))
246 }
247
248 fn taproot_leaf_depths(total_leaves: usize) -> Vec<usize> {
249 assert!(total_leaves > 0, "Taproot tree must contain at least one leaf");
250
251 let next_pow2 = total_leaves.next_power_of_two();
252 let depth = next_pow2.ilog2() as usize;
253
254 let shallow_count = next_pow2 - total_leaves;
255 let deep_count = total_leaves - shallow_count;
256
257 let mut depths = Vec::with_capacity(total_leaves);
258 depths.extend(iter::repeat_n(depth, deep_count));
259
260 if depth > 0 {
261 depths.extend(iter::repeat_n(depth - 1, shallow_count));
262 }
263
264 depths
265 }
266
267 fn taproot_spending_info(&self) -> Result<taproot::TaprootSpendInfo, ProgramError> {
268 let mut builder = taproot::TaprootBuilder::new();
269 let (script, version) = self.script_version()?;
270 let depths = Self::taproot_leaf_depths(1 + self.get_storage_len());
271
272 builder = builder
273 .add_leaf_with_ver(depths[0], script, version)
274 .expect("tap tree should be valid");
275
276 for (slot, depth) in self.get_storage().iter().zip(depths.into_iter().skip(1)) {
277 builder = builder
278 .add_hidden(depth, tap_data_hash(slot))
279 .expect("tap tree should be valid");
280 }
281
282 Ok(builder
283 .finalize(secp256k1::SECP256K1, self.pub_key)
284 .expect("tap tree should be valid"))
285 }
286
287 fn control_block(&self) -> Result<taproot::ControlBlock, ProgramError> {
288 let info = self.taproot_spending_info()?;
289 let script_ver = self.script_version()?;
290
291 Ok(info.control_block(&script_ver).expect("control block should exist"))
292 }
293}
294
295#[cfg(test)]
296mod tests {
297 use simplicityhl::{
298 Arguments,
299 elements::{AssetId, confidential, pset::Input},
300 };
301
302 use super::*;
303
304 const DUMMY_PROGRAM: &str = r#"
306 fn main() {
307 let ab: u16 = <(u8, u8)>::into((0x10, 0x01));
308 let c: u16 = 0x1001;
309 assert!(jet::eq_16(ab, c));
310 let ab: u8 = <(u4, u4)>::into((0b1011, 0b1101));
311 let c: u8 = 0b10111101;
312 assert!(jet::eq_8(ab, c));
313 }
314 "#;
315
316 #[derive(Clone)]
317 struct EmptyArguments;
318
319 impl ArgumentsTrait for EmptyArguments {
320 fn build_arguments(&self) -> Arguments {
321 Arguments::default()
322 }
323 }
324
325 fn dummy_asset_id(byte: u8) -> AssetId {
326 AssetId::from_slice(&[byte; 32]).unwrap()
327 }
328
329 fn dummy_program() -> Program {
330 Program::new(DUMMY_PROGRAM, Box::new(EmptyArguments))
331 }
332
333 fn dummy_network() -> SimplicityNetwork {
334 SimplicityNetwork::default_regtest()
335 }
336
337 fn make_pst_with_script(script: Script) -> PartiallySignedTransaction {
338 let txout = TxOut {
339 asset: confidential::Asset::Explicit(dummy_asset_id(0xAA)),
340 value: confidential::Value::Explicit(1000),
341 script_pubkey: script,
342 ..Default::default()
343 };
344 let input = Input {
345 witness_utxo: Some(txout),
346 ..Default::default()
347 };
348
349 let mut pst = PartiallySignedTransaction::new_v2();
350
351 pst.add_input(input);
352
353 pst
354 }
355
356 #[test]
357 fn test_get_env_idx() {
358 let program = dummy_program();
359 let network = dummy_network();
360
361 let correct_script = program.get_script_pubkey(&network);
362 let wrong_script = Script::new();
363
364 let mut pst = make_pst_with_script(wrong_script);
365
366 let correct_txout = TxOut {
367 asset: confidential::Asset::Explicit(dummy_asset_id(0xAA)),
368 value: confidential::Value::Explicit(1000),
369 script_pubkey: correct_script,
370 ..Default::default()
371 };
372
373 pst.add_input(Input {
374 witness_utxo: Some(correct_txout),
375 ..Default::default()
376 });
377
378 assert!(matches!(
380 program.get_env(&pst, 0, &network).unwrap_err(),
381 ProgramError::ScriptPubkeyMismatch { .. }
382 ));
383
384 assert!(program.get_env(&pst, 1, &network).is_ok());
385 }
386
387 #[test]
388 fn test_taproot_leaf_depths_known_values() {
389 let cases = [
390 (1, vec![0]),
391 (2, vec![1, 1]),
392 (3, vec![2, 2, 1]),
393 (4, vec![2, 2, 2, 2]),
394 (5, vec![3, 3, 2, 2, 2]),
395 (6, vec![3, 3, 3, 3, 2, 2]),
396 (8, vec![3, 3, 3, 3, 3, 3, 3, 3]),
397 ];
398
399 for (n, expected) in cases {
400 assert_eq!(Program::taproot_leaf_depths(n), expected, "n={n}");
401 }
402 }
403}