WIP

Author: zgrannan

Cargo.lock

@@ -5,6 +5,14 @@ authors = ["Prusti Devs <[email protected]>"]
 edition = "2021"
 build = "build.rs"
 
+[features]
+default = ["memory_profiling"]
+memory_profiling = [
+    "tikv-jemallocator",
+    "axum",
+    "tokio",
+    "jemalloc_pprof",
+]
 
 [dependencies]
 # For RustRover
@@ -34,6 +42,10 @@ shell-escape = "0.1.5"
 smallvec = { version = "^1.11", features = ["union", "const_new"] }
 tracing = "0.1"
 tracing-subscriber = "0.3"
+tikv-jemallocator = { version = "0.6", features = ["profiling", "unprefixed_malloc_on_supported_platforms"], optional = true }
+axum = { version = "0.7", features = ["multipart"], optional = true }
+tokio = { version = "1.0", features = ["full", "rt-multi-thread"], optional = true }
+jemalloc_pprof = { version = "0.6", optional = true }
 
 [dev-dependencies]
 reqwest = { version = "^0.11", features = ["blocking"] }

Cargo.toml

@@ -1,5 +1,4 @@
-# Build stage for Rust projects
-FROM rust:1.75 as rust-builder
+FROM rust:1.75 as rust-deps
 
 # Install required dependencies
 RUN apt-get update && apt-get install -y \
@@ -11,7 +10,27 @@ WORKDIR /usr/src/app
 # Copy Rust project files
 COPY . .
 
-# Build PCS
+FROM rust-deps as dev-profile
+
+# Install extra useful stuff
+RUN apt-get update && apt-get install -y \
+    golang-go tmux vim graphviz \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install latest pprof
+RUN go install github.com/google/pprof@latest
+
+# Build a debug version
+RUN cargo build
+
+# Set up Rust environment variables
+ENV PATH="/usr/local/cargo/bin:${PATH}"
+ENV RUSTUP_HOME="/usr/local/rustup"
+ENV CARGO_HOME="/usr/local/cargo"
+
+FROM rust-deps as rust-builder
+
+# Build PCG binary
 RUN cargo build --release
 
 # Build pcg-server
@@ -66,4 +85,7 @@ ENV RUST_BACKTRACE=1
 # Expose port for pcg-server
 EXPOSE 4000
 
+# Expose port for memory profiling
+EXPOSE 4444
+
 CMD ["./pcg-server"]

Dockerfile

@@ -17,6 +17,7 @@ pkgs.mkShell {
     pkg-config
     rustNightly
     libiconv
+    zlib
   ] ++ (if pkgs.stdenv.isDarwin then [
     darwin.apple_sdk.frameworks.Security
     darwin.apple_sdk.frameworks.SystemConfiguration

shell.nix

@@ -214,11 +214,11 @@ pub struct AbstractionBlockEdge<'tcx> {
 }
 
 impl<'tcx> EdgeData<'tcx> for AbstractionBlockEdge<'tcx> {
-    fn blocked_nodes(&self, repacker: PlaceRepacker<'_, 'tcx>) -> FxHashSet<PCGNode<'tcx>> {
+    fn blocked_nodes(&self, _repacker: PlaceRepacker<'_, 'tcx>) -> FxHashSet<PCGNode<'tcx>> {
         self.inputs().into_iter().map(|i| i.into()).collect()
     }
 
-    fn blocked_by_nodes(&self, repacker: PlaceRepacker<'_, 'tcx>) -> FxHashSet<LocalNode<'tcx>> {
+    fn blocked_by_nodes(&self, _repacker: PlaceRepacker<'_, 'tcx>) -> FxHashSet<LocalNode<'tcx>> {
         self.outputs()
             .into_iter()
             .map(|o| o.to_local_node())

src/borrows/edge/abstraction.rs

@@ -89,15 +89,19 @@ impl DataflowStmtPhase {
     }
 }
 
+#[derive(Clone)]
+pub (crate) struct PCGDebugData {
+    pub(crate) dot_output_dir: String,
+    pub(crate) dot_graphs: Rc<RefCell<DotGraphs>>,
+}
+
 #[derive(Clone)]
 pub struct PlaceCapabilitySummary<'a, 'tcx> {
     cgx: Rc<PCGContext<'a, 'tcx>>,
     pub(crate) block: Option<BasicBlock>,
 
     pub(crate) pcg: PCG<'a, 'tcx>,
-    dot_graphs: Option<Rc<RefCell<DotGraphs>>>,
-
-    dot_output_dir: Option<String>,
+    debug_data: Option<PCGDebugData>,
 
     join_history: FxHashSet<BasicBlock>,
 }
@@ -262,16 +266,19 @@ impl<'a, 'tcx> PlaceCapabilitySummary<'a, 'tcx> {
         self.pcg.borrow.set_block(block);
     }
 
-    pub fn set_dot_graphs(&mut self, dot_graphs: Rc<RefCell<DotGraphs>>) {
-        self.dot_graphs = Some(dot_graphs);
+    pub fn set_debug_data(&mut self, output_dir: String, dot_graphs: Rc<RefCell<DotGraphs>>) {
+        self.debug_data = Some(PCGDebugData {
+            dot_output_dir: output_dir,
+            dot_graphs,
+        });
     }
 
     pub(crate) fn block(&self) -> BasicBlock {
         self.block.unwrap()
     }
 
-    pub fn dot_graphs(&self) -> Rc<RefCell<DotGraphs>> {
-        self.dot_graphs.clone().unwrap()
+    pub fn dot_graphs(&self) -> Option<Rc<RefCell<DotGraphs>>> {
+        self.debug_data.as_ref().map(|data| data.dot_graphs.clone())
     }
 
     fn dot_filename_for(
@@ -283,9 +290,11 @@ impl<'a, 'tcx> PlaceCapabilitySummary<'a, 'tcx> {
         format!(
             "{}/{}",
             output_dir,
-            self.dot_graphs()
-                .borrow()
-                .relative_filename(phase, self.block(), statement_index)
+            self.dot_graphs().unwrap().borrow().relative_filename(
+                phase,
+                self.block(),
+                statement_index
+            )
         )
     }
 
@@ -296,19 +305,22 @@ impl<'a, 'tcx> PlaceCapabilitySummary<'a, 'tcx> {
         if self.block().as_usize() == 0 {
             assert!(!matches!(phase, DataflowStmtPhase::Join(_)));
         }
-        if let Some(output_dir) = &self.dot_output_dir {
+        if let Some(debug_data) = &self.debug_data {
             if phase == DataflowStmtPhase::Initial {
-                self.dot_graphs()
+                debug_data
+                    .dot_graphs
                     .borrow_mut()
                     .register_new_iteration(statement_index);
             }
-            let relative_filename =
-                self.dot_graphs()
-                    .borrow()
-                    .relative_filename(phase, self.block(), statement_index);
-            let filename = self.dot_filename_for(&output_dir, phase, statement_index);
-            if !self
-                .dot_graphs()
+            let relative_filename = debug_data.dot_graphs.borrow().relative_filename(
+                phase,
+                self.block(),
+                statement_index,
+            );
+            let filename =
+                self.dot_filename_for(&debug_data.dot_output_dir, phase, statement_index);
+            if !debug_data
+                .dot_graphs
                 .borrow_mut()
                 .insert(statement_index, phase, relative_filename)
             {
@@ -335,8 +347,7 @@ impl<'a, 'tcx> PlaceCapabilitySummary<'a, 'tcx> {
     pub(crate) fn new(
         cgx: Rc<PCGContext<'a, 'tcx>>,
         block: Option<BasicBlock>,
-        dot_output_dir: Option<String>,
-        dot_graphs: Option<Rc<RefCell<DotGraphs>>>,
+        debug_data: Option<PCGDebugData>,
     ) -> Self {
         let fpcs = FreePlaceCapabilitySummary::new(cgx.rp);
         let borrows =
@@ -349,8 +360,7 @@ impl<'a, 'tcx> PlaceCapabilitySummary<'a, 'tcx> {
             cgx,
             block,
             pcg,
-            dot_graphs,
-            dot_output_dir,
+            debug_data,
             join_history: FxHashSet::default(),
         }
     }
@@ -438,8 +448,10 @@ impl JoinSemiLattice for PlaceCapabilitySummary<'_, '_> {
                 }
             }
         }
-        self.dot_graphs().borrow_mut().register_new_iteration(0);
-        self.generate_dot_graph(DataflowStmtPhase::Join(other.block()), 0);
+        if let Some(debug_data) = &self.debug_data {
+            debug_data.dot_graphs.borrow_mut().register_new_iteration(0);
+            self.generate_dot_graph(DataflowStmtPhase::Join(other.block()), 0);
+        }
         fpcs || borrows
     }
 }

src/combined_pcs/domain.rs

@@ -34,7 +34,9 @@ use crate::{
     utils::PlaceRepacker,
 };
 
-use super::{domain::PlaceCapabilitySummary, DataflowStmtPhase, DotGraphs, EvalStmtPhase};
+use super::{
+    domain::PlaceCapabilitySummary, DataflowStmtPhase, DotGraphs, EvalStmtPhase, PCGDebugData,
+};
 
 #[rustversion::since(2024-10-14)]
 type OutputFacts = Box<PoloniusOutput>;
@@ -137,36 +139,59 @@ impl<'mir, 'tcx> PCGContext<'mir, 'tcx> {
     }
 }
 
+struct PCGEngineDebugData {
+    debug_output_dir: String,
+    dot_graphs: IndexVec<BasicBlock, Rc<RefCell<DotGraphs>>>,
+}
+
 pub struct PCGEngine<'a, 'tcx> {
     pub(crate) cgx: Rc<PCGContext<'a, 'tcx>>,
     pub(crate) fpcs: FpcsEngine<'a, 'tcx>,
     pub(crate) borrows: BorrowsEngine<'a, 'tcx>,
-    debug_output_dir: Option<String>,
-    dot_graphs: IndexVec<BasicBlock, Rc<RefCell<DotGraphs>>>,
+    debug_data: Option<PCGEngineDebugData>,
     curr_block: Cell<BasicBlock>,
 }
 impl<'a, 'tcx> PCGEngine<'a, 'tcx> {
+    fn dot_graphs(&self, block: BasicBlock) -> Option<Rc<RefCell<DotGraphs>>> {
+        self.debug_data
+            .as_ref()
+            .map(|data| data.dot_graphs[block].clone())
+    }
+    fn debug_output_dir(&self) -> Option<String> {
+        self.debug_data
+            .as_ref()
+            .map(|data| data.debug_output_dir.clone())
+    }
     fn initialize(&self, state: &mut PlaceCapabilitySummary<'a, 'tcx>, block: BasicBlock) {
         if let Some(existing_block) = state.block {
             assert!(existing_block == block);
             return;
         }
         state.set_block(block);
-        state.set_dot_graphs(self.dot_graphs[block].clone());
+        if let Some(debug_data) = &self.debug_data {
+            state.set_debug_data(
+                debug_data.debug_output_dir.clone(),
+                debug_data.dot_graphs[block].clone(),
+            );
+        }
         assert!(state.is_initialized());
     }
 
     pub(crate) fn new(cgx: PCGContext<'a, 'tcx>, debug_output_dir: Option<String>) -> Self {
-        if let Some(dir_path) = &debug_output_dir {
-            if std::path::Path::new(dir_path).exists() {
-                std::fs::remove_dir_all(dir_path).expect("Failed to delete directory contents");
+        let debug_data = debug_output_dir.map(|dir_path| {
+            if std::path::Path::new(&dir_path).exists() {
+                std::fs::remove_dir_all(&dir_path).expect("Failed to delete directory contents");
             }
             create_dir_all(&dir_path).expect("Failed to create directory for DOT files");
-        }
-        let dot_graphs = IndexVec::from_fn_n(
-            |_| Rc::new(RefCell::new(DotGraphs::new())),
-            cgx.rp.body().basic_blocks.len(),
-        );
+            let dot_graphs = IndexVec::from_fn_n(
+                |_| Rc::new(RefCell::new(DotGraphs::new())),
+                cgx.rp.body().basic_blocks.len(),
+            );
+            PCGEngineDebugData {
+                debug_output_dir: dir_path.clone(),
+                dot_graphs,
+            }
+        });
         let fpcs = FpcsEngine(cgx.rp);
         let borrows = BorrowsEngine::new(
             cgx.rp.tcx(),
@@ -176,10 +201,9 @@ impl<'a, 'tcx> PCGEngine<'a, 'tcx> {
         );
         Self {
             cgx: Rc::new(cgx),
-            dot_graphs,
             fpcs,
             borrows,
-            debug_output_dir,
+            debug_data,
             curr_block: Cell::new(START_BLOCK),
         }
     }
@@ -199,20 +223,19 @@ impl<'a, 'tcx> Analysis<'tcx> for PCGEngine<'a, 'tcx> {
     const NAME: &'static str = "pcs";
 
     fn bottom_value(&self, body: &Body<'tcx>) -> Self::Domain {
-        let block = self.curr_block.get();
-        let (block, dot_graphs) = if block.as_usize() < body.basic_blocks.len() {
-            self.curr_block.set(block.plus(1));
-            (Some(block), Some(self.dot_graphs[block].clone()))
+        let curr_block = self.curr_block.get();
+        let (block, debug_data) = if curr_block.as_usize() < body.basic_blocks.len() {
+            self.curr_block.set(curr_block.plus(1));
+            let debug_data = self.debug_output_dir().map(|dir| PCGDebugData {
+                dot_output_dir: dir,
+                dot_graphs: self.dot_graphs(curr_block).unwrap(),
+            });
+            (Some(curr_block), debug_data)
         } else {
-            // For results cursor, don't set block
+            // For results cursor, don't set block or consider debug data
             (None, None)
         };
-        PlaceCapabilitySummary::new(
-            self.cgx.clone(),
-            block,
-            self.debug_output_dir.clone(),
-            dot_graphs,
-        )
+        PlaceCapabilitySummary::new(self.cgx.clone(), block, debug_data)
     }
 
     fn initialize_start_block(&self, _body: &Body<'tcx>, state: &mut Self::Domain) {

src/combined_pcs/engine.rs

@@ -320,6 +320,7 @@ pub fn run_combined_pcs<'mir, 'tcx>(
                 format!("{}/block_{}_iterations.json", dir_path, block.index());
             state
                 .dot_graphs()
+                .unwrap()
                 .borrow()
                 .write_json_file(&block_iterations_json_file);
         }