Skip to content

Commit

Permalink
Creating permission steps incrementally (cognitive-engineering-lab#78)
Browse files Browse the repository at this point in the history 
* initial incremental stepper variant.

* track collection structure.

* Minor tweaks with ForLoopDesugaring & branch joins

* Minor refactor and table entry filters.

Fixed table entry filtering scheme that was causing some of the loop
examples to fail. Table entries are now first grouped by line number
which is also how they are collapsed, this seems to have broken a
different example, and the issue with closures still remains. Needed
still is documentation and addressing comments.

* Fix loop spans and ifs without and else.

Removed old algorithm alongside other useless utilities like the
DFSFinder that was unstable. No "join" segments are inserted for
branches which I will need to revisit.

* Add weird_exprs tests, updated loops.

* Updated tests and documentation.

* Update test doc for SegmentedMir.
  • Loading branch information
@gavinleroy @willcrichton
gavinleroy authored and willcrichton committed Aug 12, 2023
1 parent 3fa01c9 commit 1e741cd
Show file tree
Hide file tree
Showing 16 changed files with 3,187 additions and 1,822 deletions.
4 changes: 2 additions & 2 deletions crates/aquascope/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ publish = false
rustc_private = true

[features]
testing = ["lazy_static"]
testing = []

[dependencies]
anyhow = "1.0.0"
Expand All @@ -32,7 +32,7 @@ miri = {git = "https://github.com/rust-lang/miri.git", rev = "35d6927663065d7fde
aquascope_workspace_utils = { version = "0.2", path = "../aquascope_workspace_utils" }

# testing utils
lazy_static = { version = "1.4", optional = true }
lazy_static = { version = "1.4" }

[dev-dependencies]
insta = { version = "1.22.0", features = ["json", "yaml", "redactions"] }
Expand Down
263 changes: 135 additions & 128 deletions crates/aquascope/src/analysis/ir_mapper/body_graph.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use itertools::Itertools;
use rustc_data_structures::{fx::FxHashMap as HashMap, graph::*};
use rustc_data_structures::{captures::Captures, graph::*};
use rustc_middle::mir::{
BasicBlock, BasicBlockData, BasicBlocks, Body, Location,
BasicBlock, BasicBlockData, BasicBlocks, Body, Location, Terminator,
TerminatorKind,
};
use smallvec::SmallVec;

Expand All @@ -19,15 +19,6 @@ impl<'a, 'tcx: 'a> CleanedBody<'a, 'tcx> {
self.0
}

// TODO: cache the results
pub(crate) fn paths_from_to(
&self,
from: BasicBlock,
to: BasicBlock,
) -> Vec<Vec<BasicBlock>> {
DFSFinder::find_paths_from_to(self, from, to)
}

/// Compute the locations successor.
///
/// If the specified location lies in the middle of a `BasicBlock`,
Expand Down Expand Up @@ -56,14 +47,49 @@ impl<'a, 'tcx: 'a> CleanedBody<'a, 'tcx> {
statement_index: 0,
})
} else {
log::debug!("No Location (or too many) successor(s) found: {nexts:?}");
None
}
}
}

pub fn terminator_in_block(&self, block: BasicBlock) -> &Terminator<'tcx> {
self.body().basic_blocks[block].terminator()
}

pub fn blocks(
&self,
) -> impl Iterator<Item = BasicBlock> + Captures<'a> + Captures<'tcx> + '_ {
self
.0
.basic_blocks
.postorder()
.iter()
.filter(|bb| CleanedBody::keep_block(&self.0.basic_blocks[**bb]))
.copied()
}

pub fn is_false_edge(&self, bb: BasicBlock) -> bool {
matches!(
self.0.basic_blocks[bb].terminator().kind,
TerminatorKind::FalseEdge { .. }
)
}

fn keep_block(bb: &BasicBlockData) -> bool {
!bb.is_cleanup && !bb.is_empty_unreachable()
}

fn is_imaginary_target(
from_data: &BasicBlockData,
target: BasicBlock,
) -> bool {
let TerminatorKind::FalseEdge { imaginary_target, .. } = from_data.terminator().kind else {
return false;
};

imaginary_target == target
}
}

// -----------
Expand Down Expand Up @@ -96,7 +122,11 @@ impl<'tcx> WithSuccessors for CleanedBody<'_, 'tcx> {
node: Self::Node,
) -> <Self as GraphSuccessors<'_>>::Iter {
<BasicBlocks as WithSuccessors>::successors(&self.0.basic_blocks, node)
.filter(|bb| CleanedBody::keep_block(&self.0.basic_blocks[*bb]))
.filter(|bb| {
let from_data = &self.0.basic_blocks[*bb];
CleanedBody::keep_block(from_data)
&& !CleanedBody::is_imaginary_target(from_data, *bb)
})
.collect::<SmallVec<[BasicBlock; 4]>>()
.into_iter()
}
Expand All @@ -119,126 +149,103 @@ impl<'tcx> WithPredecessors for CleanedBody<'_, 'tcx> {
}
}

/// Finds all paths between two nodes.
///
/// This DFS will find all unique paths between two nodes. This
/// includes allowing loops to be traversed (at most once).
/// This is quite a HACK to briefly satisfy the needs of the
/// [stepper](crate::analysis::stepper::compute_permission_steps).
struct DFSFinder<'graph, G>
where
G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
{
graph: &'graph G,
paths: Vec<Vec<G::Node>>,
stack: Vec<G::Node>,
visited: HashMap<G::Node, u8>,
}

impl<'graph, G> DFSFinder<'graph, G>
where
G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
{
pub fn new(graph: &'graph G) -> Self {
Self {
graph,
paths: vec![],
stack: vec![],
visited: HashMap::default(),
}
}

pub fn find_paths_from_to(
graph: &'graph G,
from: G::Node,
to: G::Node,
) -> Vec<Vec<G::Node>> {
let mut dfs = Self::new(graph);
dfs.search(from, to);
dfs.paths.into_iter().unique().collect::<Vec<_>>()
}

fn insert(&mut self, n: G::Node) -> bool {
let v = self.visited.entry(n).or_default();
if *v >= 2 {
return false;
}
*v += 1;
true
}

fn remove(&mut self, n: G::Node) {
let v = self.visited.entry(n).or_default();
assert!(*v > 0);
*v -= 1;
}

fn search(&mut self, from: G::Node, to: G::Node) {
if !self.insert(from) {
return;
}

self.stack.push(from);

if from == to {
self.paths.push(self.stack.clone());
self.remove(to);
self.stack.pop().unwrap();
return;
}

for v in self.graph.successors(from) {
self.search(v, to);
}

self.stack.pop().unwrap();
self.remove(from);
}
}

#[cfg(test)]
mod test {
use rustc_data_structures::graph::vec_graph::VecGraph;
use rustc_utils::BodyExt;

use super::*;
use super::{super::AllPostDominators, *};
use crate::test_utils as tu;

#[test]
fn if_shape() {
// Diamond shaped IF.
let graph = VecGraph::new(6, vec![
(0u32, 1u32),
(1u32, 2u32),
(1u32, 3u32),
(2u32, 4u32),
(3u32, 4u32),
(4u32, 5u32),
]);

let paths_0_5 = vec![vec![0, 1, 2, 4, 5], vec![0, 1, 3, 4, 5]];

assert_eq!(DFSFinder::find_paths_from_to(&graph, 0, 5), paths_0_5);
}
// CleanedBody tests

#[test]
fn while_loop_shape() {
// While loop shape:
// 0 -> 1 -> 2 -> 3 -> 5
// ^ |
// | v
// |-- 4
let graph = VecGraph::new(6, vec![
(0u32, 1u32),
(1u32, 2u32),
(2u32, 3u32),
(3u32, 5u32),
(3u32, 4u32),
(4u32, 2u32),
]);

let paths_0_5 = vec![vec![0, 1, 2, 3, 5], vec![0, 1, 2, 3, 4, 2, 3, 5]];
let mut paths = DFSFinder::find_paths_from_to(&graph, 0, 5);
paths.sort_by_key(|l| l.len());

assert_eq!(paths, paths_0_5);
fn cleaned_body_simple_if() {
// EXPECTED MIR:
// -------------
// bb0: {
// StorageLive(_2);
// _2 = const 0_i32;
// FakeRead(ForLet(None), _2);
// StorageLive(_3);
// StorageLive(_4);
// _4 = const true;
// switchInt(move _4) -> [0: bb3, otherwise: bb1];
// }
//
// bb1: {
// _5 = CheckedAdd(_2, const 1_i32);
// assert(!move (_5.1: bool), <removed>) -> [success: bb2, unwind: bb5];
// }
//
// bb2: {
// _2 = move (_5.0: i32);
// _3 = const ();
// goto -> bb4;
// }
//
// bb3: {
// _3 = const ();
// goto -> bb4;
// }
//
// bb4: {
// StorageDead(_4);
// StorageDead(_3);
// _0 = _2;
// StorageDead(_2);
// return;
// }
//
// bb5 (cleanup): {
// resume;
// }

tu::compile_normal(
r#"
fn foo() -> i32 {
let mut v1 = 0;
if true {
v1 += 1;
}
return v1;
}
"#,
|tcx| {
tu::for_each_body(tcx, |_, wfacts| {
let cleaned_graph = CleanedBody(&wfacts.body);

let post_doms = AllPostDominators::<BasicBlock>::build(
&cleaned_graph,
wfacts.body.all_returns().map(|loc| loc.block),
);

let cleaned_blocks = cleaned_graph.blocks().collect::<Vec<_>>();

let bb0 = BasicBlock::from_usize(0);
let bb1 = BasicBlock::from_usize(1);
let bb2 = BasicBlock::from_usize(2);
let bb3 = BasicBlock::from_usize(3);
let bb4 = BasicBlock::from_usize(4);
let bb5 = BasicBlock::from_usize(5);

assert!(cleaned_blocks.contains(&bb0));
assert!(cleaned_blocks.contains(&bb1));
assert!(cleaned_blocks.contains(&bb2));
assert!(cleaned_blocks.contains(&bb3));
assert!(cleaned_blocks.contains(&bb4));
// Cleanup blocks
assert!(!cleaned_blocks.contains(&bb5));

for &bb in vec![bb0, bb1, bb2, bb3, bb4].iter() {
assert!(post_doms.is_postdominated_by(bb, bb4));
}

assert!(!post_doms.is_postdominated_by(bb0, bb2));
assert!(!post_doms.is_postdominated_by(bb0, bb3));
assert!(post_doms.is_postdominated_by(bb1, bb2));
assert!(!post_doms.is_postdominated_by(bb1, bb3));
})
},
);
}
}
30 changes: 15 additions & 15 deletions crates/aquascope/src/analysis/ir_mapper/mir_locations.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,24 +26,24 @@ impl MirOrderedLocations {
}

pub fn exit_location(&self) -> Option<Location> {
self.exit_block.map(|block| {
let statement_index = *self
.locations
.get(&block)
.expect("Block with no associated locations")
.last()
.unwrap();
Location {
block,
statement_index,
}
})
let block = self.exit_block?;
self.locations.get(&block).map_or_else(
// Block has no associated index then default to the start
|| Some(block.start_location()),
// Get the last associated index of the block
|vs| {
vs.last().map(|&statement_index| Location {
block,
statement_index,
})
},
)
}

pub fn get_entry_exit_locations(&self) -> Option<(Location, Location)> {
self
.entry_location()
.and_then(|mn| self.exit_location().map(|mx| (mn, mx)))
let entry = self.entry_location()?;
let exit = self.exit_location()?;
Some((entry, exit))
}

pub fn values(&self) -> impl Iterator<Item = Location> + Captures<'_> {
Expand Down
Loading

0 comments on commit 1e741cd

Please sign in to comment.