mz_timely_util/containers/stack.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file at the
// root of this repository, or online at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! A chunked columnar container based on the columnation library. It stores the local
//! portion in region-allocated data, too, which is different to the `TimelyStack` type.
use std::cell::Cell;
use timely::container::columnation::{Columnation, TimelyStack};
use timely::container::{ContainerBuilder, PushInto};
/// A Stacked container builder that keep track of container memory usage.
#[derive(Default)]
pub struct AccountedStackBuilder<CB> {
pub bytes: Cell<usize>,
pub builder: CB,
}
impl<T, CB> ContainerBuilder for AccountedStackBuilder<CB>
where
T: Clone + Columnation + 'static,
CB: ContainerBuilder<Container = TimelyStack<T>>,
{
type Container = TimelyStack<T>;
fn extract(&mut self) -> Option<&mut Self::Container> {
let container = self.builder.extract()?;
let mut new_bytes = 0;
container.heap_size(|_, cap| new_bytes += cap);
self.bytes.set(self.bytes.get() + new_bytes);
Some(container)
}
fn finish(&mut self) -> Option<&mut Self::Container> {
let container = self.builder.finish()?;
let mut new_bytes = 0;
container.heap_size(|_, cap| new_bytes += cap);
self.bytes.set(self.bytes.get() + new_bytes);
Some(container)
}
}
impl<T, CB: PushInto<T>> PushInto<T> for AccountedStackBuilder<CB> {
#[inline]
fn push_into(&mut self, item: T) {
self.builder.push_into(item);
}
}