refactor(transformer): make StatementInjectorStore methods generic over GetAddress (#6885)

`StatementInjectorStore::insert_before` etc take any `GetAddress`. This allows using them with an `Ancestor` as well as a `Statement`.

Split the implementation of each function into 2 parts - generic outer wrapper and non-generic inner function to allow compiler to make better inlining decisions, and reduced compile time.
This commit is contained in:
overlookmotel 2024-10-25 12:50:58 +00:00
parent 4618aa22d9
commit c383c34a61

View file

@ -58,18 +58,30 @@ pub struct StatementInjectorStore<'a> {
insertions: RefCell<FxHashMap<Address, Vec<AdjacentStatement<'a>>>>, insertions: RefCell<FxHashMap<Address, Vec<AdjacentStatement<'a>>>>,
} }
// Public methods
impl<'a> StatementInjectorStore<'a> { impl<'a> StatementInjectorStore<'a> {
/// Create new `StatementInjectorStore`. /// Create new `StatementInjectorStore`.
pub fn new() -> Self { pub fn new() -> Self {
Self { insertions: RefCell::new(FxHashMap::default()) } Self { insertions: RefCell::new(FxHashMap::default()) }
} }
// Each of the `insert_before` / `insert_after` functions is split into 2 parts:
//
// 1. Outer function which is generic over any `GetAddress`.
// 2. Inner function which is non-generic and takes `Address`.
//
// Outer functions are marked `#[inline]`, as `GetAddress::address` is generally only 1 or 2 instructions.
// The non-trivial inner functions are not marked `#[inline]` - compiler can decide whether to inline or not.
/// Add a statement to be inserted immediately before the target statement. /// Add a statement to be inserted immediately before the target statement.
#[expect(dead_code)] #[expect(dead_code)]
pub fn insert_before(&self, target: &Statement<'a>, stmt: Statement<'a>) { #[inline]
pub fn insert_before<A: GetAddress>(&self, target: &A, stmt: Statement<'a>) {
self.insert_before_address(target.address(), stmt);
}
fn insert_before_address(&self, target: Address, stmt: Statement<'a>) {
let mut insertions = self.insertions.borrow_mut(); let mut insertions = self.insertions.borrow_mut();
let adjacent_stmts = insertions.entry(target.address()).or_default(); let adjacent_stmts = insertions.entry(target).or_default();
let index = adjacent_stmts let index = adjacent_stmts
.iter() .iter()
.position(|s| matches!(s.direction, Direction::After)) .position(|s| matches!(s.direction, Direction::After))
@ -78,20 +90,34 @@ impl<'a> StatementInjectorStore<'a> {
} }
/// Add a statement to be inserted immediately after the target statement. /// Add a statement to be inserted immediately after the target statement.
pub fn insert_after(&self, target: &Statement<'a>, stmt: Statement<'a>) { #[inline]
pub fn insert_after<A: GetAddress>(&self, target: &A, stmt: Statement<'a>) {
self.insert_after_address(target.address(), stmt);
}
fn insert_after_address(&self, target: Address, stmt: Statement<'a>) {
let mut insertions = self.insertions.borrow_mut(); let mut insertions = self.insertions.borrow_mut();
let adjacent_stmts = insertions.entry(target.address()).or_default(); let adjacent_stmts = insertions.entry(target).or_default();
adjacent_stmts.push(AdjacentStatement { stmt, direction: Direction::After }); adjacent_stmts.push(AdjacentStatement { stmt, direction: Direction::After });
} }
/// Add multiple statements to be inserted immediately before the target statement. /// Add multiple statements to be inserted immediately before the target statement.
#[expect(dead_code)] #[expect(dead_code)]
pub fn insert_many_before<S>(&self, target: &Statement<'a>, stmts: S) #[inline]
pub fn insert_many_before<A, S>(&self, target: &A, stmts: S)
where
A: GetAddress,
S: IntoIterator<Item = Statement<'a>>,
{
self.insert_many_before_address(target.address(), stmts);
}
fn insert_many_before_address<S>(&self, target: Address, stmts: S)
where where
S: IntoIterator<Item = Statement<'a>>, S: IntoIterator<Item = Statement<'a>>,
{ {
let mut insertions = self.insertions.borrow_mut(); let mut insertions = self.insertions.borrow_mut();
let adjacent_stmts = insertions.entry(target.address()).or_default(); let adjacent_stmts = insertions.entry(target).or_default();
adjacent_stmts.splice( adjacent_stmts.splice(
0..0, 0..0,
stmts.into_iter().map(|stmt| AdjacentStatement { stmt, direction: Direction::Before }), stmts.into_iter().map(|stmt| AdjacentStatement { stmt, direction: Direction::Before }),
@ -99,19 +125,28 @@ impl<'a> StatementInjectorStore<'a> {
} }
/// Add multiple statements to be inserted immediately after the target statement. /// Add multiple statements to be inserted immediately after the target statement.
pub fn insert_many_after<S>(&self, target: &Statement<'a>, stmts: S) #[inline]
pub fn insert_many_after<A, S>(&self, target: &A, stmts: S)
where
A: GetAddress,
S: IntoIterator<Item = Statement<'a>>,
{
self.insert_many_after_address(target.address(), stmts);
}
fn insert_many_after_address<S>(&self, target: Address, stmts: S)
where where
S: IntoIterator<Item = Statement<'a>>, S: IntoIterator<Item = Statement<'a>>,
{ {
let mut insertions = self.insertions.borrow_mut(); let mut insertions = self.insertions.borrow_mut();
let adjacent_stmts = insertions.entry(target.address()).or_default(); let adjacent_stmts = insertions.entry(target).or_default();
adjacent_stmts.extend( adjacent_stmts.extend(
stmts.into_iter().map(|stmt| AdjacentStatement { stmt, direction: Direction::After }), stmts.into_iter().map(|stmt| AdjacentStatement { stmt, direction: Direction::After }),
); );
} }
/// Insert statements immediately before / after the target statement. /// Insert statements immediately before / after the target statement.
pub(self) fn insert_into_statements( fn insert_into_statements(
&self, &self,
statements: &mut AVec<'a, Statement<'a>>, statements: &mut AVec<'a, Statement<'a>>,
ctx: &mut TraverseCtx<'a>, ctx: &mut TraverseCtx<'a>,