improve statement profiler accuracy

Summary:
Prior, we kept a `BcStatementLocation` that tracked which BC instruction was the first associated with a statement. This is used for a few things, including breakpoint evaluation, statement profiling, etc.

For statement profiling, this would lead to bytecode like this for a statement misreporting all of the time it takes to `Call Bar` as time spent in "whatever the last first-statement instruction was executed during the Call Foo->Return".
```
    A: Add
    B: Call Foo
    C: Call Bar
```

Previously, only `A` here was marked in BcStatementLocations. Now we mark `A` and `C` as belonging to the statement, too. This causes extra `before_stmt` calls that can then do the additional lookups/accounting in the profiler to credit the correct statement with the time to execute the instruction at `C`.

`BcStatementLocations` still differentiates between the "first" instruction for a statement, vs all the others (the `continued` ones), so that components like the debug adaptor can still "break on next statement", rather than "next statement-annotated instruction".

Reviewed By: cjhopman

Differential Revision: D73038348

fbshipit-source-id: 6804c86ffb2a082b8848b99850dd75e0d6ba1afa

Author: Jeremy Braun

starlark/src/assert/assert.rs

@@ -289,8 +289,10 @@ impl<'a> Assert<'a> {
         }
         let loader = ReturnFileLoader { modules: &modules };
         let ast = AstModule::parse(path, program.to_owned(), &self.dialect)?;
-        let gc_always = |_span: FileSpanRef, eval: &mut Evaluator| {
-            eval.trigger_gc();
+        let gc_always = |_span: FileSpanRef, continued: bool, eval: &mut Evaluator| {
+            if !continued {
+                eval.trigger_gc();
+            }
         };
         let mut eval = Evaluator::new(module);
         eval.enable_static_typechecking(self.static_typechecking);

starlark/src/debug/adapter/implementation.rs

@@ -111,8 +111,16 @@ impl<'a, 'e: 'a> BeforeStmtFuncDyn<'a, 'e> for DapAdapterEvalHookImpl {
     fn call<'v>(
         &mut self,
         span_loc: FileSpanRef,
+        continued: bool,
         eval: &mut Evaluator<'v, 'a, 'e>,
     ) -> crate::Result<()> {
+        // The debug adapter should only break on the "initial" instruction that
+        // makes up any given statement. "Continued" instructions are part of
+        // the still-executing/previous statement, and should be ignored.
+        if continued {
+            return Ok(());
+        }
+
         let stop = if self.state.disable_breakpoints.load(Ordering::SeqCst) > 0 {
             false
         } else {

starlark/src/eval/bc/instrs.rs

@@ -239,7 +239,7 @@ impl BcInstrs {
             let loop_pad = loop_ends.len() * 2;
 
             if newline {
-                if let Some(loc) = self.stmt_locs.stmt_at(ip) {
+                if let Some((loc, _continued)) = self.stmt_locs.stmt_at(ip) {
                     writeln!(f, "{:loop_pad$} {:ip_pad$}  # {}", "", "", loc.span.span)?;
                 }
             } else {

starlark/src/eval/bc/opcode.rs

@@ -148,6 +148,25 @@ impl BcOpcode {
         BcOpcode::do_dispatch_all(handler)
     }
 
+    #[inline(always)]
+    pub(crate) fn is_call(self) -> bool {
+        match self {
+            BcOpcode::Call
+            | BcOpcode::CallPos
+            | BcOpcode::CallFrozenDef
+            | BcOpcode::CallFrozenDefPos
+            | BcOpcode::CallFrozenNative
+            | BcOpcode::CallFrozenNativePos
+            | BcOpcode::CallFrozen
+            | BcOpcode::CallFrozenPos
+            | BcOpcode::CallMethod
+            | BcOpcode::CallMethodPos
+            | BcOpcode::CallMaybeKnownMethod
+            | BcOpcode::CallMaybeKnownMethodPos => true,
+            _ => false,
+        }
+    }
+
     /// Get opcode by opcode number.
     pub(crate) fn by_number(n: u32) -> Option<BcOpcode> {
         struct ByNumber {

starlark/src/eval/bc/writer.rs

@@ -40,6 +40,7 @@ use crate::eval::bc::instr_impl::InstrMov;
 use crate::eval::bc::instr_impl::InstrStoreLocalCaptured;
 use crate::eval::bc::instrs::BcInstrsWriter;
 use crate::eval::bc::instrs::PatchAddr;
+use crate::eval::bc::opcode::BcOpcode;
 use crate::eval::bc::repr::BC_INSTR_ALIGN;
 use crate::eval::bc::slow_arg::BcInstrSlowArg;
 use crate::eval::bc::stack_ptr::BcSlot;
@@ -73,6 +74,8 @@ pub(crate) struct BcStatementLocations {
 }
 
 impl BcStatementLocations {
+    const CONTINUED_BIT: u32 = 1 << 31;
+
     pub(crate) fn new() -> Self {
         Self {
             locs: Vec::new(),
@@ -89,6 +92,7 @@ impl BcStatementLocations {
     fn push(&mut self, addr: BcAddr, span: BcStmtLoc) {
         let idx = Self::idx_for(addr);
         let stmt_idx = self.locs.len().try_into().unwrap();
+        debug_assert_eq!(0, stmt_idx & Self::CONTINUED_BIT);
         self.locs.push(span);
         while self.stmts.len() <= idx {
             self.stmts.push(u32::MAX);
@@ -97,10 +101,39 @@ impl BcStatementLocations {
         self.stmts[idx] = stmt_idx;
     }
 
-    pub(crate) fn stmt_at(&self, offset: BcAddr) -> Option<&BcStmtLoc> {
+    fn last_stmt_idx(&self) -> Option<u32> {
+        for stmt_idx in (&self.stmts).iter().rev() {
+            if *stmt_idx != u32::MAX {
+                return Some(*stmt_idx & !Self::CONTINUED_BIT);
+            }
+        }
+        None
+    }
+
+    fn push_prev(&mut self, addr: BcAddr) {
+        if let Some(stmt_idx) = self.last_stmt_idx() {
+            let idx = Self::idx_for(addr);
+            while self.stmts.len() <= idx {
+                self.stmts.push(u32::MAX);
+            }
+            // If the preceding statement ended in a call opcode, and the
+            // current opcode is the start of a new statement, it may already be
+            // marked. If so, preserve that, rather than marking it as a part of
+            // the previous statement (which it is not).
+            if self.stmts[idx] == u32::MAX {
+                self.stmts[idx] = stmt_idx | Self::CONTINUED_BIT;
+            }
+        }
+    }
+
+    pub(crate) fn stmt_at(&self, offset: BcAddr) -> Option<(&BcStmtLoc, bool)> {
         match self.stmts.get(Self::idx_for(offset)) {
             None | Some(&u32::MAX) => None,
-            Some(v) => Some(&self.locs[*v as usize]),
+            Some(v) => {
+                let continued = 0 != (v & Self::CONTINUED_BIT);
+                let idx = v & !Self::CONTINUED_BIT;
+                Some((&self.locs[idx as usize], continued))
+            }
         }
     }
 }
@@ -123,8 +156,10 @@ pub(crate) struct BcWriter<'f> {
     instrs: BcInstrsWriter,
     /// Instruction spans, used for errors.
     slow_args: Vec<(BcAddr, BcInstrSlowArg)>,
-    /// For each statement, will store the span and the BcAddr for the first instruction.
+    /// For each statement, will store the span for the first instruction and any instruction after a call.
     stmt_locs: BcStatementLocations,
+    /// The last-written opcode
+    last_opcode: BcOpcode,
     /// Current stack size.
     stack_size: u32,
     /// Local slot count.
@@ -159,6 +194,7 @@ impl<'f> BcWriter<'f> {
             instrs: BcInstrsWriter::new(),
             slow_args: Vec::new(),
             stmt_locs: BcStatementLocations::new(),
+            last_opcode: BcOpcode::End,
             stack_size: 0,
             local_names,
             definitely_assigned,
@@ -176,6 +212,7 @@ impl<'f> BcWriter<'f> {
             instrs,
             slow_args: spans,
             stmt_locs,
+            last_opcode: _,
             stack_size,
             local_names,
             definitely_assigned,
@@ -219,6 +256,15 @@ impl<'f> BcWriter<'f> {
         slow_arg: BcInstrSlowArg,
         arg: I::Arg,
     ) -> (BcAddr, *const I::Arg) {
+        // If the previously written instruction was a form of a call
+        // instruction, instrument this instruction with the current statement
+        // span so that the time this and following instructions take can be
+        // attributed to the correct statement.
+        if self.last_opcode.is_call() {
+            self.stmt_locs.push_prev(self.ip());
+        }
+        self.last_opcode = BcOpcode::for_instr::<I>();
+
         self.slow_args.push((self.ip(), slow_arg));
         self.instrs.write::<I>(arg)
     }

starlark/src/eval/runtime/before_stmt.rs

@@ -35,22 +35,23 @@ pub(crate) struct BeforeStmt<'a, 'e: 'a> {
 // TODO(cjhopman): pull DAP into the crate, and hide this function.
 #[doc(hidden)]
 pub enum BeforeStmtFunc<'a, 'e: 'a> {
-    Fn(&'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a, 'e>)),
+    Fn(&'a dyn for<'v1> Fn(FileSpanRef, bool, &mut Evaluator<'v1, 'a, 'e>)),
     Dyn(Box<dyn BeforeStmtFuncDyn<'a, 'e>>),
 }
 
 impl<'a, 'e: 'a> BeforeStmtFunc<'a, 'e> {
     pub(crate) fn call<'v>(
         &mut self,
         span: FileSpanRef,
+        continued: bool,
         eval: &mut Evaluator<'v, 'a, 'e>,
     ) -> crate::Result<()> {
         match self {
             BeforeStmtFunc::Fn(f) => {
-                f(span, eval);
+                f(span, continued, eval);
                 Ok(())
             }
-            BeforeStmtFunc::Dyn(d) => d.call(span, eval),
+            BeforeStmtFunc::Dyn(d) => d.call(span, continued, eval),
         }
     }
 }
@@ -65,6 +66,7 @@ pub trait BeforeStmtFuncDyn<'a, 'e: 'a> {
     fn call<'v>(
         &mut self,
         span: FileSpanRef,
+        continued: bool,
         eval: &mut Evaluator<'v, 'a, 'e>,
     ) -> crate::Result<()>;
 }
@@ -75,10 +77,10 @@ impl<'a, 'e: 'a> BeforeStmt<'a, 'e> {
     }
 }
 
-impl<'a, 'e: 'a> From<&'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a, 'e>)>
+impl<'a, 'e: 'a> From<&'a dyn for<'v1> Fn(FileSpanRef, bool, &mut Evaluator<'v1, 'a, 'e>)>
     for BeforeStmtFunc<'a, 'e>
 {
-    fn from(value: &'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a, 'e>)) -> Self {
+    fn from(value: &'a dyn for<'v1> Fn(FileSpanRef, bool, &mut Evaluator<'v1, 'a, 'e>)) -> Self {
         Self::Fn(value)
     }
 }

starlark/src/eval/runtime/evaluator.rs

@@ -309,7 +309,7 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
             }
             ProfileMode::Statement | ProfileMode::Coverage => {
                 self.stmt_profile.enable();
-                self.before_stmt_fn(&|span, eval| eval.stmt_profile.before_stmt(span));
+                self.before_stmt_fn(&|span, _continued, eval| eval.stmt_profile.before_stmt(span));
             }
             ProfileMode::TimeFlame => {
                 self.time_flame_profile.enable();
@@ -422,7 +422,7 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
 
     pub(crate) fn before_stmt_fn(
         &mut self,
-        f: &'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a, 'e>),
+        f: &'a dyn for<'v1> Fn(FileSpanRef, bool, &mut Evaluator<'v1, 'a, 'e>),
     ) {
         self.before_stmt(f.into())
     }
@@ -891,8 +891,8 @@ pub(crate) struct EvalCallbacksEnabled<'a> {
 impl<'a> EvalCallbacksEnabled<'a> {
     fn before_stmt(&mut self, eval: &mut Evaluator, ip: BcPtrAddr) -> crate::Result<()> {
         let offset = ip.offset_from(self.bc_start_ptr);
-        if let Some(loc) = self.stmt_locs.stmt_at(offset) {
-            before_stmt(loc.span, eval)?;
+        if let Some((loc, continued)) = self.stmt_locs.stmt_at(offset) {
+            before_stmt(loc.span, continued, eval)?;
         }
         Ok(())
     }
@@ -916,11 +916,15 @@ impl<'a> EvaluationCallbacks for EvalCallbacksEnabled<'a> {
     }
 }
 
-// This function should be called before every meaningful statement.
+// This function should be called before every meaningful statement (continued==false), and after a call that returns into a previously entered statement (continued==true).
 // The purposes are GC, profiling and debugging.
 //
 // This function is called only if `before_stmt` is set before compilation start.
-pub(crate) fn before_stmt(span: FrameSpan, eval: &mut Evaluator) -> crate::Result<()> {
+pub(crate) fn before_stmt(
+    span: FrameSpan,
+    continued: bool,
+    eval: &mut Evaluator,
+) -> crate::Result<()> {
     assert!(
         eval.eval_instrumentation.before_stmt.enabled(),
         "this code should only be called if `before_stmt` is set"
@@ -931,7 +935,7 @@ pub(crate) fn before_stmt(span: FrameSpan, eval: &mut Evaluator) -> crate::Resul
     let mut result = Ok(());
     for f in &mut fs {
         if result.is_ok() {
-            result = f.call(span.span.file_span_ref(), eval);
+            result = f.call(span.span.file_span_ref(), continued, eval);
         }
     }
     let added = eval.eval_instrumentation.change(|eval_instrumentation| {

starlark/src/eval/runtime/profile/golden/statement.csv.golden

@@ -5,18 +5,18 @@
 # ```
 
 File,Span,Duration(s),Count
-"TOTAL","",2.128,304
-"test.star","7:13-19",1.400,200
-"test.star","3:5-9:1",0.140,20
+"TOTAL","",3.864,552
+"test.star","7:13-19",2.800,400
+"test.star","3:5-9:1",0.280,40
+"test.star","13:9-27",0.280,40
 "test.star","6:9-9:1",0.140,20
 "test.star","12:9-17",0.140,20
-"test.star","13:9-27",0.140,20
+"test.star","11:5-14:1",0.056,8
 "test.star","10:5-11",0.028,4
-"test.star","11:5-14:1",0.028,4
 "test.star","14:5-13",0.028,4
+"test.star","16:1-7",0.021,3
+"test.star","17:1-7",0.021,3
+"test.star","18:1-7",0.021,3
+"test.star","20:1-11",0.021,3
 "test.star","2:1-9:1",0.014,2
 "test.star","9:1-16:1",0.014,2
-"test.star","16:1-7",0.014,2
-"test.star","17:1-7",0.014,2
-"test.star","18:1-7",0.014,2
-"test.star","20:1-11",0.014,2

starlark/src/tests/bc/golden/call.golden

@@ -26,5 +26,6 @@ Instructions:
    72: Const 40 ->&6
    96: Const 50 ->&7
   120: CallFrozenNative noop {&3..&8 2 p q r *&0 **&1} instrs.star.bzl:2:5-10:6 ->&2
+        # instrs.star.bzl:2:5-10:6
   208: ReturnConst None
   224: End

starlark/src/tests/bc/golden/def_inline_const_args_inlined.golden

@@ -19,5 +19,6 @@ Instructions:
     0: Const True ->&1
    24: Const 10 ->&2
    48: CallFrozenNativePos noop &1..&3 instrs.star.bzl:2:12-22 ->&0
+        # instrs.star.bzl:6:5-25
   104: Return &0
   112: End