Wrap String in ProfileData

Summary: Types makes code easier to read.

Reviewed By: bobyangyf

Differential Revision: D38551784

fbshipit-source-id: 2b973ef99ec3041554ff8859f4572677d991400a

Author: stepancheg

starlark/src/eval/mod.rs

@@ -35,6 +35,7 @@ pub use runtime::file_loader::ReturnFileLoader;
 pub use runtime::params::ParametersParser;
 pub use runtime::params::ParametersSpec;
 pub use runtime::params::ParametersSpecBuilder;
+pub use runtime::profile::data::ProfileData;
 pub use runtime::profile::ProfileMode;
 
 use crate::collections::symbol_map::Symbol;

starlark/src/eval/runtime/evaluator.rs

@@ -17,12 +17,10 @@
 
 use std::cell::Cell;
 use std::collections::HashSet;
-use std::fs;
 use std::mem;
 use std::mem::MaybeUninit;
 use std::path::Path;
 
-use anyhow::Context;
 use gazebo::any::AnyLifetime;
 use gazebo::cast;
 use gazebo::dupe::Dupe;
@@ -47,6 +45,7 @@ use crate::eval::runtime::call_stack::CheapCallStack;
 use crate::eval::runtime::call_stack::FrozenFileSpan;
 use crate::eval::runtime::inlined_frame::InlinedFrames;
 use crate::eval::runtime::profile::bc::BcProfile;
+use crate::eval::runtime::profile::data::ProfileData;
 use crate::eval::runtime::profile::flame::FlameProfile;
 use crate::eval::runtime::profile::heap::HeapProfile;
 use crate::eval::runtime::profile::heap::HeapProfileFormat;
@@ -301,15 +300,12 @@ impl<'v, 'a> Evaluator<'v, 'a> {
     /// Write a profile to a file.
     /// Only valid if corresponding profiler was enabled.
     pub fn write_profile<P: AsRef<Path>>(&self, filename: P) -> anyhow::Result<()> {
-        let profile = self.gen_profile()?;
-        fs::write(filename.as_ref(), profile)
-            .with_context(|| format!("writing profile to `{}`", filename.as_ref().display()))?;
-        Ok(())
+        self.gen_profile()?.write(filename.as_ref())
     }
 
     /// Generate profile for a given mode.
     /// Only valid if corresponding profiler was enabled.
-    pub fn gen_profile(&self) -> anyhow::Result<String> {
+    pub fn gen_profile(&self) -> anyhow::Result<ProfileData> {
         let mode = match &self.profile_or_instrumentation_mode {
             ProfileOrInstrumentationMode::None => {
                 return Err(EvaluatorError::ProfilingNotEnabled.into());
@@ -335,7 +331,9 @@ impl<'v, 'a> Evaluator<'v, 'a> {
                 .stmt_profile
                 .gen()
                 .ok_or_else(|| EvaluatorError::StmtProfilingNotEnabled.into()),
-            ProfileMode::Bytecode | ProfileMode::BytecodePairs => self.bc_profile.gen_csv(),
+            ProfileMode::Bytecode | ProfileMode::BytecodePairs => {
+                self.bc_profile.gen_csv().map(ProfileData::new)
+            }
             ProfileMode::TimeFlame => self
                 .flame_profile
                 .gen()

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

@@ -0,0 +1,46 @@
+/*
+ * 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::fs;
+use std::path::Path;
+
+use anyhow::Context;
+
+/// Collected profiling data.
+#[derive(Clone, Debug)]
+pub struct ProfileData {
+    /// Serialized to text (e.g. CSV or flamegraph).
+    profile: String,
+}
+
+impl ProfileData {
+    pub(crate) fn new(profile: String) -> ProfileData {
+        ProfileData { profile }
+    }
+
+    /// Generate a string with profile data (e.g. CSV or flamegraph, depending on profile type).
+    pub fn gen(&self) -> String {
+        self.profile.clone()
+    }
+
+    /// Write to a file.
+    pub fn write(&self, path: &Path) -> anyhow::Result<()> {
+        fs::write(path, &self.profile)
+            .with_context(|| format!("write profile data to `{}`", path.display()))?;
+        Ok(())
+    }
+}

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

@@ -23,6 +23,7 @@ use std::time::Instant;
 use gazebo::prelude::*;
 
 use crate as starlark;
+use crate::eval::runtime::profile::data::ProfileData;
 use crate::eval::runtime::profile::flamegraph::FlameGraphWriter;
 use crate::eval::runtime::small_duration::SmallDuration;
 use crate::values::layout::pointer::RawPointer;
@@ -167,15 +168,15 @@ impl<'v> FlameProfile<'v> {
     }
 
     // We could expose profile on the Heap, but it's an implementation detail that it works here.
-    pub(crate) fn gen(&self) -> Option<String> {
+    pub(crate) fn gen(&self) -> Option<ProfileData> {
         self.0.as_ref().map(|box x| Self::gen_profile(x))
     }
 
-    fn gen_profile(x: &FlameData) -> String {
+    fn gen_profile(x: &FlameData) -> ProfileData {
         // Need to write out lines which look like:
         // root;calls1;calls2 1
         // All the numbers at the end must be whole numbers (we use milliseconds)
         let names = x.values.map(|x| x.to_repr());
-        Stacks::new(&names, &x.frames).render()
+        ProfileData::new(Stacks::new(&names, &x.frames).render())
     }
 }

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

@@ -19,6 +19,7 @@ use std::fmt::Debug;
 
 use gazebo::dupe::Dupe;
 
+use crate::eval::runtime::profile::data::ProfileData;
 use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
 use crate::values::Heap;
 use crate::values::Value;
@@ -59,29 +60,29 @@ impl HeapProfile {
     }
 
     // We could expose profile on the Heap, but it's an implementation detail that it works here.
-    pub(crate) fn gen(&self, heap: &Heap, format: HeapProfileFormat) -> Option<String> {
+    pub(crate) fn gen(&self, heap: &Heap, format: HeapProfileFormat) -> Option<ProfileData> {
         if !self.enabled {
             None
         } else {
             Some(Self::gen_enabled(heap, format))
         }
     }
 
-    pub(crate) fn gen_enabled(heap: &Heap, format: HeapProfileFormat) -> String {
+    pub(crate) fn gen_enabled(heap: &Heap, format: HeapProfileFormat) -> ProfileData {
         match format {
             HeapProfileFormat::Summary => Self::write_summarized_heap_profile(heap),
             HeapProfileFormat::FlameGraph => Self::write_flame_heap_profile(heap),
         }
     }
 
-    fn write_flame_heap_profile(heap: &Heap) -> String {
+    fn write_flame_heap_profile(heap: &Heap) -> ProfileData {
         let stacks = AggregateHeapProfileInfo::collect(heap, None);
-        stacks.gen_flame_graph()
+        ProfileData::new(stacks.gen_flame_graph())
     }
 
-    fn write_summarized_heap_profile(heap: &Heap) -> String {
+    fn write_summarized_heap_profile(heap: &Heap) -> ProfileData {
         let stacks = AggregateHeapProfileInfo::collect(heap, None);
-        stacks.gen_summary_csv()
+        ProfileData::new(stacks.gen_summary_csv())
     }
 }
 

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

@@ -22,6 +22,7 @@ use gazebo::dupe::Dupe;
 
 pub(crate) mod bc;
 pub(crate) mod csv;
+pub(crate) mod data;
 pub(crate) mod flame;
 pub(crate) mod flamegraph;
 pub(crate) mod heap;

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

@@ -30,6 +30,7 @@ use crate::codemap::FileSpanRef;
 use crate::codemap::ResolvedFileSpan;
 use crate::codemap::Span;
 use crate::eval::runtime::profile::csv::CsvWriter;
+use crate::eval::runtime::profile::data::ProfileData;
 use crate::eval::runtime::small_duration::SmallDuration;
 
 #[derive(Debug, thiserror::Error)]
@@ -181,9 +182,11 @@ impl StmtProfile {
     }
 
     // None = not applicable because not enabled
-    pub(crate) fn gen(&self) -> Option<String> {
+    pub(crate) fn gen(&self) -> Option<ProfileData> {
         let now = Instant::now();
-        self.0.as_ref().map(|data| data.write_to_string(now))
+        self.0
+            .as_ref()
+            .map(|data| ProfileData::new(data.write_to_string(now)))
     }
 
     pub(crate) fn coverage(&self) -> anyhow::Result<HashSet<ResolvedFileSpan>> {

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

@@ -21,6 +21,7 @@ use std::time::Duration;
 
 use crate::collections::SmallMap;
 use crate::eval::runtime::profile::csv::CsvWriter;
+use crate::eval::runtime::profile::data::ProfileData;
 use crate::eval::runtime::small_duration::SmallDuration;
 use crate::values::FrozenStringValue;
 
@@ -60,11 +61,11 @@ impl TypecheckProfile {
         w.finish()
     }
 
-    pub(crate) fn gen(&self) -> Option<String> {
+    pub(crate) fn gen(&self) -> Option<ProfileData> {
         if !self.enabled {
             return None;
         }
-        Some(self.gen_csv())
+        Some(ProfileData::new(self.gen_csv()))
     }
 }