tx3_resolver/inputs/
searching.rs

1use std::collections::HashSet;
2
3use tx3_lang::{
4    backend::{UtxoPattern, UtxoStore},
5    AssetClass, UtxoRef,
6};
7
8use crate::inputs::CanonicalQuery;
9
10use super::Error;
11
12#[derive(Debug, Clone)]
13pub enum Subset {
14    All,
15    Specific(HashSet<UtxoRef>),
16}
17
18impl Subset {
19    fn count(&self) -> Option<usize> {
20        match self {
21            Self::All => None,
22            Self::Specific(s) => Some(s.len()),
23        }
24    }
25
26    #[allow(dead_code)]
27    fn union(a: Self, b: Self) -> Self {
28        match (a, b) {
29            (Self::All, _) => Self::All,
30            (_, Self::All) => Self::All,
31            (Self::Specific(s1), Self::Specific(s2)) => {
32                Self::Specific(s1.union(&s2).cloned().collect())
33            }
34        }
35    }
36
37    fn intersection(a: Self, b: Self) -> Self {
38        match (a, b) {
39            (Self::All, x) => x,
40            (x, Self::All) => x,
41            (Self::Specific(s1), Self::Specific(s2)) => {
42                Self::Specific(s1.intersection(&s2).cloned().collect())
43            }
44        }
45    }
46}
47
48impl From<HashSet<UtxoRef>> for Subset {
49    fn from(value: HashSet<UtxoRef>) -> Self {
50        Self::Specific(value)
51    }
52}
53
54#[derive(Debug, Clone)]
55pub struct SearchSpace {
56    pub union: HashSet<UtxoRef>,
57    pub intersection: HashSet<UtxoRef>,
58    pub by_address_count: Option<usize>,
59    pub by_asset_class_count: Option<usize>,
60    pub by_ref_count: Option<usize>,
61}
62
63impl SearchSpace {
64    fn new() -> Self {
65        Self {
66            union: HashSet::new(),
67            intersection: HashSet::new(),
68            by_address_count: None,
69            by_asset_class_count: None,
70            by_ref_count: None,
71        }
72    }
73
74    fn is_empty(&self) -> bool {
75        self.union.is_empty()
76    }
77
78    fn add_matches(&mut self, utxos: HashSet<UtxoRef>) {
79        self.union = self.union.union(&utxos).cloned().collect();
80        self.intersection = self.intersection.intersection(&utxos).cloned().collect();
81    }
82
83    fn add_address_matches(&mut self, subset: Subset) {
84        match subset {
85            Subset::All => (),
86            Subset::Specific(utxos) => {
87                *self.by_address_count.get_or_insert(0) += utxos.len();
88                self.add_matches(utxos);
89            }
90        }
91    }
92
93    fn add_asset_class_matches(&mut self, subset: Subset) {
94        match subset {
95            Subset::All => (),
96            Subset::Specific(utxos) => {
97                *self.by_asset_class_count.get_or_insert(0) += utxos.len();
98                self.add_matches(utxos);
99            }
100        }
101    }
102
103    fn add_ref_matches(&mut self, utxos: HashSet<UtxoRef>) {
104        *self.by_ref_count.get_or_insert(0) += utxos.len();
105        self.add_matches(utxos);
106    }
107
108    pub fn take(&self, take: Option<usize>) -> HashSet<UtxoRef> {
109        let Some(take) = take else {
110            return self.union.clone();
111        };
112
113        // first we take from the intersection which are the best matches
114        let best: HashSet<_> = self.intersection.iter().take(take).cloned().collect();
115
116        if best.len() < take {
117            let others: HashSet<_> = self.union.difference(&best).cloned().collect();
118            let remaining: HashSet<_> = others.into_iter().take(take - best.len()).collect();
119            best.union(&remaining).cloned().collect()
120        } else {
121            best
122        }
123    }
124}
125
126async fn narrow_by_asset_class<T: UtxoStore>(
127    store: &T,
128    parent: Subset,
129    class: &AssetClass,
130) -> Result<Subset, Error> {
131    // skip filtering lovelace since it's not an custom asset
132    if matches!(class, AssetClass::Naked) {
133        return Ok(Subset::All);
134    }
135
136    let AssetClass::Defined(policy, name) = class else {
137        return Ok(Subset::All);
138    };
139
140    let utxos = store
141        .narrow_refs(UtxoPattern::by_asset(policy, name))
142        .await?;
143
144    Ok(Subset::intersection(parent, Subset::Specific(utxos)))
145}
146
147pub async fn narrow_search_space<T: UtxoStore>(
148    store: &T,
149    criteria: &CanonicalQuery,
150) -> Result<SearchSpace, Error> {
151    let mut search_space = SearchSpace::new();
152
153    let parent_subset = if let Some(address) = &criteria.address {
154        let utxos = store.narrow_refs(UtxoPattern::by_address(address)).await?;
155        Subset::Specific(utxos)
156    } else {
157        Subset::All
158    };
159
160    search_space.add_address_matches(parent_subset.clone());
161
162    if let Some(assets) = &criteria.min_amount {
163        for (class, amount) in assets.iter() {
164            if *amount > 0 {
165                let subset = narrow_by_asset_class(store, parent_subset.clone(), class).await?;
166                search_space.add_asset_class_matches(subset);
167            }
168        }
169    }
170
171    if !criteria.refs.is_empty() {
172        search_space.add_ref_matches(criteria.refs.clone());
173    }
174
175    if search_space.is_empty() {
176        return Err(Error::InputQueryTooBroad);
177    }
178
179    Ok(search_space)
180}
tx3_resolver/inputs/searching.rs

tx3_resolver/inputs/
searching.rs
