Skip to content

Commit

Permalink
Don't require updating "save" as frequently.
Browse files Browse the repository at this point in the history
  • Loading branch information
@plucia-mitre
plucia-mitre committed Nov 1, 2024
1 parent bb4922b commit 132126c
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 47 deletions.
2 changes: 1 addition & 1 deletion search/treesearchlib/src/automata/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
//!
//! Most of the public modules (except for [thread] and [results]) in this
//! module are different implementations of the our NFA engines. They *should*
//! all be functionally equivalent, but they each have different performance
//! all be functionally equavalent, but they each have different performance
//! characteristics. In every case, the "main" function for each is
//! "run_program". The most performant of these modules (currently
//! [pikevm_loop_ring_rc]) is re-exported at the top level of this module for
Expand Down
42 changes: 22 additions & 20 deletions search/treesearchlib/src/automata/pikevm_loop_ring_rc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,71 +43,75 @@ pub fn run_program<Endian: BitOrder + Clone + PartialEq, S: StatesRc<ThreadRc>>(

// #[cfg_attr(test, mutate)]
pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, S: StatesRc<ThreadRc>>(
mut pc: usize,
saved: &mut Rc<SavedData>,
thread: &mut ThreadRc,
sp: usize,
prog: &Pattern<Endian>,
input: &AddressedBits,
states: &mut S,
cache: &mut LookupCache<Endian>,
) -> Option<Results> {
let mut pc = thread.pc_idx;
let saved = &mut thread.saved;
let mut start = thread.start;
loop {
match prog.steps.get(pc).unwrap() {
Op::Byte { value } => {
if input.len() > sp && *value == input[sp] {
states.add(sp + 1, ThreadRc::new(pc + 1, saved));
states.add(sp + 1, ThreadRc::new(pc + 1, saved, start));
}
break;
}
Op::MaskedByte { mask, value } => {
if input.len() > sp && *value == (input[sp] & *mask) {
states.add(sp + 1, ThreadRc::new(pc + 1, saved));
states.add(sp + 1, ThreadRc::new(pc + 1, saved, start));
}
break;
}
Op::ByteMultiNonconsuming { value } => {
if input.len() > sp && value.contains(&input[sp]) {
states.add(sp, ThreadRc::new(pc + 1, saved));
states.add(sp, ThreadRc::new(pc + 1, saved, start));
}
break;
}
Op::Match { match_number } => {
let mut final_saved = Rc::make_mut(saved).clone();
final_saved.start = start;
return Some(Results {
matched: true,
match_number: Some(*match_number),
saved: Some(Rc::make_mut(saved).clone()),
saved: Some(final_saved),
});
}
Op::Jmp { dest } => {
pc = *dest;
// states.add(sp, Thread::new(*dest, &saved));
// states.add(sp, Thread::new(*dest, &saved, start));
}
Op::Split { dest1, dest2 } => {
states.add(sp, ThreadRc::new(*dest1, saved));
// states.add(sp, Thread::new(*dest2, &saved));
states.add(sp, ThreadRc::new(*dest1, saved, start));
// states.add(sp, Thread::new(*dest2, &saved, start));
pc = *dest2;
}
Op::SplitMulti { dests } => {
if !dests.is_empty() {
pc = *dests.first().unwrap();
if dests.len() > 1 {
for dest in &dests[1..] {
states.add(sp, ThreadRc::new(*dest, saved));
states.add(sp, ThreadRc::new(*dest, saved, start));
}
}
} else {
break;
}
}
Op::SaveStart => {
Rc::make_mut(saved).start = Some(sp);
start = Some(sp);
pc += 1;
}
Op::Save { slot } => {
// let saved_ready_to_mut = Rc::make_mut(saved);

// saved_ready_to_mut.captures.insert(*slot, sp);
// // states.add(sp, Thread::new(pc + 1, &saved));
// // states.add(sp, Thread::new(pc + 1, &saved, start));
// *saved = Rc::new(saved_ready_to_mut.clone());
Rc::make_mut(saved).captures.insert(*slot, sp);
pc += 1;
Expand All @@ -120,12 +124,12 @@ pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, S: StatesRc<Threa
pc += 1;
}
Op::AnyByte => {
states.add(sp + 1, ThreadRc::new(pc + 1, saved));
states.add(sp + 1, ThreadRc::new(pc + 1, saved, start));
break;
}
Op::AnyByteSequence { min, max, interval } => {
for i in (*min..(*max + 1)).step_by(*interval) {
states.add(sp + i, ThreadRc::new(pc + 1, saved));
states.add(sp + i, ThreadRc::new(pc + 1, saved, start));
}
break;
}
Expand All @@ -138,7 +142,7 @@ pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, S: StatesRc<Threa
saved,
);
if let Ok(good_result) = result {
states.add(sp + good_result.size, ThreadRc::new(pc + 1, saved));
states.add(sp + good_result.size, ThreadRc::new(pc + 1, saved, start));
break;
}
}
Expand All @@ -156,7 +160,7 @@ pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, S: StatesRc<Threa
saved,
);
if let Ok(good_result) = result {
states.add(sp + good_result.size, ThreadRc::new(pc + 1, saved));
states.add(sp + good_result.size, ThreadRc::new(pc + 1, saved, start));
break;
}
}
Expand Down Expand Up @@ -186,7 +190,7 @@ pub fn pikevm_inner_new_rc<Endian: BitOrder + Clone + PartialEq, S: StatesRc<Thr
) -> Results {
let mut cache = LookupCache::new(max_cache_size);

states.add(0, ThreadRc::new(0, &Rc::new(SavedData::default())));
states.add(0, ThreadRc::new(0, &Rc::new(SavedData::default()), None));
let mut sp = 0;
while sp <= input.len() {
// eprintln!("{}->{}", input.len(), sp);
Expand All @@ -198,9 +202,7 @@ pub fn pikevm_inner_new_rc<Endian: BitOrder + Clone + PartialEq, S: StatesRc<Thr
}
let mut t = option_next_thread.unwrap();

if let Some(x) =
process_thread_rc(t.pc_idx, &mut t.saved, sp, prog, input, states, &mut cache)
{
if let Some(x) = process_thread_rc(&mut t, sp, prog, input, states, &mut cache) {
return x;
}
}
Expand Down
68 changes: 43 additions & 25 deletions search/treesearchlib/src/automata/pikevm_loop_ring_rc_priority.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,19 +48,36 @@ fn add_thread<Endian: BitOrder + Clone + PartialEq, State: StatesRc<ThreadRc>>(
pc: usize,
input: &AddressedBits,
saved: &mut Rc<SavedData>,
start: Option<usize>,
) {
match prog.steps.get(pc).unwrap() {
Op::ByteMultiNonconsuming { value } => {
if input.len() > sp && value.contains(&input[sp]) {
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved, start);
}
}
Op::Jmp { dest } => {
add_thread::<Endian, State>(prog, states, sp, *dest, input, saved);
add_thread::<Endian, State>(prog, states, sp, *dest, input, saved, start);
}
Op::Split { dest1, dest2 } => {
add_thread::<Endian, State>(prog, states, sp, *dest1, input, &mut Rc::clone(saved));
add_thread::<Endian, State>(prog, states, sp, *dest2, input, &mut Rc::clone(saved));
add_thread::<Endian, State>(
prog,
states,
sp,
*dest1,
input,
&mut Rc::clone(saved),
start,
);
add_thread::<Endian, State>(
prog,
states,
sp,
*dest2,
input,
&mut Rc::clone(saved),
start,
);
}
Op::SplitMulti { dests } => {
if !dests.is_empty() {
Expand All @@ -72,15 +89,15 @@ fn add_thread<Endian: BitOrder + Clone + PartialEq, State: StatesRc<ThreadRc>>(
*dest,
input,
&mut Rc::clone(saved),
start,
);
}
} else {
panic!("Cannot have an empty SplitMulti!");
}
}
Op::SaveStart => {
Rc::make_mut(saved).start = Some(sp);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved, Some(sp));
}
Op::Save { slot } => {
// let saved_ready_to_mut = Rc::make_mut(saved);
Expand All @@ -89,57 +106,61 @@ fn add_thread<Endian: BitOrder + Clone + PartialEq, State: StatesRc<ThreadRc>>(
// // states.add(sp, Thread::new(pc + 1, &saved));
// *saved = Rc::new(saved_ready_to_mut.clone());
Rc::make_mut(saved).captures.insert(*slot, sp);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved, start);
}
Op::Label { value } => {
Rc::make_mut(saved).labels.insert(
value.clone(),
TryInto::<i128>::try_into(sp).unwrap() + i128::from(input.get_base_address()),
);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp, pc + 1, input, saved, start);
}
_ => states.add(sp, ThreadRc::new(pc, saved)),
_ => states.add(sp, ThreadRc::new(pc, saved, start)),
}
}

// #[cfg_attr(test, mutate)]
pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, State: StatesRc<ThreadRc>>(
mut pc: usize,
saved: &mut Rc<SavedData>,
thread: &mut ThreadRc,
sp: usize,
prog: &Pattern<Endian>,
input: &AddressedBits,
states: &mut State, // StatesRingRcFixed<ThreadRc>,
cache: &mut LookupCache<Endian>,
) -> Option<Results> {
let mut pc = thread.pc_idx;
let saved = &mut thread.saved;
let start = thread.start;
loop {
match prog.steps.get(pc).unwrap() {
Op::Byte { value } => {
if input.len() > sp && *value == input[sp] {
add_thread::<Endian, State>(prog, states, sp + 1, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp + 1, pc + 1, input, saved, start);
}
break;
}
Op::MaskedByte { mask, value } => {
if input.len() > sp && *value == (input[sp] & *mask) {
add_thread::<Endian, State>(prog, states, sp + 1, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp + 1, pc + 1, input, saved, start);
}
break;
}
Op::Match { match_number } => {
let mut final_saved = Rc::make_mut(saved).clone();
final_saved.start = start;
return Some(Results {
matched: true,
match_number: Some(*match_number),
saved: Some(Rc::make_mut(saved).clone()),
saved: Some(final_saved),
});
}
Op::AnyByte => {
add_thread::<Endian, State>(prog, states, sp + 1, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp + 1, pc + 1, input, saved, start);
break;
}
Op::AnyByteSequence { min, max, interval } => {
for i in (*min..(*max + 1)).step_by(*interval) {
add_thread::<Endian, State>(prog, states, sp + i, pc + 1, input, saved);
add_thread::<Endian, State>(prog, states, sp + i, pc + 1, input, saved, start);
}
break;
}
Expand All @@ -159,6 +180,7 @@ pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, State: StatesRc<T
pc + 1,
input,
saved,
start,
);
break;
}
Expand All @@ -185,6 +207,7 @@ pub fn process_thread_rc<Endian: BitOrder + Clone + PartialEq, State: StatesRc<T
pc + 1,
input,
saved,
start,
);
break;
}
Expand Down Expand Up @@ -225,6 +248,7 @@ pub fn pikevm_inner_new_rc<Endian: BitOrder + Clone + PartialEq, State: StatesRc
0,
input,
&mut Rc::new(SavedData::default()),
None,
);

// add_thread::<Endian, State>(prog, states, 0, 0, &mut Rc::new(SavedData::default()));
Expand All @@ -242,15 +266,9 @@ pub fn pikevm_inner_new_rc<Endian: BitOrder + Clone + PartialEq, State: StatesRc
let mut t = option_next_thread.unwrap();
// eprintln!("Pulled thread: pc={} at sp {}", t.pc_idx, sp);

if let Some(x) = process_thread_rc::<Endian, State>(
t.pc_idx,
&mut t.saved,
sp,
prog,
input,
states,
&mut cache,
) {
if let Some(x) =
process_thread_rc::<Endian, State>(&mut t, sp, prog, input, states, &mut cache)
{
return x;
}
// eprintln!("After States: {:?}", states);
Expand Down
5 changes: 4 additions & 1 deletion search/treesearchlib/src/automata/thread.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,15 @@ impl Thread {
pub struct ThreadRc {
pub pc_idx: usize,
pub saved: Rc<SavedData>,
pub start: Option<usize>,
}

impl ThreadRc {
pub fn new(pc_idx: usize, saved: &Rc<SavedData>) -> Self {
pub fn new(pc_idx: usize, saved: &Rc<SavedData>, start: Option<usize>) -> Self {
Self {
pc_idx,
saved: Rc::clone(saved),
start,
}
}

Expand All @@ -78,6 +80,7 @@ impl Default for ThreadRc {
Self {
pc_idx: 0,
saved: Rc::new(SavedData::new()),
start: None,
}
}
}
Expand Down

0 comments on commit 132126c

Please sign in to comment.