weavatrix_graph/
filter.rs1use 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}