Skip to main content

flams_backend_types/
search.rs

1#![allow(clippy::wildcard_imports)]
2
3use ftml_ontology::narrative::documents::Document;
4use ftml_ontology::narrative::elements::paragraphs::ParagraphKind;
5use ftml_uris::Language;
6use ftml_uris::{DocumentElementUri, DocumentUri, SymbolUri};
7use smallvec::SmallVec;
8
9#[derive(Clone, Debug, PartialEq, Eq)]
10#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
11pub enum QueryFilter {
12    SymbolsOnly,
13    Any(FragmentQueryFilter),
14}
15
16impl Default for QueryFilter {
17    #[inline]
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl QueryFilter {
24    #[inline]
25    #[must_use]
26    pub const fn new() -> Self {
27        Self::Any(FragmentQueryFilter::new())
28    }
29}
30
31#[derive(Clone, Debug, PartialEq, Eq, Default)]
32#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
33pub struct FragmentQueryFilter {
34    #[cfg_attr(feature = "serde", serde(default))]
35    pub in_documents: SmallVec<DocumentUri, 2>,
36    #[cfg_attr(feature = "serde", serde(default))]
37    pub languages: SmallVec<Language, 2>,
38    #[cfg_attr(feature = "serde", serde(default))]
39    pub flags: QueryFilterFlags,
40}
41
42impl FragmentQueryFilter {
43    #[inline]
44    #[must_use]
45    pub const fn new() -> Self {
46        Self {
47            in_documents: SmallVec::new(),
48            languages: SmallVec::new(),
49            flags: QueryFilterFlags::new(),
50        }
51    }
52
53    pub fn close(&mut self, mut get: impl FnMut(&DocumentUri) -> Option<Document>) {
54        if self.in_documents.is_empty() {
55            return;
56        }
57        let mut dones = Vec::new();
58        let mut todos: Vec<_> = std::mem::take(&mut self.in_documents).into_vec();
59        while let Some(uri) = todos.pop() {
60            if dones.contains(&uri) {
61                continue;
62            }
63            let d = get(&uri);
64            dones.push(uri);
65            if let Some(d) = d {
66                use ftml_ontology::utils::RefTree;
67
68                for e in d.dfs() {
69                    use ftml_ontology::narrative::elements::DocumentElementRef;
70
71                    if let DocumentElementRef::DocumentReference { target, .. } = e {
72                        todos.push(target.clone());
73                    }
74                }
75            }
76        }
77        self.in_documents = dones.into();
78    }
79}
80
81#[derive(Copy, Clone, Debug, PartialEq, Eq)]
82#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
83pub struct QueryFilterFlags(u8);
84
85impl Default for QueryFilterFlags {
86    #[inline]
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92impl QueryFilterFlags {
93    #[inline]
94    #[must_use]
95    pub const fn new() -> Self {
96        Self(0b0111_1111)
97    }
98    #[inline]
99    #[must_use]
100    pub const fn none() -> Self {
101        Self(0)
102    }
103    #[inline]
104    #[must_use]
105    pub const fn set_allow_documents(self) -> Self {
106        Self(self.0 | 0b0000_0001)
107    }
108    #[inline]
109    #[must_use]
110    pub const fn unset_allow_documents(self) -> Self {
111        Self(self.0 & 0b1111_1110)
112    }
113    #[inline]
114    #[must_use]
115    pub const fn allow_documents(self) -> bool {
116        self.0 % 2 == 1
117    }
118    #[inline]
119    #[must_use]
120    pub const fn set_allow_paragraphs(self) -> Self {
121        Self(self.0 | 0b0000_0010)
122    }
123    #[inline]
124    #[must_use]
125    pub const fn unset_allow_paragraphs(self) -> Self {
126        Self(self.0 & 0b1111_1101)
127    }
128    #[inline]
129    #[must_use]
130    pub const fn allow_paragraphs(self) -> bool {
131        (self.0 & 0b0000_0010) >> 1 == 1
132    }
133    #[inline]
134    #[must_use]
135    pub const fn set_allow_definitions(self) -> Self {
136        Self(self.0 | 0b0000_0100)
137    }
138    #[inline]
139    #[must_use]
140    pub const fn unset_allow_definitions(self) -> Self {
141        Self(self.0 & 0b1111_1011)
142    }
143    #[inline]
144    #[must_use]
145    pub const fn allow_definitions(self) -> bool {
146        (self.0 & 0b0000_0100) >> 2 == 1
147    }
148    #[inline]
149    #[must_use]
150    pub const fn set_allow_examples(self) -> Self {
151        Self(self.0 | 0b0000_1000)
152    }
153    #[inline]
154    #[must_use]
155    pub const fn unset_allow_examples(self) -> Self {
156        Self(self.0 & 0b1111_0111)
157    }
158    #[inline]
159    #[must_use]
160    pub const fn allow_examples(self) -> bool {
161        (self.0 & 0b0000_1000) >> 3 == 1
162    }
163    #[inline]
164    #[must_use]
165    pub const fn set_allow_assertions(self) -> Self {
166        Self(self.0 | 0b0001_0000)
167    }
168    #[inline]
169    #[must_use]
170    pub const fn unset_allow_assertions(self) -> Self {
171        Self(self.0 & 0b1110_1111)
172    }
173    #[inline]
174    #[must_use]
175    pub const fn allow_assertions(self) -> bool {
176        (self.0 & 0b0001_0000) >> 4 == 1
177    }
178    #[inline]
179    #[must_use]
180    pub const fn set_allow_problems(self) -> Self {
181        Self(self.0 | 0b0010_0000)
182    }
183    #[inline]
184    #[must_use]
185    pub const fn unset_allow_problems(self) -> Self {
186        Self(self.0 & 0b1101_1111)
187    }
188    #[inline]
189    #[must_use]
190    pub const fn allow_problems(self) -> bool {
191        (self.0 & 0b0010_0000) >> 5 == 1
192    }
193    #[inline]
194    #[must_use]
195    pub const fn is_definition_like(self) -> bool {
196        self.0 > 0b1000_0000
197    }
198    #[inline]
199    #[must_use]
200    pub const fn definition_like_only() -> Self {
201        Self(0b1001_0110).set_allow_definitions()
202    }
203}
204
205#[derive(Debug, Clone)]
206#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
207#[cfg_attr(feature = "typescript", derive(tsify::Tsify))]
208#[cfg_attr(feature = "typescript", tsify(into_wasm_abi, from_wasm_abi))]
209pub enum SearchResult {
210    Document(DocumentUri),
211    Paragraph {
212        uri: DocumentElementUri,
213        fors: Vec<SymbolUri>,
214        def_like: bool,
215        kind: SearchResultKind,
216    },
217}
218
219#[derive(Copy, Clone, Debug)]
220#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
221#[cfg_attr(feature = "bincode", derive(bincode::Encode, bincode::Decode))]
222#[cfg_attr(feature = "typescript", derive(tsify::Tsify))]
223#[cfg_attr(feature = "typescript", tsify(into_wasm_abi, from_wasm_abi))]
224pub enum SearchResultKind {
225    Document = 0,
226    Paragraph = 1,
227    Definition = 2,
228    Example = 3,
229    Assertion = 4,
230    Problem = 5,
231}
232impl SearchResultKind {
233    #[must_use]
234    pub const fn as_str(&self) -> &'static str {
235        match self {
236            Self::Document => "Document",
237            Self::Paragraph => "Paragraph",
238            Self::Definition => "Definition",
239            Self::Example => "Example",
240            Self::Assertion => "Assertion",
241            Self::Problem => "Problem",
242        }
243    }
244}
245
246impl From<SearchResultKind> for u64 {
247    fn from(value: SearchResultKind) -> Self {
248        match value {
249            SearchResultKind::Document => 0,
250            SearchResultKind::Paragraph => 1,
251            SearchResultKind::Definition => 2,
252            SearchResultKind::Example => 3,
253            SearchResultKind::Assertion => 4,
254            SearchResultKind::Problem => 5,
255        }
256    }
257}
258
259impl TryFrom<u64> for SearchResultKind {
260    type Error = ();
261    fn try_from(value: u64) -> Result<Self, Self::Error> {
262        Ok(match value {
263            0 => Self::Document,
264            1 => Self::Paragraph,
265            2 => Self::Definition,
266            3 => Self::Example,
267            4 => Self::Assertion,
268            5 => Self::Problem,
269            _ => return Err(()),
270        })
271    }
272}
273impl TryFrom<ParagraphKind> for SearchResultKind {
274    type Error = ();
275    fn try_from(value: ParagraphKind) -> Result<Self, Self::Error> {
276        Ok(match value {
277            ParagraphKind::Assertion => Self::Assertion,
278            ParagraphKind::Definition => Self::Definition,
279            ParagraphKind::Example => Self::Example,
280            ParagraphKind::Paragraph => Self::Paragraph,
281            _ => return Err(()),
282        })
283    }
284}
285
286#[cfg(feature = "vectorsearch")]
287const LEN: usize = 384;
288
289#[cfg(feature = "vectorsearch")]
290#[derive(bincode::Encode, bincode::Decode, Debug, Clone)]
291pub struct Embedding(pub(crate) [f32; LEN]);
292
293#[cfg(feature = "vectorsearch")]
294impl Embedding {
295    #[must_use]
296    #[inline]
297    pub const fn zero() -> Self {
298        Self([0.0; LEN])
299    }
300
301    #[must_use]
302    #[inline]
303    pub const fn as_bytes(&self) -> &[u8; 4 * LEN] {
304        unsafe {
305            self.0
306                .as_ptr()
307                .cast::<[u8; 4 * LEN]>()
308                .as_ref()
309                .unwrap_unchecked()
310        }
311    }
312    #[inline]
313    #[must_use]
314    pub const fn new(vec: [f32; LEN]) -> Self {
315        Self(vec)
316    }
317}
318
319#[cfg(feature = "vectorsearch")]
320impl<'b> std::ops::Rem<&'b Embedding> for &Embedding {
321    type Output = f64;
322    fn rem(self, rhs: &'b Embedding) -> Self::Output {
323        #[allow(clippy::suspicious_arithmetic_impl)]
324        simsimd::SpatialSimilarity::cos(&self.0, &rhs.0).map_or(0.0, |r| 1.0 - r)
325    }
326}