Use allocative

Reviewed By: ndmitchell

Differential Revision: D41040212

fbshipit-source-id: d3b7748162bda0e630f7d47b2191f6ae0492f49c

Author: stepancheg

starlark/Cargo.toml

@@ -62,6 +62,10 @@ num-traits = "0.2"
 inventory = "0.1.10"
 clap = { version = "4.0.7", features = ["derive", "wrap_help"] }
 
+allocative.version = "0.2"
+# @oss-disable: allocative.path = "../../allocative/allocative"
+allocative.features = ["bumpalo"]
+
 [dev-dependencies]
 rand = { version = "0.8.4", features = ["small_rng"] }
 

starlark/src/environment/globals.rs

@@ -18,6 +18,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use allocative::Allocative;
 use derive_more::Display;
 use gazebo::prelude::*;
 use itertools::Itertools;
@@ -52,7 +53,7 @@ use crate::values::Heap;
 use crate::values::Value;
 
 /// The global values available during execution.
-#[derive(Clone, Dupe, Debug, Display)]
+#[derive(Clone, Dupe, Debug, Display, Allocative)]
 #[display(fmt = "globals")]
 pub struct Globals(Arc<GlobalsData>);
 
@@ -61,10 +62,12 @@ pub struct Globals(Arc<GlobalsData>);
 #[display(fmt = "methods")]
 pub struct Methods(Arc<MethodsData>);
 
-#[derive(Debug)]
+#[derive(Debug, Allocative)]
 struct GlobalsData {
     heap: FrozenHeapRef,
+    #[allocative(skip)] // TODO(nga): do not skip.
     variables: SymbolMap<FrozenValue>,
+    #[allocative(skip)]
     variable_names: Vec<FrozenStringValue>,
     docstring: Option<String>,
 }

starlark/src/environment/modules.rs

@@ -28,6 +28,7 @@ use std::sync::Arc;
 use std::time::Duration;
 use std::time::Instant;
 
+use allocative::Allocative;
 use derive_more::Display;
 use gazebo::any::ProvidesStaticType;
 use gazebo::prelude::*;
@@ -74,7 +75,7 @@ enum ModuleError {
 /// can be obtained using [`frozen_heap`](FrozenModule::frozen_heap). Be careful not to use
 /// these values after the [`FrozenModule`] has been released unless you obtain a reference
 /// to the frozen heap.
-#[derive(Debug, Clone, Dupe)]
+#[derive(Debug, Clone, Dupe, Allocative)]
 // We store the two elements separately since the FrozenHeapRef contains
 // a copy of the FrozenModuleData inside it.
 // Two Arc's should still be plenty cheap enough to qualify for `Dupe`.
@@ -89,7 +90,7 @@ pub struct FrozenModule {
     pub(crate) eval_duration: Duration,
 }
 
-#[derive(Debug, Clone, Dupe, ProvidesStaticType, Display)]
+#[derive(Debug, Clone, Dupe, ProvidesStaticType, Display, Allocative)]
 #[display(fmt = "{:?}", self)] // Type should not be user visible
 pub(crate) struct FrozenModuleRef(pub(crate) Arc<FrozenModuleData>);
 
@@ -107,7 +108,7 @@ impl FrozenModuleRef {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Allocative)]
 pub(crate) struct FrozenModuleData {
     pub(crate) names: FrozenNames,
     pub(crate) slots: FrozenSlots,

starlark/src/environment/names.rs

@@ -17,6 +17,8 @@
 
 use std::cell::RefCell;
 
+use allocative::Allocative;
+
 use crate::collections::Hashed;
 use crate::collections::SmallMap;
 use crate::environment::slots::ModuleSlotId;
@@ -43,8 +45,11 @@ use crate::values::FrozenStringValue;
 #[derive(Debug)]
 pub(crate) struct MutableNames(RefCell<SmallMap<FrozenStringValue, (ModuleSlotId, Visibility)>>);
 
-#[derive(Debug)]
-pub(crate) struct FrozenNames(SmallMap<FrozenStringValue, (ModuleSlotId, Visibility)>);
+#[derive(Debug, Allocative)]
+pub(crate) struct FrozenNames(
+    // TODO(nga): measure.
+    #[allocative(skip)] SmallMap<FrozenStringValue, (ModuleSlotId, Visibility)>,
+);
 
 impl MutableNames {
     pub fn new() -> Self {

starlark/src/environment/slots.rs

@@ -18,6 +18,7 @@
 use std::cell::RefCell;
 use std::cell::RefMut;
 
+use allocative::Allocative;
 use gazebo::prelude::*;
 
 use crate::values::Freezer;
@@ -38,7 +39,7 @@ impl ModuleSlotId {
 pub(crate) struct MutableSlots<'v>(RefCell<Vec<Option<Value<'v>>>>);
 
 // Indexed slots of a module. May contain unassigned values as `None`.
-#[derive(Debug)]
+#[derive(Debug, Allocative)]
 pub(crate) struct FrozenSlots(Vec<Option<FrozenValue>>);
 
 impl<'v> MutableSlots<'v> {

starlark/src/eval/runtime/profile/heap.rs

@@ -17,6 +17,7 @@
 
 use std::fmt::Debug;
 
+use allocative::Allocative;
 use gazebo::dupe::Dupe;
 
 use crate::eval::runtime::profile::data::ProfileData;
@@ -26,7 +27,7 @@ use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
 use crate::values::Heap;
 use crate::values::Value;
 
-#[derive(Copy, Clone, Dupe, Debug)]
+#[derive(Copy, Clone, Dupe, Debug, Allocative)]
 pub(crate) enum RetainedHeapProfileMode {
     Flame,
     Summary,

starlark/src/eval/runtime/profile/mod.rs

@@ -18,6 +18,7 @@
 use std::fmt::Display;
 use std::str::FromStr;
 
+use allocative::Allocative;
 use gazebo::dupe::Dupe;
 
 pub(crate) mod bc;
@@ -31,7 +32,7 @@ pub(crate) mod time_flame;
 pub(crate) mod typecheck;
 
 /// How to profile starlark code.
-#[derive(Debug, PartialEq, Eq, Hash, Clone, Dupe)]
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Dupe, Allocative)]
 #[non_exhaustive]
 pub enum ProfileMode {
     /// The heap profile mode provides information about the time spent in each function and allocations

starlark/src/eval/runtime/small_duration.rs

@@ -21,10 +21,11 @@ use std::ops::AddAssign;
 use std::ops::Div;
 use std::time::Duration;
 
+use allocative::Allocative;
 use gazebo::dupe::Dupe;
 
 /// Slightly faster than `Duration`.
-#[derive(Copy, Clone, Dupe, Default, Debug)]
+#[derive(Copy, Clone, Dupe, Default, Debug, Allocative)]
 pub(crate) struct SmallDuration {
     /// `u64::MAX` nanos is 500 years.
     pub(crate) nanos: u64,

starlark/src/values/layout/heap/arena.rs

@@ -36,6 +36,7 @@ use std::ptr;
 use std::slice;
 use std::time::Instant;
 
+use allocative::Allocative;
 use bumpalo::Bump;
 use either::Either;
 use gazebo::prelude::*;
@@ -74,7 +75,8 @@ pub(crate) const MIN_ALLOC: usize = {
     )
 };
 
-#[derive(Default)]
+#[derive(Default, Allocative)]
+// TODO(nga): implement `Allocative` for heap elements.
 pub(crate) struct Arena {
     /// Arena for things which don't need dropping (e.g. strings)
     non_drop: Bump,

starlark/src/values/layout/heap/heap_type.rs

@@ -36,6 +36,7 @@ use std::sync::Arc;
 use std::time::Instant;
 use std::usize;
 
+use allocative::Allocative;
 use either::Either;
 use gazebo::cast;
 use gazebo::prelude::*;
@@ -129,7 +130,7 @@ pub struct FrozenHeap {
 
 /// `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)]
+#[derive(Default, Allocative)]
 #[allow(clippy::non_send_fields_in_send_ty)]
 struct FrozenFrozenHeap {
     arena: Arena,
@@ -161,7 +162,7 @@ impl Debug for FrozenFrozenHeap {
 /// 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(Clone, Dupe, Debug)]
+#[derive(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(Arc<FrozenFrozenHeap>);