ckb_sdk/tx_builder/
mod.rs

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