Skip to main content

weavatrix_graph/
filter.rs

1use crate::{Confidence, Edge, EdgeKind, EvidenceKind};
2use crate::{String, Vec};
3
4#[derive(Debug, Clone, PartialEq, Eq, Default)]
5pub struct EdgeFilter {
6    kinds: Vec<EdgeKind>,
7    evidence: Vec<EvidenceKind>,
8    extractors: Vec<String>,
9    minimum_confidence: Option<Confidence>,
10}
11
12impl EdgeFilter {
13    #[must_use]
14    pub const fn new() -> Self {
15        Self {
16            kinds: Vec::new(),
17            evidence: Vec::new(),
18            extractors: Vec::new(),
19            minimum_confidence: None,
20        }
21    }
22
23    #[must_use]
24    pub fn with_kind(mut self, kind: EdgeKind) -> Self {
25        if !self.kinds.contains(&kind) {
26            self.kinds.push(kind);
27        }
28        self
29    }
30
31    #[must_use]
32    pub fn with_evidence(mut self, evidence: EvidenceKind) -> Self {
33        if !self.evidence.contains(&evidence) {
34            self.evidence.push(evidence);
35        }
36        self
37    }
38
39    #[must_use]
40    pub fn with_extractor(mut self, extractor: impl Into<String>) -> Self {
41        let extractor = extractor.into();
42        if !self.extractors.contains(&extractor) {
43            self.extractors.push(extractor);
44        }
45        self
46    }
47
48    #[must_use]
49    pub const fn with_minimum_confidence(mut self, confidence: Confidence) -> Self {
50        self.minimum_confidence = Some(confidence);
51        self
52    }
53
54    #[must_use]
55    pub fn matches(&self, edge: &Edge) -> bool {
56        (self.kinds.is_empty() || self.kinds.contains(&edge.kind))
57            && (self.evidence.is_empty() || self.evidence.contains(&edge.provenance.evidence))
58            && (self.extractors.is_empty() || self.extractors.contains(&edge.provenance.extractor))
59            && self.minimum_confidence.is_none_or(|minimum| {
60                confidence_rank(edge.provenance.confidence) >= confidence_rank(minimum)
61            })
62    }
63}
64
65const fn confidence_rank(confidence: Confidence) -> u8 {
66    match confidence {
67        Confidence::Exact => 4,
68        Confidence::High => 3,
69        Confidence::Medium => 2,
70        Confidence::Low => 1,
71    }
72}