Skip to main content

weavatrix_graph/
filter.rs

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