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