Skip to main content

refget_store/
memory.rs

1//! In-memory sequence store implementation.
2
3use std::collections::HashMap;
4
5use refget_model::SequenceMetadata;
6
7use crate::{SequenceStore, StoreResult, extract_subsequence};
8
9/// An in-memory sequence store suitable for testing and small datasets.
10pub struct InMemorySequenceStore {
11    /// Map from digest (MD5 or sha512t24u) to (metadata, sequence bytes).
12    index: HashMap<String, (SequenceMetadata, Vec<u8>)>,
13}
14
15impl InMemorySequenceStore {
16    /// Create a new empty store.
17    pub fn new() -> Self {
18        Self { index: HashMap::new() }
19    }
20
21    /// Add a sequence to the store.
22    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}