Repository: ergoplatform/sigma-rust (language: rust)

Issue:
[150 SigmaUSD] Fix missing `ContextExtenion` support 
### ergo-lib
- In `signing::make_context` set `Context::extension` to `TransactionContext::spending_tx::inputs(self_index)::extension`;
- In `Prover::generate_proof` add `ext:ContextExtension` parameter and set `ProverResult::extension` to it;

### Wasm API
- Add `ContextExtension::new()` and `ContextExtension::with_pair(id: u8, value: &Constant) -> ContextExtension`;
- Add `UnsignedTransaction::with_context_ext(input_id: &BoxId, ext: &ContextExtension) -> Result<UnsignedTransaction, JsValue>`; 


Discussion / hints:
I'm making a start on this one

Fix patch:
diff --git a/bindings/ergo-lib-wasm/src/context_extension.rs b/bindings/ergo-lib-wasm/src/context_extension.rs
index d5b9f04d7..6ecd3f234 100644
--- a/bindings/ergo-lib-wasm/src/context_extension.rs
+++ b/bindings/ergo-lib-wasm/src/context_extension.rs
@@ -16,6 +16,17 @@ pub struct ContextExtension(
 
 #[wasm_bindgen]
 impl ContextExtension {
+    /// Create new ContextExtension instance
+    #[wasm_bindgen(constructor)]
+    pub fn new() -> Self {
+        Self(ergo_lib::ergotree_interpreter::sigma_protocol::prover::ContextExtension::empty())
+    }
+
+    /// Set the supplied pair in the ContextExtension
+    pub fn set_pair(&mut self, id: u8, value: &Constant) {
+        self.0.values.insert(id, value.clone().into());
+    }
+
     /// Returns the number of elements in the collection
     pub fn len(&self) -> usize {
         let wrapped: ergo_lib::ergotree_interpreter::sigma_protocol::prover::ContextExtension =
diff --git a/bindings/ergo-lib-wasm/src/transaction.rs b/bindings/ergo-lib-wasm/src/transaction.rs
index 25911eeee..25e03f3bb 100644
--- a/bindings/ergo-lib-wasm/src/transaction.rs
+++ b/bindings/ergo-lib-wasm/src/transaction.rs
@@ -2,7 +2,9 @@
 
 use crate::box_coll::ErgoBoxCandidates;
 use crate::box_coll::ErgoBoxes;
+use crate::context_extension::ContextExtension;
 use crate::data_input::DataInputs;
+use crate::ergo_box::BoxId;
 use crate::error_conversion::to_js;
 use crate::input::{Inputs, UnsignedInputs};
 use crate::json::TransactionJsonEip12;
@@ -153,6 +155,25 @@ pub struct UnsignedTransaction(chain::transaction::unsigned::UnsignedTransaction
 
 #[wasm_bindgen]
 impl UnsignedTransaction {
+    /// Consumes the calling UnsignedTransaction and returns a new UnsignedTransaction containing
+    /// the ContextExtension in the provided input box id or returns an error if the input box cannot be found.
+    /// After the call the calling UnsignedTransaction will be null.
+    pub fn with_input_context_ext(
+        mut self,
+        input_id: &BoxId,
+        ext: &ContextExtension,
+    ) -> Result<UnsignedTransaction, JsValue> {
+        let mut input = self
+            .0
+            .inputs
+            .iter_mut()
+            .find(|input| input.box_id == input_id.clone().into())
+            .ok_or_else(|| JsValue::from_str("Box input id not found"))?;
+        input.extension = ext.clone().into();
+
+        Ok(self.clone())
+    }
+
     /// Get id for transaction
     pub fn id(&self) -> TxId {
         self.0.id().into()
diff --git a/ergo-lib/src/wallet/signing.rs b/ergo-lib/src/wallet/signing.rs
index f9f1a1865..afec3b177 100644
--- a/ergo-lib/src/wallet/signing.rs
+++ b/ergo-lib/src/wallet/signing.rs
@@ -14,8 +14,8 @@ use crate::chain::{
 
 use ergotree_interpreter::eval::context::{Context, TxIoVec};
 use ergotree_interpreter::eval::env::Env;
+use ergotree_interpreter::sigma_protocol::prover::Prover;
 use ergotree_interpreter::sigma_protocol::prover::ProverError;
-use ergotree_interpreter::sigma_protocol::prover::{ContextExtension, Prover};
 use thiserror::Error;
 
 /// Errors on transaction signing
@@ -142,7 +142,17 @@ pub fn make_context(
                 .map(|b| Rc::new(b.clone()))
                 .ok_or(TxSigningError::InputBoxNotFound(idx))
         })?;
-
+    let extension = tx_ctx
+        .spending_tx
+        .inputs
+        .get(self_index)
+        .ok_or_else(|| {
+            TxSigningError::ContextError(
+                "self_index not found in spending transaction inputs".to_string(),
+            )
+        })?
+        .extension
+        .clone();
     Ok(Context {
         height,
         self_box: self_box_ir,
@@ -150,7 +160,7 @@ pub fn make_context(
         data_inputs: data_inputs_ir,
         inputs: inputs_ir,
         pre_header: state_ctx.pre_header.clone(),
-        extension: ContextExtension::empty(),
+        extension,
         headers: state_ctx.headers.clone(),
     })
 }
@@ -196,17 +206,14 @@ pub fn sign_reduced_transaction(
     let tx = reduced_tx.unsigned_tx.clone();
     let message_to_sign = tx.bytes_to_sign()?;
     let signed_inputs = tx.inputs.enumerated().try_mapped(|(idx, input)| {
+        let inputs = reduced_tx.reduced_inputs();
+        let reduced_input = inputs.get(idx).unwrap();
         prover
             .generate_proof(
-                reduced_tx
-                    .reduced_inputs()
-                    .get(idx)
-                    .unwrap()
-                    .reduction_result
-                    .sigma_prop
-                    .clone(),
+                reduced_input.reduction_result.sigma_prop.clone(),
                 message_to_sign.as_slice(),
                 &HintsBag::empty(),
+                reduced_input.extension.clone(),
             )
             .map(|proof| Input::new(input.box_id.clone(), proof.into()))
             .map_err(|e| TxSigningError::ProverError(e, idx))
@@ -224,6 +231,7 @@ mod tests {
     use super::*;
     use ergotree_interpreter::sigma_protocol::private_input::DlogProverInput;
     use ergotree_interpreter::sigma_protocol::private_input::PrivateInput;
+    use ergotree_interpreter::sigma_protocol::prover::ContextExtension;
     use ergotree_interpreter::sigma_protocol::prover::TestProver;
     use ergotree_interpreter::sigma_protocol::verifier::TestVerifier;
     use ergotree_interpreter::sigma_protocol::verifier::Verifier;
diff --git a/ergotree-interpreter/src/sigma_protocol/prover.rs b/ergotree-interpreter/src/sigma_protocol/prover.rs
index 49ce27683..7a4e3c2ae 100644
--- a/ergotree-interpreter/src/sigma_protocol/prover.rs
+++ b/ergotree-interpreter/src/sigma_protocol/prover.rs
@@ -118,9 +118,10 @@ pub trait Prover {
         hints_bag: &HintsBag,
     ) -> Result<ProverResult, ProverError> {
         let expr = tree.proposition()?;
+        let ctx_ext = ctx.extension.clone();
         let reduction_result =
             reduce_to_crypto(expr.as_ref(), env, ctx).map_err(ProverError::EvalError)?;
-        self.generate_proof(reduction_result.sigma_prop, message, hints_bag)
+        self.generate_proof(reduction_result.sigma_prop, message, hints_bag, ctx_ext)
     }
 
     /// Generate proofs for the given message for the given Sigma boolean expression
@@ -129,6 +130,7 @@ pub trait Prover {
         sigmabool: SigmaBoolean,
         message: &[u8],
         hints_bag: &HintsBag,
+        ctx_ext: ContextExtension,
     ) -> Result<ProverResult, ProverError> {
         let unchecked_tree_opt = match sigmabool {
             SigmaBoolean::TrivialProp(true) => Ok(None),
@@ -145,7 +147,7 @@ pub trait Prover {
         };
         Ok(ProverResult {
             proof,
-            extension: ContextExtension::empty(),
+            extension: ctx_ext,
         })
     }
 }
diff --git a/ergotree-ir/src/chain/ergo_box.rs b/ergotree-ir/src/chain/ergo_box.rs
index bd8b2d86c..a09ccac8d 100644
--- a/ergotree-ir/src/chain/ergo_box.rs
+++ b/ergotree-ir/src/chain/ergo_box.rs
@@ -286,7 +286,7 @@ impl From<ErgoBox> for super::json::ergo_box::ErgoBoxJson {
         }
     }
 }
-/// Contains the same fields as `ErgoBox`, except if transaction id and index,
+/// Contains the same fields as `ErgoBox`, except for transaction id and index,
 /// that will be calculated after full transaction formation.
 /// Use `ErgoBoxCandidateBuilder` from ergo-lib crate to create an instance.
 #[cfg_attr(feature = "json", derive(serde::Serialize, serde::Deserialize))]


Test patch: