/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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 heap storing AValue traits. The heap is a sequence of the
//! AValue vtable, followed by the payload.
//! Every payload must be at least 1 usize large (even ZST).
//! Some elements are created using reserve, in which case they point
//! to a BlackHole until they are filled in.
//!
//! Some elements can be overwritten (typically during GC) by a usize.
//! In these cases the bottom bit of the usize as used by the heap
//! to tag it as being a usize, and the word after is the size of the
//! item it replaced.

use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem;
use std::mem::MaybeUninit;
use std::ptr;
use std::slice;
use std::time::Instant;

use allocative::Allocative;
use allocative::Visitor;
use bumpalo::Bump;
use dupe::Dupe;
use either::Either;
use starlark_map::small_map::SmallMap;

use crate::cast::transmute;
use crate::collections::StarlarkHashValue;
use crate::values::layout::aligned_size::AlignedSize;
use crate::values::layout::avalue::starlark_str;
use crate::values::layout::avalue::AValue;
use crate::values::layout::avalue::BlackHole;
use crate::values::layout::heap::allocator::api::ArenaAllocator;
use crate::values::layout::heap::allocator::api::ChunkAllocationDirection;
use crate::values::layout::heap::call_enter_exit::CallEnter;
use crate::values::layout::heap::call_enter_exit::CallExit;
use crate::values::layout::heap::call_enter_exit::NeedsDrop;
use crate::values::layout::heap::call_enter_exit::NoDrop;
use crate::values::layout::heap::heap_type::HeapKind;
use crate::values::layout::heap::profile::alloc_counts::AllocCounts;
use crate::values::layout::heap::profile::by_type::HeapSummary;
use crate::values::layout::heap::repr::AValueForward;
use crate::values::layout::heap::repr::AValueHeader;
use crate::values::layout::heap::repr::AValueOrForward;
use crate::values::layout::heap::repr::AValueRepr;
use crate::values::layout::vtable::AValueVTable;
use crate::values::string::StarlarkStr;
use crate::values::Value;
use crate::values::ValueLike;

/// Min size of allocated object including header.
/// Should be able to fit `BlackHole` or forward.
pub(crate) const MIN_ALLOC: AlignedSize = {
    const fn max(a: AlignedSize, b: AlignedSize) -> AlignedSize {
        if a.bytes() > b.bytes() { a } else { b }
    }

    max(
        AlignedSize::of::<AValueForward>(),
        AlignedSize::of::<AValueRepr<BlackHole>>(),
    )
};

#[derive(Default)]
pub(crate) struct Arena<A: ArenaAllocator> {
    /// Arena for things which don't need dropping (e.g. strings)
    non_drop: A,
    /// Arena for things which might need dropping (e.g. Vec, with memory on heap)
    drop: A,
}

/// Reservation is morally a Reservation<T>, but we treat is as an
/// existential.
/// Tied to the lifetime of the heap.
pub(crate) struct Reservation<'v, 'v2, T: AValue<'v2>> {
    pointer: *mut AValueRepr<T>, // Secretly AValueObject<T>
    phantom: PhantomData<(&'v (), &'v2 T)>,
}

impl<'v, 'v2, T: AValue<'v2>> Reservation<'v, 'v2, T> {
    pub(crate) fn fill(self, x: T) {
        unsafe {
            ptr::write(
                self.pointer,
                AValueRepr {
                    header: AValueHeader::new::<T>(),
                    payload: x,
                },
            );
        }
    }

    pub(crate) fn ptr(&self) -> &'v AValueHeader {
        unsafe { &(*self.pointer).header }
    }
}

pub(crate) trait ArenaVisitor<'v> {
    fn regular_value(&mut self, value: &'v AValueOrForward);
    fn call_enter(&mut self, function: Value<'v>, time: Instant);
    fn call_exit(&mut self, time: Instant);
}

/// Iterate over chunk contents.
struct ChunkIter<'c> {
    chunk: &'c [MaybeUninit<u8>],
}

impl<'c> Iterator for ChunkIter<'c> {
    type Item = &'c AValueOrForward;

    fn next(&mut self) -> Option<&'c AValueOrForward> {
        unsafe {
            if self.chunk.is_empty() {
                None
            } else {
                let or_forward = &*(self.chunk.as_ptr() as *const AValueOrForward);
                let n = or_forward.alloc_size();
                debug_assert!(n.bytes() as usize <= self.chunk.len());
                self.chunk = self.chunk.split_at(n.bytes() as usize).1;
                Some(or_forward)
            }
        }
    }
}

impl<A: ArenaAllocator> Arena<A> {
    pub(crate) fn is_empty(&self) -> bool {
        self.allocated_bytes() == 0
    }

    /// Number of allocated bytes plus padding size.
    pub(crate) fn allocated_bytes(&self) -> usize {
        // This overestimates the allocates size, see `Bump::allocated_bytes()`:
        // "those padding bytes will add to this methodâ€™s resulting sum".
        // https://docs.rs/bumpalo/3.11.0/bumpalo/struct.Bump.html#method.allocated_bytes
        self.drop.allocated_bytes() + self.non_drop.allocated_bytes()
    }

    pub(crate) fn available_bytes(&self) -> usize {
        self.drop.remaining_capacity() + self.non_drop.remaining_capacity()
    }

    /// Don't forget to call this function to release memory.
    pub(crate) fn finish(&mut self) {
        self.drop.finish();
        self.non_drop.finish();
    }

    fn alloc_uninit<'v, 'v2: 'v, T: AValue<'v2>>(
        bump: &'v A,
        extra_len: usize,
    ) -> (
        &'v mut MaybeUninit<AValueRepr<T>>,
        &'v mut [MaybeUninit<T::ExtraElem>],
    ) {
        assert!(
            mem::align_of::<T>() <= AValueHeader::ALIGN,
            "Unexpected alignment in Starlark arena. Type {} has alignment {}, expected <= {}",
            std::any::type_name::<T>(),
            mem::align_of::<T>(),
            AValueHeader::ALIGN,
        );

        let size = T::alloc_size_for_extra_len(extra_len);
        let p = bump.alloc(size).as_ptr();
        unsafe {
            let repr = &mut *(p as *mut MaybeUninit<AValueRepr<T>>);
            let extra = slice::from_raw_parts_mut(
                (p as *mut u8).add(AValueRepr::<T>::offset_of_extra()) as *mut _,
                extra_len,
            );
            (repr, extra)
        }
    }

    fn bump_for_type<'v, T: AValue<'v>>(&self) -> &A {
        if mem::needs_drop::<T>() {
            &self.drop
        } else {
            &self.non_drop
        }
    }

    // Reservation should really be an incremental type
    pub(crate) fn reserve_with_extra<'v, 'v2: 'v, T: AValue<'v2>>(
        &'v self,
        extra_len: usize,
    ) -> (Reservation<'v, 'v2, T>, &'v mut [MaybeUninit<T::ExtraElem>]) {
        // We don't create reservations for strings because we don't need to,
        // but also because we need to be able to reconstruct a `Pointer`
        // from `AValueHeader` (with `TAG_STR` when appropriate).
        // `BlackHole` assumes it is created for non-string, so
        // it returns `false` from `is_str`.
        assert!(!T::IS_STR);

        let (p, extra) = Self::alloc_uninit::<T>(self.bump_for_type::<T>(), extra_len);
        // If we don't have a vtable we can't skip over missing elements to drop,
        // so very important to put in a current vtable
        // We always alloc at least one pointer worth of space, so can write in a one-ST blackhole

        let x = BlackHole(T::alloc_size_for_extra_len(extra_len));
        let p = unsafe {
            transmute!(
                &mut MaybeUninit<AValueRepr<T>>,
                &mut MaybeUninit<AValueRepr<BlackHole>>,
                p
            )
        };
        let p = p.write(AValueRepr {
            header: AValueHeader(AValueVTable::new_black_hole()),
            payload: x,
        });
        let p = unsafe { transmute!(&mut AValueRepr<BlackHole>, &mut AValueRepr<T>, p) };

        (
            Reservation {
                pointer: p,
                phantom: PhantomData,
            },
            extra,
        )
    }

    /// Allocate a type `T`.
    pub(crate) fn alloc<'v, 'v2: 'v, T: AValue<'v2, ExtraElem = ()>>(
        &'v self,
        x: T,
    ) -> &'v AValueRepr<T> {
        debug_assert!(x.extra_len() == 0);
        let bump = self.bump_for_type::<T>();
        let (p, extra) = Self::alloc_uninit::<T>(bump, 0);
        debug_assert!(extra.is_empty());
        p.write(AValueRepr {
            header: AValueHeader::new::<T>(),
            payload: x,
        })
    }

    /// Allocate a type `T` plus `extra` bytes.
    ///
    /// The type `T` will never be dropped, so had better not do any memory allocation.
    pub(crate) fn alloc_extra<'v, 'v2: 'v, T: AValue<'v2>>(
        &'v self,
        x: T,
    ) -> (*mut AValueRepr<T>, &'v mut [MaybeUninit<T::ExtraElem>]) {
        let bump = self.bump_for_type::<T>();
        let (p, extra) = Self::alloc_uninit::<T>(bump, x.extra_len());
        let p = p.write(AValueRepr {
            header: AValueHeader::new::<T>(),
            payload: x,
        });
        (p, extra)
    }

    #[inline]
    pub(crate) fn alloc_str_init(
        &self,
        len: usize,
        hash: StarlarkHashValue,
        init: impl FnOnce(*mut u8),
    ) -> *mut AValueHeader {
        assert!(len > 1);
        let (v, extra) = self.alloc_extra::<_>(starlark_str(len, hash));
        debug_assert_eq!(StarlarkStr::payload_len_for_len(len), extra.len());
        unsafe {
            extra.last_mut().unwrap_unchecked().write(0usize);
        }
        init(extra.as_mut_ptr() as *mut u8);
        unsafe { &mut (*v).header }
    }

    #[inline]
    pub(crate) fn alloc_str(&self, x: &str) -> *mut AValueHeader {
        self.alloc_str_init(x.len(), StarlarkStr::UNINIT_HASH, |dest| unsafe {
            ptr::copy_nonoverlapping(x.as_ptr(), dest, x.len())
        })
    }

    fn iter_chunk<'a>(chunk: &'a [MaybeUninit<u8>]) -> ChunkIter<'a> {
        ChunkIter { chunk }
    }

    // Iterate over the values in the heap in the order they
    // were added.
    pub(crate) fn for_each_ordered<'a>(&'a mut self, mut f: impl FnMut(&'a AValueOrForward)) {
        // We get the chunks from most newest to oldest as per the bumpalo spec.
        // And within each chunk, the values are filled newest to oldest.
        // So need to do two sets of reversing.
        for bump in [&mut self.drop, &mut self.non_drop] {
            let chunks = unsafe { bump.iter_allocated_chunks_rev().collect::<Vec<_>>() };
            // Use a single buffer to reduce allocations, but clear it after use
            let mut buffer = Vec::new();
            for chunk in chunks.iter().rev() {
                match A::CHUNK_ALLOCATION_DIRECTION {
                    ChunkAllocationDirection::Down => {
                        buffer.extend(Arena::<A>::iter_chunk(chunk));
                        for x in buffer.iter().rev() {
                            f(x);
                        }
                        buffer.clear();
                    }
                    ChunkAllocationDirection::Up => {
                        for x in Arena::<A>::iter_chunk(chunk) {
                            f(x);
                        }
                    }
                }
            }
        }
    }

    pub(crate) unsafe fn visit_arena<'v>(
        &'v mut self,
        heap_kind: HeapKind,
        forward_heap_kind: HeapKind,
        visitor: &mut impl ArenaVisitor<'v>,
    ) {
        fn fix_function<'v>(function: Value<'v>, forward_heap_kind: HeapKind) -> Value<'v> {
            if let Some(function) = function.unpack_frozen() {
                return function.to_value();
            }

            unsafe {
                match function
                    .0
                    .unpack_ptr()
                    .expect("int cannot be stored in heap")
                    .unpack_forward()
                {
                    None => function,
                    Some(forward) => forward.forward_ptr().unpack_value(forward_heap_kind),
                }
            }
        }

        self.for_each_ordered(|x| match x.unpack() {
            Either::Left(header) => {
                let value = header.unpack_value(heap_kind);
                if let Some(call_enter) = value.downcast_ref::<CallEnter<NeedsDrop>>() {
                    visitor.call_enter(
                        fix_function(call_enter.function, forward_heap_kind),
                        call_enter.time,
                    );
                } else if let Some(call_enter) = value.downcast_ref::<CallEnter<NoDrop>>() {
                    visitor.call_enter(
                        fix_function(call_enter.function, forward_heap_kind),
                        call_enter.time,
                    );
                } else if let Some(call_exit) = value.downcast_ref::<CallExit<NeedsDrop>>() {
                    visitor.call_exit(call_exit.time);
                } else if let Some(call_exit) = value.downcast_ref::<CallExit<NoDrop>>() {
                    visitor.call_exit(call_exit.time);
                } else {
                    visitor.regular_value(x);
                }
            }
            Either::Right(_forward) => {
                visitor.regular_value(x);
            }
        });
    }

    // Iterate over the values in the drop bump in any order
    pub(crate) fn for_each_drop_unordered<'a>(&'a mut self, mut f: impl FnMut(&'a AValueHeader)) {
        unsafe {
            for chunk in self.drop.iter_allocated_chunks_rev() {
                for x in Arena::<A>::iter_chunk(chunk) {
                    if let Some(x) = x.unpack_header() {
                        f(x);
                    }
                }
            }
        }
    }

    fn for_each_unordered_in_bump<'a>(bump: &'a A, mut f: impl FnMut(&'a AValueHeader)) {
        // SAFETY: We're consuming the iterator immediately and not allocating from the arena during.
        unsafe {
            bump.iter_allocated_chunks_rev().for_each(|slice| {
                for x in Arena::<A>::iter_chunk(slice) {
                    if let Some(x) = x.unpack_header() {
                        f(x);
                    }
                }
            })
        }
    }

    // Iterate over the values in the both bumps in any order
    fn for_each_unordered<'a>(&'a self, mut f: impl FnMut(&'a AValueHeader)) {
        for bump in [&self.drop, &self.non_drop] {
            Self::for_each_unordered_in_bump(bump, &mut f);
        }
    }

    // For each Rust-level type (the String) report how many entries there are in the heap, and how much size they consume
    pub(crate) fn allocated_summary(&self) -> HeapSummary {
        // Record how many times each header occurs
        // We deliberately hash by the AValueHeader for higher performance, less type lookup
        let mut entries: HashMap<AValueHeader, (&'static str, AllocCounts)> = HashMap::new();
        let f = |x: &AValueHeader| {
            let v = x.unpack();
            let e = entries
                .entry(x.dupe())
                .or_insert_with(|| (v.get_type(), AllocCounts::default()));
            e.1.count += 1;
            e.1.bytes += v.total_memory()
        };
        self.for_each_unordered(f);

        // For a given type, the AValueHeader isn't always unique
        // (if they get compiled in different translation units),
        // so not just a simple map.
        let mut summary = SmallMap::new();
        for (_, (name, counts)) in entries {
            *summary.entry(name).or_insert_with(AllocCounts::default) += counts;
        }
        HeapSummary { summary }
    }

    /// Memory allocated in the arena but not used for allocation in starlark.
    pub(crate) fn unused_capacity(&self) -> usize {
        self.drop.remaining_capacity() + self.non_drop.remaining_capacity()
    }
}

impl<A: ArenaAllocator> Drop for Arena<A> {
    fn drop(&mut self) {
        self.for_each_drop_unordered(|x| {
            // Safe to convert to *mut because we are the only owner
            let value = x.payload_ptr();
            x.0.drop_in_place(value);
        });
    }
}

impl<A: ArenaAllocator> Allocative for Arena<A> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut allocative::Visitor<'b>) {
        let Arena { drop, non_drop } = self;

        fn visit_bump<'a, 'b: 'a, A: ArenaAllocator>(bump: &A, visitor: &'a mut Visitor<'b>) {
            let mut visitor =
                visitor.enter_unique(allocative::Key::new("data"), mem::size_of::<*const ()>());
            let mut allocated_visitor =
                visitor.enter(allocative::Key::new("allocated"), bump.allocated_bytes());
            Arena::for_each_unordered_in_bump(bump, |x| {
                let key = x.unpack().type_as_allocative_key();
                let size = x.alloc_size();
                let mut object_visitor = allocated_visitor.enter(key, size.bytes() as usize);
                // We visit both drop and non-drop bumps, because although
                // non-drop `Bump` cannot contain malloc pointers, it can still provide
                // useful information about headers/payload/padding.
                x.unpack().as_allocative().visit(&mut object_visitor);
                object_visitor.exit();
            });
            allocated_visitor.exit();
            visitor.visit_simple(
                allocative::Key::new("allocation_overhead"),
                bump.allocation_overhead(),
            );
            visitor.exit();
        }

        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let mut visitor = visitor.enter(allocative::Key::new("drop"), mem::size_of::<Bump>());
            visit_bump(drop, &mut visitor);
            visitor.exit();
        }
        {
            let mut visitor =
                visitor.enter(allocative::Key::new("non_drop"), mem::size_of::<Bump>());
            visit_bump(non_drop, &mut visitor);
            visitor.exit();
        }
        visitor.exit();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::values::any::StarlarkAny;
    use crate::values::layout::avalue::simple;

    fn to_repr(x: &AValueHeader) -> String {
        let mut s = String::new();
        x.unpack().collect_repr(&mut s);
        s
    }

    fn mk_str(x: &str) -> impl AValue<'static, ExtraElem = ()> {
        simple(StarlarkAny::new(x.to_owned()))
    }

    fn reserve_str<'v, T: AValue<'static>>(
        arena: &'v Arena<Bump>,
        _: &T,
    ) -> Reservation<'v, 'static, T> {
        arena.reserve_with_extra::<T>(0).0
    }

    #[test]
    fn test_trait_arena_iteration() {
        // We want iteration to proceed in the same order as allocation,
        // otherwise profiling won't work
        const LIMIT: usize = 10000;
        let mut arena = Arena::default();
        let mut reserved = Vec::new();
        for i in 0..LIMIT {
            if i % 100 == 0 {
                let r = reserve_str(&arena, &mk_str(""));
                reserved.push((r, i));
            } else {
                arena.alloc(mk_str(&i.to_string()));
            }
        }
        assert!(!reserved.is_empty());
        for (r, i) in reserved {
            r.fill(mk_str(&i.to_string()));
        }

        // Not a functional part of the test, just makes sure we go through
        // the interesting cases (last time 56 was sufficient, so 10K is plenty of margin of error)
        assert!(
            arena.drop.iter_allocated_chunks().count() > 1,
            "Didn't allocate enough to test properly"
        );
        let mut j = 0;
        arena.for_each_ordered(|i| {
            if let Some(i) = i.unpack_header() {
                assert_eq!(to_repr(i), j.to_string());
                j += 1;
            }
        });
        assert_eq!(j, LIMIT);
        j = 0;
        arena.for_each_drop_unordered(|_| j += 1);
        assert_eq!(j, LIMIT);
    }

    #[test]
    // Make sure that even if there are some blackholes when we drop, we can still walk to heap
    fn drop_with_blackhole() {
        let mut arena = Arena::default();
        arena.alloc(mk_str("test"));
        // reserve but do not fill!
        reserve_str(&arena, &mk_str(""));
        arena.alloc(mk_str("hello"));
        let mut res = Vec::new();
        arena.for_each_ordered(|x| {
            if let Some(x) = x.unpack_header() {
                res.push(x);
            }
        });
        assert_eq!(res.len(), 3);
        assert_eq!(to_repr(res[0]), "test");
        assert_eq!(to_repr(res[2]), "hello");
    }

    #[test]
    fn test_allocated_summary() {
        let arena = Arena::<Bump>::default();
        arena.alloc(mk_str("test"));
        arena.alloc(mk_str("test"));
        let res = arena.allocated_summary().summary;
        assert_eq!(res.len(), 1);
        let entry = res.values().next().unwrap();
        assert_eq!(entry.count, 2);
        // Because `arena.allocated_bytes` over-approximates.
        assert!(entry.bytes <= arena.allocated_bytes());
    }

    #[test]
    fn test_is_empty() {
        let arena = Arena::<Bump>::default();
        assert!(arena.is_empty());
        arena.alloc_str("xyz");
        assert!(!arena.is_empty());
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

//! Marker objects to track allocations.

use std::fmt::Debug;
use std::time::Instant;

use allocative::Allocative;
use starlark_derive::starlark_value;
use starlark_derive::NoSerialize;

use crate as starlark;
use crate::any::ProvidesStaticType;
use crate::values::StarlarkValue;
use crate::values::Trace;
use crate::values::Value;

/// A type which is either drop or non-drop.
pub(crate) trait MaybeDrop: Debug + Sync + Send + Allocative + 'static {}

/// Type which has `Drop`.
#[derive(ProvidesStaticType, Debug, Trace, Allocative)]
pub(crate) struct NeedsDrop;
impl Drop for NeedsDrop {
    fn drop(&mut self) {
        // Just make this type `Drop`.
        // Note `mem::needs_drop()` would return `true` for this type,
        // even if `drop` is optimized away: https://rust.godbolt.org/z/1cxKoMzdM
    }
}

/// Type which doesn't have `Drop`.
#[derive(ProvidesStaticType, Debug, Trace, Allocative)]
pub(crate) struct NoDrop;

impl MaybeDrop for NeedsDrop {}
impl MaybeDrop for NoDrop {}

#[derive(
    Trace,
    Debug,
    derive_more::Display,
    ProvidesStaticType,
    NoSerialize,
    Allocative
)]
#[display(fmt = "CallEnter")]
pub(crate) struct CallEnter<'v, D: MaybeDrop + 'static> {
    pub(crate) function: Value<'v>,
    pub(crate) time: Instant,
    pub(crate) maybe_drop: D,
}

#[starlark_value(type = "call_enter")]
impl<'v, D: MaybeDrop + Trace<'v> + 'v> StarlarkValue<'v> for CallEnter<'v, D> {}

#[derive(
    Debug,
    derive_more::Display,
    ProvidesStaticType,
    NoSerialize,
    Allocative
)]
#[display(fmt = "CallExit")]
pub(crate) struct CallExit<D: MaybeDrop + 'static> {
    pub(crate) time: Instant,
    pub(crate) maybe_drop: D,
}

#[starlark_value(type = "call_exit")]
impl<'v, D: MaybeDrop> StarlarkValue<'v> for CallExit<D> {}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::cell::Cell;
use std::cell::UnsafeCell;
use std::mem;
use std::mem::MaybeUninit;

/// Faster but less safe alternative to `RefCell<T>`:
/// all operations except `borrow` are `unsafe` and may lead to undefined behavior.
#[derive(Debug)]
pub(crate) struct FastCell<T> {
    /// The value.
    ///
    /// The idea is this: when we take the value out of cell,
    /// we write zeros to `value` field. So accessing this field will result in definite crash.
    /// This is worse than proper panic, but no runtime overhead like in `RefCell`.
    value: UnsafeCell<MaybeUninit<T>>,
    /// Whether the cell contains a value or zeros.
    init: Cell<bool>,
}

impl<T> Drop for FastCell<T> {
    fn drop(&mut self) {
        if self.init.get() {
            unsafe {
                (*self.value.get()).assume_init_drop();
            }
        }
    }
}

impl<T: Default> Default for FastCell<T> {
    fn default() -> Self {
        Self {
            value: UnsafeCell::new(MaybeUninit::new(T::default())),
            init: Cell::new(true),
        }
    }
}

impl<T> FastCell<T> {
    /// Get a reference to the value.
    ///
    /// This operation is safe under assumption that other `unsafe` operations
    /// do not leave self in invalid state.
    pub(crate) fn borrow(&self) -> &T {
        debug_assert!(self.init.get());
        unsafe { (*self.value.get()).assume_init_ref() }
    }

    pub(crate) fn try_borrow(&self) -> Option<&T> {
        if self.init.get() {
            Some(self.borrow())
        } else {
            None
        }
    }

    /// Get a mutable reference to the value.
    ///
    /// This function is unsafe because it's caller responsibility to guarantee
    /// there are no other references to the value, and nobody is going
    /// to obtain references to value while mutable reference exists.
    pub(crate) unsafe fn get_mut(&self) -> *mut T {
        debug_assert!(self.init.get());
        (*self.value.get()).as_mut_ptr()
    }

    /// Take the value out of the cell.
    pub(crate) unsafe fn take(&self) -> T {
        assert!(self.init.get());
        self.init.set(false);
        // Replace the `value` field with zeros so that accessing it will crash.
        mem::replace(&mut *self.value.get(), MaybeUninit::zeroed()).assume_init()
    }

    /// Put the value into the cell.
    pub(crate) unsafe fn set(&self, value: T) {
        assert!(!self.init.get());
        self.init.set(true);
        *self.value.get() = MaybeUninit::new(value);
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::cell::Cell;
use std::cell::RefCell;
use std::cmp;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::fmt::Formatter;
use std::hash::Hash;
use std::hash::Hasher;
use std::intrinsics::copy_nonoverlapping;
use std::marker::PhantomData;
use std::mem;
use std::mem::MaybeUninit;
use std::ops::Deref;
use std::ptr;
use std::slice;
use std::sync::Arc;
use std::time::Instant;
use std::usize;

use allocative::Allocative;
use bumpalo::Bump;
use dupe::Dupe;
use either::Either;
use starlark_map::small_set::SmallSet;

use crate::cast;
use crate::cast::transmute;
use crate::collections::maybe_uninit_backport::maybe_uninit_write_slice;
use crate::collections::maybe_uninit_backport::maybe_uninit_write_slice_cloned;
use crate::collections::Hashed;
use crate::collections::StarlarkHashValue;
use crate::eval::compiler::def::FrozenDef;
use crate::values::any::StarlarkAny;
use crate::values::array::Array;
use crate::values::layout::avalue::any_array_avalue;
use crate::values::layout::avalue::array_avalue;
use crate::values::layout::avalue::complex;
use crate::values::layout::avalue::complex_no_freeze;
use crate::values::layout::avalue::frozen_list_avalue;
use crate::values::layout::avalue::frozen_tuple_avalue;
use crate::values::layout::avalue::list_avalue;
use crate::values::layout::avalue::simple;
use crate::values::layout::avalue::tuple_avalue;
use crate::values::layout::avalue::AValue;
use crate::values::layout::avalue::VALUE_EMPTY_ARRAY;
use crate::values::layout::avalue::VALUE_EMPTY_FROZEN_LIST;
use crate::values::layout::heap::arena::Arena;
use crate::values::layout::heap::arena::ArenaVisitor;
use crate::values::layout::heap::arena::Reservation;
use crate::values::layout::heap::call_enter_exit::CallEnter;
use crate::values::layout::heap::call_enter_exit::CallExit;
use crate::values::layout::heap::call_enter_exit::NeedsDrop;
use crate::values::layout::heap::call_enter_exit::NoDrop;
use crate::values::layout::heap::fast_cell::FastCell;
use crate::values::layout::heap::maybe_uninit_slice_util::maybe_uninit_write_from_exact_size_iter;
use crate::values::layout::heap::profile::by_type::HeapSummary;
use crate::values::layout::heap::repr::AValueRepr;
use crate::values::layout::static_string::constant_string;
use crate::values::layout::typed::string::StringValueLike;
use crate::values::layout::value::FrozenValue;
use crate::values::layout::value::Value;
use crate::values::string::intern::interner::FrozenStringInterner;
use crate::values::string::StarlarkStr;
use crate::values::AllocFrozenValue;
use crate::values::AllocValue;
use crate::values::ComplexValue;
use crate::values::FrozenRef;
use crate::values::FrozenStringValue;
use crate::values::FrozenValueTyped;
use crate::values::StarlarkValue;
use crate::values::StringValue;
use crate::values::Trace;
use crate::values::UnpackValue;
use crate::values::ValueOf;
use crate::values::ValueTyped;

#[derive(Copy, Clone, Dupe)]
pub(crate) enum HeapKind {
    Unfrozen,
    Frozen,
}

/// A heap on which [`Value`]s can be allocated. The values will be annotated with the heap lifetime.
#[derive(Default)]
pub struct Heap {
    /// Peak memory seen when a garbage collection takes place (may be lower than currently allocated)
    peak_allocated: Cell<usize>,
    arena: FastCell<Arena<Bump>>,
}

impl Debug for Heap {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut x = f.debug_struct("Heap");
        x.field(
            "bytes",
            &self.arena.try_borrow().map(|x| x.allocated_bytes()),
        );
        x.finish()
    }
}

/// A heap on which [`FrozenValue`]s can be allocated.
/// Can be kept alive by a [`FrozenHeapRef`].
#[derive(Default)]
pub struct FrozenHeap {
    /// My memory.
    arena: Arena<Bump>,
    /// Memory I depend on.
    refs: RefCell<SmallSet<FrozenHeapRef>>,
    /// String interner.
    str_interner: RefCell<FrozenStringInterner>,
}

/// `FrozenHeap` when it is no longer modified and can be share between threads.
/// Although, `arena` is not safe to share between threads, but at least `refs` is.
#[derive(Default, Allocative)]
#[allow(clippy::non_send_fields_in_send_ty)]
struct FrozenFrozenHeap {
    arena: Arena<Bump>,
    refs: Box<[FrozenHeapRef]>,
}

// Safe because we never mutate the Arena other than with &mut
unsafe impl Sync for FrozenFrozenHeap {}
unsafe impl Send for FrozenFrozenHeap {}

impl Debug for FrozenHeap {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut x = f.debug_struct("FrozenHeap");
        x.field("bytes", &self.arena.allocated_bytes());
        x.field("refs", &self.refs.try_borrow().map(|x| x.len()));
        x.finish()
    }
}

impl Debug for FrozenFrozenHeap {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut x = f.debug_struct("FrozenHeap");
        x.field("bytes", &self.arena.allocated_bytes());
        x.field("refs", &self.refs.len());
        x.finish()
    }
}

/// A reference to a [`FrozenHeap`] that keeps alive all values on the underlying heap.
/// Note that the [`Hash`] is consistent for a single [`FrozenHeapRef`], but non-deterministic
/// across executions and distinct but observably identical [`FrozenHeapRef`] values.
#[derive(Default, Clone, Dupe, Debug, Allocative)]
// The Eq/Hash are by pointer rather than value, since we produce unique values
// given an underlying FrozenHeap.
pub struct FrozenHeapRef(Option<Arc<FrozenFrozenHeap>>);

fn _test_frozen_heap_ref_send_sync()
where
    FrozenHeapRef: Send + Sync,
{
}

impl Hash for FrozenHeapRef {
    fn hash<H: Hasher>(&self, state: &mut H) {
        if let Some(arc) = &self.0 {
            let x: &FrozenFrozenHeap = Deref::deref(&arc);
            ptr::hash(x, state);
        }
    }
}

impl PartialEq<FrozenHeapRef> for FrozenHeapRef {
    fn eq(&self, other: &FrozenHeapRef) -> bool {
        match (&self.0, &other.0) {
            (Some(a), Some(b)) => Arc::ptr_eq(a, b),
            (None, None) => true,
            (Some(_), None) | (None, Some(_)) => false,
        }
    }
}

impl Eq for FrozenHeapRef {}

impl FrozenHeapRef {
    /// Number of bytes allocated on this heap, not including any memory
    /// allocated outside of the starlark heap.
    pub fn allocated_bytes(&self) -> usize {
        self.0.as_ref().map_or(0, |a| a.arena.allocated_bytes())
    }

    /// Number of bytes allocated by the heap but not filled.
    /// Note that these bytes will _never_ be filled as no further allocations can
    /// be made on this heap (it has been sealed).
    pub fn available_bytes(&self) -> usize {
        self.0.as_ref().map_or(0, |a| a.arena.available_bytes())
    }

    /// Obtain a summary of how much memory is currently allocated by this heap.
    /// Doesn't include the heaps it keeps alive by reference.
    pub fn allocated_summary(&self) -> HeapSummary {
        self.0
            .as_ref()
            .map_or_else(HeapSummary::default, |a| a.arena.allocated_summary())
    }
}

impl FrozenHeap {
    /// Create a new [`FrozenHeap`].
    pub fn new() -> Self {
        Self::default()
    }

    /// After all values have been allocated, convert the [`FrozenHeap`] into a
    /// [`FrozenHeapRef`] which can be [`clone`](Clone::clone)d, shared between threads,
    /// and ensures the underlying values allocated on the [`FrozenHeap`] remain valid.
    pub fn into_ref(self) -> FrozenHeapRef {
        let FrozenHeap {
            mut arena, refs, ..
        } = self;
        arena.finish();
        let refs = refs.into_inner();
        if arena.is_empty() && refs.is_empty() {
            FrozenHeapRef::default()
        } else {
            FrozenHeapRef(Some(Arc::new(FrozenFrozenHeap {
                arena,
                refs: refs.into_iter().collect(),
            })))
        }
    }

    /// Keep the argument [`FrozenHeapRef`] alive as long as this [`FrozenHeap`]
    /// is kept alive. Used if a [`FrozenValue`] in this heap points at values in another
    /// [`FrozenHeap`].
    pub fn add_reference(&self, heap: &FrozenHeapRef) {
        if heap.0.is_none() {
            return;
        }

        let mut refs = self.refs.borrow_mut();
        if !refs.contains(heap) {
            refs.insert(heap.dupe());
        }
    }

    fn alloc_raw(&self, x: impl AValue<'static, ExtraElem = ()> + Send + Sync) -> FrozenValue {
        let v: &AValueRepr<_> = self.arena.alloc(x);
        unsafe { FrozenValue::new_repr(cast::ptr_lifetime(v)) }
    }

    #[inline]
    fn alloc_str_init(
        &self,
        len: usize,
        hash: StarlarkHashValue,
        init: impl FnOnce(*mut u8),
    ) -> FrozenStringValue {
        let v = self.arena.alloc_str_init(len, hash, init);

        unsafe {
            let value = FrozenValue::new_ptr(&*v, true);
            FrozenStringValue::new_unchecked(value)
        }
    }

    fn alloc_str_impl(&self, x: &str, hash: StarlarkHashValue) -> FrozenStringValue {
        if let Some(x) = constant_string(x) {
            x
        } else {
            self.alloc_str_init(x.len(), hash, |dest| unsafe {
                copy_nonoverlapping(x.as_ptr(), dest, x.len())
            })
        }
    }

    /// Allocate a string on this heap. Be careful about the warnings
    /// around [`FrozenValue`].
    pub fn alloc_str(&self, x: &str) -> FrozenStringValue {
        self.alloc_str_impl(x, StarlarkStr::UNINIT_HASH)
    }

    /// Intern string.
    pub(crate) fn alloc_str_intern(&self, s: &str) -> FrozenStringValue {
        if let Some(s) = constant_string(s) {
            s
        } else {
            let s = Hashed::new(s);
            self.str_interner
                .borrow_mut()
                .intern(s, || self.alloc_str_hashed(s))
        }
    }

    /// Allocate prehashed string.
    pub fn alloc_str_hashed(&self, x: Hashed<&str>) -> FrozenStringValue {
        self.alloc_str_impl(x.key(), x.hash())
    }

    /// Allocate a tuple with the given elements on this heap.
    pub(crate) fn alloc_tuple<'v>(&'v self, elems: &[FrozenValue]) -> FrozenValue {
        if elems.is_empty() {
            return FrozenValue::new_empty_tuple();
        }

        unsafe {
            let (avalue, extra) = self
                .arena
                .alloc_extra::<_>(frozen_tuple_avalue(elems.len()));
            maybe_uninit_write_slice(extra, elems);
            FrozenValue::new_repr(&*avalue)
        }
    }

    /// Allocate a tuple from iterator of elements.
    pub(crate) fn alloc_tuple_iter(
        &self,
        elems: impl IntoIterator<Item = FrozenValue>,
    ) -> FrozenValue {
        let elems = elems.into_iter();
        let (lower, upper) = elems.size_hint();
        if Some(lower) == upper {
            if lower == 0 {
                return FrozenValue::new_empty_tuple();
            }

            unsafe {
                let (avalue, extra) = self.arena.alloc_extra::<_>(frozen_tuple_avalue(lower));
                maybe_uninit_write_from_exact_size_iter(extra, elems, FrozenValue::new_none());
                FrozenValue::new_repr(&*avalue)
            }
        } else {
            self.alloc_tuple(&elems.collect::<Vec<_>>())
        }
    }

    /// Allocate a list with the given elements on this heap.
    pub(crate) fn alloc_list(&self, elems: &[FrozenValue]) -> FrozenValue {
        if elems.is_empty() {
            return FrozenValue::new_repr(&VALUE_EMPTY_FROZEN_LIST);
        }

        unsafe {
            let (avalue, elem_places) = self.arena.alloc_extra(frozen_list_avalue(elems.len()));
            maybe_uninit_write_slice(elem_places, elems);
            FrozenValue::new_repr(&*avalue)
        }
    }

    pub(crate) fn alloc_list_iter(
        &self,
        elems: impl IntoIterator<Item = FrozenValue>,
    ) -> FrozenValue {
        let elems = elems.into_iter();
        let (lower, upper) = elems.size_hint();
        if Some(lower) == upper {
            if lower == 0 {
                return FrozenValue::new_repr(&VALUE_EMPTY_FROZEN_LIST);
            }

            unsafe {
                let (avalue, elem_places) = self.arena.alloc_extra(frozen_list_avalue(lower));
                maybe_uninit_write_from_exact_size_iter(
                    elem_places,
                    elems,
                    FrozenValue::new_none(),
                );
                FrozenValue::new_repr(&*avalue)
            }
        } else {
            self.alloc_list(&elems.collect::<Vec<_>>())
        }
    }

    pub(crate) fn alloc_simple_typed<T: StarlarkValue<'static> + Send + Sync>(
        &self,
        val: T,
    ) -> FrozenValueTyped<'static, T> {
        // SAFETY: we we've just allocated `T`.
        unsafe { FrozenValueTyped::new_unchecked(self.alloc_raw(simple(val))) }
    }

    /// Allocate a simple [`StarlarkValue`] on this heap.
    ///
    /// Simple value is any starlark value which:
    /// * bound by `'static` lifetime (in particular, it cannot contain references to other `Value`s)
    /// * is not special builtin (e.g. `None`)
    pub fn alloc_simple<T: StarlarkValue<'static> + Send + Sync>(&self, val: T) -> FrozenValue {
        self.alloc_simple_typed(val).to_frozen_value()
    }

    /// Allocate a simple [`StarlarkValue`] and return `FrozenRef` to it.
    pub(crate) fn alloc_simple_frozen_ref<T: StarlarkValue<'static> + Send + Sync>(
        &self,
        value: T,
    ) -> FrozenRef<'static, T> {
        self.alloc_simple_typed(value).as_frozen_ref()
    }

    /// Allocate any value in the frozen heap.
    pub(crate) fn alloc_any<T: Debug + Display + Send + Sync>(
        &self,
        value: T,
    ) -> FrozenRef<'static, T> {
        let value = self.alloc_simple_frozen_ref(StarlarkAny::new(value));
        value.map(|r| &r.0)
    }

    /// Allocate any value, use `Debug` implementation for `Display`.
    pub fn alloc_any_display_from_debug<T: Debug + Send + Sync>(
        &self,
        value: T,
    ) -> FrozenRef<'static, T> {
        #[derive(derive_more::Display, Debug)]
        #[display(fmt = "{:?}", _0)]
        struct Wrapper<T: Debug + Send + Sync>(T);
        self.alloc_any(Wrapper(value)).map(|r| &r.0)
    }

    pub(crate) fn alloc_any_display_from_type_name<T: Debug + Send + Sync>(
        &self,
        value: T,
    ) -> FrozenRef<'static, T> {
        #[derive(derive_more::Display, Debug)]
        #[display(fmt = "{}", "std::any::type_name::<T>()")]
        struct Wrapper<T: Debug + Send + Sync>(T);
        self.alloc_any(Wrapper(value)).map(|r| &r.0)
    }

    fn do_alloc_any_slice_display_from_debug<T: Debug + Send + Sync + Clone>(
        &self,
        values: &[T],
    ) -> FrozenRef<'static, [T]> {
        let (_any_array, content) = self.arena.alloc_extra(any_array_avalue(values.len()));
        // Drop lifetime.
        let content = unsafe { transmute!(&mut [MaybeUninit<T>], &mut [MaybeUninit<T>], content) };
        FrozenRef::new(&*maybe_uninit_write_slice_cloned(content, values))
    }

    /// Allocate a slice in the frozen heap.
    pub(crate) fn alloc_any_slice_display_from_debug<T: Debug + Send + Sync + Clone>(
        &self,
        values: &[T],
    ) -> FrozenRef<'static, [T]> {
        if values.is_empty() {
            FrozenRef::new(&[])
        } else if values.len() == 1 {
            self.alloc_any_display_from_debug(values[0].clone())
                .map(|r| slice::from_ref(r))
        } else {
            self.do_alloc_any_slice_display_from_debug(values)
        }
    }

    /// Allocate a new value on a [`FrozenHeap`].
    pub fn alloc<T: AllocFrozenValue>(&self, val: T) -> FrozenValue {
        val.alloc_frozen_value(self)
    }

    /// Number of bytes allocated on this heap, not including any memory
    /// allocated outside of the starlark heap.
    pub fn allocated_bytes(&self) -> usize {
        self.arena.allocated_bytes()
    }

    /// Number of bytes allocated by the heap but not yet filled.
    pub fn available_bytes(&self) -> usize {
        self.arena.available_bytes()
    }

    /// Obtain a summary of how much memory is currently allocated by this heap.
    pub fn allocated_summary(&self) -> HeapSummary {
        self.arena.allocated_summary()
    }

    /// Memory allocated in the arena, but not used for allocation of starlark values.
    pub(crate) fn unused_capacity(&self) -> usize {
        self.arena.unused_capacity()
    }
}

/// Used to `freeze` values by [`Freeze::freeze`](crate::values::Freeze::freeze).
// A freezer is a pair of the FrozenHeap and a "magic" value,
// which we happen to use for the slots (see `FrozenSlotsRef`)
// but could be used for anything.
pub struct Freezer {
    /// Freezing into this heap.
    pub(crate) heap: FrozenHeap,
    /// Defs frozen by this freezer.
    pub(crate) frozen_defs: RefCell<Vec<FrozenRef<'static, FrozenDef>>>,
}

impl Freezer {
    pub(crate) fn new(heap: FrozenHeap) -> Self {
        Freezer {
            heap,
            frozen_defs: RefCell::new(Vec::new()),
        }
    }

    pub(crate) fn into_ref(self) -> FrozenHeapRef {
        self.heap.into_ref()
    }

    /// Allocate a new value while freezing. Usually not a great idea.
    pub fn alloc<'v, T: AllocFrozenValue>(&'v self, val: T) -> FrozenValue {
        val.alloc_frozen_value(&self.heap)
    }

    pub(crate) fn reserve<'v, 'v2: 'v, T: AValue<'v2, ExtraElem = ()>>(
        &'v self,
    ) -> (FrozenValue, Reservation<'v, 'v2, T>) {
        let (fv, r, extra) = self.reserve_with_extra::<T>(0);
        debug_assert!(extra.is_empty());
        (fv, r)
    }

    pub(crate) fn reserve_with_extra<'v, 'v2: 'v, T: AValue<'v2>>(
        &'v self,
        extra_len: usize,
    ) -> (
        FrozenValue,
        Reservation<'v, 'v2, T>,
        &'v mut [MaybeUninit<T::ExtraElem>],
    ) {
        let (r, extra) = self.heap.arena.reserve_with_extra::<T>(extra_len);
        let fv = FrozenValue::new_ptr(unsafe { cast::ptr_lifetime(r.ptr()) }, false);
        (fv, r, extra)
    }

    /// Freeze a nested value while freezing yourself.
    pub fn freeze(&self, value: Value) -> anyhow::Result<FrozenValue> {
        // Case 1: We have our value encoded in our pointer
        if let Some(x) = value.unpack_frozen() {
            return Ok(x);
        }

        // Case 2: We have already been replaced with a forwarding, or need to freeze
        let value = value.0.unpack_ptr().unwrap();
        match value.unpack_overwrite() {
            Either::Left(x) => Ok(unsafe { x.unpack_frozen_value() }),
            Either::Right(v) => unsafe { v.heap_freeze(self) },
        }
    }

    /// Frozen heap where the values are frozen to.
    ///
    /// Can be used to allocate additional values while freezing.
    pub fn frozen_heap(&self) -> &FrozenHeap {
        &self.heap
    }
}

impl Heap {
    /// Create a new [`Heap`].
    pub fn new() -> Self {
        Self::default()
    }

    /// Number of bytes allocated on this heap, not including any memory
    /// allocated outside of the starlark heap.
    pub fn allocated_bytes(&self) -> usize {
        self.arena.borrow().allocated_bytes()
    }

    /// Peak memory allocated to this heap, even if the value is now lower
    /// as a result of a subsequent garbage collection.
    pub fn peak_allocated_bytes(&self) -> usize {
        cmp::max(self.allocated_bytes(), self.peak_allocated.get())
    }

    /// Number of bytes allocated by the heap but not yet filled.
    pub fn available_bytes(&self) -> usize {
        self.arena.borrow().available_bytes()
    }

    fn alloc_raw<'v, 'v2: 'v2>(&'v self, x: impl AValue<'v2, ExtraElem = ()>) -> Value<'v> {
        let arena = self.arena.borrow();
        let v: &AValueRepr<_> = arena.alloc(x);

        // We have an arena inside a RefCell which stores ValueMem<'v>
        // However, we promise not to clear the RefCell other than for GC
        // so we can make the `arena` available longer
        unsafe {
            let value = Value::new_repr(cast::ptr_lifetime(v));
            transmute!(Value, Value, value)
        }
    }

    fn alloc_raw_typed<'v, A: AValue<'v, ExtraElem = ()>>(
        &'v self,
        x: A,
    ) -> ValueTyped<'v, A::StarlarkValue> {
        unsafe { ValueTyped::new_unchecked(self.alloc_raw(x)) }
    }

    pub(crate) fn alloc_str_init<'v>(
        &'v self,
        len: usize,
        init: impl FnOnce(*mut u8),
    ) -> StringValue<'v> {
        let arena = self.arena.borrow();
        let v = arena.alloc_str_init(len, StarlarkStr::UNINIT_HASH, init);

        // We have an arena inside a RefCell which stores ValueMem<'v>
        // However, we promise not to clear the RefCell other than for GC
        // so we can make the `arena` available longer
        unsafe {
            let value = Value::new_ptr(&*v, true);
            StringValue::new_unchecked(value)
        }
    }

    /// Allocate a string on the heap.
    pub fn alloc_str<'v>(&'v self, x: &str) -> StringValue<'v> {
        if let Some(x) = constant_string(x) {
            x.to_string_value()
        } else {
            self.alloc_str_init(x.len(), |dest| unsafe {
                copy_nonoverlapping(x.as_ptr(), dest, x.len())
            })
        }
    }

    /// Allocate a string on the heap, based on two concatenated strings.
    pub fn alloc_str_concat<'v>(&'v self, x: &str, y: &str) -> StringValue<'v> {
        if x.is_empty() {
            self.alloc_str(y)
        } else if y.is_empty() {
            self.alloc_str(x)
        } else {
            self.alloc_str_init(x.len() + y.len(), |dest| unsafe {
                copy_nonoverlapping(x.as_ptr(), dest, x.len());
                copy_nonoverlapping(y.as_ptr(), dest.add(x.len()), y.len())
            })
        }
    }

    /// Allocate a string on the heap, based on three concatenated strings.
    pub fn alloc_str_concat3<'v>(&'v self, x: &str, y: &str, z: &str) -> StringValue<'v> {
        if x.is_empty() {
            self.alloc_str_concat(y, z)
        } else if y.is_empty() {
            self.alloc_str_concat(x, z)
        } else if z.is_empty() {
            self.alloc_str_concat(x, y)
        } else {
            self.alloc_str_init(x.len() + y.len() + z.len(), |dest| unsafe {
                copy_nonoverlapping(x.as_ptr(), dest, x.len());
                let dest = dest.add(x.len());
                copy_nonoverlapping(y.as_ptr(), dest, y.len());
                let dest = dest.add(y.len());
                copy_nonoverlapping(z.as_ptr(), dest, z.len());
            })
        }
    }

    /// Allocate a tuple with the given elements.
    pub(crate) fn alloc_tuple<'v>(&'v self, elems: &[Value<'v>]) -> Value<'v> {
        if elems.is_empty() {
            return Value::new_empty_tuple();
        }

        unsafe {
            let arena = self.arena.borrow();
            let (avalue, extra) = arena.alloc_extra(tuple_avalue(elems.len()));
            maybe_uninit_write_slice(extra, elems);
            Value::new_repr(&*avalue)
        }
    }

    pub(crate) fn alloc_tuple_iter<'v>(
        &'v self,
        elems: impl IntoIterator<Item = Value<'v>>,
    ) -> Value<'v> {
        let elems = elems.into_iter();
        let (lower, upper) = elems.size_hint();
        if Some(lower) == upper {
            if lower == 0 {
                return Value::new_empty_tuple();
            }

            unsafe {
                let arena = self.arena.borrow();
                let (avalue, extra) = arena.alloc_extra(tuple_avalue(lower));
                maybe_uninit_write_from_exact_size_iter(extra, elems, Value::new_none());
                Value::new_repr(&*avalue)
            }
        } else {
            self.alloc_tuple(&elems.collect::<Vec<_>>())
        }
    }

    pub(crate) fn alloc_array<'v>(&'v self, cap: usize) -> ValueTyped<'v, Array<'v>> {
        if cap == 0 {
            return FrozenValueTyped::new_repr(VALUE_EMPTY_ARRAY.repr()).to_value_typed();
        }

        let cap: u32 = cap.try_into().expect("capacity overflows u32::MAX");

        unsafe {
            let (avalue, _) = self.arena.borrow().alloc_extra(array_avalue(cap));
            ValueTyped::new_repr(&*avalue)
        }
    }

    /// Allocate a list with the given elements.
    pub(crate) fn alloc_list<'v>(&'v self, elems: &[Value<'v>]) -> Value<'v> {
        let array = self.alloc_array(elems.len());
        array.extend_from_slice(elems);
        self.alloc_raw(list_avalue(array))
    }

    /// Allocate a list with the given elements.
    pub(crate) fn alloc_list_iter<'v>(
        &'v self,
        elems: impl IntoIterator<Item = Value<'v>>,
    ) -> Value<'v> {
        let elems = elems.into_iter();
        let array = self.alloc_array(0);
        let list = self.alloc_raw_typed(list_avalue(array));
        list.0.extend(elems, self);
        list.to_value()
    }

    /// Allocate a list by concatenating two slices.
    pub(crate) fn alloc_list_concat<'v>(&'v self, a: &[Value<'v>], b: &[Value<'v>]) -> Value<'v> {
        let array = self.alloc_array(a.len() + b.len());
        array.extend_from_slice(a);
        array.extend_from_slice(b);
        self.alloc_raw(list_avalue(array))
    }

    pub(crate) fn alloc_char<'v>(&'v self, x: char) -> StringValue<'v> {
        let mut dst = [0; 4];
        let res = x.encode_utf8(&mut dst);
        self.alloc_str(res)
    }

    /// Allocate a simple [`StarlarkValue`] on this heap.
    ///
    /// Simple value is any starlark value which:
    /// * bound by `'static` lifetime (in particular, it cannot contain references to other `Value`s)
    /// * is not special builtin (e.g. `None`)
    ///
    /// Must be [`Send`] and [`Sync`] because it will be reused in frozen values.
    pub fn alloc_simple<'v, T: StarlarkValue<'static> + Send + Sync>(&'v self, x: T) -> Value<'v> {
        self.alloc_raw(simple(x))
    }

    /// Allocate a [`ComplexValue`] on the [`Heap`].
    pub fn alloc_complex<'v, T>(&'v self, x: T) -> Value<'v>
    where
        T: ComplexValue<'v>,
        T::Frozen: StarlarkValue<'static>,
    {
        self.alloc_raw(complex(x))
    }

    /// Allocate a value which can be traced (garbage collected), but cannot be frozen.
    pub fn alloc_complex_no_freeze<'v, T>(&'v self, x: T) -> Value<'v>
    where
        T: StarlarkValue<'v> + Trace<'v>,
    {
        // When specializations are stable, we can have single `alloc_complex` function,
        // which enables or not enables freezing depending on whether `T` implements `Freeze`.
        self.alloc_raw(complex_no_freeze(x))
    }

    /// Allocate a new value on a [`Heap`].
    pub fn alloc<'v, T: AllocValue<'v>>(&'v self, x: T) -> Value<'v> {
        x.alloc_value(self)
    }

    /// Allocate a value and return [`ValueTyped`] of it.
    /// Can fail if the [`AllocValue`] trait generates a different type on the heap.
    pub fn alloc_typed<'v, T: AllocValue<'v> + StarlarkValue<'v>>(
        &'v self,
        x: T,
    ) -> ValueTyped<'v, T> {
        ValueTyped::new(self.alloc(x)).expect("just allocated value must have the right type")
    }

    /// Allocate a value and return [`ValueOf`] of it.
    pub fn alloc_value_of<'v, T>(&'v self, x: T) -> ValueOf<'v, &'v T>
    where
        T: AllocValue<'v>,
        &'v T: UnpackValue<'v>,
    {
        let value = self.alloc(x);
        ValueOf::unpack_value(value)
            .expect("just allocate value must be unpackable to the type of value")
    }

    pub(crate) unsafe fn visit_arena<'v>(
        &'v self,
        forward_heap_kind: HeapKind,
        v: &mut impl ArenaVisitor<'v>,
    ) {
        (*self.arena.get_mut()).visit_arena(HeapKind::Unfrozen, forward_heap_kind, v)
    }

    /// Garbage collect any values that are unused. This function is _unsafe_ in
    /// the sense that any `Value<'v>` not returned by `Tracer` _will become
    /// invalid_. Furthermore, any references to values, e.g `&'v str` will
    /// also become invalid.
    pub(crate) unsafe fn garbage_collect<'v>(&'v self, f: impl FnOnce(&Tracer<'v>)) {
        // Record the highest peak, so it never decreases
        self.peak_allocated.set(self.peak_allocated_bytes());
        self.garbage_collect_internal(f)
    }

    unsafe fn garbage_collect_internal<'v>(&'v self, f: impl FnOnce(&Tracer<'v>)) {
        // Must rewrite all Value's so they point at the new heap.
        // Take the arena out of the heap to make sure nobody allocates in it,
        // but hold the reference until the GC is done.
        let _arena = self.arena.take();

        let tracer = Tracer::<'v> {
            arena: Arena::default(),
            phantom: PhantomData,
        };
        f(&tracer);
        self.arena.set(tracer.arena);
    }

    /// Obtain a summary of how much memory is currently allocated by this heap.
    pub fn allocated_summary(&self) -> HeapSummary {
        self.arena.borrow().allocated_summary()
    }

    pub(crate) fn record_call_enter<'v>(&'v self, function: Value<'v>) {
        let time = Instant::now();
        assert!(mem::needs_drop::<CallEnter<NeedsDrop>>());
        assert!(!mem::needs_drop::<CallEnter<NoDrop>>());
        self.alloc_complex_no_freeze(CallEnter {
            function,
            time,
            maybe_drop: NeedsDrop,
        });
        self.alloc_complex_no_freeze(CallEnter {
            function,
            time,
            maybe_drop: NoDrop,
        });
    }

    pub(crate) fn record_call_exit<'v>(&'v self) {
        let time = Instant::now();
        assert!(mem::needs_drop::<CallExit<NeedsDrop>>());
        assert!(!mem::needs_drop::<CallExit<NoDrop>>());
        self.alloc_simple(CallExit {
            time,
            maybe_drop: NeedsDrop,
        });
        self.alloc_simple(CallExit {
            time,
            maybe_drop: NoDrop,
        });
    }

    /// Memory allocated in the arena, but not used for allocation of starlark values.
    pub(crate) fn unused_capacity(&self) -> usize {
        self.arena.borrow().unused_capacity()
    }
}

/// Used to perform garbage collection by [`Trace::trace`](crate::values::Trace::trace).
pub struct Tracer<'v> {
    arena: Arena<Bump>,
    phantom: PhantomData<&'v ()>,
}

impl<'v> Tracer<'v> {
    /// Walk over a value during garbage collection.
    pub fn trace(&self, value: &mut Value<'v>) {
        *value = self.adjust(*value)
    }

    /// Helper function to annotate that this field has been considered for tracing,
    /// but is not relevant because it has a static lifetime containing no relevant values.
    /// Does nothing.
    pub fn trace_static<T: 'static>(&self, value: &mut T) {
        // Nothing to do because T can't contain the lifetime 'v
        let _ = value;
    }

    pub(crate) fn reserve<'a, 'v2: 'v + 'a, T: AValue<'v2, ExtraElem = ()>>(
        &'a self,
    ) -> (Value<'v>, Reservation<'a, 'v2, T>) {
        let (v, r, extra) = self.reserve_with_extra::<T>(0);
        debug_assert!(extra.is_empty());
        (v, r)
    }

    pub(crate) fn reserve_with_extra<'a, 'v2: 'v + 'a, T: AValue<'v2>>(
        &'a self,
        extra_len: usize,
    ) -> (
        Value<'v>,
        Reservation<'a, 'v2, T>,
        &'a mut [MaybeUninit<T::ExtraElem>],
    ) {
        assert!(!T::IS_STR, "strings cannot be reserved");
        let (r, extra) = self.arena.reserve_with_extra::<T>(extra_len);
        let v = Value::new_ptr(unsafe { cast::ptr_lifetime(r.ptr()) }, false);
        (v, r, extra)
    }

    pub(crate) fn alloc_str(&self, x: &str) -> Value<'v> {
        let v = self.arena.alloc_str(x);
        unsafe { Value::new_ptr(&*v, true) }
    }

    fn adjust(&self, value: Value<'v>) -> Value<'v> {
        // Case 1, doesn't point at the old arena
        if !value.0.is_unfrozen() {
            return value;
        }
        let old_val = value.0.unpack_ptr().unwrap();

        // Case 2: We have already been replaced with a forwarding, or need to freeze
        let res = match old_val.unpack_overwrite() {
            Either::Left(x) => unsafe { x.unpack_unfrozen_value() },
            Either::Right(v) => unsafe { v.heap_copy(self) },
        };

        res
    }
}

#[test]
fn test_send_sync()
where
    FrozenHeapRef: Send + Sync,
{
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::mem::MaybeUninit;
use std::slice;

/// Populate a slice with values from iterator, handle panic by writing provided fallback value.
///
/// buck2 is compiled without unwinding, but other users of starlark-rust may.
/// It's hard to say how dangerous is to leave for example a list containing
/// uninitialized memory in heap in case of panic. So let's just initialize
/// the slice with valid values on panic.
#[inline]
pub(crate) fn maybe_uninit_write_from_exact_size_iter<T: Copy>(
    slice: &mut [MaybeUninit<T>],
    iter: impl IntoIterator<Item = T>,
    write_on_panic: T,
) {
    /// On drop, populate the iterator with the provided value.
    struct WriteRemOnDrop<'a, T: Copy>(slice::IterMut<'a, MaybeUninit<T>>, T);

    impl<'a, T: Copy> Drop for WriteRemOnDrop<'a, T> {
        fn drop(&mut self) {
            for elem in &mut self.0 {
                elem.write(self.1);
            }
        }
    }

    let mut slice = WriteRemOnDrop(slice.iter_mut(), write_on_panic);
    let mut iter = iter.into_iter();

    for (elem_place, elem) in slice.0.by_ref().zip(iter.by_ref()) {
        elem_place.write(elem);
    }

    assert!(
        slice.0.next().is_none() && iter.next().is_none(),
        "iterator provided size_hint incorrectly",
    );

    // Should live to this point.
    drop(slice);
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

//! Starlark heap implementation.

pub(crate) mod allocator;
pub(crate) mod arena;
pub(crate) mod call_enter_exit;
mod fast_cell;
pub(crate) mod heap_type;
pub(crate) mod maybe_uninit_slice_util;
pub(crate) mod profile;
pub(crate) mod repr;


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::hash::Hash;
use std::mem;
use std::mem::ManuallyDrop;
use std::ptr;

use dupe::Dupe;
use either::Either;

use crate::any::AnyLifetime;
use crate::cast;
use crate::values::layout::avalue::AValue;
use crate::values::layout::heap::heap_type::HeapKind;
use crate::values::layout::value_alloc_size::ValueAllocSize;
use crate::values::layout::vtable::AValueDyn;
use crate::values::layout::vtable::AValueVTable;
use crate::values::layout::vtable::StarlarkValueRawPtr;
use crate::values::FrozenValue;
use crate::values::StarlarkValue;
use crate::values::Value;

#[derive(Clone)]
#[repr(C)]
pub(crate) struct AValueHeader(pub(crate) &'static AValueVTable);

impl Hash for AValueHeader {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        ptr::hash(self.0, state);
    }
}

impl PartialEq for AValueHeader {
    fn eq(&self, other: &Self) -> bool {
        ptr::eq(self.0, other.0)
    }
}

impl Eq for AValueHeader {}

// Implements Copy so this is fine
impl Dupe for AValueHeader {}

/// How object is represented in arena.
#[repr(C, align(8))]
pub(crate) struct AValueRepr<T> {
    pub(crate) header: AValueHeader,
    /// Payload of the object, i.e. `StarlarkValue`.
    /// Note that `T` may have larger alignment that `AValueHeader`,
    /// so we cannot add fixed offset to `self` to get to the payload.
    pub(crate) payload: T,
}

/// "Forward" pointer (pointer to another heap during GC).
///
/// This pointer has `TAG_STR` bit set if it points to a string.
///
/// Lower bit (which is the same bit as `TAG_UNFROZEN`) is always unset
/// regardless of whether it points to frozen or unfrozen value.
/// User of this struct must set this bit explicitly if needed.
#[derive(Copy, Clone, Dupe)]
pub(crate) struct ForwardPtr(usize);

impl ForwardPtr {
    pub(crate) fn new(ptr: usize) -> ForwardPtr {
        debug_assert!(ptr & 1 == 0);
        ForwardPtr(ptr)
    }

    /// It's caller responsibility to ensure that forward pointer points to a frozen value.
    pub(crate) unsafe fn unpack_frozen_value(self) -> FrozenValue {
        FrozenValue::new_ptr_usize_with_str_tag(self.0)
    }

    /// It's caller responsibility to ensure that forward pointer points to an unfrozen value.
    pub(crate) unsafe fn unpack_unfrozen_value<'v>(self) -> Value<'v> {
        Value::new_ptr_usize_with_str_tag(self.0)
    }

    pub(crate) unsafe fn unpack_value<'v>(self, heap_kind: HeapKind) -> Value<'v> {
        match heap_kind {
            HeapKind::Unfrozen => self.unpack_unfrozen_value(),
            HeapKind::Frozen => self.unpack_frozen_value().to_value(),
        }
    }
}

/// This is object written over [`AValueRepr`] during GC.
#[repr(C)]
pub(crate) struct AValueForward {
    /// Moved object pointer with lowest bit set.
    forward_ptr: usize,
    /// Size of `AValueRepr<T>` including extra.
    object_size: ValueAllocSize,
}

impl AValueForward {
    pub(crate) fn new(forward_ptr: ForwardPtr, object_size: ValueAllocSize) -> AValueForward {
        AValueForward {
            forward_ptr: forward_ptr.0 | 1,
            object_size,
        }
    }

    /// Unpack forward pointer.
    pub(crate) fn forward_ptr(&self) -> ForwardPtr {
        debug_assert!((self.forward_ptr & 1) != 0);
        ForwardPtr(self.forward_ptr & !1)
    }
}

/// Object on the heap, either a real object or a forward.
#[repr(C)]
pub(crate) union AValueOrForward {
    // We intentionally do not implement `Copy` for these types
    // to avoid accidentally copying them.
    pub(crate) header: ManuallyDrop<AValueHeader>,
    pub(crate) forward: ManuallyDrop<AValueForward>,
    pub(crate) flags: usize,
}

impl AValueOrForward {
    /// Is this pointer a value or forward?
    #[inline]
    fn is_forward(&self) -> bool {
        unsafe { (self.flags & 1) != 0 }
    }

    pub(crate) fn unpack(&self) -> Either<&AValueHeader, &AValueForward> {
        if self.is_forward() {
            Either::Right(unsafe { &self.forward })
        } else {
            Either::Left(unsafe { &self.header })
        }
    }

    /// Unpack something that might have been overwritten.
    pub(crate) fn unpack_overwrite<'v>(&'v self) -> Either<ForwardPtr, AValueDyn<'v>> {
        match self.unpack() {
            Either::Left(header) => Either::Right(header.unpack()),
            Either::Right(forward) => Either::Left(forward.forward_ptr()),
        }
    }

    #[inline]
    pub(crate) unsafe fn unpack_header_unchecked(&self) -> &AValueHeader {
        debug_assert!(!self.is_forward());
        &self.header
    }

    pub(crate) fn unpack_header(&self) -> Option<&AValueHeader> {
        self.unpack().left()
    }

    pub(crate) fn unpack_forward(&self) -> Option<&AValueForward> {
        self.unpack().right()
    }

    /// Size of allocation for this object:
    /// following object is allocated at `self + alloc_size + align up`.
    pub(crate) fn alloc_size(&self) -> ValueAllocSize {
        match self.unpack() {
            Either::Left(ptr) => ptr.unpack().memory_size(),
            Either::Right(forward) => {
                // Overwritten, so the next word will be the size of the memory
                forward.object_size
            }
        }
    }
}

impl AValueForward {
    pub(crate) fn assert_does_not_overwrite_extra<'v, T: AValue<'v>>() {
        assert!(mem::size_of::<AValueForward>() <= AValueRepr::<T>::offset_of_extra());
    }
}

impl AValueHeader {
    /// Alignment of objects in Starlark heap.
    /// We must use 8 byte alignment because we use three lowest bits for tags.
    /// Note the alignment of `AValueHeader` may be smaller than this.
    pub(crate) const ALIGN: usize = 8;

    pub(crate) fn new<'v, T: AValue<'v>>() -> AValueHeader {
        let header = AValueHeader::new_const::<T>();

        let vtable_ptr = header.0 as *const AValueVTable as usize;
        // Check that the LSB is not set, as we reuse that for overwrite
        debug_assert!(vtable_ptr & 1 == 0);

        header
    }

    pub(crate) const fn new_const<'v, T: AValue<'v>>() -> AValueHeader {
        AValueHeader(AValueVTable::new::<T>())
    }

    #[inline]
    pub(crate) fn payload_ptr(&self) -> StarlarkValueRawPtr {
        StarlarkValueRawPtr::new_header(self)
    }

    pub(crate) unsafe fn payload<'v, T: StarlarkValue<'v>>(&self) -> &T {
        debug_assert_eq!(self.unpack().static_type_of_value(), T::static_type_id());
        &*self.payload_ptr().value_ptr::<T>()
    }

    pub(crate) unsafe fn unpack_value<'v>(&'v self, heap_kind: HeapKind) -> Value<'v> {
        match heap_kind {
            HeapKind::Unfrozen => Value::new_ptr_query_is_str(self),
            HeapKind::Frozen => {
                FrozenValue::new_ptr_query_is_str(cast::ptr_lifetime(self)).to_value()
            }
        }
    }

    pub(crate) fn unpack<'v>(&'v self) -> AValueDyn<'v> {
        unsafe {
            // TODO: this assertion does not belong here.
            //   Instead, `Value` should be a `Pointer<AValueOrForward>`
            //   instead of `Pointer<AValueHeader>`,
            //   and assertion should be where we unpack the pointer.
            debug_assert!(
                !(*(self as *const AValueHeader as *const AValueOrForward)).is_forward(),
                "value is a forward pointer; value cannot be unpacked during GC or freeze"
            );
        }
        unsafe { AValueDyn::new(self.payload_ptr(), self.0) }
    }

    /// After performing the overwrite any existing pointers to this value
    /// are corrupted.
    pub unsafe fn overwrite_with_forward<'v, T: AValue<'v>>(
        me: *mut AValueRepr<T>,
        forward_ptr: ForwardPtr,
    ) -> T {
        // TODO(nga): we don't need to do virtual call to obtain memory size
        let sz = (*me).header.unpack().memory_size();
        let p = me as *const AValueRepr<T>;
        let res = ptr::read(p).payload;
        let p = me as *mut AValueForward;
        *p = AValueForward::new(forward_ptr, sz);
        res
    }

    /// Cast header pointer to repr pointer.
    #[inline]
    pub(crate) unsafe fn as_repr<'v, A: AValue<'v>>(&self) -> &AValueRepr<A> {
        &*(self.as_repr_v::<A::StarlarkValue>() as *const _ as *const _)
    }

    #[inline]
    pub(crate) unsafe fn as_repr_v<'v, T: StarlarkValue<'v>>(&self) -> &AValueRepr<T> {
        debug_assert_eq!(T::static_type_id(), self.unpack().static_type_of_value());
        &*(self as *const AValueHeader as *const AValueRepr<T>)
    }

    fn as_avalue_or_header(&self) -> &AValueOrForward {
        unsafe { &*(self as *const AValueHeader as *const AValueOrForward) }
    }

    /// Size of allocation for this object: following object is allocated at `self + alloc_size`.
    pub(crate) fn alloc_size(&self) -> ValueAllocSize {
        self.as_avalue_or_header().alloc_size()
    }
}

impl<T> AValueRepr<T> {
    const _ASSERTIONS: () = {
        assert!(mem::align_of::<Self>() == AValueHeader::ALIGN);
    };

    pub(crate) const fn with_metadata(
        metadata: &'static AValueVTable,
        payload: T,
    ) -> AValueRepr<T> {
        AValueRepr {
            header: AValueHeader(metadata),
            payload,
        }
    }

    pub(crate) fn offset_of_payload() -> usize {
        memoffset::offset_of!(Self, payload)
    }

    /// Padding between header and payload.
    /// Non-zero when alignment of payload is larger than alignment of pointer.
    pub(crate) fn padding_after_header() -> usize {
        Self::offset_of_payload() - mem::size_of::<AValueHeader>()
    }

    /// Offset of value extra content relative to `AValueRepr` start.
    pub(crate) fn offset_of_extra<'v>() -> usize
    where
        T: AValue<'v>,
    {
        Self::offset_of_payload() + T::offset_of_extra()
    }

    pub(crate) fn from_payload_ptr_mut(payload_ptr: *mut T) -> *mut AValueRepr<T> {
        let payload_ptr = payload_ptr as usize;
        let header_ptr = payload_ptr - Self::offset_of_payload();
        header_ptr as *mut AValueRepr<T>
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::mem::MaybeUninit;
use std::ptr::NonNull;

use crate::values::layout::value_alloc_size::ValueAllocSize;

pub(crate) enum ChunkAllocationDirection {
    /// Next allocation in the chunk has higher address than the previous one.
    #[allow(dead_code)] // TODO(nga): used in the following diff D40738710.
    Up,
    /// Next allocation in the chunk has lower address than the previous one.
    Down,
}

/// Fast memory allocator for the heap.
pub(crate) trait ArenaAllocator {
    /// Number of bytes allocated by this allocator.
    ///
    /// That is:
    /// * space occupied by allocated values
    /// * padding
    /// * reserved but not yet allocated space
    /// * does not include metadata
    fn allocated_bytes(&self) -> usize;
    /// Number of bytes reserved but not yet allocated by this allocator.
    fn remaining_capacity(&self) -> usize;
    /// Estimate the size of allocated metadata.
    fn allocation_overhead(&self) -> usize;
    /// Allocate given number of words.
    fn alloc(&self, size: ValueAllocSize) -> NonNull<u8>;

    /// This allocator chunk allocation direction.
    const CHUNK_ALLOCATION_DIRECTION: ChunkAllocationDirection;

    type ChunkRevIterator<'a>: Iterator<Item = &'a [MaybeUninit<u8>]>
    where
        Self: 'a;

    /// Iterate allocated chunks in the reverse order.
    unsafe fn iter_allocated_chunks_rev(&self) -> Self::ChunkRevIterator<'_>;

    /// No more allocation, reclaim memory if possible.
    fn finish(&mut self);
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::mem::MaybeUninit;
use std::ptr::NonNull;
use std::slice;

use bumpalo::Bump;
use bumpalo::ChunkRawIter;

use crate::values::layout::heap::allocator::api::ArenaAllocator;
use crate::values::layout::heap::allocator::api::ChunkAllocationDirection;
use crate::values::layout::value_alloc_size::ValueAllocSize;

pub(crate) struct ChunkIteratorWrapper<'a> {
    iter: ChunkRawIter<'a>,
}

impl<'a> Iterator for ChunkIteratorWrapper<'a> {
    type Item = &'a [MaybeUninit<u8>];

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.iter
            .next()
            .map(|(ptr, len)| unsafe { slice::from_raw_parts(ptr as *mut MaybeUninit<u8>, len) })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl ArenaAllocator for Bump {
    fn allocated_bytes(&self) -> usize {
        Bump::allocated_bytes(self)
    }

    fn remaining_capacity(&self) -> usize {
        Bump::chunk_capacity(self)
    }

    fn allocation_overhead(&self) -> usize {
        // Need this to estimate correctly:
        // https://github.com/fitzgen/bumpalo/pull/184
        // For now we can either ignore it or hardcode the bumpalo footer size.
        0
    }

    #[inline]
    fn alloc(&self, size: ValueAllocSize) -> NonNull<u8> {
        Bump::alloc_layout(self, size.layout())
    }

    const CHUNK_ALLOCATION_DIRECTION: ChunkAllocationDirection = ChunkAllocationDirection::Down;

    type ChunkRevIterator<'a> = ChunkIteratorWrapper<'a>;

    unsafe fn iter_allocated_chunks_rev(&self) -> Self::ChunkRevIterator<'_> {
        ChunkIteratorWrapper {
            iter: Bump::iter_allocated_chunks_raw(self),
        }
    }

    fn finish(&mut self) {}
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

pub(crate) mod api;
pub(crate) mod bumpalo;


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::cell::RefCell;
use std::collections::hash_map;
use std::collections::HashMap;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::rc::Rc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Instant;

use allocative::Allocative;
use dupe::Dupe;
use either::Either;
use starlark_map::small_map::SmallMap;

use crate::eval::runtime::profile::data::ProfileDataImpl;
use crate::eval::runtime::profile::flamegraph::FlameGraphData;
use crate::eval::runtime::profile::flamegraph::FlameGraphNode;
use crate::eval::runtime::profile::heap::RetainedHeapProfileMode;
use crate::eval::runtime::small_duration::SmallDuration;
use crate::eval::ProfileData;
use crate::values::layout::heap::arena::ArenaVisitor;
use crate::values::layout::heap::heap_type::HeapKind;
use crate::values::layout::heap::profile::alloc_counts::AllocCounts;
use crate::values::layout::heap::profile::arc_str::ArcStr;
use crate::values::layout::heap::profile::by_type::HeapSummary;
use crate::values::layout::heap::profile::string_index::StringId;
use crate::values::layout::heap::profile::string_index::StringIndex;
use crate::values::layout::heap::profile::summary_by_function::HeapSummaryByFunction;
use crate::values::layout::heap::repr::AValueOrForward;
use crate::values::layout::pointer::RawPointer;
use crate::values::Heap;
use crate::values::Value;

/// A mapping from function Value to FunctionId, which must be continuous
#[derive(Default)]
struct FunctionIds {
    values: HashMap<RawPointer, StringId>,
    strings: StringIndex,
}

impl FunctionIds {
    fn get_value(&mut self, x: Value) -> StringId {
        match self.values.entry(x.ptr_value()) {
            hash_map::Entry::Occupied(v) => *v.get(),
            hash_map::Entry::Vacant(outer) => {
                let function_id = self.strings.index(&x.to_str());
                outer.insert(function_id);
                function_id
            }
        }
    }
}

/// A stack frame, its caller and the functions it called, and the allocations it made itself.
struct StackFrameData {
    callees: SmallMap<StringId, StackFrameBuilder>,
    allocs: HeapSummary,
    /// Time spent in this frame excluding callees.
    /// Double, because enter/exit are recorded twice, in drop and non-drop heaps.
    time_x2: SmallDuration,
    /// How many times this function was called (with this stack).
    /// Double.
    calls_x2: u32,
}

#[derive(Clone, Dupe)]
struct StackFrameBuilder(Rc<RefCell<StackFrameData>>);

impl StackFrameBuilder {
    fn new() -> Self {
        Self(Rc::new(RefCell::new(StackFrameData {
            callees: Default::default(),
            allocs: Default::default(),
            time_x2: SmallDuration::default(),
            calls_x2: 0,
        })))
    }

    /// Enter a new stack frame.
    fn push(&self, function: StringId) -> Self {
        let mut this = self.0.borrow_mut();

        let callee = this
            .callees
            .entry(function)
            .or_insert_with(StackFrameBuilder::new);

        callee.dupe()
    }

    fn build(&self) -> StackFrame {
        StackFrame {
            callees: self
                .0
                .borrow()
                .callees
                .iter()
                .map(|(f, s)| (*f, s.build()))
                .collect(),
            allocs: self.0.borrow().allocs.clone(),
            time_x2: self.0.borrow().time_x2,
            calls_x2: self.0.borrow().calls_x2,
        }
    }
}

/// An accumulator for stack frames that lets us visit the heap.
pub(crate) struct StackCollector {
    /// Timestamp of last call enter or exit.
    last_time: Option<Instant>,
    ids: FunctionIds,
    current: Vec<StackFrameBuilder>,
    /// What we are collecting.
    /// When unset, we are collecting allocated memory (not retained).
    /// When set, must be set to correct heap type (unfrozen or frozen), we are traversing.
    retained: Option<HeapKind>,
}

impl StackCollector {
    pub(crate) fn new(retained: Option<HeapKind>) -> Self {
        Self {
            ids: FunctionIds::default(),
            current: vec![StackFrameBuilder::new()],
            last_time: None,
            retained,
        }
    }
}

impl<'v> ArenaVisitor<'v> for StackCollector {
    fn regular_value(&mut self, value: &'v AValueOrForward) {
        let value = match (value.unpack(), self.retained) {
            (Either::Left(header), None) => unsafe { header.unpack_value(HeapKind::Unfrozen) },
            (Either::Right(forward), Some(retained)) => unsafe {
                forward.forward_ptr().unpack_value(retained)
            },
            _ => return,
        };

        let frame = match self.current.last() {
            Some(frame) => frame,
            None => return,
        };

        // Value allocated in this frame, record it!
        let typ = value.get_ref().get_type();
        let mut frame = frame.0.borrow_mut();
        frame.allocs.add(
            typ,
            AllocCounts {
                count: 1,
                bytes: value.get_ref().total_memory(),
            },
        );
    }

    fn call_enter(&mut self, function: Value<'v>, time: Instant) {
        if let Some(last_time) = self.last_time {
            self.current.last_mut().unwrap().0.borrow_mut().time_x2 +=
                time.saturating_duration_since(last_time);
            self.current.last_mut().unwrap().0.borrow_mut().calls_x2 += 1;
        }

        let frame = match self.current.last() {
            Some(frame) => frame,
            None => return,
        };

        // New frame, enter it.
        let id = self.ids.get_value(function);
        let new_frame = frame.push(id);
        self.current.push(new_frame);

        self.last_time = Some(time)
    }

    fn call_exit(&mut self, time: Instant) {
        if let Some(last_time) = self.last_time {
            self.current.last_mut().unwrap().0.borrow_mut().time_x2 +=
                time.saturating_duration_since(last_time);
        }
        self.current.pop().unwrap();
        self.last_time = Some(time);
    }
}

/// Aggregated stack frame data.
#[derive(Clone, Default, Allocative)]
pub(crate) struct StackFrame {
    /// Aggregated callees.
    pub(crate) callees: SmallMap<StringId, StackFrame>,
    /// Aggregated allocations in this frame, without callees.
    pub(crate) allocs: HeapSummary,
    /// Time spend in this frame excluding callees.
    /// `x2` because enter/exit are recorded twice, in drop and non-drop heaps.
    pub(crate) time_x2: SmallDuration,
    /// How many times this frame was called with the same callers.
    /// `x2` because enter/exit are recorded twice, in drop and non-drop heaps.
    pub(crate) calls_x2: u32,
}

impl StackFrame {
    fn merge_callees<'a>(
        frames: &'a [StackFrameWithContext<'a>],
        strings: &mut StringIndex,
    ) -> SmallMap<StringId, StackFrame> {
        let mut group_by_callee: SmallMap<&str, Vec<StackFrameWithContext>> = SmallMap::new();
        for frame in frames {
            for (name, callee) in frame.callees() {
                group_by_callee
                    .entry(name)
                    .or_insert_with(Vec::new)
                    .push(callee);
            }
        }
        group_by_callee
            .into_iter()
            .map(|(name, frames)| {
                let name = strings.index(name);
                (name, StackFrame::merge(frames, strings))
            })
            .collect()
    }

    fn merge<'a>(
        frames: impl IntoIterator<Item = StackFrameWithContext<'a>>,
        strings: &mut StringIndex,
    ) -> StackFrame {
        let frames = Vec::from_iter(frames);
        let callees = StackFrame::merge_callees(&frames, strings);
        let allocs = HeapSummary::merge(frames.iter().map(|f| &f.frame.allocs));
        let time_x2 = frames.iter().map(|f| f.frame.time_x2).sum();
        let calls_x2 = frames.iter().map(|f| f.frame.calls_x2).sum();
        StackFrame {
            callees,
            allocs,
            time_x2,
            calls_x2,
        }
    }
}

struct StackFrameWithContext<'c> {
    frame: &'c StackFrame,
    strings: &'c StringIndex,
}

impl<'c> StackFrameWithContext<'c> {
    fn callees(&self) -> impl Iterator<Item = (&'c ArcStr, StackFrameWithContext<'c>)> + '_ {
        self.frame.callees.iter().map(move |(id, callee)| {
            (
                self.strings.get(*id),
                StackFrameWithContext {
                    frame: callee,
                    strings: self.strings,
                },
            )
        })
    }

    /// Write this stack frame's data to a file in flamegraph.pl format.
    fn write_flame_graph(&self, node: &mut FlameGraphNode) {
        for (k, v) in &self.frame.allocs.summary {
            node.child((*k).into()).add(v.bytes as u64);
        }

        for (id, frame) in self.callees() {
            let child_node = node.child(id.dupe());
            frame.write_flame_graph(child_node);
        }
    }
}

/// `Clone` wrapper.
#[derive(Default, Allocative)]
pub(crate) struct UnusedCapacity(AtomicUsize);

impl Clone for UnusedCapacity {
    fn clone(&self) -> Self {
        UnusedCapacity(AtomicUsize::new(self.0.load(Ordering::Relaxed)))
    }
}

impl UnusedCapacity {
    pub(crate) fn new(value: usize) -> UnusedCapacity {
        UnusedCapacity(AtomicUsize::new(value))
    }

    pub(crate) fn get(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }

    pub(crate) fn set(&self, value: usize) {
        self.0.store(value, Ordering::Relaxed);
    }
}

/// Aggregated heap profiling data when heap profiling is enabled.
///
/// Can be:
/// * written as CSV or flamegraph
/// * merged with another data
#[derive(Clone, Allocative)]
pub struct AggregateHeapProfileInfo {
    pub(crate) strings: StringIndex,
    pub(crate) root: StackFrame,
    /// Memory allocated in bump, but unused.
    pub(crate) unused_capacity: UnusedCapacity,
}

impl Debug for AggregateHeapProfileInfo {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("AggregateHeapProfileInfo")
            .finish_non_exhaustive()
    }
}

impl Default for AggregateHeapProfileInfo {
    fn default() -> AggregateHeapProfileInfo {
        let strings = StringIndex::default();
        AggregateHeapProfileInfo {
            root: StackFrame::default(),
            strings,
            unused_capacity: UnusedCapacity::default(),
        }
    }
}

impl AggregateHeapProfileInfo {
    pub(crate) fn collect(heap: &Heap, retained: Option<HeapKind>) -> AggregateHeapProfileInfo {
        let mut collector = StackCollector::new(retained);
        unsafe {
            heap.visit_arena(HeapKind::Unfrozen, &mut collector);
        }
        assert_eq!(1, collector.current.len());
        let unused_capacity = if retained.is_some() {
            // Filled later.
            UnusedCapacity::default()
        } else {
            UnusedCapacity::new(heap.unused_capacity())
        };
        AggregateHeapProfileInfo {
            strings: collector.ids.strings,
            root: collector.current.pop().unwrap().build(),
            unused_capacity,
        }
    }

    fn root(&self) -> StackFrameWithContext {
        StackFrameWithContext {
            frame: &self.root,
            strings: &self.strings,
        }
    }

    /// Merge aggregated heap profile from multiple sources (e.g. from several runs).
    pub fn merge<'a>(
        profiles: impl IntoIterator<Item = &'a AggregateHeapProfileInfo>,
    ) -> AggregateHeapProfileInfo {
        let profiles: Vec<_> = Vec::from_iter(profiles);

        let mut strings = StringIndex::default();
        let unused_capacity =
            UnusedCapacity::new(profiles.iter().map(|p| p.unused_capacity.get()).sum());
        let roots = profiles.into_iter().map(|p| p.root());
        let root = StackFrame::merge(roots, &mut strings);
        AggregateHeapProfileInfo {
            strings,
            root,
            unused_capacity,
        }
    }

    /// Write this out recursively to a file.
    pub fn gen_flame_graph(&self) -> String {
        let mut data = FlameGraphData::default();
        self.root().write_flame_graph(data.root());
        data.root()
            .child(ArcStr::new_static("unused_capacity"))
            .add(self.unused_capacity.get() as u64);
        data.write()
    }

    /// Write per-function summary in CSV format.
    pub fn gen_summary_csv(&self) -> String {
        HeapSummaryByFunction::init(self).gen_csv()
    }
}

#[derive(Debug, Allocative)]
pub(crate) struct RetainedHeapProfile {
    pub(crate) info: AggregateHeapProfileInfo,
    pub(crate) mode: RetainedHeapProfileMode,
}

impl RetainedHeapProfile {
    pub(crate) fn to_profile(&self) -> ProfileData {
        ProfileData {
            profile: ProfileDataImpl::AggregateHeapProfileInfo(Box::new(self.info.clone())),
            profile_mode: self.mode.to_profile_mode(),
        }
    }
}

#[cfg(test)]
mod tests {
    use dupe::Dupe;

    use crate::const_frozen_string;
    use crate::values::layout::heap::heap_type::HeapKind;
    use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
    use crate::values::layout::heap::profile::aggregated::StackFrame;
    use crate::values::layout::heap::profile::summary_by_function::HeapSummaryByFunction;
    use crate::values::Freezer;
    use crate::values::FrozenHeap;
    use crate::values::Heap;

    fn total_alloc_count(frame: &StackFrame) -> usize {
        frame.allocs.total().count
            + frame
                .callees
                .values()
                .map(|c| total_alloc_count(c.dupe()))
                .sum::<usize>()
    }

    #[test]
    fn test_stacks_collect() {
        let heap = Heap::new();
        heap.record_call_enter(const_frozen_string!("enter").to_value());
        heap.alloc_str("xxyy");
        heap.alloc_str("zzww");
        heap.record_call_exit();

        let stacks = AggregateHeapProfileInfo::collect(&heap, None);
        assert!(stacks.root.allocs.summary.is_empty());
        assert_eq!(1, stacks.root.callees.len());
        assert_eq!(2, total_alloc_count(&stacks.root));
    }

    #[test]
    fn test_stacks_collect_retained() {
        let heap = Heap::new();
        heap.record_call_enter(const_frozen_string!("enter").to_value());
        let s0 = heap.alloc_str("xxyy");
        let s1 = heap.alloc_str("zzww");
        heap.alloc_str("rrtt");
        heap.record_call_exit();

        let freezer = Freezer::new(FrozenHeap::new());
        freezer.freeze(s0.to_value()).unwrap();
        freezer.freeze(s1.to_value()).unwrap();

        let stacks = AggregateHeapProfileInfo::collect(&heap, Some(HeapKind::Frozen));
        assert!(stacks.root.allocs.summary.is_empty());
        assert_eq!(1, stacks.root.callees.len());
        // 3 allocated, 2 retained.
        assert_eq!(
            2,
            stacks
                .root
                .callees
                .values()
                .next()
                .unwrap()
                .allocs
                .summary
                .get("string")
                .unwrap()
                .count
        );
        assert_eq!(2, total_alloc_count(&stacks.root));
    }

    #[test]
    fn test_merge() {
        fn make() -> AggregateHeapProfileInfo {
            let heap = Heap::new();
            heap.record_call_enter(const_frozen_string!("xx").to_value());
            let s = heap.alloc_str("abc");
            heap.record_call_exit();
            let freezer = Freezer::new(FrozenHeap::new());
            freezer.freeze(s.to_value()).unwrap();

            AggregateHeapProfileInfo::collect(&heap, Some(HeapKind::Frozen))
        }

        let merge = AggregateHeapProfileInfo::merge([&make(), &make(), &make()]);
        let summary = HeapSummaryByFunction::init(&merge);
        assert_eq!(1, summary.info().len());
        let (xx_id, xx_info) = summary.info()[0];
        assert_eq!("xx", &**xx_id);
        assert_eq!(3, xx_info.alloc.get("string").unwrap().count);
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::iter::Sum;
use std::ops::Add;
use std::ops::AddAssign;

use allocative::Allocative;
use dupe::Dupe;

/// Allocations counters.
#[derive(Default, Copy, Clone, Dupe, Debug, Allocative)]
pub(crate) struct AllocCounts {
    pub(crate) bytes: usize,
    pub(crate) count: usize,
}

impl AddAssign for AllocCounts {
    fn add_assign(&mut self, other: AllocCounts) {
        self.bytes += other.bytes;
        self.count += other.count;
    }
}

impl Add for AllocCounts {
    type Output = AllocCounts;

    fn add(self, other: AllocCounts) -> AllocCounts {
        AllocCounts {
            bytes: self.bytes + other.bytes,
            count: self.count + other.count,
        }
    }
}

impl<'a> Sum<&'a AllocCounts> for AllocCounts {
    fn sum<I>(iter: I) -> AllocCounts
    where
        I: Iterator<Item = &'a AllocCounts>,
    {
        iter.fold(AllocCounts::default(), |acc, x| acc + *x)
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 at
 *
 *     https://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.
 */

use std::borrow::Borrow;
use std::hash::Hash;
use std::hash::Hasher;
use std::ops::Deref;
use std::sync::Arc;

use allocative::Allocative;
use dupe::Dupe;

#[derive(Clone, Dupe, Debug, Allocative)]
enum Inner {
    Arc(Arc<str>),
    Static(&'static str),
}

/// Wrapper for `Arc<str>`.
#[derive(Clone, Dupe, Debug, derive_more::Display, Allocative)]
#[display(fmt = "{}", "&**self")]
pub(crate) struct ArcStr(Inner);

impl ArcStr {
    pub(crate) fn new_static(s: &'static str) -> ArcStr {
        ArcStr(Inner::Static(s))
    }

    pub(crate) fn as_str(&self) -> &str {
        match &self.0 {
            Inner::Arc(s) => s,
            Inner::Static(s) => s,
        }
    }
}

impl PartialEq for ArcStr {
    fn eq(&self, other: &Self) -> bool {
        self.as_str() == other.as_str()
    }
}

impl Eq for ArcStr {}

impl Hash for ArcStr {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.as_str().hash(state)
    }
}

impl Deref for ArcStr {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl Borrow<str> for ArcStr {
    fn borrow(&self) -> &str {
        self
    }
}

impl<'a> From<&'a str> for ArcStr {
    fn from(s: &'a str) -> Self {
        if s.is_empty() {
            ArcStr(Inner::Static(""))
        } else {
            ArcStr(Inner::Arc(Arc::from(s)))
        }
    }
}