#![forbid(unsafe_code)]
mod callgraph;
mod calls;
mod classdefs;
mod insn_locator;
mod locators;
mod persist;
mod query;
mod scanner;
mod table;
pub use callgraph::CallGraph;
pub use calls::{CallConstraint, Certainty, ResultWant, Slot, Want};
pub use insn_locator::{InsnLocator, MethodRef};
pub use locators::{FieldLocator, MethodLocator, StringLocator, TypeLocator};
pub use persist::FORMAT_VERSION as PERSIST_FORMAT_VERSION;
pub use query::{CompiledQuery, Compound};
pub use scanner::{CodeScanner, FieldAccess, IdxSet, RefKind};
pub use table::Backing;
use asc_dex::Dex;
use query::Parts;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::sync::{Mutex, OnceLock, PoisonError};
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum FindError {
#[error(transparent)]
Dex(#[from] asc_dex::DexError),
#[error("invalid pattern: {0}")]
Pattern(#[from] regex::Error),
#[error("invalid query: {0}")]
Invalid(String),
}
pub type Result<T> = std::result::Result<T, FindError>;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct ClassQuery {
pub pattern: String,
pub precise: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct MemberQuery {
pub class: Option<ClassQuery>,
pub name: Option<String>,
pub exact_name: bool,
pub signature: Option<MemberSignature>,
}
impl ClassQuery {
pub fn exact(descriptor: impl Into<String>) -> ClassQuery {
ClassQuery { pattern: descriptor.into(), precise: true }
}
pub fn fuzzy(pattern: impl Into<String>) -> ClassQuery {
ClassQuery { pattern: pattern.into(), precise: false }
}
}
impl MemberQuery {
#[must_use]
pub fn with_class(self, class: Option<ClassQuery>) -> MemberQuery {
MemberQuery { class, ..self }
}
#[must_use]
pub fn with_name(self, name: Option<String>) -> MemberQuery {
MemberQuery { name, ..self }
}
#[must_use]
pub fn with_exact_name(self, exact_name: bool) -> MemberQuery {
MemberQuery { exact_name, ..self }
}
#[must_use]
pub fn with_signature(self, signature: Option<MemberSignature>) -> MemberQuery {
MemberQuery { signature, ..self }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MemberSignature {
Method {
params: Vec<TypeMatch>,
ret: Option<TypeMatch>,
},
Field(TypeMatch),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TypeMatch {
Descriptor(String),
Simple {
name: String,
dims: usize,
},
}
impl TypeMatch {
pub fn matches(&self, descriptor: &str) -> bool {
query::CompiledType::new(self).matches(&asc_dex::mutf8::encode(descriptor))
}
}
impl std::fmt::Display for TypeMatch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
TypeMatch::Descriptor(d) => f.write_str(d),
TypeMatch::Simple { name, dims } => write!(f, "{name}{}", "[]".repeat(*dims)),
}
}
}
impl std::fmt::Display for MemberSignature {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MemberSignature::Method { params, ret } => {
f.write_str("(")?;
for (i, p) in params.iter().enumerate() {
let simple = |t: &TypeMatch| matches!(t, TypeMatch::Simple { .. });
if i > 0 && (simple(p) || simple(¶ms[i - 1])) {
f.write_str(", ")?;
}
write!(f, "{p}")?;
}
f.write_str(")")?;
ret.as_ref().map_or(Ok(()), |r| write!(f, "{r}"))
}
MemberSignature::Field(t) => write!(f, ":{t}"),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Query {
String(String),
StringLiteral(String),
Type(String),
Method(MemberQuery),
MethodSet(Vec<MemberQuery>),
Nested(Box<Query>),
Virtual {
target: Box<Query>,
via: BTreeMap<String, String>,
fan_out: Option<usize>,
},
Field(MemberQuery),
FieldAccess {
field: MemberQuery,
access: FieldAccess,
},
Call {
target: Box<Query>,
constraints: Vec<CallConstraint>,
},
Extends(String),
Implements(String),
Defined(Defined),
All {
all: Vec<Query>,
none: Vec<Query>,
},
}
impl Query {
pub fn kind(&self) -> RefKind {
match self {
Query::String(_) | Query::StringLiteral(_) => RefKind::String,
Query::Type(_) | Query::Extends(_) | Query::Implements(_) => RefKind::Type,
Query::Method(_) | Query::MethodSet(_) | Query::Nested(_) | Query::Virtual { .. } | Query::Call { .. } | Query::Defined(_) => {
RefKind::Method
}
Query::Field(_) | Query::FieldAccess { .. } => RefKind::Field,
Query::All { all, .. } => all.first().map_or(RefKind::String, Query::kind),
}
}
pub fn relation(&self) -> Option<ClassRelation> {
match self {
Query::Extends(_) => Some(ClassRelation::Extends),
Query::Implements(_) => Some(ClassRelation::Implements),
_ => None,
}
}
pub fn field_access(&self) -> Option<FieldAccess> {
match self {
Query::FieldAccess { access, .. } => Some(*access),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Defined {
Flags(u32),
Annotation(String),
Members {
classes: Vec<String>,
name: Option<String>,
proto: Option<String>,
},
Override(MemberQuery),
In(Box<Query>),
}
impl Defined {
pub fn label(&self) -> String {
match self {
Defined::Flags(_) => "flag".into(),
Defined::Annotation(a) => format!("@{a}"),
Defined::Members { name: Some(n), .. } => format!("override {n}"),
Defined::Members { .. } => "in".into(),
Defined::Override(m) => format!("override {}", m.name.as_deref().unwrap_or("")),
Defined::In(_) => "in".into(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ClassRelation {
Extends,
Implements,
}
impl ClassRelation {
pub fn verb(self) -> &'static str {
match self {
ClassRelation::Extends => "extends",
ClassRelation::Implements => "implements",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct Hit {
pub offset: u32,
pub matched_idx: u32,
pub methods: Option<MethodRef>,
pub field_access: Option<FieldAccess>,
pub note: Option<String>,
}
#[derive(Default)]
pub struct SearchState {
strings: OnceLock<StringLocator>,
types: OnceLock<TypeLocator>,
methods: OnceLock<MethodLocator>,
fields: OnceLock<FieldLocator>,
insns: OnceLock<InsnLocator>,
calls: OnceLock<CallGraph>,
building: [Mutex<()>; 6],
}
fn built<'c, T>(cell: &'c OnceLock<T>, lock: &Mutex<()>, build: impl FnOnce() -> Result<T>) -> Result<&'c T> {
if let Some(v) = cell.get() {
return Ok(v);
}
let _building = lock.lock().unwrap_or_else(PoisonError::into_inner);
if let Some(v) = cell.get() {
return Ok(v);
}
let v = build()?;
Ok(cell.get_or_init(|| v))
}
impl SearchState {
pub fn new() -> Self {
Self::default()
}
pub fn is_warm(&self) -> bool {
self.insns.get().is_some()
}
pub fn heap_bytes(&self) -> usize {
self.strings.get().map_or(0, StringLocator::heap_bytes)
+ self.types.get().map_or(0, TypeLocator::heap_bytes)
+ self.methods.get().map_or(0, MethodLocator::heap_bytes)
+ self.fields.get().map_or(0, FieldLocator::heap_bytes)
+ self.insns.get().map_or(0, InsnLocator::heap_bytes)
+ self.calls.get().map_or(0, CallGraph::heap_bytes)
}
pub fn call_graph(&self) -> Option<&CallGraph> {
self.calls.get()
}
pub fn set_call_graph(&self, graph: CallGraph) -> bool {
self.calls.set(graph).is_ok()
}
}
enum StateRef<'s> {
Owned(Box<SearchState>),
Shared(&'s SearchState),
}
pub struct FindRefs<'d, 'a, 's> {
dex: &'d Dex<'a>,
state: StateRef<'s>,
}
impl<'d, 'a, 's> FindRefs<'d, 'a, 's> {
pub fn new(dex: &'d Dex<'a>) -> Self {
FindRefs { dex, state: StateRef::Owned(Box::default()) }
}
pub fn with_state(dex: &'d Dex<'a>, state: &'s SearchState) -> Self {
FindRefs { dex, state: StateRef::Shared(state) }
}
fn state(&self) -> &SearchState {
match &self.state {
StateRef::Owned(s) => s,
StateRef::Shared(s) => s,
}
}
pub fn dex(&self) -> &'d Dex<'a> {
self.dex
}
pub fn string_locator(&self) -> Result<&StringLocator> {
let st = self.state();
built(&st.strings, &st.building[0], || StringLocator::build(self.dex))
}
pub fn type_locator(&self) -> Result<&TypeLocator> {
let st = self.state();
built(&st.types, &st.building[1], || TypeLocator::build(self.dex))
}
pub fn method_locator(&self) -> Result<&MethodLocator> {
let st = self.state();
built(&st.methods, &st.building[2], || MethodLocator::build(self.dex))
}
pub fn field_locator(&self) -> Result<&FieldLocator> {
let st = self.state();
built(&st.fields, &st.building[3], || FieldLocator::build(self.dex))
}
pub fn insn_locator(&self) -> Result<&InsnLocator> {
let st = self.state();
built(&st.insns, &st.building[4], || Ok(InsnLocator::build(self.dex)?))
}
pub fn call_graph(&self) -> Result<&CallGraph> {
let st = self.state();
built(&st.calls, &st.building[5], || Ok(CallGraph::build(self.dex.buf(), self.insn_locator()?)))
}
pub fn locate(&self, query: &Query) -> Result<BTreeSet<u32>> {
self.locate_compiled(&CompiledQuery::new(query)?)
}
pub fn locate_compiled(&self, query: &CompiledQuery) -> Result<BTreeSet<u32>> {
let buf = self.dex.buf();
match (query.kind(), query.parts()) {
(_, Parts::Pattern(None)) | (_, Parts::Literal(None)) => Ok(BTreeSet::new()),
(RefKind::String, Parts::Pattern(Some(re))) => Ok(self.string_locator()?.locate(buf, re)),
(RefKind::String, Parts::Literal(Some(f))) => Ok(self.string_locator()?.locate_literal(buf, f)),
(_, Parts::Relation(..)) => Ok(self.class_relation(query)?.unwrap_or_default().into_iter().collect()),
(RefKind::Type, Parts::Pattern(Some(re))) => Ok(self.type_locator()?.locate(buf, self.string_locator()?, re)),
(_, Parts::Call(c)) => self.locate_compiled(&c.target),
(_, Parts::Virtual(target, _)) => self.locate_compiled(target),
(RefKind::Method, Parts::Member(_) | Parts::MemberSet(_)) => {
self.method_locator()?.locate_compiled(self.dex, self.string_locator()?, self.type_locator()?, query)
}
(RefKind::Field, Parts::Member(_)) => {
self.field_locator()?.locate_compiled(self.dex, self.string_locator()?, self.type_locator()?, query)
}
_ => Ok(BTreeSet::new()),
}
}
pub fn cannot_match(&self, query: &CompiledQuery) -> Result<bool> {
if let Some(c) = query.compound() {
for part in &c.all {
if self.cannot_match(part)? {
return Ok(true);
}
}
return Ok(false);
}
match query.parts() {
Parts::Call(c) => self.cannot_match(&c.target),
Parts::Virtual(target, _) => self.cannot_match(target),
Parts::Relation(_, descriptor) => Ok(self.dex.find_type_idx(descriptor)?.is_none()),
Parts::Defined(Defined::Flags(_)) => Ok(false),
Parts::Defined(Defined::Annotation(a)) => Ok(self.dex.find_type_idx(&asc_dex::mutf8::encode(a))?.is_none()),
Parts::Defined(Defined::Members { classes, .. }) => {
for c in classes {
if self.dex.find_type_idx(&asc_dex::mutf8::encode(c))?.is_some() {
return Ok(false);
}
}
Ok(true)
}
Parts::Defined(Defined::Override(_) | Defined::In(_)) => Ok(false),
Parts::Member(query::CompiledMember { class: Some(query::CompiledClass::Precise(descriptor)), .. })
if self.dex.find_type_idx(descriptor)?.is_none() =>
{
Ok(true)
}
_ => Ok(self.locate_compiled(query)?.is_empty()),
}
}
pub fn find_ref(&self, query: &Query) -> Result<Vec<Hit>> {
self.find_ref_compiled(&CompiledQuery::new(query)?)
}
pub fn find_ref_compiled(&self, query: &CompiledQuery) -> Result<Vec<Hit>> {
if query.relation().is_some() || query.compound().is_some() || query.defined().is_some() {
return Ok(Vec::new());
}
let idxs = self.locate_compiled(query)?;
if idxs.is_empty() {
return Ok(Vec::new());
}
let mut hits = self.scan_idxs(query.kind(), &idxs)?;
if let Some(access) = query.query().field_access() {
hits.retain(|hit| hit.field_access == Some(access));
}
self.refine(query, hits)
}
fn refine(&self, query: &CompiledQuery, hits: Vec<Hit>) -> Result<Vec<Hit>> {
match query.parts() {
Parts::Call(c) => {
let hits = self.refine(&c.target, hits)?;
self.constrain(c, hits)
}
Parts::Virtual(_, via) => self.label_dispatch(via, hits),
_ => Ok(hits),
}
}
fn label_dispatch(&self, via: &BTreeMap<String, String>, mut hits: Vec<Hit>) -> Result<Vec<Hit>> {
if via.is_empty() {
return Ok(hits);
}
let mut labels: HashMap<u32, Option<&str>> = HashMap::new();
for hit in &mut hits {
let label = match labels.get(&hit.matched_idx) {
Some(&l) => l,
None => {
let m = self.dex.method_id(hit.matched_idx)?;
let name = self.dex.method_display(hit.matched_idx)?;
let proto = self.dex.proto_id(u32::from(m.proto_idx))?;
let mut sig = format!("{name}(");
for p in self.dex.type_list(proto.parameters_off)? {
sig.push_str(&self.dex.type_descriptor(u32::from(p))?);
}
sig.push(')');
sig.push_str(&self.dex.type_descriptor(proto.return_type_idx)?);
let l = via.get(&sig).or_else(|| via.get(&name)).map(String::as_str);
labels.insert(hit.matched_idx, l);
l
}
};
if let Some(label) = label {
hit.note = Some(format!("(dispatch {label})"));
}
}
Ok(hits)
}
pub fn find_many(&self, queries: &[&CompiledQuery]) -> Result<Vec<Vec<Hit>>> {
let mut out: Vec<Vec<Hit>> = (0..queries.len()).map(|_| Vec::new()).collect();
let mut located: Vec<Option<BTreeSet<u32>>> = Vec::with_capacity(queries.len());
for q in queries {
let skip = q.relation().is_some() || q.compound().is_some();
located.push(if skip { None } else { Some(self.locate_compiled(q)?).filter(|s| !s.is_empty()) });
}
for kind in [RefKind::String, RefKind::Type, RefKind::Method, RefKind::Field] {
let mine: Vec<usize> = (0..queries.len()).filter(|&i| queries[i].kind() == kind && located[i].is_some()).collect();
let union: BTreeSet<u32> = match mine.as_slice() {
[] => continue,
[only] => located[*only].clone().unwrap_or_default(),
_ => mine.iter().flat_map(|&i| located[i].iter().flatten().copied()).collect(),
};
for hit in self.scan_idxs(kind, &union)? {
for &i in &mine {
if located[i].as_ref().is_some_and(|s| s.contains(&hit.matched_idx)) {
let access = queries[i].query().field_access();
if access.is_none() || access == hit.field_access {
out[i].push(hit.clone());
}
}
}
}
}
for (q, hits) in queries.iter().zip(out.iter_mut()) {
*hits = self.refine(q, std::mem::take(hits))?;
}
Ok(out)
}
pub fn class_relation(&self, query: &CompiledQuery) -> Result<Option<Vec<u32>>> {
match query.parts() {
Parts::Relation(rel, descriptor) => Ok(Some(classdefs::related_classes(self.dex, *rel, descriptor)?)),
_ => Ok(None),
}
}
pub fn scan_idxs(&self, kind: RefKind, idxs: &BTreeSet<u32>) -> Result<Vec<Hit>> {
let buf = self.dex.buf();
let insn = self.insn_locator()?;
let scanner = CodeScanner::new(insn);
let max_idx = (buf.len() / 4) as u32;
let set = IdxSet::from_iter(idxs.iter().copied().filter(|&i| i < max_idx));
let matches = scanner.scan_instruction_candidates(buf, kind, &set);
let offsets: Vec<u32> = matches.iter().map(|m| m.0).collect();
let methods = insn.locate(buf, &offsets);
Ok(matches
.into_iter()
.zip(methods)
.map(|((offset, matched_idx), methods)| Hit {
offset,
matched_idx,
methods,
field_access: (kind == RefKind::Field).then(|| FieldAccess::from_opcode(buf[offset as usize])).flatten(),
note: None,
})
.collect())
}
}