StringIndexMap

Summary:
Another abstraction, map `StringId -> V`.

Before this diff:

```
pub(crate) struct HeapSummaryByFunction {
    info: Vec<FuncInfo>,
}
```

Now field type is changed to `StringIndexMap<FuncInfo>`.

Also, `StringIndexMap` is sparse, which is easier to work with.

Reviewed By: bobyangyf

Differential Revision: D38518061

fbshipit-source-id: f0fb716945cbc99d47489e8784c0af06fa96fc30

Author: stepancheg

starlark/src/values/layout/heap/profile/string_index.rs

@@ -16,6 +16,7 @@
  */
 
 use gazebo::dupe::Dupe;
+use gazebo::prelude::*;
 use starlark_map::small_set::SmallSet;
 
 /// Map strings to integers 0, 1, 2, ...
@@ -45,3 +46,32 @@ impl StringIndex {
         self.strings.get_index(id.0).expect("invalid string id")
     }
 }
+
+#[derive(Debug, Default_)]
+pub(crate) struct StringIndexMap<V> {
+    values: Vec<Option<V>>,
+}
+
+impl<V> StringIndexMap<V> {
+    pub(crate) fn or_insert(&mut self, id: StringId) -> &mut V
+    where
+        V: Default,
+    {
+        // Resize.
+        while self.values.get(id.0).is_none() {
+            self.values.push(None);
+        }
+        self.values[id.0].get_or_insert_with(Default::default)
+    }
+
+    pub(crate) fn iter(&self) -> impl Iterator<Item = (StringId, &V)> {
+        self.values
+            .iter()
+            .enumerate()
+            .filter_map(|(id, value)| value.as_ref().map(|value| (StringId(id), value)))
+    }
+
+    pub(crate) fn values(&self) -> impl Iterator<Item = &V> {
+        self.values.iter().filter_map(Option::as_ref)
+    }
+}

starlark/src/values/layout/heap/profile/summary.rs

@@ -25,6 +25,7 @@ use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
 use crate::values::layout::heap::profile::aggregated::StackFrame;
 use crate::values::layout::heap::profile::alloc_counts::AllocCounts;
 use crate::values::layout::heap::profile::string_index::StringId;
+use crate::values::layout::heap::profile::string_index::StringIndexMap;
 
 /// Information relating to a function.
 #[derive(Default, Debug, Clone)]
@@ -65,19 +66,14 @@ impl FuncInfo {
 /// so pull out top_stack/top_info as a cache.
 pub(crate) struct HeapSummaryByFunction {
     /// Information about all functions. Map from `StringId`.
-    info: Vec<FuncInfo>,
+    info: StringIndexMap<FuncInfo>,
 }
 
 impl HeapSummaryByFunction {
-    fn ensure(&mut self, x: StringId) -> &mut FuncInfo {
-        if self.info.len() <= x.0 {
-            self.info.resize(x.0 + 1, FuncInfo::default());
-        }
-        &mut self.info[x.0]
-    }
-
     pub(crate) fn init(stacks: &AggregateHeapProfileInfo) -> HeapSummaryByFunction {
-        let mut info = HeapSummaryByFunction { info: Vec::new() };
+        let mut info = HeapSummaryByFunction {
+            info: StringIndexMap::default(),
+        };
         info.init_children(&stacks.root, stacks.root_id);
         info
     }
@@ -96,28 +92,24 @@ impl HeapSummaryByFunction {
         frame: &StackFrame,
         caller: StringId,
     ) -> SmallDuration {
-        self.ensure(func).time += frame.time_x2;
-        self.ensure(func).calls += frame.calls_x2 as usize;
-        *self.ensure(func).callers.entry(caller).or_insert(0) += 1;
+        self.info.or_insert(func).time += frame.time_x2;
+        self.info.or_insert(func).calls += frame.calls_x2 as usize;
+        *self.info.or_insert(func).callers.entry(caller).or_insert(0) += 1;
         for (t, allocs) in &frame.allocs.summary {
-            *self.ensure(func).alloc.entry(t).or_default() += *allocs;
+            *self.info.or_insert(func).alloc.entry(t).or_default() += *allocs;
         }
 
         let time_rec = frame.time_x2 + self.init_children(frame, func);
-        self.ensure(func).time_rec += time_rec;
+        self.info.or_insert(func).time_rec += time_rec;
         time_rec
     }
 
     fn totals(&self) -> FuncInfo {
-        FuncInfo::merge(self.info.iter())
+        FuncInfo::merge(self.info.values())
     }
 
     fn info(&self) -> Vec<(StringId, &FuncInfo)> {
-        self.info
-            .iter()
-            .enumerate()
-            .map(|(id, x)| (StringId(id), x))
-            .collect::<Vec<_>>()
+        self.info.iter().collect::<Vec<_>>()
     }
 
     pub(crate) fn gen_csv(&self, stacks: &AggregateHeapProfileInfo) -> String {
@@ -216,7 +208,7 @@ _ignore = str([1])     # allocate a string in non_drop
 
         let info = HeapSummaryByFunction::init(&stacks);
 
-        let total = FuncInfo::merge(info.info.iter());
+        let total = FuncInfo::merge(info.info.values());
         // from non-drop heap
         assert_eq!(total.alloc.get("string").unwrap().count, 1);
         // from drop heap