tx3_resolver/inputs/
searching.rs

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

tx3_resolver/inputs/
searching.rs
