Introduce flame_guard in rustpython-stdlib too (#6722)

Cargo.lock

@@ -3160,6 +3160,7 @@ dependencies = [
  "digest",
  "dns-lookup",
  "dyn-clone",
+ "flame",
  "flate2",
  "foreign-types-shared",
  "gethostname",

Cargo.toml

@@ -15,7 +15,7 @@ importlib = ["rustpython-vm/importlib"]
 encodings = ["rustpython-vm/encodings"]
 stdio = ["rustpython-vm/stdio"]
 stdlib = ["rustpython-stdlib", "rustpython-pylib", "encodings"]
-flame-it = ["rustpython-vm/flame-it", "flame", "flamescope"]
+flame-it = ["rustpython-vm/flame-it", "rustpython-stdlib/flame-it", "flame", "flamescope"]
 freeze-stdlib = ["stdlib", "rustpython-vm/freeze-stdlib", "rustpython-pylib?/freeze-stdlib"]
 jit = ["rustpython-vm/jit"]
 threading = ["rustpython-vm/threading", "rustpython-stdlib/threading"]

crates/stdlib/Cargo.toml

@@ -22,6 +22,7 @@ ssl-rustls-fips = ["ssl-rustls", "aws-lc-rs/fips"]
 ssl-openssl = ["ssl", "openssl", "openssl-sys", "foreign-types-shared", "openssl-probe"]
 ssl-vendor = ["ssl-openssl", "openssl/vendored"]
 tkinter = ["dep:tk-sys", "dep:tcl-sys", "dep:widestring"]
+flame-it = ["flame"]
 
 [dependencies]
 # rustpython crates
@@ -33,6 +34,7 @@ ahash = { workspace = true }
 ascii = { workspace = true }
 cfg-if = { workspace = true }
 crossbeam-utils = { workspace = true }
+flame = { workspace = true, optional = true }
 hex = { workspace = true }
 itertools = { workspace = true }
 indexmap = { workspace = true }

crates/stdlib/src/json.rs

@@ -74,6 +74,7 @@ mod _json {
             scan_once: PyObjectRef,
             vm: &VirtualMachine,
         ) -> PyResult<PyIterReturn> {
+            flame_guard!("JsonScanner::parse");
             let c = match s.chars().next() {
                 Some(c) => c,
                 None => {
@@ -153,6 +154,7 @@ mod _json {
         }
 
         fn parse_number(&self, s: &str, vm: &VirtualMachine) -> Option<(PyResult, usize)> {
+            flame_guard!("JsonScanner::parse_number");
             let mut has_neg = false;
             let mut has_decimal = false;
             let mut has_exponent = false;
@@ -213,6 +215,7 @@ mod _json {
     }
 
     fn encode_string(s: &str, ascii_only: bool) -> String {
+        flame_guard!("_json::encode_string");
         let mut buf = Vec::<u8>::with_capacity(s.len() + 2);
         machinery::write_json_string(s, ascii_only, &mut buf)
             // SAFETY: writing to a vec can't fail
@@ -253,6 +256,7 @@ mod _json {
         strict: OptionalArg<bool>,
         vm: &VirtualMachine,
     ) -> PyResult<(Wtf8Buf, usize)> {
+        flame_guard!("_json::scanstring");
         machinery::scanstring(s.as_wtf8(), end, strict.unwrap_or(true))
             .map_err(|e| py_decode_error(e, s, vm))
     }

crates/stdlib/src/lib.rs

@@ -8,6 +8,9 @@
 extern crate rustpython_derive;
 extern crate alloc;
 
+#[macro_use]
+pub(crate) mod macros;
+
 pub mod array;
 mod binascii;
 mod bisect;

crates/stdlib/src/macros.rs

@@ -0,0 +1,7 @@
+#[macro_export]
+macro_rules! flame_guard {
+    ($name:expr) => {
+        #[cfg(feature = "flame-it")]
+        let _guard = ::flame::start_guard($name);
+    };
+}