1use std::{
11 collections::{BTreeMap, BTreeSet},
12 ops::Deref,
13 sync::Arc,
14};
15
16use object_rainbow::{
17 Address, ByteNode, FailFuture, Fetch, FetchBytes, Hash, ListHashes, Node, Object, Output,
18 Parse, ParseInput, ParseSliceExtra, PointInput, Resolve, Singular, Tagged, ToOutput,
19 Topological, Traversible,
20};
21use object_rainbow_fetchall::fetchall;
22use object_rainbow_local_map::LocalMap;
23
24#[derive(Clone)]
25struct MarshalledInner {
26 data: Arc<[u8]>,
27 root: Hash,
28 at: usize,
29}
30
31impl MarshalledInner {
32 fn read_usize(&self, at: usize) -> object_rainbow::Result<usize> {
33 u64::from_le_bytes(
34 self.data
35 .get(
36 at..at
37 .checked_add(8)
38 .ok_or(object_rainbow::Error::UnsupportedLength)?,
39 )
40 .ok_or(object_rainbow::Error::UnsupportedLength)?
41 .try_into()
42 .unwrap(),
43 )
44 .try_into()
45 .map_err(|_| object_rainbow::Error::UnsupportedLength)
46 }
47
48 fn data_len(&self) -> object_rainbow::Result<usize> {
49 self.read_usize(self.at)
50 }
51
52 fn data_begin(&self) -> object_rainbow::Result<usize> {
53 self.at
54 .checked_add(8)
55 .ok_or(object_rainbow::Error::UnsupportedLength)
56 }
57
58 fn data(&self) -> object_rainbow::Result<&[u8]> {
59 Ok(&self.data[self.data_begin()?..self.data_end()?])
60 }
61
62 fn data_end(&self) -> object_rainbow::Result<usize> {
63 self.data_begin()?
64 .checked_add(self.data_len()?)
65 .ok_or(object_rainbow::Error::UnsupportedLength)
66 }
67
68 fn reference_at(&self, index: usize) -> object_rainbow::Result<usize> {
69 self.read_usize(
70 self.data_end()?
71 .checked_add(
72 index
73 .checked_mul(8)
74 .ok_or(object_rainbow::Error::UnsupportedLength)?,
75 )
76 .ok_or(object_rainbow::Error::UnsupportedLength)?,
77 )
78 }
79
80 fn data_vec(&self) -> object_rainbow::Result<Vec<u8>> {
81 Ok(self.data()?.into())
82 }
83
84 fn resolve_node(&self, address: Address) -> object_rainbow::Result<ByteNode> {
85 let referenced = MarshalledInner {
86 data: self.data.clone(),
87 root: address.hash,
88 at: self.reference_at(address.index)?,
89 };
90 Ok((referenced.data_vec()?, referenced.to_resolve()))
91 }
92
93 fn to_resolve(&self) -> Arc<dyn Resolve> {
94 Arc::new(self.clone())
95 }
96}
97
98impl FetchBytes for MarshalledInner {
99 fn fetch_bytes(&'_ self) -> FailFuture<'_, ByteNode> {
100 Box::pin(async move { Ok((self.data_vec()?, self.to_resolve())) })
101 }
102
103 fn fetch_data(&'_ self) -> FailFuture<'_, Vec<u8>> {
104 Box::pin(async move { self.data_vec() })
105 }
106
107 fn fetch_bytes_local(&self) -> object_rainbow::Result<Option<ByteNode>> {
108 Ok(Some((self.data_vec()?, self.to_resolve())))
109 }
110}
111
112impl Resolve for MarshalledInner {
113 fn resolve(&'_ self, address: Address, _: &Arc<dyn Resolve>) -> FailFuture<'_, ByteNode> {
114 Box::pin(async move { self.resolve_node(address) })
115 }
116
117 fn resolve_data(&'_ self, address: Address) -> FailFuture<'_, Vec<u8>> {
118 Box::pin(async move {
119 let (data, _) = self.resolve_node(address)?;
120 Ok(data)
121 })
122 }
123
124 fn try_resolve_local(
125 &self,
126 address: Address,
127 _: &Arc<dyn Resolve>,
128 ) -> object_rainbow::Result<Option<ByteNode>> {
129 self.resolve_node(address).map(Some)
130 }
131}
132
133impl Singular for MarshalledInner {
134 fn hash(&self) -> Hash {
135 self.root
136 }
137}
138
139enum Action {
140 WriteLocation { at: usize, of: Hash },
141 SaveFull { hash: Hash },
142 FinishLocation { at: usize, of: Hash },
143}
144
145trait ToBytes: Copy {
146 fn to_bytes(self) -> [u8; 8];
147}
148
149impl ToBytes for u64 {
150 fn to_bytes(self) -> [u8; 8] {
151 self.to_le_bytes()
152 }
153}
154
155impl ToBytes for usize {
156 fn to_bytes(self) -> [u8; 8] {
157 (self as u64).to_bytes()
158 }
159}
160
161#[derive(Clone)]
162pub struct MarshalledRoot {
163 marshalled: MarshalledInner,
164}
165
166impl FetchBytes for MarshalledRoot {
167 fn fetch_bytes(&'_ self) -> FailFuture<'_, ByteNode> {
168 self.marshalled.fetch_bytes()
169 }
170
171 fn fetch_data(&'_ self) -> FailFuture<'_, Vec<u8>> {
172 self.marshalled.fetch_data()
173 }
174
175 fn fetch_bytes_local(&self) -> object_rainbow::Result<Option<ByteNode>> {
176 self.marshalled.fetch_bytes_local()
177 }
178}
179
180impl Singular for MarshalledRoot {
181 fn hash(&self) -> Hash {
182 self.marshalled.hash()
183 }
184}
185
186pub fn marshall(map: &LocalMap, root: Hash) -> MarshalledRoot {
187 let mut data = Vec::<u8>::new();
188 let mut locations = BTreeMap::<Hash, usize>::new();
189 let mut started = BTreeSet::<Hash>::new();
190 let mut stack = vec![Action::SaveFull { hash: root }];
191 while let Some(action) = stack.pop() {
192 match action {
193 Action::WriteLocation { at, of } => {
194 data[at..at + 8].copy_from_slice(&locations.get(&of).unwrap().to_bytes());
195 }
196 Action::SaveFull { hash } => {
197 if locations.contains_key(&hash) {
198 continue;
199 }
200 assert!(started.insert(hash));
201 let (references, d) = map.get(hash).unwrap();
202 stack.push(Action::FinishLocation {
203 at: data.len(),
204 of: hash,
205 });
206 data.extend_from_slice(&d.len().to_bytes());
207 data.extend_from_slice(d);
208 for &hash in references {
209 stack.push(Action::WriteLocation {
210 at: data.len(),
211 of: hash,
212 });
213 data.extend_from_slice(&u64::MAX.to_bytes());
214 stack.push(Action::SaveFull { hash });
215 }
216 }
217 Action::FinishLocation { at, of } => {
218 assert!(started.contains(&of));
219 assert!(locations.insert(of, at).is_none());
220 }
221 }
222 }
223 assert_eq!(*locations.get(&root).unwrap(), 0);
224 let data = Arc::from(data);
225 let marshalled = MarshalledInner { data, root, at: 0 };
226 MarshalledRoot { marshalled }
227}
228
229impl ToOutput for MarshalledRoot {
230 fn to_output(&self, output: &mut impl Output) {
231 self.marshalled.root.to_output(output);
232 self.marshalled.data.to_output(output);
233 }
234}
235
236impl Tagged for MarshalledRoot {}
237impl ListHashes for MarshalledRoot {}
238impl Topological for MarshalledRoot {}
239
240impl<I: ParseInput> Parse<I> for MarshalledRoot {
241 fn parse(mut input: I) -> object_rainbow::Result<Self> {
242 let root = input.parse_inline()?;
243 let data = Arc::<[u8]>::from(input.parse_all()?.as_ref());
244 let marshalled = MarshalledInner { data, root, at: 0 };
245 Ok(Self { marshalled })
246 }
247}
248
249#[derive(Tagged)]
250pub struct Marshalled<T> {
251 root: MarshalledRoot,
252 object: T,
253}
254
255impl<T: ToOutput> ToOutput for Marshalled<T> {
256 fn to_output(&self, output: &mut impl Output) {
257 if output.is_mangling() {
258 self.object.to_output(output);
259 }
260 self.root.to_output(output);
261 }
262}
263
264impl<T> ListHashes for Marshalled<T> {}
265impl<T> Topological for Marshalled<T> {}
266
267impl<I: PointInput, T: Object<I::Extra>> Parse<I> for Marshalled<T> {
268 fn parse(input: I) -> object_rainbow::Result<Self> {
269 let extra = input.extra().clone();
270 let root = input.parse::<MarshalledRoot>()?;
271 let object = T::parse_slice_extra(
272 root.marshalled.data()?,
273 &root.marshalled.to_resolve(),
274 &extra,
275 )?;
276 if object.full_hash() != root.hash() {
277 return Err(object_rainbow::Error::FullHashMismatch);
278 }
279 Ok(Self { root, object })
280 }
281}
282
283impl<T: ToOutput> FetchBytes for Marshalled<T> {
284 fn fetch_bytes(&'_ self) -> FailFuture<'_, ByteNode> {
285 self.root.fetch_bytes()
286 }
287
288 fn fetch_data(&'_ self) -> FailFuture<'_, Vec<u8>> {
289 self.root.fetch_data()
290 }
291
292 fn fetch_bytes_local(&self) -> object_rainbow::Result<Option<ByteNode>> {
293 self.root.fetch_bytes_local()
294 }
295
296 fn fetch_data_local(&self) -> Option<Vec<u8>> {
297 Some(self.object.output())
298 }
299}
300
301impl<T: Send + Sync + Clone + ToOutput> Fetch for Marshalled<T> {
302 type T = T;
303
304 fn fetch_full(&'_ self) -> FailFuture<'_, Node<Self::T>> {
305 Box::pin(async move { Ok((self.object.clone(), self.root.marshalled.to_resolve())) })
306 }
307
308 fn fetch(&'_ self) -> FailFuture<'_, Self::T> {
309 Box::pin(async move { Ok(self.object.clone()) })
310 }
311
312 fn try_fetch_local(&self) -> object_rainbow::Result<Option<Node<Self::T>>> {
313 Ok(Some((
314 self.object.clone(),
315 self.root.marshalled.to_resolve(),
316 )))
317 }
318
319 fn fetch_local(&self) -> Option<Self::T> {
320 Some(self.object.clone())
321 }
322
323 fn try_unwrap(self: Arc<Self>) -> Option<Self::T> {
324 Arc::try_unwrap(self).ok().map(|Self { object, .. }| object)
325 }
326}
327
328impl<T: Send + Sync + ToOutput> Singular for Marshalled<T> {
329 fn hash(&self) -> Hash {
330 self.root.hash()
331 }
332}
333
334impl<T: Traversible> Marshalled<T> {
335 pub async fn new(object: T) -> object_rainbow::Result<Self> {
336 let map = fetchall(&object).await?;
337 let root = marshall(&map, object.full_hash());
338 Ok(Self { root, object })
339 }
340}
341
342impl<T> Deref for Marshalled<T> {
343 type Target = T;
344
345 fn deref(&self) -> &Self::Target {
346 &self.object
347 }
348}