tx3_resolver/inputs/assign/
mod.rs

1//! Assignment stage: given all queries with their ranked candidates, find an
2//! allocation that satisfies everything simultaneously.
3
4use std::collections::HashSet;
5
6use crate::job::{QueryResolution, ResolveJob};
7use crate::Error;
8use tx3_tir::model::{
9    assets::CanonicalAssets,
10    core::{Utxo, UtxoRef, UtxoSet},
11};
12
13#[cfg(test)]
14mod tests;
15
16/// How tightly constrained a query is, from most to least specific.
17/// Lower values get priority during assignment so that the most constrained
18/// queries pick UTxOs first.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
20enum Specificity {
21    SingleCollateral = 0,
22    Single = 1,
23    ManyCollateral = 2,
24    Many = 3,
25}
26
27impl QueryResolution {
28    fn specificity(&self) -> Specificity {
29        match (self.query.support_many, self.query.collateral) {
30            (false, true) => Specificity::SingleCollateral,
31            (false, false) => Specificity::Single,
32            (true, true) => Specificity::ManyCollateral,
33            (true, false) => Specificity::Many,
34        }
35    }
36
37    fn constraint_tightness(&self) -> (usize, Specificity, &str) {
38        (self.candidates.len(), self.specificity(), &self.name)
39    }
40
41    fn pick(&mut self, used: &HashSet<UtxoRef>) {
42        let available: Vec<Utxo> = self
43            .candidates
44            .drain(..)
45            .filter(|utxo| !used.contains(&utxo.r#ref))
46            .collect();
47
48        let target = self
49            .query
50            .min_amount
51            .clone()
52            .unwrap_or(CanonicalAssets::empty());
53
54        let selection = if self.query.support_many {
55            pick_many(available, &target)
56        } else {
57            pick_single(available, &target)
58        };
59
60        self.selection = Some(selection);
61    }
62}
63
64/// Select the first candidate that fully covers the target.
65fn pick_single(candidates: Vec<Utxo>, target: &CanonicalAssets) -> UtxoSet {
66    candidates
67        .into_iter()
68        .find(|utxo| utxo.assets.contains_total(target))
69        .into_iter()
70        .collect()
71}
72
73/// Accumulate candidates greedily until the target is fully covered,
74/// then trim any excess UTxOs that aren't needed.
75fn pick_many(candidates: Vec<Utxo>, target: &CanonicalAssets) -> UtxoSet {
76    let mut matched = UtxoSet::new();
77    let mut pending = target.clone();
78
79    for candidate in candidates {
80        if candidate.assets.contains_some(&pending) {
81            matched.insert(candidate.clone());
82            let to_include = candidate.assets.clone();
83            pending = pending - to_include;
84        }
85
86        if pending.is_empty_or_negative() {
87            break;
88        }
89    }
90
91    if !pending.is_empty_or_negative() {
92        return UtxoSet::new();
93    }
94
95    while let Some(utxo) = find_first_excess_utxo(&matched, target) {
96        matched.remove(&utxo);
97    }
98
99    matched
100}
101
102fn find_first_excess_utxo(utxos: &UtxoSet, target: &CanonicalAssets) -> Option<Utxo> {
103    if utxos.len() == 1 {
104        return None;
105    }
106
107    let available = utxos.total_assets();
108
109    let excess = available - target.clone();
110
111    if excess.is_empty_or_negative() {
112        return None;
113    }
114
115    for utxo in utxos.iter_sorted_by_ref() {
116        if excess.contains_total(&utxo.assets) {
117            return Some(utxo.clone());
118        }
119    }
120
121    None
122}
123
124impl ResolveJob {
125    /// Given queries with their ranked candidates (from the approximation
126    /// stage), find an allocation that satisfies all queries simultaneously
127    /// using greedy assignment. Returns an error if any query cannot be resolved.
128    pub fn assign_all(&mut self) -> Result<(), Error> {
129        // Sort indices by constraint tightness so tightest queries pick first.
130        let mut indices: Vec<usize> = (0..self.input_queries.len()).collect();
131        indices.sort_by(|&a, &b| {
132            self.input_queries[a]
133                .constraint_tightness()
134                .cmp(&self.input_queries[b].constraint_tightness())
135        });
136
137        let mut used: HashSet<UtxoRef> = HashSet::new();
138
139        for idx in indices {
140            self.input_queries[idx].pick(&used);
141
142            if let Some(selection) = &self.input_queries[idx].selection {
143                for utxo in selection.iter() {
144                    used.insert(utxo.r#ref.clone());
145                }
146            }
147        }
148
149        let pool_refs = self.pool_refs();
150        for qr in &self.input_queries {
151            qr.ensure_resolved(&pool_refs)?;
152        }
153
154        Ok(())
155    }
156}
tx3_resolver/inputs/assign/mod.rs

tx3_resolver/inputs/assign/
mod.rs
