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