Skip to main content

ckb_sdk/tx_builder/
mod.rs

1#![allow(clippy::doc_overindented_list_items)]
2
3pub mod acp;
4pub mod cheque;
5pub mod dao;
6pub mod omni_lock;
7pub mod transfer;
8pub mod udt;
9
10use std::collections::{HashMap, HashSet};
11#[cfg(not(target_arch = "wasm32"))]
12use std::sync::Arc;
13
14use anyhow::anyhow;
15#[cfg(not(target_arch = "wasm32"))]
16use ckb_chain_spec::consensus::Consensus;
17#[cfg(not(target_arch = "wasm32"))]
18use ckb_script::{TransactionScriptsVerifier, TxVerifyEnv};
19#[cfg(not(target_arch = "wasm32"))]
20use ckb_traits::{CellDataProvider, ExtensionProvider, HeaderProvider};
21#[cfg(not(target_arch = "wasm32"))]
22use ckb_types::core::cell::{CellProvider, HeaderChecker};
23#[cfg(not(target_arch = "wasm32"))]
24use ckb_types::core::{cell::resolve_transaction, HeaderView};
25use ckb_types::{
26    core::{error::OutPointError, Capacity, CapacityError, FeeRate, TransactionView},
27    packed::{Byte32, CellInput, CellOutput, Script, WitnessArgs},
28    prelude::*,
29};
30use thiserror::Error;
31
32use crate::types::ScriptGroup;
33use crate::types::{HumanCapacity, ScriptId};
34use crate::unlock::{ScriptUnlocker, UnlockError};
35use crate::util::calculate_dao_maximum_withdraw4;
36use crate::{constants::DAO_TYPE_HASH, NetworkType};
37use crate::{
38    traits::{
39        CellCollector, CellCollectorError, CellDepResolver, CellQueryOptions, HeaderDepResolver,
40        TransactionDependencyError, TransactionDependencyProvider, ValueRangeOption,
41    },
42    RpcError,
43};
44
45/// Transaction builder errors
46#[derive(Error, Debug)]
47pub enum TxBuilderError {
48    #[error("invalid parameter: `{0}`")]
49    InvalidParameter(anyhow::Error),
50
51    #[error("transaction dependency provider error: `{0}`")]
52    TxDep(#[from] TransactionDependencyError),
53    #[error("ChangeIndex alread set: `{0}`")]
54    ChangeIndex(usize),
55
56    #[error("cell collector error: `{0}`")]
57    CellCollector(#[from] CellCollectorError),
58
59    #[error("balance capacity error: `{0}`")]
60    BalanceCapacity(#[from] BalanceTxCapacityError),
61
62    #[error("resolve cell dep failed: `{0}`")]
63    ResolveCellDepFailed(Script),
64
65    #[error("resolve header dep by transaction hash failed: `{0}`")]
66    ResolveHeaderDepByTxHashFailed(Byte32),
67
68    #[error("resolve header dep by block number failed: `{0}`")]
69    ResolveHeaderDepByNumberFailed(u64),
70
71    #[error("unlock error: `{0}`")]
72    Unlock(#[from] UnlockError),
73
74    #[error("build_balance_unlocked exceed max loop times, current is: `{0}`")]
75    ExceedCycleMaxLoopTimes(u32),
76    #[error("witness idx `{0}` is out of bound `{1}")]
77    WitnessOutOfBound(usize, usize),
78    #[error("unsupported networktype `{0}")]
79    UnsupportedNetworkType(NetworkType),
80    #[error("can not find specifed output to put small change")]
81    NoOutputForSmallChange,
82
83    #[error("other error: `{0}`")]
84    Other(anyhow::Error),
85}
86
87/// Transaction Builder interface
88#[async_trait::async_trait]
89pub trait TxBuilder: Send + Sync {
90    /// Build base transaction
91    async fn build_base_async(
92        &self,
93        cell_collector: &mut dyn CellCollector,
94        cell_dep_resolver: &dyn CellDepResolver,
95        header_dep_resolver: &dyn HeaderDepResolver,
96        tx_dep_provider: &dyn TransactionDependencyProvider,
97    ) -> Result<TransactionView, TxBuilderError>;
98    #[cfg(not(target_arch = "wasm32"))]
99    /// Build base transaction
100    fn build_base(
101        &self,
102        cell_collector: &mut dyn CellCollector,
103        cell_dep_resolver: &dyn CellDepResolver,
104        header_dep_resolver: &dyn HeaderDepResolver,
105        tx_dep_provider: &dyn TransactionDependencyProvider,
106    ) -> Result<TransactionView, TxBuilderError> {
107        crate::rpc::block_on(self.build_base_async(
108            cell_collector,
109            cell_dep_resolver,
110            header_dep_resolver,
111            tx_dep_provider,
112        ))
113    }
114
115    /// Build balanced transaction that ready to sign:
116    ///  * Build base transaction
117    ///  * Fill placeholder witness for lock script
118    ///  * balance the capacity
119    async fn build_balanced_async(
120        &self,
121        cell_collector: &mut dyn CellCollector,
122        cell_dep_resolver: &dyn CellDepResolver,
123        header_dep_resolver: &dyn HeaderDepResolver,
124        tx_dep_provider: &dyn TransactionDependencyProvider,
125        balancer: &CapacityBalancer,
126        unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
127    ) -> Result<TransactionView, TxBuilderError> {
128        let base_tx = self
129            .build_base_async(
130                cell_collector,
131                cell_dep_resolver,
132                header_dep_resolver,
133                tx_dep_provider,
134            )
135            .await?;
136        let (tx_filled_witnesses, _) =
137            fill_placeholder_witnesses_async(base_tx, tx_dep_provider, unlockers).await?;
138        Ok(balance_tx_capacity_async(
139            &tx_filled_witnesses,
140            balancer,
141            cell_collector,
142            tx_dep_provider,
143            cell_dep_resolver,
144            header_dep_resolver,
145        )
146        .await?)
147    }
148    #[cfg(not(target_arch = "wasm32"))]
149    fn build_balanced(
150        &self,
151        cell_collector: &mut dyn CellCollector,
152        cell_dep_resolver: &dyn CellDepResolver,
153        header_dep_resolver: &dyn HeaderDepResolver,
154        tx_dep_provider: &dyn TransactionDependencyProvider,
155        balancer: &CapacityBalancer,
156        unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
157    ) -> Result<TransactionView, TxBuilderError> {
158        crate::rpc::block_on(self.build_balanced_async(
159            cell_collector,
160            cell_dep_resolver,
161            header_dep_resolver,
162            tx_dep_provider,
163            balancer,
164            unlockers,
165        ))
166    }
167
168    /// Build unlocked transaction that ready to send or for further unlock:
169    ///   * build base transaction
170    ///   * balance the capacity
171    ///   * unlock(sign) the transaction
172    ///
173    /// Return value:
174    ///   * The built transaction
175    ///   * The script groups that not unlocked by given `unlockers`
176    ///
177    #[cfg(not(target_arch = "wasm32"))]
178    fn build_unlocked(
179        &self,
180        cell_collector: &mut dyn CellCollector,
181        cell_dep_resolver: &dyn CellDepResolver,
182        header_dep_resolver: &dyn HeaderDepResolver,
183        tx_dep_provider: &dyn TransactionDependencyProvider,
184        balancer: &CapacityBalancer,
185        unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
186    ) -> Result<(TransactionView, Vec<ScriptGroup>), TxBuilderError> {
187        let balanced_tx = self.build_balanced(
188            cell_collector,
189            cell_dep_resolver,
190            header_dep_resolver,
191            tx_dep_provider,
192            balancer,
193            unlockers,
194        )?;
195        Ok(unlock_tx(balanced_tx, tx_dep_provider, unlockers)?)
196    }
197
198    #[cfg(target_arch = "wasm32")]
199    async fn build_unlocked_async(
200        &self,
201        cell_collector: &mut dyn CellCollector,
202        cell_dep_resolver: &dyn CellDepResolver,
203        header_dep_resolver: &dyn HeaderDepResolver,
204        tx_dep_provider: &dyn TransactionDependencyProvider,
205        balancer: &CapacityBalancer,
206        unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
207    ) -> Result<(TransactionView, Vec<ScriptGroup>), TxBuilderError> {
208        let balanced_tx = self
209            .build_balanced_async(
210                cell_collector,
211                cell_dep_resolver,
212                header_dep_resolver,
213                tx_dep_provider,
214                balancer,
215                unlockers,
216            )
217            .await?;
218        Ok(unlock_tx_async(balanced_tx, tx_dep_provider, unlockers).await?)
219    }
220
221    /// Build unlocked transaction that ready to send or for further unlock, it's similar to `build_unlocked`,
222    /// except it will try to check the consumed cycles limitation:
223    /// If all input unlocked, and transaction fee can not meet the required transaction fee rate because of a big estimated cycles,
224    /// it will tweak the change cell capacity or collect more cells to balance the transaction.
225    ///
226    /// Return value:
227    ///   * The built transaction
228    ///   * The script groups that not unlocked by given `unlockers`
229    #[cfg(not(target_arch = "wasm32"))]
230    fn build_balance_unlocked(
231        &self,
232        cell_collector: &mut dyn CellCollector,
233        cell_dep_resolver: &dyn CellDepResolver,
234        header_dep_resolver: &dyn HeaderDepResolver,
235        tx_dep_provider: &'static dyn TransactionDependencyProvider,
236        balancer: &CapacityBalancer,
237        unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
238    ) -> Result<(TransactionView, Vec<ScriptGroup>), TxBuilderError> {
239        let base_tx = self.build_base(
240            cell_collector,
241            cell_dep_resolver,
242            header_dep_resolver,
243            tx_dep_provider,
244        )?;
245        let (tx_filled_witnesses, _) =
246            fill_placeholder_witnesses(base_tx, tx_dep_provider, unlockers)?;
247        let (balanced_tx, mut change_idx) = rebalance_tx_capacity(
248            &tx_filled_witnesses,
249            balancer,
250            cell_collector,
251            tx_dep_provider,
252            cell_dep_resolver,
253            header_dep_resolver,
254            0,
255            None,
256        )?;
257        let (mut tx, unlocked_group) = unlock_tx(balanced_tx, tx_dep_provider, unlockers)?;
258        if unlocked_group.is_empty() {
259            let mut ready = false;
260            const MAX_LOOP_TIMES: u32 = 16;
261            let mut n = 0;
262            while !ready && n < MAX_LOOP_TIMES {
263                n += 1;
264
265                let (new_tx, new_change_idx, ok) = balancer.check_cycle_fee(
266                    tx,
267                    cell_collector,
268                    tx_dep_provider,
269                    cell_dep_resolver,
270                    header_dep_resolver,
271                    change_idx,
272                )?;
273                tx = new_tx;
274                ready = ok;
275                change_idx = new_change_idx;
276                if !ready {
277                    let (new_tx, _) = unlock_tx(tx, tx_dep_provider, unlockers)?;
278                    tx = new_tx
279                }
280            }
281            if !ready && n >= MAX_LOOP_TIMES {
282                return Err(TxBuilderError::ExceedCycleMaxLoopTimes(n));
283            }
284        }
285        Ok((tx, unlocked_group))
286    }
287    // #[cfg(target_arch = "wasm32")]
288    #[cfg(not(target_arch = "wasm32"))]
289    async fn build_balance_unlocked_async(
290        &self,
291        cell_collector: &mut dyn CellCollector,
292        cell_dep_resolver: &dyn CellDepResolver,
293        header_dep_resolver: &dyn HeaderDepResolver,
294        tx_dep_provider: &'static dyn TransactionDependencyProvider,
295        balancer: &CapacityBalancer,
296        unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
297    ) -> Result<(TransactionView, Vec<ScriptGroup>), TxBuilderError> {
298        let base_tx = self
299            .build_base_async(
300                cell_collector,
301                cell_dep_resolver,
302                header_dep_resolver,
303                tx_dep_provider,
304            )
305            .await?;
306        let (tx_filled_witnesses, _) =
307            fill_placeholder_witnesses_async(base_tx, tx_dep_provider, unlockers).await?;
308        let (balanced_tx, mut change_idx) = rebalance_tx_capacity_async(
309            &tx_filled_witnesses,
310            balancer,
311            cell_collector,
312            tx_dep_provider,
313            cell_dep_resolver,
314            header_dep_resolver,
315            0,
316            None,
317        )
318        .await?;
319        let (mut tx, unlocked_group) =
320            unlock_tx_async(balanced_tx, tx_dep_provider, unlockers).await?;
321        if unlocked_group.is_empty() {
322            let mut ready = false;
323            const MAX_LOOP_TIMES: u32 = 16;
324            let mut n = 0;
325            while !ready && n < MAX_LOOP_TIMES {
326                n += 1;
327
328                let (new_tx, new_change_idx, ok) = balancer
329                    .check_cycle_fee_async(
330                        tx,
331                        cell_collector,
332                        tx_dep_provider,
333                        cell_dep_resolver,
334                        header_dep_resolver,
335                        change_idx,
336                    )
337                    .await?;
338                tx = new_tx;
339                ready = ok;
340                change_idx = new_change_idx;
341                if !ready {
342                    let (new_tx, _) = unlock_tx_async(tx, tx_dep_provider, unlockers).await?;
343                    tx = new_tx
344                }
345            }
346            if !ready && n >= MAX_LOOP_TIMES {
347                return Err(TxBuilderError::ExceedCycleMaxLoopTimes(n));
348            }
349        }
350        Ok((tx, unlocked_group))
351    }
352}
353
354#[derive(Debug, Eq, PartialEq, Hash, Clone)]
355pub enum TransferAction {
356    /// This action will crate a new cell, typecial lock script: cheque, sighash, multisig
357    Create,
358    /// This action will query the exists cell and update the amount, typecial lock script: acp
359    Update,
360}
361
362#[derive(Error, Debug)]
363pub enum TransactionFeeError {
364    #[error("transaction dependency provider error: `{0}`")]
365    TxDep(#[from] TransactionDependencyError),
366
367    #[error("header dependency provider error: `{0}`")]
368    HeaderDep(#[from] anyhow::Error),
369
370    #[error("out point error: `{0}`")]
371    OutPoint(#[from] OutPointError),
372
373    #[error("unexpected dao withdraw cell in inputs")]
374    UnexpectedDaoWithdrawInput,
375
376    #[error("capacity error: `{0}`")]
377    CapacityError(#[from] CapacityError),
378
379    #[error("capacity sub overflow, delta: `{0}`")]
380    CapacityOverflow(u64),
381}
382
383/// Calculate the actual transaction fee of the transaction, include dao
384/// withdraw capacity.
385#[allow(clippy::unnecessary_lazy_evaluations)]
386#[cfg(not(target_arch = "wasm32"))]
387pub fn tx_fee(
388    tx: TransactionView,
389    tx_dep_provider: &dyn TransactionDependencyProvider,
390    header_dep_resolver: &dyn HeaderDepResolver,
391) -> Result<u64, TransactionFeeError> {
392    let mut input_total: u64 = 0;
393    for input in tx.inputs() {
394        let mut is_withdraw = false;
395        let since: u64 = input.since().unpack();
396        let cell = tx_dep_provider.get_cell(&input.previous_output())?;
397        if since != 0 {
398            if let Some(type_script) = cell.type_().to_opt() {
399                if type_script.code_hash().as_slice() == DAO_TYPE_HASH.as_bytes() {
400                    is_withdraw = true;
401                }
402            }
403        }
404        let capacity: u64 = if is_withdraw {
405            let tx_hash = input.previous_output().tx_hash();
406            let prepare_header = header_dep_resolver
407                .resolve_by_tx(&tx_hash)
408                .map_err(TransactionFeeError::HeaderDep)?
409                .ok_or_else(|| {
410                    TransactionFeeError::HeaderDep(anyhow!(
411                        "resolve prepare header by transaction hash failed: {}",
412                        tx_hash
413                    ))
414                })?;
415            let data = tx_dep_provider.get_cell_data(&input.previous_output())?;
416            assert_eq!(data.len(), 8);
417            let deposit_number = {
418                let mut number_bytes = [0u8; 8];
419                number_bytes.copy_from_slice(data.as_ref());
420                u64::from_le_bytes(number_bytes)
421            };
422            let deposit_header = header_dep_resolver
423                .resolve_by_number(deposit_number)
424                .map_err(TransactionFeeError::HeaderDep)?
425                .ok_or_else(|| {
426                    TransactionFeeError::HeaderDep(anyhow!(
427                        "resolve deposit header by block number failed: {}",
428                        deposit_number
429                    ))
430                })?;
431            let occupied_capacity = cell
432                .occupied_capacity(Capacity::bytes(data.len()).unwrap())
433                .unwrap();
434            calculate_dao_maximum_withdraw4(
435                &deposit_header,
436                &prepare_header,
437                &cell,
438                occupied_capacity.as_u64(),
439            )
440        } else {
441            cell.capacity().unpack()
442        };
443        input_total += capacity;
444    }
445    let output_total = tx.outputs_capacity()?.as_u64();
446    #[allow(clippy::unnecessary_lazy_evaluations)]
447    input_total
448        .checked_sub(output_total)
449        .ok_or_else(|| TransactionFeeError::CapacityOverflow(output_total - input_total))
450}
451
452/// Calculate the actual transaction fee of the transaction, include dao
453/// withdraw capacity.
454#[allow(clippy::unnecessary_lazy_evaluations)]
455pub async fn tx_fee_async(
456    tx: TransactionView,
457    tx_dep_provider: &dyn TransactionDependencyProvider,
458    header_dep_resolver: &dyn HeaderDepResolver,
459) -> Result<u64, TransactionFeeError> {
460    let mut input_total: u64 = 0;
461    for input in tx.inputs() {
462        let mut is_withdraw = false;
463        let since: u64 = input.since().unpack();
464        let cell = tx_dep_provider
465            .get_cell_async(&input.previous_output())
466            .await?;
467        if since != 0 {
468            if let Some(type_script) = cell.type_().to_opt() {
469                if type_script.code_hash().as_slice() == DAO_TYPE_HASH.as_bytes() {
470                    is_withdraw = true;
471                }
472            }
473        }
474        let capacity: u64 = if is_withdraw {
475            let tx_hash = input.previous_output().tx_hash();
476            let prepare_header = header_dep_resolver
477                .resolve_by_tx_async(&tx_hash)
478                .await
479                .map_err(TransactionFeeError::HeaderDep)?
480                .ok_or_else(|| {
481                    TransactionFeeError::HeaderDep(anyhow!(
482                        "resolve prepare header by transaction hash failed: {}",
483                        tx_hash
484                    ))
485                })?;
486            let data = tx_dep_provider
487                .get_cell_data_async(&input.previous_output())
488                .await?;
489            assert_eq!(data.len(), 8);
490            let deposit_number = {
491                let mut number_bytes = [0u8; 8];
492                number_bytes.copy_from_slice(data.as_ref());
493                u64::from_le_bytes(number_bytes)
494            };
495            let deposit_header = header_dep_resolver
496                .resolve_by_number_async(deposit_number)
497                .await
498                .map_err(TransactionFeeError::HeaderDep)?
499                .ok_or_else(|| {
500                    TransactionFeeError::HeaderDep(anyhow!(
501                        "resolve deposit header by block number failed: {}",
502                        deposit_number
503                    ))
504                })?;
505            let occupied_capacity = cell
506                .occupied_capacity(Capacity::bytes(data.len()).unwrap())
507                .unwrap();
508            calculate_dao_maximum_withdraw4(
509                &deposit_header,
510                &prepare_header,
511                &cell,
512                occupied_capacity.as_u64(),
513            )
514        } else {
515            cell.capacity().unpack()
516        };
517        input_total += capacity;
518    }
519    let output_total = tx.outputs_capacity()?.as_u64();
520    #[allow(clippy::unnecessary_lazy_evaluations)]
521    input_total
522        .checked_sub(output_total)
523        .ok_or_else(|| TransactionFeeError::CapacityOverflow(output_total - input_total))
524}
525
526#[derive(Debug, Clone)]
527pub enum SinceSource {
528    /// The vaule in the tuple is offset of the args, and the `since` is stored in `lock.args[offset..offset+8]`
529    LockArgs(usize),
530    /// raw since value
531    Value(u64),
532}
533
534impl Default for SinceSource {
535    fn default() -> SinceSource {
536        SinceSource::Value(0)
537    }
538}
539
540/// Provide capacity locked by a list of lock scripts.
541///
542/// The cells collected by `lock_script` will filter out those have type script
543/// or data length is not `0` or is not mature.
544#[derive(Debug, Clone)]
545pub struct CapacityProvider {
546    /// The lock scripts provider capacity. The second field of the tuple is the
547    /// placeholder witness of the lock script.
548    pub lock_scripts: Vec<(Script, WitnessArgs, SinceSource)>,
549}
550
551impl CapacityProvider {
552    /// create a new capacity provider.
553    pub fn new(lock_scripts: Vec<(Script, WitnessArgs, SinceSource)>) -> CapacityProvider {
554        CapacityProvider { lock_scripts }
555    }
556
557    /// create a new capacity provider with the default since source.
558    pub fn new_simple(lock_scripts: Vec<(Script, WitnessArgs)>) -> CapacityProvider {
559        let lock_scripts = lock_scripts
560            .into_iter()
561            .map(|(script, witness)| (script, witness, SinceSource::default()))
562            .collect();
563        CapacityProvider { lock_scripts }
564    }
565}
566
567#[derive(Error, Debug)]
568pub enum BalanceTxCapacityError {
569    #[error("calculate transaction fee error: `{0}`")]
570    TxFee(#[from] TransactionFeeError),
571
572    #[error("transaction dependency provider error: `{0}`")]
573    TxDep(#[from] TransactionDependencyError),
574
575    #[error("capacity not enough: `{0}`")]
576    CapacityNotEnough(String),
577
578    #[error("Force small change as fee failed, fee: `{0}`")]
579    ForceSmallChangeAsFeeFailed(u64),
580
581    #[error("empty capacity provider")]
582    EmptyCapacityProvider,
583
584    #[error("cell collector error: `{0}`")]
585    CellCollector(#[from] CellCollectorError),
586
587    #[error("resolve cell dep failed: `{0}`")]
588    ResolveCellDepFailed(Script),
589
590    #[error("invalid witness args: `{0}`")]
591    InvalidWitnessArgs(anyhow::Error),
592
593    #[error("Fail to parse since value from args, offset: `{0}`, args length: `{1}`")]
594    InvalidSinceValue(usize, usize),
595
596    #[error("change index not found at given index: `{0}`")]
597    ChangeIndexNotFound(usize),
598
599    #[error("Fail to estimate_cycles: `{0}`")]
600    FailEstimateCycles(#[from] RpcError),
601
602    #[error("verify script error: {0}")]
603    VerifyScript(String),
604
605    #[error("should not try to rebalance, orignal fee {0}, required fee: {1},")]
606    AlreadyBalance(u64, u64),
607}
608
609/// Transaction capacity balancer config.
610///
611/// CapacityBalancer will try to balance the transaction capacity by adding inputs from CapacityProvider.
612#[derive(Debug, Clone)]
613pub struct CapacityBalancer {
614    pub fee_rate: FeeRate,
615
616    /// Search cell by this lock script and filter out cells with data or with
617    /// type script or not mature.
618    pub capacity_provider: CapacityProvider,
619
620    /// Change cell's lock script if `None` use capacity_provider's first lock script
621    pub change_lock_script: Option<Script>,
622
623    /// When there is no more inputs for create a change cell to balance the
624    /// transaction capacity, force the addition capacity as fee, the value is
625    /// actual maximum transaction fee.
626    pub force_small_change_as_fee: Option<u64>,
627}
628
629impl CapacityBalancer {
630    /// Create a new balancer.
631    ///
632    /// # Arguments
633    ///
634    /// * `capacity_provider` - Use live cells with this lock script as capacity provider.
635    /// * `placeholder_witness` - The witness used as a placeholder when adding new inputs.
636    ///     This placeholder ensures that the transaction size does not increase after signing,
637    ///     thus maintaining the validity of fee calculation.
638    /// * `fee_rate` - The fee rate used to calculate the transaction fee.
639    pub fn new_simple(
640        capacity_provider: Script,
641        placeholder_witness: WitnessArgs,
642        fee_rate: u64,
643    ) -> CapacityBalancer {
644        CapacityBalancer {
645            fee_rate: FeeRate::from_u64(fee_rate),
646            capacity_provider: CapacityProvider::new_simple(vec![(
647                capacity_provider,
648                placeholder_witness,
649            )]),
650            change_lock_script: None,
651            force_small_change_as_fee: None,
652        }
653    }
654
655    /// Create new simple capacity balancer with since source.
656    pub fn new_simple_with_since(
657        capacity_provider: Script,
658        placeholder_witness: WitnessArgs,
659        since_source: SinceSource,
660        fee_rate: u64,
661    ) -> CapacityBalancer {
662        CapacityBalancer {
663            fee_rate: FeeRate::from_u64(fee_rate),
664            capacity_provider: CapacityProvider::new(vec![(
665                capacity_provider,
666                placeholder_witness,
667                since_source,
668            )]),
669            change_lock_script: None,
670            force_small_change_as_fee: None,
671        }
672    }
673
674    pub fn new_with_provider(fee_rate: u64, capacity_provider: CapacityProvider) -> Self {
675        CapacityBalancer {
676            fee_rate: FeeRate::from_u64(fee_rate),
677            capacity_provider,
678            change_lock_script: None,
679            force_small_change_as_fee: None,
680        }
681    }
682
683    /// Set or clear the force_small_change_as_fee
684    pub fn set_max_fee(&mut self, max_fee: Option<u64>) {
685        self.force_small_change_as_fee = max_fee;
686    }
687    #[cfg(not(target_arch = "wasm32"))]
688    pub fn balance_tx_capacity(
689        &mut self,
690        tx: &TransactionView,
691        cell_collector: &mut dyn CellCollector,
692        tx_dep_provider: &dyn TransactionDependencyProvider,
693        cell_dep_resolver: &dyn CellDepResolver,
694        header_dep_resolver: &dyn HeaderDepResolver,
695    ) -> Result<TransactionView, BalanceTxCapacityError> {
696        balance_tx_capacity(
697            tx,
698            self,
699            cell_collector,
700            tx_dep_provider,
701            cell_dep_resolver,
702            header_dep_resolver,
703        )
704    }
705
706    pub async fn balance_tx_capacity_async(
707        &mut self,
708        tx: &TransactionView,
709        cell_collector: &mut dyn CellCollector,
710        tx_dep_provider: &dyn TransactionDependencyProvider,
711        cell_dep_resolver: &dyn CellDepResolver,
712        header_dep_resolver: &dyn HeaderDepResolver,
713    ) -> Result<TransactionView, BalanceTxCapacityError> {
714        balance_tx_capacity_async(
715            tx,
716            self,
717            cell_collector,
718            tx_dep_provider,
719            cell_dep_resolver,
720            header_dep_resolver,
721        )
722        .await
723    }
724
725    #[allow(clippy::too_many_arguments)]
726    #[cfg(not(target_arch = "wasm32"))]
727    pub fn rebalance_tx_capacity(
728        &self,
729        tx: &TransactionView,
730        cell_collector: &mut dyn CellCollector,
731        tx_dep_provider: &dyn TransactionDependencyProvider,
732        cell_dep_resolver: &dyn CellDepResolver,
733        header_dep_resolver: &dyn HeaderDepResolver,
734        accepted_min_fee: u64,
735        change_index: Option<usize>,
736    ) -> Result<(TransactionView, Option<usize>), BalanceTxCapacityError> {
737        if let Some(idx) = change_index {
738            let output = tx
739                .outputs()
740                .get(idx)
741                .ok_or(BalanceTxCapacityError::ChangeIndexNotFound(idx))?;
742            let base_change_occupied_capacity = output
743                .occupied_capacity(Capacity::zero())
744                .expect("init change occupied capacity")
745                .as_u64();
746            let output_header_extra = 4 + 4 + 4;
747            // NOTE: extra_min_fee +1 is for `FeeRate::fee` round
748            let extra_min_fee = self
749                .fee_rate
750                .fee(output.as_slice().len() as u64 + output_header_extra)
751                .as_u64()
752                + 1;
753            let original_fee = tx_fee(tx.clone(), tx_dep_provider, header_dep_resolver)?;
754            if original_fee >= accepted_min_fee {
755                return Err(BalanceTxCapacityError::AlreadyBalance(
756                    original_fee,
757                    accepted_min_fee,
758                ));
759            }
760            let extra_fee = accepted_min_fee - original_fee;
761            // The extra capacity (delta - extra_min_fee) is enough to hold the change cell.
762            let original_capacity: u64 = output.capacity().unpack();
763            if original_capacity >= base_change_occupied_capacity + extra_min_fee + extra_fee {
764                let output = output
765                    .as_builder()
766                    .capacity(original_capacity - extra_fee)
767                    .build();
768                let mut outputs: Vec<_> = tx.outputs().into_iter().collect();
769                outputs[idx] = output;
770                let tx = tx.as_advanced_builder().set_outputs(outputs).build();
771                return Ok((tx, change_index));
772            };
773        }
774
775        rebalance_tx_capacity(
776            tx,
777            self,
778            cell_collector,
779            tx_dep_provider,
780            cell_dep_resolver,
781            header_dep_resolver,
782            accepted_min_fee,
783            change_index,
784        )
785    }
786    #[allow(clippy::too_many_arguments)]
787    pub async fn rebalance_tx_capacity_async(
788        &self,
789        tx: &TransactionView,
790        cell_collector: &mut dyn CellCollector,
791        tx_dep_provider: &dyn TransactionDependencyProvider,
792        cell_dep_resolver: &dyn CellDepResolver,
793        header_dep_resolver: &dyn HeaderDepResolver,
794        accepted_min_fee: u64,
795        change_index: Option<usize>,
796    ) -> Result<(TransactionView, Option<usize>), BalanceTxCapacityError> {
797        if let Some(idx) = change_index {
798            let output = tx
799                .outputs()
800                .get(idx)
801                .ok_or(BalanceTxCapacityError::ChangeIndexNotFound(idx))?;
802            let base_change_occupied_capacity = output
803                .occupied_capacity(Capacity::zero())
804                .expect("init change occupied capacity")
805                .as_u64();
806            let output_header_extra = 4 + 4 + 4;
807            // NOTE: extra_min_fee +1 is for `FeeRate::fee` round
808            let extra_min_fee = self
809                .fee_rate
810                .fee(output.as_slice().len() as u64 + output_header_extra)
811                .as_u64()
812                + 1;
813            let original_fee =
814                tx_fee_async(tx.clone(), tx_dep_provider, header_dep_resolver).await?;
815            if original_fee >= accepted_min_fee {
816                return Err(BalanceTxCapacityError::AlreadyBalance(
817                    original_fee,
818                    accepted_min_fee,
819                ));
820            }
821            let extra_fee = accepted_min_fee - original_fee;
822            // The extra capacity (delta - extra_min_fee) is enough to hold the change cell.
823            let original_capacity: u64 = output.capacity().unpack();
824            if original_capacity >= base_change_occupied_capacity + extra_min_fee + extra_fee {
825                let output = output
826                    .as_builder()
827                    .capacity(original_capacity - extra_fee)
828                    .build();
829                let mut outputs: Vec<_> = tx.outputs().into_iter().collect();
830                outputs[idx] = output;
831                let tx = tx.as_advanced_builder().set_outputs(outputs).build();
832                return Ok((tx, change_index));
833            };
834        }
835
836        rebalance_tx_capacity_async(
837            tx,
838            self,
839            cell_collector,
840            tx_dep_provider,
841            cell_dep_resolver,
842            header_dep_resolver,
843            accepted_min_fee,
844            change_index,
845        )
846        .await
847    }
848    #[cfg(not(target_arch = "wasm32"))]
849    pub fn check_cycle_fee(
850        &self,
851        tx: TransactionView,
852        cell_collector: &mut dyn CellCollector,
853        tx_dep_provider: &'static dyn TransactionDependencyProvider,
854        cell_dep_resolver: &dyn CellDepResolver,
855        header_dep_resolver: &dyn HeaderDepResolver,
856        change_index: Option<usize>,
857    ) -> Result<(TransactionView, Option<usize>, bool), BalanceTxCapacityError> {
858        let cycle_resolver = CycleResolver::new(tx_dep_provider);
859        let cycle = cycle_resolver.estimate_cycles(&tx)?;
860        let cycle_size = (cycle as f64 * bytes_per_cycle()) as usize;
861        let serialized_size = tx.data().as_reader().serialized_size_in_block();
862        if serialized_size >= cycle_size {
863            return Ok((tx, None, true));
864        }
865        let fee = tx_fee(tx.clone(), tx_dep_provider, header_dep_resolver).unwrap();
866        let cycle_fee = self.fee_rate.fee(cycle_size as u64).as_u64();
867
868        if fee >= cycle_fee {
869            return Ok((tx, None, true));
870        }
871
872        let (tx, idx) = self.rebalance_tx_capacity(
873            &tx,
874            cell_collector,
875            tx_dep_provider,
876            cell_dep_resolver,
877            header_dep_resolver,
878            cycle_fee,
879            change_index,
880        )?;
881        Ok((tx, idx, false))
882    }
883    #[cfg(not(target_arch = "wasm32"))]
884    pub async fn check_cycle_fee_async(
885        &self,
886        tx: TransactionView,
887        cell_collector: &mut dyn CellCollector,
888        tx_dep_provider: &'static dyn TransactionDependencyProvider,
889        cell_dep_resolver: &dyn CellDepResolver,
890        header_dep_resolver: &dyn HeaderDepResolver,
891        change_index: Option<usize>,
892    ) -> Result<(TransactionView, Option<usize>, bool), BalanceTxCapacityError> {
893        let cycle_resolver = CycleResolver::new(tx_dep_provider);
894        let cycle = cycle_resolver.estimate_cycles(&tx)?;
895        let cycle_size = (cycle as f64 * bytes_per_cycle()) as usize;
896        let serialized_size = tx.data().as_reader().serialized_size_in_block();
897        if serialized_size >= cycle_size {
898            return Ok((tx, None, true));
899        }
900        let fee = tx_fee_async(tx.clone(), tx_dep_provider, header_dep_resolver)
901            .await
902            .unwrap();
903        let cycle_fee = self.fee_rate.fee(cycle_size as u64).as_u64();
904
905        if fee >= cycle_fee {
906            return Ok((tx, None, true));
907        }
908
909        let (tx, idx) = self
910            .rebalance_tx_capacity_async(
911                &tx,
912                cell_collector,
913                tx_dep_provider,
914                cell_dep_resolver,
915                header_dep_resolver,
916                cycle_fee,
917                change_index,
918            )
919            .await?;
920        Ok((tx, idx, false))
921    }
922}
923
924const DEFAULT_BYTES_PER_CYCLE: f64 = 0.000_170_571_4;
925pub const fn bytes_per_cycle() -> f64 {
926    DEFAULT_BYTES_PER_CYCLE
927}
928#[cfg(not(target_arch = "wasm32"))]
929pub struct CycleResolver<DL> {
930    tx_dep_provider: DL,
931    tip_header: HeaderView,
932    consensus: Arc<Consensus>,
933}
934#[cfg(not(target_arch = "wasm32"))]
935impl<
936        DL: CellDataProvider
937            + HeaderProvider
938            + ExtensionProvider
939            + CellProvider
940            + HeaderChecker
941            + Send
942            + Sync
943            + Clone
944            + 'static,
945    > CycleResolver<DL>
946{
947    pub fn new(tx_dep_provider: DL) -> Self {
948        CycleResolver {
949            tx_dep_provider,
950            tip_header: HeaderView::new_advanced_builder().build(), // TODO
951            consensus: Default::default(),                          // TODO
952        }
953    }
954
955    fn estimate_cycles(&self, tx: &TransactionView) -> Result<u64, BalanceTxCapacityError> {
956        let rtx = resolve_transaction(
957            tx.clone(),
958            &mut HashSet::new(),
959            &self.tx_dep_provider,
960            &self.tx_dep_provider,
961        )
962        .map_err(|err| {
963            BalanceTxCapacityError::VerifyScript(format!("Resolve transaction error: {:?}", err))
964        })?;
965
966        let verifier = TransactionScriptsVerifier::new(
967            Arc::new(rtx),
968            self.tx_dep_provider.clone(),
969            Arc::clone(&self.consensus),
970            Arc::new(TxVerifyEnv::new_submit(&self.tip_header)),
971        );
972        verifier.verify(u64::MAX).map_err(|err| {
973            BalanceTxCapacityError::VerifyScript(format!("Verify script error : {:?}", err))
974        })
975    }
976}
977
978/// Fill more inputs to balance the transaction capacity
979#[cfg(not(target_arch = "wasm32"))]
980pub fn balance_tx_capacity(
981    tx: &TransactionView,
982    balancer: &CapacityBalancer,
983    cell_collector: &mut dyn CellCollector,
984    tx_dep_provider: &dyn TransactionDependencyProvider,
985    cell_dep_resolver: &dyn CellDepResolver,
986    header_dep_resolver: &dyn HeaderDepResolver,
987) -> Result<TransactionView, BalanceTxCapacityError> {
988    let (tx, _change_idx) = rebalance_tx_capacity(
989        tx,
990        balancer,
991        cell_collector,
992        tx_dep_provider,
993        cell_dep_resolver,
994        header_dep_resolver,
995        0,
996        None,
997    )?;
998    Ok(tx)
999}
1000pub async fn balance_tx_capacity_async(
1001    tx: &TransactionView,
1002    balancer: &CapacityBalancer,
1003    cell_collector: &mut dyn CellCollector,
1004    tx_dep_provider: &dyn TransactionDependencyProvider,
1005    cell_dep_resolver: &dyn CellDepResolver,
1006    header_dep_resolver: &dyn HeaderDepResolver,
1007) -> Result<TransactionView, BalanceTxCapacityError> {
1008    let (tx, _change_idx) = rebalance_tx_capacity_async(
1009        tx,
1010        balancer,
1011        cell_collector,
1012        tx_dep_provider,
1013        cell_dep_resolver,
1014        header_dep_resolver,
1015        0,
1016        None,
1017    )
1018    .await?;
1019    Ok(tx)
1020}
1021
1022#[allow(clippy::too_many_arguments)]
1023#[cfg(not(target_arch = "wasm32"))]
1024fn rebalance_tx_capacity(
1025    tx: &TransactionView,
1026    balancer: &CapacityBalancer,
1027    cell_collector: &mut dyn CellCollector,
1028    tx_dep_provider: &dyn TransactionDependencyProvider,
1029    cell_dep_resolver: &dyn CellDepResolver,
1030    header_dep_resolver: &dyn HeaderDepResolver,
1031    accepted_min_fee: u64,
1032    change_index: Option<usize>,
1033) -> Result<(TransactionView, Option<usize>), BalanceTxCapacityError> {
1034    let capacity_provider = &balancer.capacity_provider;
1035    if capacity_provider.lock_scripts.is_empty() {
1036        return Err(BalanceTxCapacityError::EmptyCapacityProvider);
1037    }
1038    let change_lock_script = balancer
1039        .change_lock_script
1040        .clone()
1041        .unwrap_or_else(|| capacity_provider.lock_scripts[0].0.clone());
1042    let (tx, base_change_output, base_change_occupied_capacity) = if let Some(idx) = change_index {
1043        let outputs = tx.outputs();
1044        let output = tx
1045            .outputs()
1046            .get(idx)
1047            .ok_or(BalanceTxCapacityError::ChangeIndexNotFound(idx))?;
1048
1049        // remove change output
1050        let outputs: Vec<_> = outputs
1051            .into_iter()
1052            .enumerate()
1053            .filter_map(|(i, output)| if idx == i { None } else { Some(output) })
1054            .collect();
1055        let base_change_occupied_capacity = output
1056            .occupied_capacity(Capacity::zero())
1057            .expect("init change occupied capacity")
1058            .as_u64();
1059        let tx = tx.data().as_advanced_builder().set_outputs(outputs).build();
1060        (tx, output, base_change_occupied_capacity)
1061    } else {
1062        let base_change_output = CellOutput::new_builder().lock(change_lock_script).build();
1063        let base_change_occupied_capacity = base_change_output
1064            .occupied_capacity(Capacity::zero())
1065            .expect("init change occupied capacity")
1066            .as_u64();
1067        (
1068            tx.clone(),
1069            base_change_output,
1070            base_change_occupied_capacity,
1071        )
1072    };
1073
1074    let mut lock_scripts = Vec::new();
1075    // remove duplicated lock script
1076    for (script, placeholder, since_source) in &capacity_provider.lock_scripts {
1077        if lock_scripts.iter().all(|(target, _, _)| target != script) {
1078            lock_scripts.push((script.clone(), placeholder.clone(), since_source.clone()));
1079        }
1080    }
1081    let mut lock_script_idx = 0;
1082    let mut cell_deps = Vec::new();
1083    #[allow(clippy::mutable_key_type)]
1084    let mut resolved_scripts = HashSet::new();
1085    let mut inputs = Vec::new();
1086    let mut change_output: Option<CellOutput> = if change_index.is_some() {
1087        Some(base_change_output.clone())
1088    } else {
1089        None
1090    };
1091    let mut changed_witnesses: HashMap<usize, WitnessArgs> = HashMap::default();
1092    let mut witnesses = Vec::new();
1093    loop {
1094        let (lock_script, placeholder_witness, since_source) = &lock_scripts[lock_script_idx];
1095        let base_query = {
1096            let mut query = CellQueryOptions::new_lock(lock_script.clone());
1097            query.secondary_script_len_range = Some(ValueRangeOption::new_exact(0));
1098            query.data_len_range = Some(ValueRangeOption::new_exact(0));
1099            query
1100        };
1101        // check if capacity provider lock script already in inputs
1102        let mut has_provider = false;
1103        for input in tx.inputs().into_iter().chain(inputs.clone().into_iter()) {
1104            let cell = tx_dep_provider.get_cell(&input.previous_output())?;
1105            if cell.lock() == *lock_script {
1106                has_provider = true;
1107            }
1108        }
1109        while tx.witnesses().item_count() + witnesses.len()
1110            < tx.inputs().item_count() + inputs.len()
1111        {
1112            witnesses.push(Default::default());
1113        }
1114        let mut ret_change_index = None;
1115        let new_tx = {
1116            let mut all_witnesses = tx.witnesses().into_iter().collect::<Vec<_>>();
1117            for (idx, witness_args) in &changed_witnesses {
1118                all_witnesses[*idx] = witness_args.as_bytes().pack();
1119            }
1120            all_witnesses.extend(witnesses.clone());
1121            let output_len = tx.outputs().len();
1122            let mut builder = tx
1123                .data()
1124                .as_advanced_builder()
1125                .cell_deps(cell_deps.clone())
1126                .inputs(inputs.clone())
1127                .set_witnesses(all_witnesses);
1128            if let Some(output) = change_output.clone() {
1129                ret_change_index = Some(output_len);
1130                builder = builder
1131                    .output(output)
1132                    .output_data(ckb_types::packed::Bytes::default());
1133            }
1134            builder.build()
1135        };
1136        let tx_size = new_tx.data().as_reader().serialized_size_in_block();
1137        let min_fee = accepted_min_fee.max(balancer.fee_rate.fee(tx_size as u64).as_u64());
1138        let mut need_more_capacity = 1;
1139        let fee_result: Result<u64, TransactionFeeError> =
1140            tx_fee(new_tx.clone(), tx_dep_provider, header_dep_resolver);
1141        match fee_result {
1142            Ok(fee) if fee == min_fee => {
1143                return Ok((new_tx, ret_change_index));
1144            }
1145            Ok(fee) if fee > min_fee => {
1146                let delta = fee - min_fee;
1147                if let Some(output) = change_output.take() {
1148                    // If change cell already exits, just change the capacity field
1149                    let old_capacity: u64 = output.capacity().unpack();
1150                    let new_capacity = old_capacity
1151                        .checked_add(delta)
1152                        .expect("change cell capacity add overflow");
1153                    // next loop round must return new_tx;
1154                    change_output = Some(output.as_builder().capacity(new_capacity).build());
1155                    need_more_capacity = 0;
1156                } else {
1157                    // If change cell not exists, add a change cell.
1158
1159                    // The output extra header size is for:
1160                    //   * first 4 bytes is for output data header (the length)
1161                    //   * second 4 bytes if for output data offset
1162                    //   * third 4 bytes is for output offset
1163                    let output_header_extra = 4 + 4 + 4;
1164                    // NOTE: extra_min_fee +1 is for `FeeRate::fee` round
1165                    let extra_min_fee = balancer
1166                        .fee_rate
1167                        .fee(base_change_output.as_slice().len() as u64 + output_header_extra)
1168                        .as_u64()
1169                        + 1;
1170                    // The extra capacity (delta - extra_min_fee) is enough to hold the change cell.
1171                    if delta >= base_change_occupied_capacity + extra_min_fee {
1172                        // next loop round must return new_tx;
1173                        change_output = Some(
1174                            base_change_output
1175                                .clone()
1176                                .as_builder()
1177                                .capacity(delta - extra_min_fee)
1178                                .build(),
1179                        );
1180                        need_more_capacity = 0;
1181                    } else {
1182                        // peek if there is more live cell owned by this capacity provider
1183                        let (more_cells, _more_capacity) =
1184                            cell_collector.collect_live_cells(&base_query, false)?;
1185                        if more_cells.is_empty() {
1186                            if let Some(capacity) = balancer.force_small_change_as_fee {
1187                                if fee > capacity {
1188                                    return Err(
1189                                        BalanceTxCapacityError::ForceSmallChangeAsFeeFailed(fee),
1190                                    );
1191                                } else {
1192                                    return Ok((new_tx, ret_change_index));
1193                                }
1194                            } else if lock_script_idx + 1 == lock_scripts.len() {
1195                                return Err(BalanceTxCapacityError::CapacityNotEnough(format!(
1196                                    "can not create change cell, left capacity={}",
1197                                    HumanCapacity(delta)
1198                                )));
1199                            } else {
1200                                lock_script_idx += 1;
1201                                continue;
1202                            }
1203                        } else {
1204                            // need more input to balance the capacity
1205                            change_output = Some(
1206                                base_change_output
1207                                    .clone()
1208                                    .as_builder()
1209                                    .capacity(base_change_occupied_capacity)
1210                                    .build(),
1211                            );
1212                        }
1213                    }
1214                }
1215            }
1216            // fee is positive and `fee < min_fee`
1217            Ok(fee) => {
1218                need_more_capacity = min_fee - fee;
1219            }
1220            Err(TransactionFeeError::CapacityOverflow(delta)) => {
1221                need_more_capacity = delta.checked_add(min_fee).ok_or_else(|| {
1222                    BalanceTxCapacityError::CapacityNotEnough(format!(
1223                        "need more capacity, value={}",
1224                        HumanCapacity(delta)
1225                    ))
1226                })?;
1227            }
1228            Err(err) => {
1229                return Err(err.into());
1230            }
1231        }
1232        if need_more_capacity > 0 {
1233            let query = {
1234                let mut query = base_query.clone();
1235                query.min_total_capacity = need_more_capacity;
1236                query
1237            };
1238            let (more_cells, _more_capacity) = cell_collector.collect_live_cells(&query, true)?;
1239            if more_cells.is_empty() {
1240                if lock_script_idx + 1 == lock_scripts.len() {
1241                    return Err(BalanceTxCapacityError::CapacityNotEnough(format!(
1242                        "need more capacity, value={}",
1243                        HumanCapacity(need_more_capacity)
1244                    )));
1245                } else {
1246                    lock_script_idx += 1;
1247                    continue;
1248                }
1249            }
1250            if !resolved_scripts.contains(lock_script) {
1251                let provider_cell_dep =
1252                    cell_dep_resolver.resolve(lock_script).ok_or_else(|| {
1253                        BalanceTxCapacityError::ResolveCellDepFailed(lock_script.clone())
1254                    })?;
1255                if tx
1256                    .cell_deps()
1257                    .into_iter()
1258                    .all(|cell_dep| cell_dep != provider_cell_dep)
1259                {
1260                    cell_deps.push(provider_cell_dep);
1261                    resolved_scripts.insert(lock_script);
1262                }
1263            }
1264            if !has_provider {
1265                if tx.witnesses().item_count() > tx.inputs().item_count() + inputs.len() {
1266                    let idx = tx.inputs().item_count() + inputs.len();
1267                    let witness_data = tx.witnesses().get(idx).expect("get witness").raw_data();
1268                    // in case witness filled before balance tx
1269                    let mut witness = if witness_data.is_empty() {
1270                        WitnessArgs::default()
1271                    } else {
1272                        WitnessArgs::from_slice(witness_data.as_ref())
1273                            .map_err(|err| BalanceTxCapacityError::InvalidWitnessArgs(err.into()))?
1274                    };
1275                    if let Some(data) = placeholder_witness.input_type().to_opt() {
1276                        witness = witness
1277                            .as_builder()
1278                            .input_type(Some(data.raw_data()).pack())
1279                            .build();
1280                    }
1281                    if let Some(data) = placeholder_witness.output_type().to_opt() {
1282                        witness = witness
1283                            .as_builder()
1284                            .output_type(Some(data.raw_data()).pack())
1285                            .build();
1286                    }
1287                    if let Some(data) = placeholder_witness.lock().to_opt() {
1288                        witness = witness
1289                            .as_builder()
1290                            .lock(Some(data.raw_data()).pack())
1291                            .build();
1292                    }
1293                    changed_witnesses.insert(idx, witness);
1294                } else {
1295                    witnesses.push(placeholder_witness.as_bytes().pack());
1296                }
1297            }
1298            let since = match since_source {
1299                SinceSource::LockArgs(offset) => {
1300                    let lock_arg = lock_script.args().raw_data();
1301                    if lock_arg.len() < offset + 8 {
1302                        return Err(BalanceTxCapacityError::InvalidSinceValue(
1303                            *offset,
1304                            lock_arg.len(),
1305                        ));
1306                    }
1307                    let mut since_bytes = [0u8; 8];
1308                    since_bytes.copy_from_slice(&lock_arg[*offset..*offset + 8]);
1309                    u64::from_le_bytes(since_bytes)
1310                }
1311                SinceSource::Value(since_value) => *since_value,
1312            };
1313            inputs.extend(
1314                more_cells
1315                    .into_iter()
1316                    .map(|cell| CellInput::new(cell.out_point, since)),
1317            );
1318        }
1319    }
1320}
1321#[allow(clippy::too_many_arguments)]
1322async fn rebalance_tx_capacity_async(
1323    tx: &TransactionView,
1324    balancer: &CapacityBalancer,
1325    cell_collector: &mut dyn CellCollector,
1326    tx_dep_provider: &dyn TransactionDependencyProvider,
1327    cell_dep_resolver: &dyn CellDepResolver,
1328    header_dep_resolver: &dyn HeaderDepResolver,
1329    accepted_min_fee: u64,
1330    change_index: Option<usize>,
1331) -> Result<(TransactionView, Option<usize>), BalanceTxCapacityError> {
1332    let capacity_provider = &balancer.capacity_provider;
1333    if capacity_provider.lock_scripts.is_empty() {
1334        return Err(BalanceTxCapacityError::EmptyCapacityProvider);
1335    }
1336    let change_lock_script = balancer
1337        .change_lock_script
1338        .clone()
1339        .unwrap_or_else(|| capacity_provider.lock_scripts[0].0.clone());
1340    let (tx, base_change_output, base_change_occupied_capacity) = if let Some(idx) = change_index {
1341        let outputs = tx.outputs();
1342        let output = tx
1343            .outputs()
1344            .get(idx)
1345            .ok_or(BalanceTxCapacityError::ChangeIndexNotFound(idx))?;
1346
1347        // remove change output
1348        let outputs: Vec<_> = outputs
1349            .into_iter()
1350            .enumerate()
1351            .filter_map(|(i, output)| if idx == i { None } else { Some(output) })
1352            .collect();
1353        let base_change_occupied_capacity = output
1354            .occupied_capacity(Capacity::zero())
1355            .expect("init change occupied capacity")
1356            .as_u64();
1357        let tx = tx.data().as_advanced_builder().set_outputs(outputs).build();
1358        (tx, output, base_change_occupied_capacity)
1359    } else {
1360        let base_change_output = CellOutput::new_builder().lock(change_lock_script).build();
1361        let base_change_occupied_capacity = base_change_output
1362            .occupied_capacity(Capacity::zero())
1363            .expect("init change occupied capacity")
1364            .as_u64();
1365        (
1366            tx.clone(),
1367            base_change_output,
1368            base_change_occupied_capacity,
1369        )
1370    };
1371
1372    let mut lock_scripts = Vec::new();
1373    // remove duplicated lock script
1374    for (script, placeholder, since_source) in &capacity_provider.lock_scripts {
1375        if lock_scripts.iter().all(|(target, _, _)| target != script) {
1376            lock_scripts.push((script.clone(), placeholder.clone(), since_source.clone()));
1377        }
1378    }
1379    let mut lock_script_idx = 0;
1380    let mut cell_deps = Vec::new();
1381    #[allow(clippy::mutable_key_type)]
1382    let mut resolved_scripts = HashSet::new();
1383    let mut inputs = Vec::new();
1384    let mut change_output: Option<CellOutput> = if change_index.is_some() {
1385        Some(base_change_output.clone())
1386    } else {
1387        None
1388    };
1389    let mut changed_witnesses: HashMap<usize, WitnessArgs> = HashMap::default();
1390    let mut witnesses = Vec::new();
1391    loop {
1392        let (lock_script, placeholder_witness, since_source) = &lock_scripts[lock_script_idx];
1393        let base_query = {
1394            let mut query = CellQueryOptions::new_lock(lock_script.clone());
1395            query.secondary_script_len_range = Some(ValueRangeOption::new_exact(0));
1396            query.data_len_range = Some(ValueRangeOption::new_exact(0));
1397            query
1398        };
1399        // check if capacity provider lock script already in inputs
1400        let mut has_provider = false;
1401        for input in tx.inputs().into_iter().chain(inputs.clone().into_iter()) {
1402            let cell = tx_dep_provider
1403                .get_cell_async(&input.previous_output())
1404                .await?;
1405            if cell.lock() == *lock_script {
1406                has_provider = true;
1407            }
1408        }
1409        while tx.witnesses().item_count() + witnesses.len()
1410            < tx.inputs().item_count() + inputs.len()
1411        {
1412            witnesses.push(Default::default());
1413        }
1414        let mut ret_change_index = None;
1415        let new_tx = {
1416            let mut all_witnesses = tx.witnesses().into_iter().collect::<Vec<_>>();
1417            for (idx, witness_args) in &changed_witnesses {
1418                all_witnesses[*idx] = witness_args.as_bytes().pack();
1419            }
1420            all_witnesses.extend(witnesses.clone());
1421            let output_len = tx.outputs().len();
1422            let mut builder = tx
1423                .data()
1424                .as_advanced_builder()
1425                .cell_deps(cell_deps.clone())
1426                .inputs(inputs.clone())
1427                .set_witnesses(all_witnesses);
1428            if let Some(output) = change_output.clone() {
1429                ret_change_index = Some(output_len);
1430                builder = builder
1431                    .output(output)
1432                    .output_data(ckb_types::packed::Bytes::default());
1433            }
1434            builder.build()
1435        };
1436        let tx_size = new_tx.data().as_reader().serialized_size_in_block();
1437        let min_fee = accepted_min_fee.max(balancer.fee_rate.fee(tx_size as u64).as_u64());
1438        let mut need_more_capacity = 1;
1439        let fee_result: Result<u64, TransactionFeeError> =
1440            tx_fee_async(new_tx.clone(), tx_dep_provider, header_dep_resolver).await;
1441        match fee_result {
1442            Ok(fee) if fee == min_fee => {
1443                return Ok((new_tx, ret_change_index));
1444            }
1445            Ok(fee) if fee > min_fee => {
1446                let delta = fee - min_fee;
1447                if let Some(output) = change_output.take() {
1448                    // If change cell already exits, just change the capacity field
1449                    let old_capacity: u64 = output.capacity().unpack();
1450                    let new_capacity = old_capacity
1451                        .checked_add(delta)
1452                        .expect("change cell capacity add overflow");
1453                    // next loop round must return new_tx;
1454                    change_output = Some(output.as_builder().capacity(new_capacity).build());
1455                    need_more_capacity = 0;
1456                } else {
1457                    // If change cell not exists, add a change cell.
1458
1459                    // The output extra header size is for:
1460                    //   * first 4 bytes is for output data header (the length)
1461                    //   * second 4 bytes if for output data offset
1462                    //   * third 4 bytes is for output offset
1463                    let output_header_extra = 4 + 4 + 4;
1464                    // NOTE: extra_min_fee +1 is for `FeeRate::fee` round
1465                    let extra_min_fee = balancer
1466                        .fee_rate
1467                        .fee(base_change_output.as_slice().len() as u64 + output_header_extra)
1468                        .as_u64()
1469                        + 1;
1470                    // The extra capacity (delta - extra_min_fee) is enough to hold the change cell.
1471                    if delta >= base_change_occupied_capacity + extra_min_fee {
1472                        // next loop round must return new_tx;
1473                        change_output = Some(
1474                            base_change_output
1475                                .clone()
1476                                .as_builder()
1477                                .capacity(delta - extra_min_fee)
1478                                .build(),
1479                        );
1480                        need_more_capacity = 0;
1481                    } else {
1482                        // peek if there is more live cell owned by this capacity provider
1483                        let (more_cells, _more_capacity) = cell_collector
1484                            .collect_live_cells_async(&base_query, false)
1485                            .await?;
1486                        if more_cells.is_empty() {
1487                            if let Some(capacity) = balancer.force_small_change_as_fee {
1488                                if fee > capacity {
1489                                    return Err(
1490                                        BalanceTxCapacityError::ForceSmallChangeAsFeeFailed(fee),
1491                                    );
1492                                } else {
1493                                    return Ok((new_tx, ret_change_index));
1494                                }
1495                            } else if lock_script_idx + 1 == lock_scripts.len() {
1496                                return Err(BalanceTxCapacityError::CapacityNotEnough(format!(
1497                                    "can not create change cell, left capacity={}",
1498                                    HumanCapacity(delta)
1499                                )));
1500                            } else {
1501                                lock_script_idx += 1;
1502                                continue;
1503                            }
1504                        } else {
1505                            // need more input to balance the capacity
1506                            change_output = Some(
1507                                base_change_output
1508                                    .clone()
1509                                    .as_builder()
1510                                    .capacity(base_change_occupied_capacity)
1511                                    .build(),
1512                            );
1513                        }
1514                    }
1515                }
1516            }
1517            // fee is positive and `fee < min_fee`
1518            Ok(fee) => {
1519                need_more_capacity = min_fee - fee;
1520            }
1521            Err(TransactionFeeError::CapacityOverflow(delta)) => {
1522                need_more_capacity = delta.checked_add(min_fee).ok_or_else(|| {
1523                    BalanceTxCapacityError::CapacityNotEnough(format!(
1524                        "need more capacity, value={}",
1525                        HumanCapacity(delta)
1526                    ))
1527                })?;
1528            }
1529            Err(err) => {
1530                return Err(err.into());
1531            }
1532        }
1533        if need_more_capacity > 0 {
1534            let query = {
1535                let mut query = base_query.clone();
1536                query.min_total_capacity = need_more_capacity;
1537                query
1538            };
1539            let (more_cells, _more_capacity) = cell_collector
1540                .collect_live_cells_async(&query, true)
1541                .await?;
1542            if more_cells.is_empty() {
1543                if lock_script_idx + 1 == lock_scripts.len() {
1544                    return Err(BalanceTxCapacityError::CapacityNotEnough(format!(
1545                        "need more capacity, value={}",
1546                        HumanCapacity(need_more_capacity)
1547                    )));
1548                } else {
1549                    lock_script_idx += 1;
1550                    continue;
1551                }
1552            }
1553            if !resolved_scripts.contains(lock_script) {
1554                let provider_cell_dep =
1555                    cell_dep_resolver.resolve(lock_script).ok_or_else(|| {
1556                        BalanceTxCapacityError::ResolveCellDepFailed(lock_script.clone())
1557                    })?;
1558                if tx
1559                    .cell_deps()
1560                    .into_iter()
1561                    .all(|cell_dep| cell_dep != provider_cell_dep)
1562                {
1563                    cell_deps.push(provider_cell_dep);
1564                    resolved_scripts.insert(lock_script);
1565                }
1566            }
1567            if !has_provider {
1568                if tx.witnesses().item_count() > tx.inputs().item_count() + inputs.len() {
1569                    let idx = tx.inputs().item_count() + inputs.len();
1570                    let witness_data = tx.witnesses().get(idx).expect("get witness").raw_data();
1571                    // in case witness filled before balance tx
1572                    let mut witness = if witness_data.is_empty() {
1573                        WitnessArgs::default()
1574                    } else {
1575                        WitnessArgs::from_slice(witness_data.as_ref())
1576                            .map_err(|err| BalanceTxCapacityError::InvalidWitnessArgs(err.into()))?
1577                    };
1578                    if let Some(data) = placeholder_witness.input_type().to_opt() {
1579                        witness = witness
1580                            .as_builder()
1581                            .input_type(Some(data.raw_data()).pack())
1582                            .build();
1583                    }
1584                    if let Some(data) = placeholder_witness.output_type().to_opt() {
1585                        witness = witness
1586                            .as_builder()
1587                            .output_type(Some(data.raw_data()).pack())
1588                            .build();
1589                    }
1590                    if let Some(data) = placeholder_witness.lock().to_opt() {
1591                        witness = witness
1592                            .as_builder()
1593                            .lock(Some(data.raw_data()).pack())
1594                            .build();
1595                    }
1596                    changed_witnesses.insert(idx, witness);
1597                } else {
1598                    witnesses.push(placeholder_witness.as_bytes().pack());
1599                }
1600            }
1601            let since = match since_source {
1602                SinceSource::LockArgs(offset) => {
1603                    let lock_arg = lock_script.args().raw_data();
1604                    if lock_arg.len() < offset + 8 {
1605                        return Err(BalanceTxCapacityError::InvalidSinceValue(
1606                            *offset,
1607                            lock_arg.len(),
1608                        ));
1609                    }
1610                    let mut since_bytes = [0u8; 8];
1611                    since_bytes.copy_from_slice(&lock_arg[*offset..*offset + 8]);
1612                    u64::from_le_bytes(since_bytes)
1613                }
1614                SinceSource::Value(since_value) => *since_value,
1615            };
1616            inputs.extend(
1617                more_cells
1618                    .into_iter()
1619                    .map(|cell| CellInput::new(cell.out_point, since)),
1620            );
1621        }
1622    }
1623}
1624
1625pub struct ScriptGroups {
1626    pub lock_groups: HashMap<Byte32, ScriptGroup>,
1627    pub type_groups: HashMap<Byte32, ScriptGroup>,
1628}
1629#[cfg(not(target_arch = "wasm32"))]
1630pub fn gen_script_groups(
1631    tx: &TransactionView,
1632    tx_dep_provider: &dyn TransactionDependencyProvider,
1633) -> Result<ScriptGroups, TransactionDependencyError> {
1634    #[allow(clippy::mutable_key_type)]
1635    let mut lock_groups: HashMap<Byte32, ScriptGroup> = HashMap::default();
1636    #[allow(clippy::mutable_key_type)]
1637    let mut type_groups: HashMap<Byte32, ScriptGroup> = HashMap::default();
1638    for (i, input) in tx.inputs().into_iter().enumerate() {
1639        let output = tx_dep_provider.get_cell(&input.previous_output())?;
1640        let lock_group_entry = lock_groups
1641            .entry(output.calc_lock_hash())
1642            .or_insert_with(|| ScriptGroup::from_lock_script(&output.lock()));
1643        lock_group_entry.input_indices.push(i);
1644        if let Some(t) = &output.type_().to_opt() {
1645            let type_group_entry = type_groups
1646                .entry(t.calc_script_hash())
1647                .or_insert_with(|| ScriptGroup::from_type_script(t));
1648            type_group_entry.input_indices.push(i);
1649        }
1650    }
1651    for (i, output) in tx.outputs().into_iter().enumerate() {
1652        if let Some(t) = &output.type_().to_opt() {
1653            let type_group_entry = type_groups
1654                .entry(t.calc_script_hash())
1655                .or_insert_with(|| ScriptGroup::from_type_script(t));
1656            type_group_entry.output_indices.push(i);
1657        }
1658    }
1659    Ok(ScriptGroups {
1660        lock_groups,
1661        type_groups,
1662    })
1663}
1664pub async fn gen_script_groups_async(
1665    tx: &TransactionView,
1666    tx_dep_provider: &dyn TransactionDependencyProvider,
1667) -> Result<ScriptGroups, TransactionDependencyError> {
1668    #[allow(clippy::mutable_key_type)]
1669    let mut lock_groups: HashMap<Byte32, ScriptGroup> = HashMap::default();
1670    #[allow(clippy::mutable_key_type)]
1671    let mut type_groups: HashMap<Byte32, ScriptGroup> = HashMap::default();
1672    for (i, input) in tx.inputs().into_iter().enumerate() {
1673        let output = tx_dep_provider
1674            .get_cell_async(&input.previous_output())
1675            .await?;
1676        let lock_group_entry = lock_groups
1677            .entry(output.calc_lock_hash())
1678            .or_insert_with(|| ScriptGroup::from_lock_script(&output.lock()));
1679        lock_group_entry.input_indices.push(i);
1680        if let Some(t) = &output.type_().to_opt() {
1681            let type_group_entry = type_groups
1682                .entry(t.calc_script_hash())
1683                .or_insert_with(|| ScriptGroup::from_type_script(t));
1684            type_group_entry.input_indices.push(i);
1685        }
1686    }
1687    for (i, output) in tx.outputs().into_iter().enumerate() {
1688        if let Some(t) = &output.type_().to_opt() {
1689            let type_group_entry = type_groups
1690                .entry(t.calc_script_hash())
1691                .or_insert_with(|| ScriptGroup::from_type_script(t));
1692            type_group_entry.output_indices.push(i);
1693        }
1694    }
1695    Ok(ScriptGroups {
1696        lock_groups,
1697        type_groups,
1698    })
1699}
1700
1701/// Fill placeholder lock script witnesses
1702///
1703/// Return value:
1704///   * The updated transaction
1705///   * The script groups that not matched by given `unlockers`
1706#[cfg(not(target_arch = "wasm32"))]
1707pub fn fill_placeholder_witnesses(
1708    balanced_tx: TransactionView,
1709    tx_dep_provider: &dyn TransactionDependencyProvider,
1710    unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
1711) -> Result<(TransactionView, Vec<ScriptGroup>), UnlockError> {
1712    let ScriptGroups { lock_groups, .. } = gen_script_groups(&balanced_tx, tx_dep_provider)?;
1713    let mut tx = balanced_tx;
1714    let mut not_matched = Vec::new();
1715    for script_group in lock_groups.values() {
1716        let script_id = ScriptId::from(&script_group.script);
1717        let script_args = script_group.script.args().raw_data();
1718        if let Some(unlocker) = unlockers.get(&script_id) {
1719            if !unlocker.is_unlocked(&tx, script_group, tx_dep_provider)? {
1720                if unlocker.match_args(script_args.as_ref()) {
1721                    tx = unlocker.fill_placeholder_witness(&tx, script_group, tx_dep_provider)?;
1722                } else {
1723                    not_matched.push(script_group.clone());
1724                }
1725            }
1726        } else {
1727            not_matched.push(script_group.clone());
1728        }
1729    }
1730    Ok((tx, not_matched))
1731}
1732
1733/// Fill placeholder lock script witnesses
1734///
1735/// Return value:
1736///   * The updated transaction
1737///   * The script groups that not matched by given `unlockers`
1738pub async fn fill_placeholder_witnesses_async(
1739    balanced_tx: TransactionView,
1740    tx_dep_provider: &dyn TransactionDependencyProvider,
1741    unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
1742) -> Result<(TransactionView, Vec<ScriptGroup>), UnlockError> {
1743    let ScriptGroups { lock_groups, .. } =
1744        gen_script_groups_async(&balanced_tx, tx_dep_provider).await?;
1745    let mut tx = balanced_tx;
1746    let mut not_matched = Vec::new();
1747    for script_group in lock_groups.values() {
1748        let script_id = ScriptId::from(&script_group.script);
1749        let script_args = script_group.script.args().raw_data();
1750        if let Some(unlocker) = unlockers.get(&script_id) {
1751            if !unlocker
1752                .is_unlocked_async(&tx, script_group, tx_dep_provider)
1753                .await?
1754            {
1755                if unlocker.match_args(script_args.as_ref()) {
1756                    tx = unlocker
1757                        .fill_placeholder_witness_async(&tx, script_group, tx_dep_provider)
1758                        .await?;
1759                } else {
1760                    not_matched.push(script_group.clone());
1761                }
1762            }
1763        } else {
1764            not_matched.push(script_group.clone());
1765        }
1766    }
1767    Ok((tx, not_matched))
1768}
1769
1770/// Build unlocked transaction that ready to send or for further unlock.
1771///
1772/// Return value:
1773///   * The built transaction
1774///   * The script groups that not unlocked by given `unlockers`
1775#[cfg(not(target_arch = "wasm32"))]
1776pub fn unlock_tx(
1777    balanced_tx: TransactionView,
1778    tx_dep_provider: &dyn TransactionDependencyProvider,
1779    unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
1780) -> Result<(TransactionView, Vec<ScriptGroup>), UnlockError> {
1781    let ScriptGroups { lock_groups, .. } = gen_script_groups(&balanced_tx, tx_dep_provider)?;
1782    let mut tx = balanced_tx;
1783    let mut not_unlocked = Vec::new();
1784    for script_group in lock_groups.values() {
1785        let script_id = ScriptId::from(&script_group.script);
1786        let script_args = script_group.script.args().raw_data();
1787        if let Some(unlocker) = unlockers.get(&script_id) {
1788            if unlocker.is_unlocked(&tx, script_group, tx_dep_provider)? {
1789                tx = unlocker.clear_placeholder_witness(&tx, script_group)?;
1790            } else if unlocker.match_args(script_args.as_ref()) {
1791                tx = unlocker.unlock(&tx, script_group, tx_dep_provider)?;
1792            } else {
1793                not_unlocked.push(script_group.clone());
1794            }
1795        } else {
1796            not_unlocked.push(script_group.clone());
1797        }
1798    }
1799    Ok((tx, not_unlocked))
1800}
1801
1802/// Build unlocked transaction that ready to send or for further unlock.
1803///
1804/// Return value:
1805///   * The built transaction
1806///   * The script groups that not unlocked by given `unlockers`
1807pub async fn unlock_tx_async(
1808    balanced_tx: TransactionView,
1809    tx_dep_provider: &dyn TransactionDependencyProvider,
1810    unlockers: &HashMap<ScriptId, Box<dyn ScriptUnlocker>>,
1811) -> Result<(TransactionView, Vec<ScriptGroup>), UnlockError> {
1812    let ScriptGroups { lock_groups, .. } =
1813        gen_script_groups_async(&balanced_tx, tx_dep_provider).await?;
1814    let mut tx = balanced_tx;
1815    let mut not_unlocked = Vec::new();
1816    for script_group in lock_groups.values() {
1817        let script_id = ScriptId::from(&script_group.script);
1818        let script_args = script_group.script.args().raw_data();
1819        if let Some(unlocker) = unlockers.get(&script_id) {
1820            if unlocker
1821                .is_unlocked_async(&tx, script_group, tx_dep_provider)
1822                .await?
1823            {
1824                tx = unlocker.clear_placeholder_witness(&tx, script_group)?;
1825            } else if unlocker.match_args(script_args.as_ref()) {
1826                tx = unlocker
1827                    .unlock_async(&tx, script_group, tx_dep_provider)
1828                    .await?;
1829            } else {
1830                not_unlocked.push(script_group.clone());
1831            }
1832        } else {
1833            not_unlocked.push(script_group.clone());
1834        }
1835    }
1836    Ok((tx, not_unlocked))
1837}
1838
1839#[cfg(test)]
1840mod anyhow_tests {
1841    use anyhow::anyhow;
1842    #[test]
1843    fn test_signer_error() {
1844        use super::TxBuilderError;
1845        let error = TxBuilderError::ResolveHeaderDepByNumberFailed(0);
1846        let error = anyhow!(error);
1847        assert_eq!(
1848            "resolve header dep by block number failed: `0`",
1849            error.to_string()
1850        );
1851    }
1852
1853    #[test]
1854    fn test_transaction_fee_error() {
1855        let error = super::TransactionFeeError::CapacityOverflow(0);
1856        let error = anyhow!(error);
1857        assert_eq!("capacity sub overflow, delta: `0`", error.to_string());
1858    }
1859
1860    #[test]
1861    fn test_balance_tx_capacity_error() {
1862        let eror = super::BalanceTxCapacityError::EmptyCapacityProvider;
1863        let error = anyhow!(eror);
1864        assert_eq!("empty capacity provider", error.to_string())
1865    }
1866}