1pub mod tuple;
11
12use crate::error::PrivacyError;
13use crate::types::PrivacyTuple;
14use sha2::{Digest, Sha256};
15
16extern crate alloc;
17use alloc::{collections::BTreeMap, string::{String, ToString}, vec::Vec};
18
19#[derive(Debug, Default)]
24pub struct TupleChain {
25 tuples: BTreeMap<String, PrivacyTuple>,
26 anchors: BTreeMap<String, String>,
28 total_count: u64,
30}
31
32impl TupleChain {
33 pub fn new() -> Self {
34 Self::default()
35 }
36
37 pub fn insert(&mut self, tuple: PrivacyTuple) -> String {
41 let id = self.tuple_id(&tuple);
42 self.tuples.insert(id.clone(), tuple);
43 self.total_count += 1;
44 id
45 }
46
47 pub fn query_by_subject(&self, subject: &str, now_ms: u64) -> Vec<&PrivacyTuple> {
51 self.tuples
52 .values()
53 .filter(|t| t.subject == subject && !t.is_expired(now_ms))
54 .collect()
55 }
56
57 pub fn query_by_predicate(&self, predicate: &str, now_ms: u64) -> Vec<&PrivacyTuple> {
59 self.tuples
60 .values()
61 .filter(|t| t.predicate == predicate && !t.is_expired(now_ms))
62 .collect()
63 }
64
65 pub fn get(&self, id: &str) -> Option<&PrivacyTuple> {
67 self.tuples.get(id)
68 }
69
70 pub fn prune_expired(&mut self, now_ms: u64) -> usize {
74 let expired: Vec<String> = self
75 .tuples
76 .iter()
77 .filter(|(_, t)| t.is_expired(now_ms))
78 .map(|(id, _)| id.clone())
79 .collect();
80 let count = expired.len();
81 for id in expired {
82 self.tuples.remove(&id);
83 }
84 count
85 }
86
87 pub fn anchor_all(&mut self) -> String {
91 let mut hasher = Sha256::new();
92 for (id, tuple) in &mut self.tuples {
93 if tuple.anchor.is_none() {
94 let anchor = tuple::anchor_tuple(tuple);
95 self.anchors.insert(id.clone(), anchor.clone());
96 hasher.update(anchor.as_bytes());
97 }
98 }
99 hex::encode(hasher.finalize())
100 }
101
102 pub fn homomorphic_add(
107 &self,
108 id_a: &str,
109 id_b: &str,
110 ) -> Result<Vec<u8>, PrivacyError> {
111 let a = self.tuples.get(id_a).ok_or_else(|| PrivacyError::TupleNotFound(id_a.into()))?;
112 let b = self.tuples.get(id_b).ok_or_else(|| PrivacyError::TupleNotFound(id_b.into()))?;
113
114 let len = a.object.len().min(b.object.len());
115 let result: Vec<u8> = a.object[..len]
116 .iter()
117 .zip(b.object[..len].iter())
118 .map(|(x, y)| x ^ y)
119 .collect();
120
121 Ok(result)
122 }
123
124 pub fn active_count(&self, now_ms: u64) -> usize {
126 self.tuples.values().filter(|t| !t.is_expired(now_ms)).count()
127 }
128
129 pub fn total_count(&self) -> u64 {
131 self.total_count
132 }
133
134 fn tuple_id(&self, tuple: &PrivacyTuple) -> String {
137 let mut hasher = Sha256::new();
138 hasher.update(tuple.subject.as_bytes());
139 hasher.update(tuple.predicate.as_bytes());
140 hasher.update(tuple.expiry_ms.to_le_bytes());
141 hasher.update(self.total_count.to_le_bytes());
142 hex::encode(hasher.finalize())[..16].to_string()
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149 use crate::ledger::tuple::build_tuple;
150
151 #[test]
152 fn test_insert_query() {
153 let mut chain = TupleChain::new();
154 let seed = [0u8; 32];
155 let t = build_tuple("alice", "owns", b"data".to_vec(), 9_999_999_999, 1e-6, &seed).unwrap();
156 let id = chain.insert(t);
157 assert!(!id.is_empty());
158 let results = chain.query_by_subject("alice", 0);
159 assert_eq!(results.len(), 1);
160 }
161
162 #[test]
163 fn test_prune_expired() {
164 let mut chain = TupleChain::new();
165 let seed = [0u8; 32];
166 let t = build_tuple("bob", "had", b"old".to_vec(), 1, 1e-6, &seed).unwrap();
168 chain.insert(t);
169 let pruned = chain.prune_expired(1_000_000_000);
170 assert!(pruned >= 1);
171 }
172}