brk_interface/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::{collections::BTreeMap, sync::OnceLock};
4
5use brk_computer::Computer;
6use brk_error::{Error, Result};
7use brk_indexer::Indexer;
8use brk_structs::Height;
9use nucleo_matcher::{
10    Config, Matcher,
11    pattern::{AtomKind, CaseMatching, Normalization, Pattern},
12};
13use quick_cache::sync::Cache;
14use tabled::settings::Style;
15use vecdb::{AnyCollectableVec, AnyStoredVec};
16
17mod deser;
18mod format;
19mod ids;
20mod index;
21mod output;
22mod pagination;
23mod params;
24mod table;
25mod vecs;
26
27pub use format::Format;
28pub use index::Index;
29pub use output::{Output, Value};
30pub use pagination::{PaginatedIndexParam, PaginationParam};
31pub use params::{IdParam, Params, ParamsOpt};
32pub use table::Tabled;
33use vecs::Vecs;
34
35use crate::vecs::{IdToVec, IndexToVec};
36
37pub fn cached_errors() -> &'static Cache<String, String> {
38    static CACHE: OnceLock<Cache<String, String>> = OnceLock::new();
39    CACHE.get_or_init(|| Cache::new(1000))
40}
41
42#[allow(dead_code)]
43pub struct Interface<'a> {
44    vecs: Vecs<'a>,
45    indexer: &'a Indexer,
46    computer: &'a Computer,
47}
48
49impl<'a> Interface<'a> {
50    pub fn build(indexer: &Indexer, computer: &Computer) -> Self {
51        let indexer = indexer.static_clone();
52        let computer = computer.static_clone();
53        let vecs = Vecs::build(indexer, computer);
54
55        Self {
56            vecs,
57            indexer,
58            computer,
59        }
60    }
61
62    pub fn get_height(&self) -> Height {
63        Height::from(self.indexer.vecs.height_to_blockhash.stamp())
64    }
65
66    pub fn search(&self, params: &Params) -> Result<Vec<(String, &&dyn AnyCollectableVec)>> {
67        let ids = &params.ids;
68        let index = params.index;
69
70        let ids_to_vec = self
71            .vecs
72            .index_to_id_to_vec
73            .get(&index)
74            .ok_or(Error::String(format!(
75                "Index \"{}\" isn't a valid index",
76                index
77            )))?;
78
79        ids.iter()
80            .map(|id| {
81                let vec = ids_to_vec.get(id.as_str()).ok_or_else(|| {
82                    let cached_errors = cached_errors();
83
84                    if let Some(message) = cached_errors.get(id) {
85                        return Error::String(message)
86                    }
87
88                    let mut message = format!(
89                        "No vec named \"{}\" indexed by \"{}\" found.\n",
90                        id,
91                        index
92                    );
93
94                    let mut matcher = Matcher::new(Config::DEFAULT);
95
96                    let matches = Pattern::new(
97                        id.as_str(),
98                        CaseMatching::Ignore,
99                        Normalization::Smart,
100                        AtomKind::Fuzzy,
101                    )
102                    .match_list(ids_to_vec.keys(), &mut matcher)
103                    .into_iter()
104                    .take(10)
105                    .map(|(s, _)| s)
106                    .collect::<Vec<_>>();
107
108                    if !matches.is_empty() {
109                        message +=
110                            &format!("\nMaybe you meant one of the following: {matches:#?} ?\n");
111                    }
112
113                    if let Some(index_to_vec) = self.id_to_index_to_vec().get(id.as_str()) {
114                        message += &format!("\nBut there is a vec named {id} which supports the following indexes: {:#?}\n", index_to_vec.keys());
115                    }
116
117                    cached_errors.insert(id.clone(), message.clone());
118
119                    Error::String(message)
120                });
121                vec.map(|vec| (id.clone(), vec))
122            })
123            .collect::<Result<Vec<_>>>()
124    }
125
126    pub fn format(
127        &self,
128        vecs: Vec<(String, &&dyn AnyCollectableVec)>,
129        params: &ParamsOpt,
130    ) -> Result<Output> {
131        let from = params.from().map(|from| {
132            vecs.iter()
133                .map(|(_, v)| v.i64_to_usize(from))
134                .min()
135                .unwrap_or_default()
136        });
137
138        let to = params.to().map(|to| {
139            vecs.iter()
140                .map(|(_, v)| v.i64_to_usize(to))
141                .min()
142                .unwrap_or_default()
143        });
144
145        let mut values = vecs
146            .iter()
147            .map(|(_, vec)| -> Result<Vec<serde_json::Value>> {
148                Ok(vec.collect_range_serde_json(from, to)?)
149            })
150            .collect::<Result<Vec<_>>>()?;
151
152        let format = params.format();
153
154        if values.is_empty() {
155            return Ok(Output::default(format));
156        }
157
158        let ids_last_i = vecs.len() - 1;
159
160        Ok(match format {
161            Some(Format::CSV) | Some(Format::TSV) => {
162                let delimiter = if format == Some(Format::CSV) {
163                    ','
164                } else {
165                    '\t'
166                };
167
168                let mut text = vecs
169                    .iter()
170                    .map(|(id, _)| id.to_owned())
171                    .collect::<Vec<_>>()
172                    .join(&delimiter.to_string());
173
174                text.push('\n');
175
176                let values_len = values.first().unwrap().len();
177
178                (0..values_len).for_each(|i| {
179                    let mut line = "".to_string();
180                    values.iter().enumerate().for_each(|(id_i, v)| {
181                        line += &v.get(i).unwrap().to_string();
182                        if id_i == ids_last_i {
183                            line.push('\n');
184                        } else {
185                            line.push(delimiter);
186                        }
187                    });
188                    text += &line;
189                });
190
191                if format == Some(Format::CSV) {
192                    Output::CSV(text)
193                } else {
194                    Output::TSV(text)
195                }
196            }
197            Some(Format::MD) => {
198                let mut table =
199                    values.to_table(vecs.iter().map(|(s, _)| s.to_owned()).collect::<Vec<_>>());
200
201                table.with(Style::markdown());
202
203                Output::MD(table.to_string())
204            }
205            Some(Format::JSON) | None => {
206                if values.len() == 1 {
207                    let mut values = values.pop().unwrap();
208                    if values.len() == 1 {
209                        let value = values.pop().unwrap();
210                        Output::Json(Value::Single(value))
211                    } else {
212                        Output::Json(Value::List(values))
213                    }
214                } else {
215                    Output::Json(Value::Matrix(values))
216                }
217            }
218        })
219    }
220
221    pub fn search_and_format(&self, params: Params) -> Result<Output> {
222        self.format(self.search(&params)?, &params.rest)
223    }
224
225    pub fn id_to_index_to_vec(&self) -> &BTreeMap<&str, IndexToVec<'_>> {
226        &self.vecs.id_to_index_to_vec
227    }
228
229    pub fn index_to_id_to_vec(&self) -> &BTreeMap<Index, IdToVec<'_>> {
230        &self.vecs.index_to_id_to_vec
231    }
232
233    pub fn get_vecid_count(&self) -> usize {
234        self.vecs.id_count
235    }
236
237    pub fn get_index_count(&self) -> usize {
238        self.vecs.index_count
239    }
240
241    pub fn get_vec_count(&self) -> usize {
242        self.vecs.vec_count
243    }
244
245    pub fn get_indexes(&self) -> &[&'static str] {
246        &self.vecs.indexes
247    }
248
249    pub fn get_accepted_indexes(&self) -> &BTreeMap<&'static str, &'static [&'static str]> {
250        &self.vecs.accepted_indexes
251    }
252
253    pub fn get_vecids(&self, pagination: PaginationParam) -> &[&str] {
254        self.vecs.ids(pagination)
255    }
256
257    pub fn get_index_to_vecids(&self, paginated_index: PaginatedIndexParam) -> Vec<&str> {
258        self.vecs.index_to_ids(paginated_index)
259    }
260
261    pub fn get_vecid_to_indexes(&self, id: String) -> Option<&Vec<&'static str>> {
262        self.vecs.id_to_indexes(id)
263    }
264}
brk_interface/lib.rs

brk_interface/
lib.rs
