Nested witness signature path injection

crates/sdk/src/program/core.rs

@@ -41,6 +41,8 @@ pub trait ProgramTrait: DynClone {
         input_index: usize,
         network: &SimplicityNetwork,
     ) -> Result<Vec<Vec<u8>>, ProgramError>;
+
+    fn load(&self) -> Result<CompiledProgram, ProgramError>;
 }
 
 #[derive(Clone)]
@@ -142,6 +144,10 @@ impl ProgramTrait for Program {
             self.control_block()?.serialize(),
         ])
     }
+
+    fn load(&self) -> Result<CompiledProgram, ProgramError> {
+        self.load()
+    }
 }
 
 impl Program {
@@ -173,7 +179,7 @@ impl Program {
         hash_script(&self.get_script_pubkey(network))
     }
 
-    fn load(&self) -> Result<CompiledProgram, ProgramError> {
+    pub fn load(&self) -> Result<CompiledProgram, ProgramError> {
         let compiled = CompiledProgram::new(self.source, self.arguments.build_arguments(), true)
             .map_err(ProgramError::Compilation)?;
         Ok(compiled)

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,24 @@ 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(wnts_name) => Some((wnts_name, &[])),
+                    RequiredSignature::WitnessWithPath(wnts_name, sig_path) => Some((wnts_name, sig_path)),
+                    _ => None,
+                };
+
+                let signed_witness: Result<WitnessValues, SignerError> = match signing_info {
+                    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()),
+                    None => Ok(program_input.witness.build_witness()),
                 };
+
                 let pruned_witness = program_input
                     .program
                     .finalize(&pst, &signed_witness.unwrap(), index, &self.network)
@@ -455,20 +463,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)?;
 
+        // put signature right after wtns field name if path is not provided
+        let sig_val = if !sig_path.is_empty() {
+            let wtns_injector = WtnsInjector::new(sig_path)?;
+
+            let compiled = program.load().map_err(SignerError::Program)?;
+
+            let abi_meta = compiled.generate_abi_meta().map_err(SignerError::ProgramGenAbiMeta)?;
+
+            let witness_types = abi_meta
+                .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(),
+            );
+
+            wtns_injector.inject_value(&local_wtns, witness_types, 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

@@ -53,4 +53,25 @@ pub enum SignerError {
 
     #[error("Failed to construct a wpkh address: {0}")]
     WpkhAddressConstruction(#[from] elements_miniscript::Error),
+
+    #[error("Failed to obtain program witness types: {0}")]
+    ProgramGenAbiMeta(String),
+
+    #[error("Missing such witness field: {0}")]
+    WtnsFieldNotFound(String),
+
+    #[error(transparent)]
+    WtnsInjectError(#[from] WtnsWrappingError),
+}
+
+#[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),
 }

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;

crates/sdk/src/signer/wtns_injector.rs

@@ -0,0 +1,212 @@
+use std::sync::Arc;
+
+use simplicityhl::{
+    ResolvedType, Value,
+    types::TypeInner,
+    value::{ValueConstructible, ValueInner},
+};
+
+use crate::signer::error::WtnsWrappingError;
+
+/// Struct for injecting specific value by given path into witness value
+#[derive(Clone)]
+pub struct WtnsInjector {
+    path: Vec<WtnsPathRoute>,
+}
+
+impl WtnsInjector {
+    /// ## Usage
+    /// ```rust,ignore
+    /// // .simf script
+    /// match witness::SOMETHING {
+    ///     Left(x: u64) => ...,
+    ///     Right([y, z]: [u64, u64]) => ...
+    /// }
+    /// // path for each variable
+    /// vec!["Left"] // for x
+    /// vec!["Right", "0"] // for y
+    /// vec!["Right", "1"] // for z
+    /// ```
+    pub fn new(path: &[String]) -> Result<Self, WtnsWrappingError> {
+        let parsed_path = path
+            .iter()
+            .map(|route| match route.as_str() {
+                "Left" => Ok(WtnsPathRoute::Either(EitherRoute::Left)),
+                "Right" => Ok(WtnsPathRoute::Either(EitherRoute::Right)),
+                s => s
+                    .parse::<usize>()
+                    .map(|n| WtnsPathRoute::Enumerable(EnumerableRoute(n)))
+                    .map_err(|_| WtnsWrappingError::ParsingError),
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        Ok(Self { path: parsed_path })
+    }
+
+    /// Constructs new value by intjecting given value into witness at the position described by `path`.
+    /// Consistency between `witness` and `witness_types` should be guaranteed by caller.
+    pub fn inject_value(
+        &self,
+        witness: &Arc<Value>,
+        witness_types: &ResolvedType,
+        value: Value,
+    ) -> Result<Value, WtnsWrappingError> {
+        enum StackItem {
+            Either(EitherRoute, Arc<ResolvedType>),
+            Array(EnumerableRoute, Arc<ResolvedType>, Arc<[Value]>),
+            Tuple(EnumerableRoute, Arc<[Value]>),
+        }
+
+        // invocations of these functions below determined from types during traversal
+        // matches! guard at top of loop guarantees that types and routes are consistent
+        fn downcast_either(val: &Value, direction: EitherRoute) -> Arc<Value> {
+            match (direction, val.inner()) {
+                (EitherRoute::Left, ValueInner::Either(either)) => Arc::clone(either.as_ref().unwrap_left()),
+                (EitherRoute::Right, ValueInner::Either(either)) => Arc::clone(either.as_ref().unwrap_right()),
+                _ => unreachable!(),
+            }
+        }
+
+        fn downcast_array(val: &Value) -> Arc<[Value]> {
+            match val.inner() {
+                ValueInner::Array(arr) => Arc::clone(arr),
+                _ => unreachable!(),
+            }
+        }
+
+        fn downcast_tuple(val: &Value) -> Arc<[Value]> {
+            match val.inner() {
+                ValueInner::Tuple(arr) => Arc::clone(arr),
+                _ => unreachable!(),
+            }
+        }
+
+        let mut stack = Vec::new();
+        let mut current_val = Arc::clone(witness);
+        let mut current_ty = witness_types;
+
+        for route in self.path.iter() {
+            if !matches!(
+                (route, current_ty.as_inner()),
+                (WtnsPathRoute::Enumerable(_), TypeInner::Array(_, _))
+                    | (WtnsPathRoute::Enumerable(_), TypeInner::Tuple(_))
+                    | (WtnsPathRoute::Either(_), TypeInner::Either(_, _))
+            ) {
+                return Err(WtnsWrappingError::UnsupportedPathType(current_ty.to_string()));
+            }
+
+            match current_ty.as_inner() {
+                TypeInner::Either(left_ty, right_ty) => {
+                    let direction: EitherRoute = (*route).try_into().expect("Checked in matches! above");
+                    match direction {
+                        EitherRoute::Left => {
+                            stack.push(StackItem::Either(direction, Arc::clone(right_ty)));
+                            current_ty = left_ty;
+                        }
+                        EitherRoute::Right => {
+                            stack.push(StackItem::Either(direction, Arc::clone(left_ty)));
+                            current_ty = right_ty;
+                        }
+                    }
+                    current_val = downcast_either(&current_val, direction);
+                }
+                TypeInner::Array(ty, len) => {
+                    let idx: EnumerableRoute = (*route).try_into().expect("Checked in matches! above");
+
+                    if idx.0 >= *len {
+                        return Err(WtnsWrappingError::IdxOutOfBounds(*len, idx.0));
+                    }
+
+                    let arr_val = downcast_array(&current_val);
+
+                    stack.push(StackItem::Array(idx, Arc::clone(ty), Arc::clone(&arr_val)));
+
+                    current_ty = ty;
+                    current_val = Arc::new(arr_val[idx.0].clone());
+                }
+                TypeInner::Tuple(tuple) => {
+                    let idx: EnumerableRoute = (*route).try_into().expect("Checked in matches! above");
+
+                    if idx.0 >= tuple.len() {
+                        return Err(WtnsWrappingError::IdxOutOfBounds(tuple.len(), idx.0));
+                    }
+
+                    let tuple_val = downcast_tuple(&current_val);
+
+                    stack.push(StackItem::Tuple(idx, Arc::clone(&tuple_val)));
+
+                    current_ty = &tuple[idx.0];
+                    current_val = Arc::new(tuple_val[idx.0].clone());
+                }
+                _ => unreachable!("checked at the top of loop"),
+            }
+        }
+
+        if value.ty() != current_ty {
+            return Err(WtnsWrappingError::RootTypeMismatch(
+                current_ty.to_string(),
+                value.ty().to_string(),
+            ));
+        }
+
+        let mut value = value;
+
+        for item in stack.into_iter().rev() {
+            value = match item {
+                StackItem::Either(direction, sibling_ty) => match direction {
+                    EitherRoute::Left => Value::left(value, (*sibling_ty).clone()),
+                    EitherRoute::Right => Value::right((*sibling_ty).clone(), value),
+                },
+                StackItem::Array(idx, elem_ty, arr) => {
+                    let mut elements = arr.to_vec();
+                    elements[idx.0] = value;
+                    Value::array(elements, (*elem_ty).clone())
+                }
+                StackItem::Tuple(idx, tuple_vals) => {
+                    let mut elements = tuple_vals.to_vec();
+                    elements[idx.0] = value;
+                    Value::tuple(elements)
+                }
+            };
+        }
+
+        Ok(value)
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum WtnsPathRoute {
+    Either(EitherRoute),
+    Enumerable(EnumerableRoute),
+}
+
+impl TryInto<EitherRoute> for WtnsPathRoute {
+    type Error = WtnsPathRoute;
+
+    fn try_into(self) -> Result<EitherRoute, Self::Error> {
+        match self {
+            Self::Either(direction) => Ok(direction),
+            _ => Err(self),
+        }
+    }
+}
+
+impl TryInto<EnumerableRoute> for WtnsPathRoute {
+    type Error = WtnsPathRoute;
+
+    fn try_into(self) -> Result<EnumerableRoute, Self::Error> {
+        match self {
+            Self::Enumerable(tuple) => Ok(tuple),
+            _ => Err(self),
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum EitherRoute {
+    Left,
+    Right,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct EnumerableRoute(usize);

crates/sdk/src/transaction/final_transaction.rs

@@ -38,9 +38,12 @@ impl FinalTransaction {
     }
 
     pub fn add_input(&mut self, partial_input: PartialInput, required_sig: RequiredSignature) {
-        if let RequiredSignature::Witness(_) = required_sig {
-            panic!("Requested signature is not NativeEcdsa or None");
-        }
+        match required_sig {
+            RequiredSignature::Witness(_) | RequiredSignature::WitnessWithPath(_, _) => {
+                panic!("Requested signature is not NativeEcdsa or None")
+            }
+            _ => {}
+        };
 
         self.inputs.push(FinalInput {
             partial_input,
@@ -74,9 +77,12 @@ impl FinalTransaction {
         issuance_input: IssuanceInput,
         required_sig: RequiredSignature,
     ) -> AssetId {
-        if let RequiredSignature::Witness(_) = required_sig {
-            panic!("Requested signature is not NativeEcdsa or None");
-        }
+        match required_sig {
+            RequiredSignature::Witness(_) | RequiredSignature::WitnessWithPath(_, _) => {
+                panic!("Requested signature is not NativeEcdsa or None")
+            }
+            _ => {}
+        };
 
         let asset_id = AssetId::from_entropy(asset_entropy(&partial_input.outpoint(), issuance_input.asset_entropy));