mz_timely_util/containers/stack.rs
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Container builder wrapper that carries a byte counter for fuel-based yielding.
17
18use std::cell::Cell;
19
20use timely::container::{ContainerBuilder, PushInto};
21
22/// A [`ContainerBuilder`] wrapper carrying a byte counter that
23/// [`AsyncOutputHandle::give_fueled`][gf] increments when emitting items. The
24/// wrapper itself does no accounting; it exists so `give_fueled` can find a
25/// counter associated with the output handle. The caller supplies the size of
26/// each pushed item, which lets the source choose whatever estimate is
27/// cheapest for its data without needing the container to introspect items.
28///
29/// [gf]: crate::builder_async::AsyncOutputHandle::give_fueled
30#[derive(Default)]
31pub struct FueledBuilder<CB> {
32 pub bytes: Cell<usize>,
33 pub builder: CB,
34}
35
36impl<CB: ContainerBuilder> ContainerBuilder for FueledBuilder<CB> {
37 type Container = CB::Container;
38
39 fn extract(&mut self) -> Option<&mut Self::Container> {
40 self.builder.extract()
41 }
42
43 fn finish(&mut self) -> Option<&mut Self::Container> {
44 self.builder.finish()
45 }
46}
47
48impl<T, CB: PushInto<T>> PushInto<T> for FueledBuilder<CB> {
49 #[inline]
50 fn push_into(&mut self, item: T) {
51 self.builder.push_into(item);
52 }
53}