Nested witness signature path injection

crates/sdk/src/program/core.rs

@@ -4,14 +4,14 @@ use std::sync::Arc;
 use dyn_clone::DynClone;
 
 use simplicityhl::CompiledProgram;
-use simplicityhl::WitnessValues;
 use simplicityhl::elements::pset::PartiallySignedTransaction;
 use simplicityhl::elements::{Address, Script, Transaction, TxOut, taproot};
 use simplicityhl::simplicity::bitcoin::{XOnlyPublicKey, secp256k1};
 use simplicityhl::simplicity::jet::Elements;
 use simplicityhl::simplicity::jet::elements::{ElementsEnv, ElementsUtxo};
 use simplicityhl::simplicity::{BitMachine, RedeemNode, Value, leaf_version};
 use simplicityhl::tracker::{DefaultTracker, TrackerLogLevel};
+use simplicityhl::{Parameters, WitnessTypes, WitnessValues};
 
 use super::arguments::ArgumentsTrait;
 use super::error::ProgramError;
@@ -20,6 +20,10 @@ use crate::provider::SimplicityNetwork;
 use crate::utils::{hash_script, tap_data_hash, tr_unspendable_key};
 
 pub trait ProgramTrait: DynClone {
+    fn get_argument_types(&self) -> Result<Parameters, ProgramError>;
+
+    fn get_witness_types(&self) -> Result<WitnessTypes, ProgramError>;
+
     fn get_env(
         &self,
         pst: &PartiallySignedTransaction,
@@ -55,6 +59,14 @@ pub struct Program {
 dyn_clone::clone_trait_object!(ProgramTrait);
 
 impl ProgramTrait for Program {
+    fn get_argument_types(&self) -> Result<Parameters, ProgramError> {
+        self.get_argument_types()
+    }
+
+    fn get_witness_types(&self) -> Result<WitnessTypes, ProgramError> {
+        self.get_witness_types()
+    }
+
     fn get_env(
         &self,
         pst: &PartiallySignedTransaction,
@@ -206,6 +218,20 @@ impl Program {
         hash_script(&self.get_script_pubkey(network))
     }
 
+    pub fn get_argument_types(&self) -> Result<Parameters, ProgramError> {
+        let compiled = self.load()?;
+        let abi_meta = compiled.generate_abi_meta().map_err(ProgramError::ProgramGenAbiMeta)?;
+
+        Ok(abi_meta.param_types)
+    }
+
+    pub fn get_witness_types(&self) -> Result<WitnessTypes, ProgramError> {
+        let compiled = self.load()?;
+        let abi_meta = compiled.generate_abi_meta().map_err(ProgramError::ProgramGenAbiMeta)?;
+
+        Ok(abi_meta.witness_types)
+    }
+
     fn load(&self) -> Result<CompiledProgram, ProgramError> {
         let compiled = CompiledProgram::new(self.source, self.arguments.build_arguments(), true)
             .map_err(ProgramError::Compilation)?;

crates/sdk/src/program/error.rs

@@ -26,4 +26,7 @@ pub enum ProgramError {
 
     #[error("Input index exceeds u32 maximum: {0}")]
     InputIndexOverflow(#[from] std::num::TryFromIntError),
+
+    #[error("Failed to obtain program witness types: {0}")]
+    ProgramGenAbiMeta(String),
 }

crates/sdk/src/signer/core.rs

@@ -1,5 +1,6 @@
 use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
+use std::sync::Arc;
 
 use simplicityhl::Value;
 use simplicityhl::WitnessValues;
@@ -31,6 +32,7 @@ use crate::constants::MIN_FEE;
 use crate::program::ProgramTrait;
 use crate::provider::ProviderTrait;
 use crate::provider::SimplicityNetwork;
+use crate::signer::wtns_injector::WtnsInjector;
 use crate::transaction::{FinalTransaction, PartialInput, PartialOutput, RequiredSignature, UTXO};
 
 use super::error::SignerError;
@@ -418,18 +420,26 @@ impl Signer {
         for (index, input_i) in inputs.iter().enumerate() {
             // we need to prune the program
             if let Some(program_input) = &input_i.program_input {
-                let signed_witness: Result<WitnessValues, SignerError> = match &input_i.required_sig {
-                    // sign the program and insert the signature into the witness
-                    RequiredSignature::Witness(witness_name) => Ok(self.get_signed_program_witness(
+                let signing_info: Option<(&String, &[String])> = match &input_i.required_sig {
+                    RequiredSignature::Witness(wtns_name) => Some((wtns_name, &[])),
+                    RequiredSignature::WitnessWithPath(wtns_name, sig_path) => Some((wtns_name, sig_path)),
+                    _ => None,
+                };
+
+                let signed_witness: Result<WitnessValues, SignerError> = match signing_info {
+                    // sign the program and inject the signature into the witness
+                    Some((witness_name, sig_path)) => Ok(self.get_signed_program_witness(
                         &pst,
                         program_input.program.as_ref(),
                         &program_input.witness.build_witness(),
                         witness_name,
+                        sig_path,
                         index,
                     )?),
-                    // just build the passed witness
-                    _ => Ok(program_input.witness.build_witness()),
+                    // just build the witness
+                    None => Ok(program_input.witness.build_witness()),
                 };
+
                 let pruned_witness = program_input
                     .program
                     .finalize(&pst, &signed_witness.unwrap(), index, &self.network)
@@ -455,20 +465,42 @@ impl Signer {
         program: &dyn ProgramTrait,
         witness: &WitnessValues,
         witness_name: &str,
+        sig_path: &[String],
         index: usize,
     ) -> Result<WitnessValues, SignerError> {
         let signature = self.sign_program(pst, program, index, &self.network)?;
 
+        // inject the signature into the wtns name directly if the path is not provided
+        let sig_val = if !sig_path.is_empty() {
+            let witness_types = program.get_witness_types()?;
+            let witness_type = witness_types
+                .get(&WitnessName::from_str_unchecked(witness_name))
+                .ok_or(SignerError::WtnsFieldNotFound(witness_name.to_string()))?;
+
+            let local_wtns = Arc::new(
+                witness
+                    .get(&WitnessName::from_str_unchecked(witness_name))
+                    .expect("checked above")
+                    .clone(),
+            );
+
+            WtnsInjector::inject_value(
+                &local_wtns,
+                witness_type,
+                sig_path,
+                Value::byte_array(signature.serialize()),
+            )?
+        } else {
+            Value::byte_array(signature.serialize())
+        };
+
         let mut hm = HashMap::new();
 
         witness.iter().for_each(|el| {
             hm.insert(el.0.clone(), el.1.clone());
         });
 
-        hm.insert(
-            WitnessName::from_str_unchecked(witness_name),
-            Value::byte_array(signature.serialize()),
-        );
+        hm.insert(WitnessName::from_str_unchecked(witness_name), sig_val);
 
         Ok(WitnessValues::from(hm))
     }

crates/sdk/src/signer/error.rs

@@ -9,6 +9,9 @@ pub enum SignerError {
     #[error(transparent)]
     Provider(#[from] ProviderError),
 
+    #[error(transparent)]
+    WtnsInjectError(#[from] WtnsWrappingError),
+
     #[error("Failed to parse a mnemonic: {0}")]
     Mnemonic(String),
 
@@ -53,4 +56,25 @@ pub enum SignerError {
 
     #[error("Failed to construct a wpkh address: {0}")]
     WpkhAddressConstruction(#[from] elements_miniscript::Error),
+
+    #[error("Missing such witness field: {0}")]
+    WtnsFieldNotFound(String),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum WtnsWrappingError {
+    #[error("Failed to parse path")]
+    ParsingError,
+
+    #[error("Unsupported path type: {0}")]
+    UnsupportedPathType(String),
+
+    #[error("Path index out of bounds: len is {0}, got {1}")]
+    IdxOutOfBounds(usize, usize),
+
+    #[error("Root type mismatch: expected {0}, got {1}")]
+    RootTypeMismatch(String, String),
+
+    #[error("Path reached undefined branch of Either")]
+    EitherBranchMismatch,
 }

crates/sdk/src/signer/mod.rs

@@ -1,5 +1,6 @@
 pub mod core;
 pub mod error;
+mod wtns_injector;
 
 pub use core::{Signer, SignerTrait};
 pub use error::SignerError;