1use std::collections::HashMap;
4
5use refget_model::SequenceMetadata;
6
7use crate::{SequenceStore, StoreResult, extract_subsequence};
8
9pub struct InMemorySequenceStore {
11 index: HashMap<String, (SequenceMetadata, Vec<u8>)>,
13}
14
15impl InMemorySequenceStore {
16 pub fn new() -> Self {
18 Self { index: HashMap::new() }
19 }
20
21 pub fn add(&mut self, metadata: SequenceMetadata, sequence: Vec<u8>) {
23 self.index.insert(metadata.md5.clone(), (metadata.clone(), sequence.clone()));
24 self.index.insert(metadata.sha512t24u.clone(), (metadata, sequence));
25 }
26}
27
28impl Default for InMemorySequenceStore {
29 fn default() -> Self {
30 Self::new()
31 }
32}
33
34impl SequenceStore for InMemorySequenceStore {
35 fn get_sequence(
36 &self,
37 digest: &str,
38 start: Option<u64>,
39 end: Option<u64>,
40 ) -> StoreResult<Option<Vec<u8>>> {
41 let Some((_, seq)) = self.index.get(digest) else {
42 return Ok(None);
43 };
44 Ok(Some(extract_subsequence(seq, start, end)))
45 }
46
47 fn get_metadata(&self, digest: &str) -> StoreResult<Option<SequenceMetadata>> {
48 Ok(self.index.get(digest).map(|(meta, _)| meta.clone()))
49 }
50
51 fn get_length(&self, digest: &str) -> StoreResult<Option<u64>> {
52 Ok(self.index.get(digest).map(|(meta, _)| meta.length))
53 }
54}
55
56#[cfg(test)]
57mod tests {
58 use super::*;
59 use refget_model::SequenceMetadata;
60
61 fn test_metadata() -> SequenceMetadata {
62 SequenceMetadata {
63 md5: "abc123".to_string(),
64 sha512t24u: "SQ.xyz789".to_string(),
65 length: 4,
66 aliases: vec![],
67 circular: false,
68 }
69 }
70
71 #[test]
72 fn test_add_and_get() {
73 let mut store = InMemorySequenceStore::new();
74 store.add(test_metadata(), b"ACGT".to_vec());
75
76 let seq = store.get_sequence("abc123", None, None).unwrap().unwrap();
77 assert_eq!(seq, b"ACGT");
78
79 let seq = store.get_sequence("SQ.xyz789", None, None).unwrap().unwrap();
80 assert_eq!(seq, b"ACGT");
81 }
82
83 #[test]
84 fn test_subsequence() {
85 let mut store = InMemorySequenceStore::new();
86 store.add(test_metadata(), b"ACGT".to_vec());
87
88 let seq = store.get_sequence("abc123", Some(1), Some(3)).unwrap().unwrap();
89 assert_eq!(seq, b"CG");
90 }
91
92 #[test]
93 fn test_not_found() {
94 let store = InMemorySequenceStore::new();
95 assert!(store.get_sequence("missing", None, None).unwrap().is_none());
96 }
97
98 #[test]
99 fn test_metadata_lookup() {
100 let mut store = InMemorySequenceStore::new();
101 let meta = test_metadata();
102 store.add(meta.clone(), b"ACGT".to_vec());
103
104 let found = store.get_metadata("abc123").unwrap().unwrap();
105 assert_eq!(found, meta);
106 }
107
108 #[test]
109 fn test_get_sequence_start_at_length_returns_empty() {
110 let mut store = InMemorySequenceStore::new();
111 store.add(test_metadata(), b"ACGT".to_vec());
112
113 let seq = store.get_sequence("abc123", Some(4), None).unwrap().unwrap();
114 assert!(seq.is_empty());
115 }
116
117 #[test]
118 fn test_get_sequence_start_beyond_length_returns_empty() {
119 let mut store = InMemorySequenceStore::new();
120 store.add(test_metadata(), b"ACGT".to_vec());
121
122 let seq = store.get_sequence("abc123", Some(100), None).unwrap().unwrap();
123 assert!(seq.is_empty());
124 }
125
126 #[test]
127 fn test_get_sequence_end_beyond_length_clamps() {
128 let mut store = InMemorySequenceStore::new();
129 store.add(test_metadata(), b"ACGT".to_vec());
130
131 let seq = store.get_sequence("abc123", Some(2), Some(100)).unwrap().unwrap();
132 assert_eq!(seq, b"GT");
133 }
134
135 #[test]
136 fn test_get_metadata_non_existent_returns_none() {
137 let store = InMemorySequenceStore::new();
138 let result = store.get_metadata("no_such_digest").unwrap();
139 assert!(result.is_none());
140 }
141
142 #[test]
143 fn test_add_same_sequence_twice_overwrites() {
144 let mut store = InMemorySequenceStore::new();
145 let meta = test_metadata();
146 store.add(meta.clone(), b"ACGT".to_vec());
147 store.add(meta.clone(), b"TTTT".to_vec());
148
149 let seq = store.get_sequence("abc123", None, None).unwrap().unwrap();
150 assert_eq!(seq, b"TTTT");
151
152 let seq = store.get_sequence("SQ.xyz789", None, None).unwrap().unwrap();
153 assert_eq!(seq, b"TTTT");
154 }
155}