sark_http/payload/owned/
segments.rs

1use std::{iter::FusedIterator, ops};
2
3use o3::{
4    buffer::{
5        Residency as _,
6        bytes::scope,
7        hold, resident, segments,
8        segments::{Nonempty, Shape, assembly},
9        storage,
10    },
11    mem,
12};
13
14use crate::{
15    owner::{handoff, handoff::stage},
16    payload::{
17        owned::{Error, Segment, iter},
18        response,
19    },
20    request::body,
21};
22
23/// Retained payload ownership preserving transport segmentation.
24#[repr(transparent)]
25#[derive(Debug)]
26pub struct Segments<'d> {
27    pub(in crate::payload::owned) storage: Shape<resident::Shared<'d>>,
28}
29
30impl<'d> Segments<'d> {
31    pub const fn new() -> Self {
32        Self {
33            storage: Shape::Empty,
34        }
35    }
36
37    pub fn try_push(
38        self,
39        segment: Segment<'d>,
40    ) -> Result<Self, mem::Rejected<(Self, Segment<'d>), assembly::Error>> {
41        let ops::ControlFlow::Continue(item) = segments::Item::new(resident::Shared::from(segment))
42        else {
43            return Ok(self);
44        };
45        match self.storage {
46            Shape::Empty => Ok(Self {
47                storage: Shape::One(item),
48            }),
49            Shape::One(first) => Self::try_merge((first, item), Segment::from_item),
50            Shape::Many(chain) => Self::try_merge((chain, item), Segment::from_item),
51        }
52    }
53
54    pub fn try_extend(
55        self,
56        source: Self,
57    ) -> Result<Self, mem::Rejected<(Self, Self), assembly::Error>> {
58        match (self.storage, source.storage) {
59            (storage, Shape::Empty) | (Shape::Empty, storage) => Ok(Self { storage }),
60            (Shape::One(first), Shape::One(second)) => Self::try_merge((first, second), Self::from),
61            (Shape::Many(first), Shape::One(second)) => {
62                Self::try_merge((first, second), Self::from)
63            }
64            (Shape::One(first), Shape::Many(second)) => {
65                Self::try_merge((first, second), Self::from)
66            }
67            (Shape::Many(first), Shape::Many(second)) => {
68                Self::try_merge((first, second), Self::from)
69            }
70        }
71    }
72
73    fn try_merge<A, B, T>(
74        parts: (A, B),
75        restore: impl FnOnce(B) -> T,
76    ) -> Result<Self, mem::Rejected<(Self, T), assembly::Error>>
77    where
78        Self: From<A>,
79        segments::Many<resident::Shared<'d>>:
80            TryFrom<(A, B), Error = mem::Rejected<(A, B), assembly::Error>>,
81    {
82        segments::Many::try_from(parts)
83            .map(|chain| Self {
84                storage: Shape::Many(chain),
85            })
86            .map_err(|rejected| {
87                rejected.map(|(first, second)| (Self::from(first), restore(second)))
88            })
89    }
90
91    #[inline]
92    pub fn charge(&self) -> hold::Charge {
93        match &self.storage {
94            Shape::Empty => hold::Charge::default(),
95            Shape::One(owner) => hold::Charge::from_owner(owner.len(), owner.resident()),
96            Shape::Many(owners) => hold::Charge::from(owners.as_chain()),
97        }
98    }
99
100    pub fn segment(&self, index: usize) -> Option<&[u8]> {
101        match &self.storage {
102            Shape::Empty => None,
103            Shape::One(segment) => (index == 0).then(|| segment.as_slice()),
104            Shape::Many(segments) => segments.as_chain().segments().nth(index).map(AsRef::as_ref),
105        }
106    }
107
108    #[doc(hidden)]
109    pub fn backings(&self) -> iter::Backings<'_, 'd> {
110        iter::Backings {
111            owners: self.storage.owners(),
112        }
113    }
114
115    pub fn into_shape(self) -> Shape<resident::Shared<'d>> {
116        self.storage
117    }
118
119    pub(crate) const fn storage(&self) -> &Shape<resident::Shared<'d>> {
120        &self.storage
121    }
122
123    pub fn take(self) -> Option<Nonempty<resident::Shared<'d>>> {
124        match self.storage {
125            Shape::Empty => None,
126            Shape::One(owner) => Some(Nonempty::One(owner)),
127            Shape::Many(owners) => Some(Nonempty::Many(owners)),
128        }
129    }
130
131    pub fn try_clone(&self) -> Result<Self, Error> {
132        let storage = match &self.storage {
133            Shape::Empty => Shape::Empty,
134            Shape::One(owner) => Shape::One(owner.try_clone().map_err(Error::Retain)?),
135            Shape::Many(source) => Shape::Many(
136                segments::Exact::try_new(source)
137                    .map_err(Error::Metadata)?
138                    .try_collect()
139                    .map_err(Error::Retain)?,
140            ),
141        };
142        Ok(Self { storage })
143    }
144}
145
146impl<'d> body::Source<'d> for Segments<'d> {
147    type Body<'a>
148        = &'a Self
149    where
150        Self: 'a,
151        'd: 'a;
152
153    fn try_project<V, T>(
154        &mut self,
155        value: V,
156        project: impl for<'a> FnOnce(V, body::View<'a, 'd>) -> T,
157    ) -> Result<T, (V, body::Error)> {
158        let plan = match &self.storage {
159            Shape::Empty => return Ok(project(value, body::View::Slice(&[]))),
160            Shape::One(owner) => {
161                return Ok(project(value, body::View::Resident(owner.as_ref())));
162            }
163            Shape::Many(chain) => chain.as_chain().try_flatten(),
164        };
165        match plan {
166            Ok(store) => {
167                let owner: resident::Shared<'d> = store.into();
168                let output = project(value, body::View::Resident(&owner));
169                *self = Self::from(owner);
170                Ok(output)
171            }
172            Err(error) => Err((value, body::Error::Owned(Error::Storage(error)))),
173        }
174    }
175
176    fn body<'a>(&'a self) -> Self::Body<'a>
177    where
178        'd: 'a,
179    {
180        self
181    }
182}
183
184impl<'d> response::Body<'d> for Segments<'d> {
185    fn len(&self) -> usize {
186        self.storage.len()
187    }
188
189    fn as_slice(&self) -> Option<&[u8]> {
190        self.storage.contiguous()
191    }
192
193    fn segments(&self) -> impl ExactSizeIterator<Item = &[u8]> + FusedIterator {
194        self.storage.segments()
195    }
196
197    fn try_resolve<'req, H: scope::Source, O: stage::Body<'req, 'd>>(
198        self,
199        slot: &handoff::Slot<'req, 'd, H, O>,
200    ) -> Result<handoff::Plan<'req, 'd, H::Selection, O>, Error> {
201        Ok(handoff::Plan::new(slot.origin().segmented(self)))
202    }
203}
204
205impl<'d> IntoIterator for Segments<'d> {
206    type Item = Segment<'d>;
207    type IntoIter = iter::Segments<'d>;
208
209    fn into_iter(self) -> Self::IntoIter {
210        iter::Segments { state: self }
211    }
212}
213
214impl<'d> From<segments::Chain<resident::Shared<'d>>> for Segments<'d> {
215    fn from(chain: segments::Chain<resident::Shared<'d>>) -> Self {
216        Self {
217            storage: chain.into_shape(),
218        }
219    }
220}
221
222impl<'d> From<segments::Many<resident::Shared<'d>>> for Segments<'d> {
223    fn from(chain: segments::Many<resident::Shared<'d>>) -> Self {
224        Self {
225            storage: Shape::Many(chain),
226        }
227    }
228}
229
230impl<'d> From<Shape<resident::Shared<'d>>> for Segments<'d> {
231    fn from(storage: Shape<resident::Shared<'d>>) -> Self {
232        Self { storage }
233    }
234}
235
236impl<'d> From<Nonempty<resident::Shared<'d>>> for Segments<'d> {
237    fn from(owners: Nonempty<resident::Shared<'d>>) -> Self {
238        Self {
239            storage: owners.into_shape(),
240        }
241    }
242}
243
244impl<'d> ops::Deref for Segments<'d> {
245    type Target = Shape<resident::Shared<'d>>;
246    fn deref(&self) -> &Self::Target {
247        &self.storage
248    }
249}
250
251impl<'d> From<segments::Item<resident::Shared<'d>>> for Segments<'d> {
252    fn from(owner: segments::Item<resident::Shared<'d>>) -> Self {
253        Self {
254            storage: Shape::One(owner),
255        }
256    }
257}
258
259impl<'d> From<storage::Shared> for Segments<'d> {
260    fn from(data: storage::Shared) -> Self {
261        Self::from(resident::Shared::from(data))
262    }
263}
264
265impl<'d> From<Segment<'d>> for Segments<'d> {
266    fn from(data: Segment<'d>) -> Self {
267        Self::from(resident::Shared::from(data))
268    }
269}
270
271impl<'d> From<resident::Shared<'d>> for Segments<'d> {
272    fn from(data: resident::Shared<'d>) -> Self {
273        Self {
274            storage: match segments::Item::new(data) {
275                ops::ControlFlow::Continue(owner) => Shape::One(owner),
276                ops::ControlFlow::Break(_) => Shape::Empty,
277            },
278        }
279    }
280}
281
282impl Default for Segments<'_> {
283    fn default() -> Self {
284        Self::new()
285    }
286}
287
288const _: () = {
289    assert!(
290        size_of::<<Segments<'static> as body::Source<'static>>::Body<'static>>()
291            == size_of::<&Segments<'static>>()
292    );
293    assert!(size_of::<Segments<'static>>() == size_of::<Shape<resident::Shared<'static>>>());
294    assert!(align_of::<Segments<'static>>() == align_of::<Shape<resident::Shared<'static>>>());
295    assert!(
296        size_of::<Option<Segments<'static>>>()
297            == size_of::<Option<Shape<resident::Shared<'static>>>>()
298    );
299};
sark_http/payload/owned/segments.rs

sark_http/payload/owned/
segments.rs
