1use objects::store::ObjectStore;
3
4use crate::{ObjectId, ObjectInfo, ObjectType, Result};
5
6#[derive(Debug, Clone, Default, PartialEq, Eq)]
7pub struct ObjectAvailabilityPlan {
8 pub have_objects: Vec<ObjectId>,
9 pub want_objects: Vec<ObjectId>,
10 pub present_objects: Vec<ObjectId>,
11 pub missing_objects: Vec<ObjectId>,
12 pub resumable_objects: Vec<ObjectId>,
13 pub lazy_objects: Vec<ObjectId>,
14 pub partial_fetch_allowed: bool,
15}
16
17pub fn has_object(store: &impl ObjectStore, info: &ObjectInfo) -> Result<bool> {
18 match (&info.id, info.obj_type) {
19 (ObjectId::Hash(hash), ObjectType::Blob) => Ok(store.has_blob(hash)?),
20 (ObjectId::Hash(hash), ObjectType::Tree) => Ok(store.has_tree(hash)?),
21 (ObjectId::StateId(state_id), ObjectType::State) => Ok(store.has_state(state_id)?),
22 (ObjectId::StateAttachment { state, id, kind: _ }, ObjectType::StateAttachment) => {
23 Ok(store.get_state_attachment(state, id)?.is_some())
24 }
25 (ObjectId::Hash(_), ObjectType::Redaction) => Ok(false),
32 (ObjectId::StateId(_), ObjectType::StateVisibility) => Ok(false),
36 _ => Ok(false),
37 }
38}
39
40pub fn plan_object_availability(
41 store: &impl ObjectStore,
42 objects: &[ObjectInfo],
43) -> Result<ObjectAvailabilityPlan> {
44 let mut plan = ObjectAvailabilityPlan::default();
45
46 for info in objects {
47 if has_object(store, info)? {
48 plan.have_objects.push(info.id.clone());
49 plan.present_objects.push(info.id.clone());
50 } else {
51 plan.want_objects.push(info.id.clone());
52 plan.missing_objects.push(info.id.clone());
53 }
54 }
55
56 Ok(plan)
57}
58
59impl ObjectAvailabilityPlan {
60 pub fn with_partial_fetch_allowed(mut self, allowed: bool) -> Self {
61 self.partial_fetch_allowed = allowed;
62 self
63 }
64
65 pub fn is_complete(&self) -> bool {
66 self.want_objects.is_empty()
67 && self.missing_objects.is_empty()
68 && self.resumable_objects.is_empty()
69 && self.lazy_objects.is_empty()
70 }
71
72 pub fn has_partial_fetch_candidates(&self) -> bool {
73 !self.resumable_objects.is_empty() || !self.lazy_objects.is_empty()
74 }
75}
76
77#[cfg(test)]
78mod tests {
79 use objects::{
80 object::{Blob, ContentHash, StateId, Tree},
81 store::{ObjectStore, Result as StoreResult},
82 };
83
84 use super::*;
85
86 #[derive(Default)]
87 struct DummyStore {
88 blob: Option<ContentHash>,
89 state: Option<StateId>,
90 }
91
92 impl ObjectStore for DummyStore {
93 fn get_blob(&self, _hash: &ContentHash) -> StoreResult<Option<Blob>> {
94 Ok(None)
95 }
96
97 fn put_blob(&self, _blob: &Blob) -> StoreResult<ContentHash> {
98 unreachable!("not used in test")
99 }
100
101 fn has_blob(&self, hash: &ContentHash) -> StoreResult<bool> {
102 Ok(self.blob == Some(*hash))
103 }
104
105 fn get_tree(&self, _hash: &ContentHash) -> StoreResult<Option<Tree>> {
106 Ok(None)
107 }
108
109 fn put_tree(&self, _tree: &Tree) -> StoreResult<ContentHash> {
110 unreachable!("not used in test")
111 }
112
113 fn has_tree(&self, _hash: &ContentHash) -> StoreResult<bool> {
114 Ok(false)
115 }
116
117 fn get_state(&self, _id: &StateId) -> StoreResult<Option<objects::object::State>> {
118 Ok(None)
119 }
120
121 fn put_state(&self, _state: &objects::object::State) -> StoreResult<()> {
122 unreachable!("not used in test")
123 }
124
125 fn has_state(&self, id: &StateId) -> StoreResult<bool> {
126 Ok(self.state == Some(*id))
127 }
128
129 fn list_states(&self) -> StoreResult<Vec<StateId>> {
130 Ok(vec![])
131 }
132
133 fn get_action(
134 &self,
135 _id: &objects::object::ActionId,
136 ) -> StoreResult<Option<objects::object::Action>> {
137 Ok(None)
138 }
139
140 fn put_action(
141 &self,
142 _action: &mut objects::object::Action,
143 ) -> StoreResult<objects::object::ActionId> {
144 unreachable!("not used in test")
145 }
146
147 fn list_actions(&self) -> StoreResult<Vec<objects::object::ActionId>> {
148 Ok(vec![])
149 }
150
151 fn list_blobs(&self) -> StoreResult<Vec<ContentHash>> {
152 Ok(vec![])
153 }
154
155 fn list_trees(&self) -> StoreResult<Vec<ContentHash>> {
156 Ok(vec![])
157 }
158 }
159
160 #[test]
161 fn test_plan_tracks_present_and_missing_objects() {
162 let blob = Blob::new(b"hello".to_vec());
163 let blob_hash = blob.hash();
164 let store = DummyStore {
165 blob: Some(blob_hash),
166 state: None,
167 };
168 let missing_hash = ContentHash::from_bytes([7; 32]);
169 let objects = vec![
170 ObjectInfo {
171 id: ObjectId::Hash(blob_hash),
172 obj_type: ObjectType::Blob,
173 size: blob.size() as u64,
174 delta_base: None,
175 },
176 ObjectInfo {
177 id: ObjectId::Hash(missing_hash),
178 obj_type: ObjectType::Tree,
179 size: 0,
180 delta_base: None,
181 },
182 ];
183
184 let plan = plan_object_availability(&store, &objects).unwrap();
185
186 assert_eq!(plan.have_objects.len(), 1);
187 assert_eq!(plan.want_objects.len(), 1);
188 assert_eq!(plan.present_objects.len(), 1);
189 assert_eq!(plan.missing_objects.len(), 1);
190 assert!(!plan.is_complete());
191 }
192
193 #[test]
194 fn missing_state_objects_are_requested() {
195 let store = DummyStore::default();
196 let state = StateId::from_bytes([9; 32]);
197 let objects = vec![ObjectInfo {
198 id: ObjectId::StateId(state),
199 obj_type: ObjectType::State,
200 size: 0,
201 delta_base: None,
202 }];
203
204 let plan = plan_object_availability(&store, &objects).unwrap();
205
206 assert!(plan.have_objects.is_empty());
207 assert_eq!(plan.want_objects, vec![ObjectId::StateId(state)]);
208 assert_eq!(plan.missing_objects, vec![ObjectId::StateId(state)]);
209 }
210
211 #[test]
212 fn immutable_state_objects_are_not_requested_when_present() {
213 let state = StateId::from_bytes([9; 32]);
214 let store = DummyStore {
215 state: Some(state),
216 ..DummyStore::default()
217 };
218 let objects = vec![ObjectInfo {
219 id: ObjectId::StateId(state),
220 obj_type: ObjectType::State,
221 size: 0,
222 delta_base: None,
223 }];
224
225 let plan = plan_object_availability(&store, &objects).unwrap();
226
227 assert_eq!(plan.have_objects, vec![ObjectId::StateId(state)]);
228 assert!(plan.want_objects.is_empty());
229 }
230
231 #[test]
232 fn test_partial_fetch_flag_helpers() {
233 let plan = ObjectAvailabilityPlan::default().with_partial_fetch_allowed(true);
234
235 assert!(plan.partial_fetch_allowed);
236 assert!(!plan.has_partial_fetch_candidates());
237 assert!(plan.is_complete());
238 }
239}