From 4d12d84bb22ae1e9dcaf3556890a9e7bcb80ee58 Mon Sep 17 00:00:00 2001
From: jjy <jjyruby@gmail.com>
Date: Thu, 18 Apr 2024 16:00:59 +0800
Subject: [PATCH] feat: support read unlock witness from other index position

---
 contracts/rgbpp-lock/src/main.rs | 30 ++++++++++++++++++
 tests/src/tests.rs               | 54 ++++++++++++++++++++++++++++++--
 2 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/contracts/rgbpp-lock/src/main.rs b/contracts/rgbpp-lock/src/main.rs
index 08d1d48..3ed91c7 100644
--- a/contracts/rgbpp-lock/src/main.rs
+++ b/contracts/rgbpp-lock/src/main.rs
@@ -97,10 +97,40 @@ fn verify_outputs(config: &RGBPPConfig, btc_tx: &BTCTx) -> Result<(), Error> {
     Ok(())
 }
 
+fn load_unlock(index: usize, source: Source) -> Result<RGBPPUnlock, Error> {
+    let witness_args = load_witness_args(index, source)?;
+    match witness_args.lock().to_opt() {
+        Some(args) => {
+            let unlock =
+                RGBPPUnlock::from_slice(&args.raw_data()).map_err(|_| Error::BadRGBPPUnlock)?;
+            Ok(unlock)
+        }
+        None => Err(Error::ItemMissing),
+    }
+}
+
+/// Fetch unlock
+///
+/// In most cases, the RGBPPUnlock is located at group_input[0].lock.
+///
+/// For RGB++ cells which seal UTXOs in the same BTC transaction, the RGBPPUnlock is also the same.
+/// Therefore, to reduce duplicated witness, we can pass an index to group_input[0].lock.
+/// In such a situation, we load RGBPPUnlock from the index position.
 fn fetch_unlock_from_witness() -> Result<RGBPPUnlock, Error> {
     let witness_args = load_witness_args(0, Source::GroupInput)?;
     match witness_args.lock().to_opt() {
+        Some(args) if args.len() == 4 => {
+            // we assume args represents index of witness if len is 4
+            let index = u32::from_le_bytes(
+                args.raw_data()[..]
+                    .try_into()
+                    .map_err(|_| Error::BadRGBPPUnlock)?,
+            );
+            // load unlock from index
+            load_unlock(index as usize, Source::Input)
+        }
         Some(args) => {
+            // parse unlock
             let unlock =
                 RGBPPUnlock::from_slice(&args.raw_data()).map_err(|_| Error::BadRGBPPUnlock)?;
             Ok(unlock)
diff --git a/tests/src/tests.rs b/tests/src/tests.rs
index e00f3f5..9999b93 100644
--- a/tests/src/tests.rs
+++ b/tests/src/tests.rs
@@ -4,13 +4,15 @@ use crate::rgbpp::{
 };
 use crate::utils::TestScripts;
 use crate::{verify_and_dump_failed_tx, Loader};
-use ckb_testtool::ckb_types::prelude::Unpack;
 use ckb_testtool::context::Context;
 use rand::random;
 use rgbpp_core::bitcoin::MIN_BTC_TIME_LOCK_AFTER;
 use rgbpp_core::error::Error as RgbppError;
-use rgbpp_core::schemas::blockchain::Script;
-use rgbpp_core::schemas::{blockchain::CellOutput, ckb_gen_types::prelude::*};
+use rgbpp_core::schemas::blockchain::{BytesOpt, Script, WitnessArgs};
+use rgbpp_core::schemas::{
+    blockchain::CellOutput,
+    ckb_gen_types::{bytes::Bytes, prelude::*},
+};
 
 const MAX_CYCLES: u64 = 10_000_000;
 
@@ -299,3 +301,49 @@ fn test_btc_time_lock_with_incorrect_output() {
     let err = verify_and_dump_failed_tx(&context, &tx, MAX_CYCLES).expect_err("fail");
     assert_script_error(err, RgbppError::OutputCellMismatch);
 }
+
+#[test]
+fn test_rgbpp_unlock_with_index_witness() {
+    let loader = Loader::default();
+    let mut context = Context::default();
+
+    let scripts = TestScripts::setup(&loader, &mut context);
+
+    let tx = build_rgbpp_tx(
+        &mut context,
+        &scripts,
+        1000,
+        vec![
+            OutputDesc {
+                lock: LockDesc::Rgbpp,
+                amount: 300,
+            },
+            OutputDesc {
+                lock: LockDesc::Rgbpp,
+                amount: 700,
+            },
+        ],
+    );
+    let tx = context.complete_tx(tx);
+    let mut witnesses: Vec<_> = tx.witnesses().unpack();
+    assert!(!witnesses[0].is_empty(), "unlock witness must isn't empty");
+
+    // append unlock witness to the last
+    let index: u32 = witnesses.len() as u32;
+    witnesses.push(witnesses[0].clone());
+
+    witnesses[0] = {
+        let index = Bytes::copy_from_slice(&index.to_le_bytes());
+        WitnessArgs::new_builder()
+            .lock(BytesOpt::new_builder().set(Some(index.pack())).build())
+            .build()
+            .as_bytes()
+    };
+
+    let tx = tx
+        .as_advanced_builder()
+        .set_witnesses(witnesses.into_iter().map(|b| b.pack()).collect())
+        .build();
+
+    verify_and_dump_failed_tx(&context, &tx, MAX_CYCLES).expect("pass");
+}