1use std::collections::BTreeMap;
15
16use crate::{DocId, DocSet};
17
18pub trait Semiring: Clone {
24 fn zero() -> Self;
26
27 fn one() -> Self;
29
30 fn plus(&self, other: &Self) -> Self;
32
33 fn times(&self, other: &Self) -> Self;
35
36 fn is_zero(&self) -> bool;
39}
40
41impl Semiring for bool {
42 fn zero() -> Self {
43 false
44 }
45
46 fn one() -> Self {
47 true
48 }
49
50 fn plus(&self, other: &Self) -> Self {
51 *self || *other
52 }
53
54 fn times(&self, other: &Self) -> Self {
55 *self && *other
56 }
57
58 fn is_zero(&self) -> bool {
59 !*self
60 }
61}
62
63#[derive(Debug, Clone, Copy, PartialEq)]
70pub struct LogSemiring(f64);
71
72impl LogSemiring {
73 pub fn from_log(value: f64) -> Option<Self> {
75 (!value.is_nan()).then_some(Self(value))
76 }
77
78 pub fn from_weight(weight: f64) -> Option<Self> {
80 if weight.is_nan() || weight < 0.0 {
81 return None;
82 }
83 if weight == 0.0 {
84 return Some(Self::zero());
85 }
86 Some(Self(weight.ln()))
87 }
88
89 pub fn log_value(self) -> f64 {
91 self.0
92 }
93
94 pub fn weight(self) -> f64 {
96 self.0.exp()
97 }
98}
99
100impl Semiring for LogSemiring {
101 fn zero() -> Self {
102 Self(f64::NEG_INFINITY)
103 }
104
105 fn one() -> Self {
106 Self(0.0)
107 }
108
109 fn plus(&self, other: &Self) -> Self {
110 if self.is_zero() {
111 return *other;
112 }
113 if other.is_zero() {
114 return *self;
115 }
116 if self.0 == f64::INFINITY || other.0 == f64::INFINITY {
117 return Self(f64::INFINITY);
118 }
119
120 let maximum = self.0.max(other.0);
121 Self(maximum + ((self.0 - maximum).exp() + (other.0 - maximum).exp()).ln())
122 }
123
124 fn times(&self, other: &Self) -> Self {
125 if self.is_zero() || other.is_zero() {
126 Self::zero()
127 } else {
128 Self(self.0 + other.0)
129 }
130 }
131
132 fn is_zero(&self) -> bool {
133 self.0 == f64::NEG_INFINITY
134 }
135}
136
137#[derive(Debug, Clone, PartialEq)]
139pub struct RelationEntry<K> {
140 pub doc_id: DocId,
141 pub value: K,
142}
143
144impl<K> RelationEntry<K> {
145 pub fn new(doc_id: DocId, value: K) -> Self {
146 Self { doc_id, value }
147 }
148}
149
150#[derive(Debug, Clone, Default, PartialEq)]
155pub struct Relation<K> {
156 entries: Vec<RelationEntry<K>>,
157}
158
159impl<K> Relation<K> {
160 pub fn new() -> Self {
162 Self {
163 entries: Vec::new(),
164 }
165 }
166
167 pub fn entries(&self) -> &[RelationEntry<K>] {
169 &self.entries
170 }
171
172 pub fn get(&self, doc_id: DocId) -> Option<&K> {
174 self.entries
175 .binary_search_by_key(&doc_id, |entry| entry.doc_id)
176 .ok()
177 .map(|index| &self.entries[index].value)
178 }
179
180 pub fn support(&self) -> DocSet {
182 DocSet::from_sorted_unchecked(self.entries.iter().map(|entry| entry.doc_id).collect())
183 }
184
185 pub fn len(&self) -> usize {
187 self.entries.len()
188 }
189
190 pub fn is_empty(&self) -> bool {
192 self.entries.is_empty()
193 }
194
195 pub fn iter(&self) -> std::slice::Iter<'_, RelationEntry<K>> {
197 self.entries.iter()
198 }
199}
200
201impl<K: Semiring> Relation<K> {
202 pub fn from_support(support: &DocSet) -> Self {
205 Self::from_terms(
206 support
207 .iter()
208 .map(|doc_id| RelationEntry::new(doc_id, K::one())),
209 )
210 }
211
212 pub fn from_terms<I>(terms: I) -> Self
215 where
216 I: IntoIterator<Item = RelationEntry<K>>,
217 {
218 let mut values = BTreeMap::<DocId, K>::new();
219 for term in terms {
220 if term.value.is_zero() {
221 continue;
222 }
223 values
224 .entry(term.doc_id)
225 .and_modify(|value| *value = value.plus(&term.value))
226 .or_insert(term.value);
227 }
228
229 let entries = values
230 .into_iter()
231 .filter_map(|(doc_id, value)| {
232 (!value.is_zero()).then_some(RelationEntry { doc_id, value })
233 })
234 .collect();
235 Self { entries }
236 }
237
238 pub fn singleton(doc_id: DocId, value: K) -> Self {
241 Self::from_terms([RelationEntry::new(doc_id, value)])
242 }
243
244 pub fn plus(&self, other: &Self) -> Self {
246 let mut entries = Vec::with_capacity(self.len() + other.len());
247 let (mut left, mut right) = (0, 0);
248
249 while left < self.len() && right < other.len() {
250 match self.entries[left].doc_id.cmp(&other.entries[right].doc_id) {
251 std::cmp::Ordering::Less => {
252 entries.push(self.entries[left].clone());
253 left += 1;
254 }
255 std::cmp::Ordering::Equal => {
256 let value = self.entries[left].value.plus(&other.entries[right].value);
257 if !value.is_zero() {
258 entries.push(RelationEntry::new(self.entries[left].doc_id, value));
259 }
260 left += 1;
261 right += 1;
262 }
263 std::cmp::Ordering::Greater => {
264 entries.push(other.entries[right].clone());
265 right += 1;
266 }
267 }
268 }
269
270 entries.extend_from_slice(&self.entries[left..]);
271 entries.extend_from_slice(&other.entries[right..]);
272 Self { entries }
273 }
274
275 pub fn times(&self, other: &Self) -> Self {
277 let mut entries = Vec::with_capacity(self.len().min(other.len()));
278 let (mut left, mut right) = (0, 0);
279
280 while left < self.len() && right < other.len() {
281 match self.entries[left].doc_id.cmp(&other.entries[right].doc_id) {
282 std::cmp::Ordering::Less => left += 1,
283 std::cmp::Ordering::Equal => {
284 let value = self.entries[left].value.times(&other.entries[right].value);
285 if !value.is_zero() {
286 entries.push(RelationEntry::new(self.entries[left].doc_id, value));
287 }
288 left += 1;
289 right += 1;
290 }
291 std::cmp::Ordering::Greater => right += 1,
292 }
293 }
294
295 Self { entries }
296 }
297}
298
299impl From<&DocSet> for Relation<bool> {
300 fn from(support: &DocSet) -> Self {
301 Self::from_support(support)
302 }
303}
304
305impl From<DocSet> for Relation<bool> {
306 fn from(support: DocSet) -> Self {
307 Self::from_support(&support)
308 }
309}
310
311impl<K> IntoIterator for Relation<K> {
312 type Item = RelationEntry<K>;
313 type IntoIter = std::vec::IntoIter<RelationEntry<K>>;
314
315 fn into_iter(self) -> Self::IntoIter {
316 self.entries.into_iter()
317 }
318}
319
320impl<'a, K> IntoIterator for &'a Relation<K> {
321 type Item = &'a RelationEntry<K>;
322 type IntoIter = std::slice::Iter<'a, RelationEntry<K>>;
323
324 fn into_iter(self) -> Self::IntoIter {
325 self.entries.iter()
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::{LogSemiring, Relation, RelationEntry, Semiring};
332 use crate::DocSet;
333
334 #[test]
335 fn boolean_relation_lifts_set_union_and_intersection() {
336 let left = Relation::<bool>::from_support(&DocSet::from(vec![1, 3]));
337 let right = Relation::<bool>::from_support(&DocSet::from(vec![2, 3]));
338
339 assert_eq!(left.plus(&right).support(), DocSet::from(vec![1, 2, 3]));
340 assert_eq!(left.times(&right).support(), DocSet::from(vec![3]));
341 }
342
343 #[test]
344 fn duplicate_terms_are_combined_and_zero_is_not_stored() {
345 let relation = Relation::from_terms([
346 RelationEntry::new(1, false),
347 RelationEntry::new(2, true),
348 RelationEntry::new(2, true),
349 ]);
350
351 assert_eq!(relation.support(), DocSet::from(vec![2]));
352 assert_eq!(relation.get(2), Some(&true));
353 }
354
355 #[test]
356 fn log_semiring_uses_log_sum_exp_and_log_space_multiplication() {
357 let point_two = LogSemiring::from_weight(0.2).unwrap();
358 let point_three = LogSemiring::from_weight(0.3).unwrap();
359
360 let sum = point_two.plus(&point_three);
361 let product = point_two.times(&point_three);
362
363 assert!((sum.weight() - 0.5).abs() < 1e-12);
364 assert!((product.weight() - 0.06).abs() < 1e-12);
365 assert!(LogSemiring::from_log(f64::NAN).is_none());
366 }
367}