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    pub fn charge(&self) -> hold::Charge {
92        match &self.storage {
93            Shape::Empty => hold::Charge::default(),
94            Shape::One(owner) => hold::Charge::from_owner(owner.len(), owner.resident()),
95            Shape::Many(owners) => hold::Charge::from(owners.as_chain()),
96        }
97    }
98
99    pub fn segment(&self, index: usize) -> Option<&[u8]> {
100        match &self.storage {
101            Shape::Empty => None,
102            Shape::One(segment) => (index == 0).then(|| segment.as_slice()),
103            Shape::Many(segments) => segments.as_chain().segments().nth(index).map(AsRef::as_ref),
104        }
105    }
106
107    #[doc(hidden)]
108    pub fn backings(&self) -> iter::Backings<'_, 'd> {
109        iter::Backings {
110            owners: self.storage.owners(),
111        }
112    }
113
114    pub fn into_shape(self) -> Shape<resident::Shared<'d>> {
115        self.storage
116    }
117
118    pub(crate) const fn storage(&self) -> &Shape<resident::Shared<'d>> {
119        &self.storage
120    }
121
122    pub fn take(self) -> Option<Nonempty<resident::Shared<'d>>> {
123        match self.storage {
124            Shape::Empty => None,
125            Shape::One(owner) => Some(Nonempty::One(owner)),
126            Shape::Many(owners) => Some(Nonempty::Many(owners)),
127        }
128    }
129
130    pub fn try_clone(&self) -> Result<Self, Error> {
131        let storage = match &self.storage {
132            Shape::Empty => Shape::Empty,
133            Shape::One(owner) => Shape::One(owner.try_clone().map_err(Error::Retain)?),
134            Shape::Many(source) => Shape::Many(
135                segments::Exact::try_new(source)
136                    .map_err(Error::Metadata)?
137                    .try_collect()
138                    .map_err(Error::Retain)?,
139            ),
140        };
141        Ok(Self { storage })
142    }
143}
144
145impl<'d> body::Source<'d> for Segments<'d> {
146    type Body<'a>
147        = &'a Self
148    where
149        Self: 'a,
150        'd: 'a;
151
152    fn try_project<V, T>(
153        &mut self,
154        value: V,
155        project: impl for<'a> FnOnce(V, body::View<'a, 'd>) -> T,
156    ) -> Result<T, (V, body::Error)> {
157        let plan = match &self.storage {
158            Shape::Empty => return Ok(project(value, body::View::Slice(&[]))),
159            Shape::One(owner) => {
160                return Ok(project(value, body::View::Resident(owner.as_ref())));
161            }
162            Shape::Many(chain) => chain.as_chain().try_flatten(),
163        };
164        match plan {
165            Ok(store) => {
166                let owner: resident::Shared<'d> = store.into();
167                let output = project(value, body::View::Resident(&owner));
168                *self = Self::from(owner);
169                Ok(output)
170            }
171            Err(error) => Err((value, body::Error::Owned(Error::Storage(error)))),
172        }
173    }
174
175    fn body<'a>(&'a self) -> Self::Body<'a>
176    where
177        'd: 'a,
178    {
179        self
180    }
181}
182
183impl<'d> response::Body<'d> for Segments<'d> {
184    fn len(&self) -> usize {
185        self.storage.len()
186    }
187
188    fn as_slice(&self) -> Option<&[u8]> {
189        self.storage.contiguous()
190    }
191
192    fn segments(&self) -> impl ExactSizeIterator<Item = &[u8]> + FusedIterator {
193        self.storage.segments()
194    }
195
196    fn try_resolve<'req, H: scope::Source, O: stage::Body<'req, 'd>>(
197        self,
198        slot: &handoff::Slot<'req, 'd, H, O>,
199    ) -> Result<handoff::Plan<'req, 'd, H::Selection, O>, Error> {
200        Ok(handoff::Plan::new(slot.origin().segmented(self)))
201    }
202}
203
204impl<'d> IntoIterator for Segments<'d> {
205    type Item = Segment<'d>;
206    type IntoIter = iter::Segments<'d>;
207
208    fn into_iter(self) -> Self::IntoIter {
209        iter::Segments { state: self }
210    }
211}
212
213impl<'d> From<segments::Chain<resident::Shared<'d>>> for Segments<'d> {
214    fn from(chain: segments::Chain<resident::Shared<'d>>) -> Self {
215        Self {
216            storage: chain.into_shape(),
217        }
218    }
219}
220
221impl<'d> From<segments::Many<resident::Shared<'d>>> for Segments<'d> {
222    fn from(chain: segments::Many<resident::Shared<'d>>) -> Self {
223        Self {
224            storage: Shape::Many(chain),
225        }
226    }
227}
228
229impl<'d> From<Shape<resident::Shared<'d>>> for Segments<'d> {
230    fn from(storage: Shape<resident::Shared<'d>>) -> Self {
231        Self { storage }
232    }
233}
234
235impl<'d> From<Nonempty<resident::Shared<'d>>> for Segments<'d> {
236    fn from(owners: Nonempty<resident::Shared<'d>>) -> Self {
237        Self {
238            storage: owners.into_shape(),
239        }
240    }
241}
242
243impl<'d> ops::Deref for Segments<'d> {
244    type Target = Shape<resident::Shared<'d>>;
245    fn deref(&self) -> &Self::Target {
246        &self.storage
247    }
248}
249
250impl<'d> From<segments::Item<resident::Shared<'d>>> for Segments<'d> {
251    fn from(owner: segments::Item<resident::Shared<'d>>) -> Self {
252        Self {
253            storage: Shape::One(owner),
254        }
255    }
256}
257
258impl<'d> From<storage::Shared> for Segments<'d> {
259    fn from(data: storage::Shared) -> Self {
260        Self::from(resident::Shared::from(data))
261    }
262}
263
264impl<'d> From<Segment<'d>> for Segments<'d> {
265    fn from(data: Segment<'d>) -> Self {
266        Self::from(resident::Shared::from(data))
267    }
268}
269
270impl<'d> From<resident::Shared<'d>> for Segments<'d> {
271    fn from(data: resident::Shared<'d>) -> Self {
272        Self {
273            storage: match segments::Item::new(data) {
274                ops::ControlFlow::Continue(owner) => Shape::One(owner),
275                ops::ControlFlow::Break(_) => Shape::Empty,
276            },
277        }
278    }
279}
280
281impl Default for Segments<'_> {
282    fn default() -> Self {
283        Self::new()
284    }
285}
286
287const _: () = {
288    assert!(
289        size_of::<<Segments<'static> as body::Source<'static>>::Body<'static>>()
290            == size_of::<&Segments<'static>>()
291    );
292    assert!(size_of::<Segments<'static>>() == size_of::<Shape<resident::Shared<'static>>>());
293    assert!(align_of::<Segments<'static>>() == align_of::<Shape<resident::Shared<'static>>>());
294    assert!(
295        size_of::<Option<Segments<'static>>>()
296            == size_of::<Option<Shape<resident::Shared<'static>>>>()
297    );
298};
sark_http/payload/owned/segments.rs

sark_http/payload/owned/
segments.rs
