Unblock graph vis

src/borrow_pcg/unblock_graph.rs

@@ -115,6 +115,10 @@ impl<'tcx> UnblockGraph<'tcx> {
         self.edges.is_empty()
     }
 
+    pub(crate) fn edges(&self) -> impl Iterator<Item = &UnblockEdge<'tcx>> {
+        self.edges.iter()
+    }
+
     fn add_dependency(&mut self, unblock_edge: UnblockEdge<'tcx>) -> bool {
         self.edges.insert(unblock_edge)
     }

src/lib.rs

@@ -56,8 +56,8 @@ pub mod utils;
 pub mod visualization;
 
 use borrow_checker::BorrowCheckerInterface;
-use borrow_pcg::graph::borrows_imgcat_debug;
-use pcg::{CapabilityKind, PcgEngine};
+use borrow_pcg::graph::{borrows_imgcat_debug, loop_abstraction::MaybeRemoteCurrentPlace};
+use pcg::{CapabilityKind, EvalStmtPhase, PcgEngine};
 use rustc_interface::{
     borrowck::{self, BorrowSet, LocationTable, PoloniusInput, RegionInferenceContext},
     dataflow::{AnalysisEngine, compute_fixpoint},
@@ -269,6 +269,8 @@ use utils::eval_stmt_data::EvalStmtData;
 struct PcgStmtVisualizationData {
     actions: EvalStmtData<Vec<AppliedActionDebugRepr>>,
     graphs: visualization::stmt_graphs::StmtGraphs,
+    return_unblock_graph: Option<String>,
+    return_unblock_actions: Option<Vec<String>>,
 }
 
 #[derive(Serialize)]
@@ -546,6 +548,49 @@ pub fn run_pcg<'a, 'tcx>(pcg_ctxt: &'a PcgCtxt<'_, 'tcx>) -> PcgOutput<'a, 'tcx>
                 (None, Vec::new())
             };
 
+            let terminator_stmt_index = body[block].statements.len();
+            let is_return = matches!(&body[block].terminator().kind, mir::TerminatorKind::Return);
+            let return_unblock = if is_return {
+                let filename = format!("{block:?}_remote_lifetime_projection_unblock.dot");
+                let post_main = &pcg_block.terminator.states[EvalStmtPhase::PostMain];
+                let remote_lifetime_projections = post_main
+                    .borrow
+                    .graph()
+                    .nodes(pcg_ctxt.compiler_ctxt)
+                    .into_iter()
+                    .filter(|node| {
+                        node.try_into_lifetime_projection().is_ok_and(|rp| {
+                            rp.base()
+                                .maybe_remote_current_place()
+                                .is_some_and(MaybeRemoteCurrentPlace::is_remote)
+                        })
+                    })
+                    .collect::<Vec<_>>();
+                let unblock_graph = borrow_pcg::unblock_graph::UnblockGraph::for_nodes(
+                    remote_lifetime_projections,
+                    &post_main.borrow,
+                    pcg_ctxt.compiler_ctxt,
+                );
+                let unblock_actions = unblock_graph
+                    .clone()
+                    .actions(pcg_ctxt.compiler_ctxt)
+                    .expect("Failed to generate return unblock actions")
+                    .into_iter()
+                    .map(|action| action.edge().display_string(pcg_ctxt.compiler_ctxt))
+                    .collect::<Vec<_>>();
+                let dot_graph = visualization::generate_unblock_dot_graph(
+                    pcg_ctxt.compiler_ctxt,
+                    &post_main.place_capabilities,
+                    &unblock_graph,
+                )
+                .expect("Failed to generate return unblock graph");
+                std::fs::write(dir_path.join(&filename), dot_graph)
+                    .expect("Failed to write return unblock graph");
+                Some((filename, unblock_actions))
+            } else {
+                None
+            };
+
             let statements = pcg_block
                 .statements()
                 .map(|stmt| {
@@ -557,6 +602,22 @@ pub fn run_pcg<'a, 'tcx>(pcg_ctxt: &'a PcgCtxt<'_, 'tcx>) -> PcgOutput<'a, 'tcx>
                             .get(stmt.location.statement_index)
                             .cloned()
                             .unwrap_or_default(),
+                        return_unblock_graph: if stmt.location.statement_index
+                            == terminator_stmt_index
+                        {
+                            return_unblock
+                                .as_ref()
+                                .map(|(filename, _)| filename.clone())
+                        } else {
+                            None
+                        },
+                        return_unblock_actions: if stmt.location.statement_index
+                            == terminator_stmt_index
+                        {
+                            return_unblock.as_ref().map(|(_, actions)| actions.clone())
+                        } else {
+                            None
+                        },
                     }
                 })
                 .collect();

src/results/mod.rs

@@ -127,6 +127,7 @@ impl<'a, 'tcx: 'a> PcgAnalysisResults<'a, 'tcx> {
         let state = self.cursor.get().expect_results_or_error()?;
         let from_pcg = &state.data.pcg;
         let from_post_main = from_pcg.states[EvalStmtPhase::PostMain].clone();
+        let states = state.data.pcg.states.to_owned();
         let self_abstraction_edges = from_post_main
             .borrow
             .graph()
@@ -194,7 +195,7 @@ impl<'a, 'tcx: 'a> PcgAnalysisResults<'a, 'tcx> {
                 ))
             })
             .collect::<Result<Vec<_>, PcgError>>()?;
-        Ok(PcgTerminator { succs })
+        Ok(PcgTerminator { succs, states })
     }
 
     /// Obtains the results of the dataflow analysis for all blocks.
@@ -455,4 +456,5 @@ impl<'a, 'tcx: 'a> PcgLocation<'a, 'tcx> {
 #[derive(Debug)]
 pub struct PcgTerminator<'a, 'tcx> {
     pub succs: Vec<PcgSuccessor<'a, 'tcx>>,
+    pub(crate) states: DomainDataStates<Pcg<'a, 'tcx>>,
 }

src/visualization/mod.rs

@@ -25,7 +25,7 @@ pub use mir_graph::SourcePos;
 
 use crate::{
     borrow_pcg::{
-        edge::borrow_flow::BorrowFlowEdgeKind, graph::BorrowsGraph,
+        edge::borrow_flow::BorrowFlowEdgeKind, graph::BorrowsGraph, unblock_graph::UnblockGraph,
         validity_conditions::ValidityConditions,
     },
     pcg::{CapabilityKind, PcgRef, place_capabilities::PlaceCapabilitiesReader},
@@ -425,6 +425,23 @@ pub(crate) fn generate_borrows_dot_graph<'a, 'tcx: 'a>(
     Ok(String::from_utf8(buf).unwrap())
 }
 
+pub(crate) fn generate_unblock_dot_graph<'a, 'tcx: 'a>(
+    ctxt: CompilerCtxt<'a, 'tcx>,
+    capabilities: &impl PlaceCapabilitiesReader<'tcx>,
+    unblock_graph: &UnblockGraph<'tcx>,
+) -> io::Result<String> {
+    let mut borrows_graph = BorrowsGraph::default();
+    for edge in unblock_graph.edges() {
+        borrows_graph.insert(edge.clone(), ctxt.bc_ctxt());
+    }
+    let constructor = BorrowsGraphConstructor::new(&borrows_graph, capabilities, ctxt.bc_ctxt());
+    let graph = constructor.construct_graph();
+    let mut buf = vec![];
+    let drawer = GraphDrawer::new(&mut buf, None);
+    drawer.draw(&graph, ctxt)?;
+    Ok(String::from_utf8(buf).unwrap())
+}
+
 pub(crate) fn generate_pcg_dot_graph<'a, 'tcx: 'a>(
     pcg: PcgRef<'a, 'tcx>,
     ctxt: impl HasBorrowCheckerCtxt<'a, 'tcx>,

test-files/227_axel.rs

@@ -0,0 +1,9 @@
+fn fn_arg_reborrow(x: &i32) -> &i32 {
+  let y = &(*x);
+  y                       // MIR: actually `result = &(*y); return`
+}
+
+fn fn_arg_copy(x: &i32) -> &i32 {
+  let y = x; // MIR: y = copy x
+  y               // MIR: actually `result = &(*y); return`
+}

visualization/src/components/PCGNavigator.tsx

@@ -42,6 +42,17 @@ export const NAVIGATOR_MIN_WIDTH_NUM = 40;
 export const NAVIGATOR_MAX_WIDTH = "200px";
 export const NAVIGATOR_MIN_WIDTH = "40px";
 
+const navigatorActionButtonStyle: React.CSSProperties = {
+  margin: "10px",
+  padding: "8px",
+  cursor: "pointer",
+  backgroundColor: "#007acc",
+  color: "white",
+  border: "none",
+  borderRadius: "4px",
+  fontSize: "12px",
+};
+
 function DotGraphButton({
   label,
   dotContent,
@@ -257,6 +268,36 @@ const getPCGDotGraphFilename = (
   return `data/${selectedFunction}/${filename}`;
 };
 
+const getReturnUnblockGraphFilename = (
+  currentPoint: CurrentPoint,
+  selectedFunction: string,
+  graphs: PcgBlockVisualizationData
+): string | null => {
+  if (
+    currentPoint.type !== "stmt" ||
+    graphs.statements.length <= currentPoint.stmt
+  ) {
+    return null;
+  }
+
+  const filename = graphs.statements[currentPoint.stmt].return_unblock_graph;
+  return filename ? `data/${selectedFunction}/${filename}` : null;
+};
+
+const getReturnUnblockActions = (
+  currentPoint: CurrentPoint,
+  graphs: PcgBlockVisualizationData
+): string[] => {
+  if (
+    currentPoint.type !== "stmt" ||
+    graphs.statements.length <= currentPoint.stmt
+  ) {
+    return [];
+  }
+
+  return graphs.statements[currentPoint.stmt].return_unblock_actions ?? [];
+};
+
 const formatCurrentPointTitle = (currentPoint: CurrentPoint): string => {
   if (currentPoint.type === "stmt") {
     const navPointName =
@@ -345,6 +386,12 @@ export default function PCGNavigator({
   };
 
   const navigationItems = buildNavigationItems();
+  const returnUnblockGraphFilename = getReturnUnblockGraphFilename(
+    currentPoint,
+    selectedFunction,
+    pcgData
+  );
+  const returnUnblockActions = getReturnUnblockActions(currentPoint, pcgData);
 
   // Resize handlers
   const handleResizeStart = (event: React.MouseEvent) => {
@@ -626,16 +673,7 @@ export default function PCGNavigator({
             {renderItems()}
           </div>
           <button
-            style={{
-              margin: "10px",
-              padding: "8px",
-              cursor: "pointer",
-              backgroundColor: "#007acc",
-              color: "white",
-              border: "none",
-              borderRadius: "4px",
-              fontSize: "12px",
-            }}
+            style={navigatorActionButtonStyle}
             onClick={async () => {
               const dotFilePath = getPCGDotGraphFilename(
                 currentPoint,
@@ -650,6 +688,24 @@ export default function PCGNavigator({
           >
             Open Current PCG in New Window
           </button>
+          {returnUnblockGraphFilename ? (
+            <button
+              style={navigatorActionButtonStyle}
+              onClick={() => {
+                openDotGraphInNewWindow(
+                  api,
+                  returnUnblockGraphFilename,
+                  "Remote Lifetime Projection Unblock Graph",
+                  {
+                    title: "Unblock Actions",
+                    items: returnUnblockActions,
+                  }
+                );
+              }}
+            >
+              Open Remote Lifetime Projection Unblock Graph in New Window
+            </button>
+          ) : null}
           <div
             style={{
               padding: "10px",

visualization/src/dot_graph.ts

@@ -1,7 +1,12 @@
 import * as Viz from "@viz-js/viz";
 import { Api } from "./api";
 
-function renderDotInPopup(dotData: string, title: string) {
+export type DotGraphSidebar = {
+    title: string;
+    items: string[];
+};
+
+function renderDotInPopup(dotData: string, title: string, sidebar?: DotGraphSidebar) {
     Viz.instance().then((viz) => {
       const svgElement = viz.renderSVGElement(dotData);
       const popup = window.open(
@@ -16,21 +21,93 @@ function renderDotInPopup(dotData: string, title: string) {
       popup.document.head.innerHTML = `
         <title>${title}</title>
         <style>
-          body { margin: 0; }
-          svg {
+          body { margin: 0; font-family: sans-serif; }
+          .dot-popup {
+            display: flex;
             width: 100vw;
             height: 100vh;
+            overflow: hidden;
+          }
+          .dot-sidebar {
+            box-sizing: border-box;
+            width: 360px;
+            flex: 0 0 360px;
+            padding: 16px;
+            border-right: 1px solid #ddd;
+            overflow: auto;
+            background: #f8f8f8;
+          }
+          .dot-sidebar h2 {
+            margin: 0 0 12px;
+            font-size: 16px;
+          }
+          .dot-sidebar ol {
+            margin: 0;
+            padding-left: 24px;
+          }
+          .dot-sidebar li {
+            margin-bottom: 8px;
+            font-family: monospace;
+            white-space: pre-wrap;
+            overflow-wrap: anywhere;
+          }
+          .dot-graph {
+            flex: 1;
+            min-width: 0;
+            overflow: hidden;
+          }
+          svg {
+            width: 100%;
+            height: 100%;
             display: block;
           }
         </style>
       `;
-      popup.document.body.appendChild(svgElement);
+      const container = document.createElement("div");
+      container.className = "dot-popup";
+
+      if (sidebar) {
+        const sidebarElement = document.createElement("div");
+        sidebarElement.className = "dot-sidebar";
+
+        const heading = document.createElement("h2");
+        heading.textContent = sidebar.title;
+        sidebarElement.appendChild(heading);
+
+        if (sidebar.items.length === 0) {
+          const empty = document.createElement("p");
+          empty.textContent = "No actions.";
+          sidebarElement.appendChild(empty);
+        } else {
+          const list = document.createElement("ol");
+          sidebar.items.forEach((item) => {
+            const listItem = document.createElement("li");
+            listItem.textContent = item;
+            list.appendChild(listItem);
+          });
+          sidebarElement.appendChild(list);
+        }
+
+        container.appendChild(sidebarElement);
+      }
+
+      const graphContainer = document.createElement("div");
+      graphContainer.className = "dot-graph";
+      graphContainer.appendChild(svgElement);
+      container.appendChild(graphContainer);
+
+      popup.document.body.appendChild(container);
     });
 }
 
-export async function openDotGraphInNewWindow(api: Api, filename: string, title?: string) {
+export async function openDotGraphInNewWindow(
+    api: Api,
+    filename: string,
+    title?: string,
+    sidebar?: DotGraphSidebar
+) {
     const dotData = await api.fetchDotFile(filename);
-    renderDotInPopup(dotData, title || `Dot Graph - ${filename}`);
+    renderDotInPopup(dotData, title || `Dot Graph - ${filename}`, sidebar);
 }
 
 export function openDotStringInNewWindow(dotData: string, title?: string) {

visualization/src/generated_types/PcgStmtVisualizationData.ts

@@ -3,4 +3,4 @@ import type { AppliedAction } from "./AppliedAction";
 import type { EvalStmtData } from "./EvalStmtData";
 import type { StmtGraphs } from "./StmtGraphs";
 
-export type PcgStmtVisualizationData = { actions: EvalStmtData<Array<AppliedAction>>, graphs: StmtGraphs, };
+export type PcgStmtVisualizationData = { actions: EvalStmtData<Array<AppliedAction>>, graphs: StmtGraphs, return_unblock_graph: string | null, return_unblock_actions: Array<string> | null, };