Fix liveness computation for RP graph

src/borrow_pcg/borrow_checker/impl.rs

@@ -5,8 +5,7 @@ use crate::borrow_pcg::coupling_graph_constructor::BorrowCheckerInterface;
 use crate::borrow_pcg::region_projection::PCGRegion;
 use crate::combined_pcs::PCGNode;
 use crate::rustc_interface::borrowck::{
-    BorrowData, BorrowIndex, BorrowSet, LocationTable, RegionInferenceContext,
-    PoloniusOutput,
+    BorrowData, BorrowIndex, BorrowSet, LocationTable, PoloniusOutput, RegionInferenceContext,
 };
 use crate::rustc_interface::data_structures::fx::FxIndexMap;
 use crate::rustc_interface::dataflow::compute_fixpoint;
@@ -136,7 +135,7 @@ impl<'mir, 'tcx> BorrowCheckerInterface<'mir, 'tcx> for BorrowCheckerImpl<'mir,
                 if let Some(local) = rp.local() {
                     local
                 } else {
-                    todo!()
+                    return true; // e.g. from a constant or a remote place
                 }
             }
             PCGNode::Place(MaybeRemotePlace::Local(p)) => p.local(),

src/borrow_pcg/coupling_graph_constructor.rs

@@ -207,17 +207,17 @@ impl<T> DebugRecursiveCallHistory<T> {
     fn add(&mut self, _action: T) {}
 }
 
-pub(crate) struct CouplingGraphConstructor<'mir, 'tcx> {
+pub(crate) struct RegionProjectionAbstractionConstructor<'mir, 'tcx> {
     repacker: PlaceRepacker<'mir, 'tcx>,
     #[allow(unused)]
     block: BasicBlock,
-    coupling_graph: coupling::DisjointSetGraph<CGNode<'tcx>>,
+    graph: coupling::DisjointSetGraph<CGNode<'tcx>>,
 }
 
 #[derive(Clone, Eq, PartialEq)]
 struct AddEdgeHistory<'a, 'tcx> {
-    bottom_connect: &'a BTreeSet<CGNode<'tcx>>,
-    upper_candidate: &'a BTreeSet<CGNode<'tcx>>,
+    bottom_connect: &'a Coupled<CGNode<'tcx>>,
+    upper_candidate: &'a Coupled<CGNode<'tcx>>,
 }
 
 impl std::fmt::Display for AddEdgeHistory<'_, '_> {
@@ -239,77 +239,92 @@ impl std::fmt::Display for AddEdgeHistory<'_, '_> {
     }
 }
 
-impl<'mir, 'tcx> CouplingGraphConstructor<'mir, 'tcx> {
-    pub(crate) fn new(
-        repacker: PlaceRepacker<'mir, 'tcx>,
-        block: BasicBlock,
-    ) -> Self {
+impl<'mir, 'tcx> RegionProjectionAbstractionConstructor<'mir, 'tcx> {
+    pub(crate) fn new(repacker: PlaceRepacker<'mir, 'tcx>, block: BasicBlock) -> Self {
         Self {
             repacker,
             block,
-            coupling_graph: coupling::DisjointSetGraph::new(),
+            graph: coupling::DisjointSetGraph::new(),
         }
     }
 
     fn add_edges_from<'a>(
         &mut self,
         bg: &coupling::DisjointSetGraph<CGNode<'tcx>>,
-        bottom_connect: &'a BTreeSet<CGNode<'tcx>>,
-        upper_candidate: &'a BTreeSet<CGNode<'tcx>>,
+        bottom_connect: &'a Coupled<CGNode<'tcx>>,
+        upper_candidate: &'a Coupled<CGNode<'tcx>>,
+        borrow_checker: &dyn BorrowCheckerInterface<'mir, 'tcx>,
         mut history: DebugRecursiveCallHistory<AddEdgeHistory<'a, 'tcx>>,
     ) {
         history.add(AddEdgeHistory {
             bottom_connect,
             upper_candidate,
         });
-        let upper_candidate = upper_candidate.clone().into();
+        let upper_candidate = upper_candidate.clone();
         let endpoints = bg.endpoints_pointing_to(&upper_candidate);
         for coupled in endpoints {
             pcg_validity_assert!(
                 coupled != upper_candidate,
                 "Coupling graph should be acyclic"
             );
-            let should_include = coupled
-                .iter()
-                .any(|n| /* self.liveness.is_live(*n, self.block) && */ !n.is_old());
-            let coupled_set: BTreeSet<_> = coupled.clone().into_iter().collect();
+            let is_root = bg.is_root(&coupled) && !coupled.iter().any(|n| n.is_old());
+            let should_include = is_root
+                || coupled.iter().any(|n| {
+                    borrow_checker.is_live(
+                        (*n).into(),
+                        Location {
+                            block: self.block,
+                            statement_index: 0,
+                        },
+                    ) && !n.is_old()
+                });
+            tracing::info!(
+                "connect {} to {} ? (is_root: {:?}): {}",
+                bottom_connect.to_short_string(self.repacker),
+                coupled.to_short_string(self.repacker),
+                is_root,
+                should_include
+            );
             if !should_include {
-                self.add_edges_from(bg, bottom_connect, &coupled_set, history.clone());
-            } else {
-                self.coupling_graph.add_edge(
+                self.add_edges_from(
+                    bg,
+                    bottom_connect,
                     &coupled,
-                    &bottom_connect.clone().into(),
-                    self.repacker,
+                    borrow_checker,
+                    history.clone(),
                 );
-                self.add_edges_from(bg, &coupled_set, &coupled_set, history.clone());
+            } else {
+                self.graph
+                    .add_edge(&coupled, &bottom_connect.clone(), self.repacker);
+                self.add_edges_from(bg, &coupled, &coupled, borrow_checker, history.clone());
             }
         }
     }
 
     pub(crate) fn construct_region_projection_abstraction(
         mut self,
         bg: &BorrowsGraph<'tcx>,
+        borrow_checker: &dyn BorrowCheckerInterface<'mir, 'tcx>,
     ) -> coupling::DisjointSetGraph<CGNode<'tcx>> {
         tracing::debug!("Construct coupling graph start");
-        let full_graph = bg.base_coupling_graph(self.repacker);
+        let full_graph = bg.base_rp_graph(self.repacker);
         if coupling_imgcat_debug() {
             full_graph.render_with_imgcat(self.repacker, "Base coupling graph");
         }
         let leaf_nodes = full_graph.leaf_nodes();
         let num_leaf_nodes = leaf_nodes.len();
         for (i, node) in leaf_nodes.into_iter().enumerate() {
             tracing::debug!("Inserting leaf node {} / {}", i, num_leaf_nodes);
-            self.coupling_graph
-                .insert_endpoint(node.clone(), self.repacker);
-            let node_set: BTreeSet<_> = node.into_iter().collect();
+            self.graph.insert_endpoint(node.clone(), self.repacker);
             self.add_edges_from(
                 &full_graph,
-                &node_set,
-                &node_set,
+                &node,
+                &node,
+                borrow_checker,
                 DebugRecursiveCallHistory::new(),
             );
         }
         tracing::debug!("Construct coupling graph end");
-        self.coupling_graph
+        self.graph
     }
 }

src/borrow_pcg/graph/mod.rs

@@ -28,7 +28,9 @@ use super::{
     borrow_pcg_edge::{
         BlockedNode, BorrowPCGEdge, BorrowPCGEdgeLike, BorrowPCGEdgeRef, LocalNode, ToBorrowsEdge,
     },
-    coupling_graph_constructor::{BorrowCheckerInterface, CGNode, CouplingGraphConstructor},
+    coupling_graph_constructor::{
+        BorrowCheckerInterface, CGNode, RegionProjectionAbstractionConstructor,
+    },
     edge::borrow::LocalBorrow,
     edge_data::EdgeData,
     has_pcs_elem::{HasPcgElems, MakePlaceOld},
@@ -150,7 +152,7 @@ impl<'tcx> BorrowsGraph<'tcx> {
             .map(|(kind, conditions)| BorrowPCGEdgeRef { kind, conditions })
     }
 
-    pub(crate) fn base_coupling_graph(
+    pub(crate) fn base_rp_graph(
         &self,
         repacker: PlaceRepacker<'_, 'tcx>,
     ) -> coupling::DisjointSetGraph<CGNode<'tcx>> {
@@ -213,7 +215,6 @@ impl<'tcx> BorrowsGraph<'tcx> {
                 continue;
             }
             seen.insert(ef);
-            tracing::debug!("Exploring from {}", ef.current());
             let edges_blocking = blocking_map.get_edges_blocking(ef.current(), repacker);
             for edge in edges_blocking.iter() {
                 match edge.kind() {
@@ -232,10 +233,6 @@ impl<'tcx> BorrowsGraph<'tcx> {
                         graph.add_edge(&inputs, &outputs, repacker);
                     }
                     _ => {
-                        // if let BorrowPCGEdgeKind::Outlives(outlives) = edge.kind() {
-                        //     graph.insert_endpoint(Coupled::singleton(outlives.long().into()));
-                        //     graph.insert_endpoint(Coupled::singleton(outlives.short().into()));
-                        // }
                         for node in edge.blocked_by_nodes(repacker) {
                             if let LocalNode::RegionProjection(rp) = node {
                                 if let Some(source) = ef.connect()
@@ -421,13 +418,12 @@ impl<'tcx> BorrowsGraph<'tcx> {
 
     fn construct_region_projection_abstraction<'mir>(
         &self,
-        #[allow(unused)]
         borrow_checker: &dyn BorrowCheckerInterface<'mir, 'tcx>,
         repacker: PlaceRepacker<'mir, 'tcx>,
         block: BasicBlock,
     ) -> coupling::DisjointSetGraph<CGNode<'tcx>> {
-        let constructor = CouplingGraphConstructor::new(repacker, block);
-        constructor.construct_region_projection_abstraction(self)
+        let constructor = RegionProjectionAbstractionConstructor::new(repacker, block);
+        constructor.construct_region_projection_abstraction(self, borrow_checker)
     }
 
     fn join_loop<'mir>(

src/coupling/mod.rs

@@ -41,6 +41,15 @@ impl<'tcx, N: Copy + Ord + Clone + DisplayWithRepacker<'tcx> + Hash> DisjointSet
             .collect()
     }
 
+    pub(crate) fn is_root(&self, node: &Coupled<N>) -> bool {
+        self.lookup(*node.iter().next().unwrap()).is_some_and(|idx| {
+            self.inner
+                .neighbors_directed(idx, Direction::Incoming)
+                .count()
+                == 0
+        })
+    }
+
     pub(crate) fn edges(&self) -> impl Iterator<Item = (Coupled<N>, Coupled<N>)> + '_ {
         self.inner.edge_references().map(|e| {
             let source = self.inner.node_weight(e.source()).unwrap();

src/utils/visitor.rs

@@ -42,6 +42,7 @@ pub(crate) trait FallableVisitor<'tcx> {
         Ok(())
     }
 
+    #[allow(unreachable_patterns)]
     fn super_statement_fallable(
         &mut self,
         statement: &mir::Statement<'tcx>,