1#![forbid(unsafe_code)]
24
25mod callgraph;
26mod calls;
27mod classdefs;
28mod insn_locator;
29mod locators;
30mod persist;
31mod query;
32mod scanner;
33mod table;
34
35pub use callgraph::CallGraph;
36pub use calls::{CallConstraint, Certainty, ResultWant, Slot, Want};
37pub use insn_locator::{InsnLocator, MethodRef};
38pub use locators::{FieldLocator, MethodLocator, StringLocator, TypeLocator};
39pub use persist::FORMAT_VERSION as PERSIST_FORMAT_VERSION;
40pub use query::{CompiledQuery, Compound};
41pub use scanner::{CodeScanner, FieldAccess, IdxSet, RefKind};
42pub use table::Backing;
43
44use asc_dex::Dex;
45use query::Parts;
46use std::collections::{BTreeMap, BTreeSet, HashMap};
47use std::sync::{Mutex, OnceLock, PoisonError};
48
49#[derive(Debug, thiserror::Error)]
51pub enum FindError {
52 #[error(transparent)]
53 Dex(#[from] asc_dex::DexError),
54 #[error("invalid pattern: {0}")]
55 Pattern(#[from] regex::Error),
56 #[error("invalid query: {0}")]
57 Invalid(String),
58}
59
60pub type Result<T> = std::result::Result<T, FindError>;
61
62#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct ClassQuery {
65 pub pattern: String,
67 pub precise: bool,
68}
69
70#[derive(Debug, Clone, PartialEq, Eq, Default)]
72pub struct MemberQuery {
73 pub class: Option<ClassQuery>,
74 pub name: Option<String>,
77 pub exact_name: bool,
80 pub signature: Option<MemberSignature>,
84}
85
86#[derive(Debug, Clone, PartialEq, Eq)]
88pub enum MemberSignature {
89 Method { params: Vec<TypeMatch>, ret: Option<TypeMatch> },
91 Field(TypeMatch),
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
97pub enum TypeMatch {
98 Descriptor(String),
100 Simple { name: String, dims: usize },
103}
104
105impl TypeMatch {
106 pub fn matches(&self, descriptor: &str) -> bool {
108 query::CompiledType::new(self).matches(&asc_dex::mutf8::encode(descriptor))
109 }
110}
111
112impl std::fmt::Display for TypeMatch {
113 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114 match self {
115 TypeMatch::Descriptor(d) => f.write_str(d),
116 TypeMatch::Simple { name, dims } => write!(f, "{name}{}", "[]".repeat(*dims)),
117 }
118 }
119}
120
121impl std::fmt::Display for MemberSignature {
122 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 match self {
125 MemberSignature::Method { params, ret } => {
126 f.write_str("(")?;
127 for (i, p) in params.iter().enumerate() {
128 let simple = |t: &TypeMatch| matches!(t, TypeMatch::Simple { .. });
131 if i > 0 && (simple(p) || simple(¶ms[i - 1])) {
132 f.write_str(", ")?;
133 }
134 write!(f, "{p}")?;
135 }
136 f.write_str(")")?;
137 ret.as_ref().map_or(Ok(()), |r| write!(f, "{r}"))
138 }
139 MemberSignature::Field(t) => write!(f, ":{t}"),
140 }
141 }
142}
143
144#[derive(Debug, Clone, PartialEq, Eq)]
146pub enum Query {
147 String(String),
149 StringLiteral(String),
151 Type(String),
152 Method(MemberQuery),
153 MethodSet(Vec<MemberQuery>),
156 Nested(Box<Query>),
161 Virtual {
170 target: Box<Query>,
171 via: BTreeMap<String, String>,
172 fan_out: Option<usize>,
177 },
178 Field(MemberQuery),
179 FieldAccess {
182 field: MemberQuery,
183 access: FieldAccess,
184 },
185 Call {
189 target: Box<Query>,
190 constraints: Vec<CallConstraint>,
191 },
192 Extends(String),
195 Implements(String),
197 Defined(Defined),
201 All {
205 all: Vec<Query>,
206 none: Vec<Query>,
207 },
208}
209
210impl Query {
211 pub fn kind(&self) -> RefKind {
214 match self {
215 Query::String(_) | Query::StringLiteral(_) => RefKind::String,
216 Query::Type(_) | Query::Extends(_) | Query::Implements(_) => RefKind::Type,
217 Query::Method(_) | Query::MethodSet(_) | Query::Nested(_) | Query::Virtual { .. } | Query::Call { .. } | Query::Defined(_) => {
218 RefKind::Method
219 }
220 Query::Field(_) | Query::FieldAccess { .. } => RefKind::Field,
221 Query::All { all, .. } => all.first().map_or(RefKind::String, Query::kind),
223 }
224 }
225
226 pub fn relation(&self) -> Option<ClassRelation> {
228 match self {
229 Query::Extends(_) => Some(ClassRelation::Extends),
230 Query::Implements(_) => Some(ClassRelation::Implements),
231 _ => None,
232 }
233 }
234
235 pub fn field_access(&self) -> Option<FieldAccess> {
237 match self {
238 Query::FieldAccess { access, .. } => Some(*access),
239 _ => None,
240 }
241 }
242}
243
244#[derive(Debug, Clone, PartialEq, Eq)]
252pub enum Defined {
253 Flags(u32),
255 Annotation(String),
257 Members { classes: Vec<String>, name: Option<String>, proto: Option<String> },
261 Override(MemberQuery),
264 In(Box<Query>),
267}
268
269impl Defined {
270 pub fn label(&self) -> String {
272 match self {
273 Defined::Flags(_) => "flag".into(),
274 Defined::Annotation(a) => format!("@{a}"),
275 Defined::Members { name: Some(n), .. } => format!("override {n}"),
276 Defined::Members { .. } => "in".into(),
277 Defined::Override(m) => format!("override {}", m.name.as_deref().unwrap_or("")),
278 Defined::In(_) => "in".into(),
279 }
280 }
281}
282
283#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
285pub enum ClassRelation {
286 Extends,
288 Implements,
290}
291
292impl ClassRelation {
293 pub fn verb(self) -> &'static str {
295 match self {
296 ClassRelation::Extends => "extends",
297 ClassRelation::Implements => "implements",
298 }
299 }
300}
301
302#[derive(Debug, Clone, PartialEq, Eq)]
304pub struct Hit {
305 pub offset: u32,
307 pub matched_idx: u32,
309 pub methods: Option<MethodRef>,
311 pub field_access: Option<FieldAccess>,
314 pub note: Option<String>,
319}
320
321#[derive(Default)]
333pub struct SearchState {
334 strings: OnceLock<StringLocator>,
335 types: OnceLock<TypeLocator>,
336 methods: OnceLock<MethodLocator>,
337 fields: OnceLock<FieldLocator>,
338 insns: OnceLock<InsnLocator>,
339 calls: OnceLock<CallGraph>,
340 building: [Mutex<()>; 6],
342}
343
344fn built<'c, T>(cell: &'c OnceLock<T>, lock: &Mutex<()>, build: impl FnOnce() -> Result<T>) -> Result<&'c T> {
346 if let Some(v) = cell.get() {
347 return Ok(v);
348 }
349 let _building = lock.lock().unwrap_or_else(PoisonError::into_inner);
350 if let Some(v) = cell.get() {
351 return Ok(v);
352 }
353 let v = build()?;
354 Ok(cell.get_or_init(|| v))
355}
356
357impl SearchState {
358 pub fn new() -> Self {
359 Self::default()
360 }
361
362 pub fn is_warm(&self) -> bool {
364 self.insns.get().is_some()
365 }
366
367 pub fn heap_bytes(&self) -> usize {
370 self.strings.get().map_or(0, StringLocator::heap_bytes)
371 + self.types.get().map_or(0, TypeLocator::heap_bytes)
372 + self.methods.get().map_or(0, MethodLocator::heap_bytes)
373 + self.fields.get().map_or(0, FieldLocator::heap_bytes)
374 + self.insns.get().map_or(0, InsnLocator::heap_bytes)
375 + self.calls.get().map_or(0, CallGraph::heap_bytes)
376 }
377
378 pub fn call_graph(&self) -> Option<&CallGraph> {
380 self.calls.get()
381 }
382
383 pub fn set_call_graph(&self, graph: CallGraph) -> bool {
387 self.calls.set(graph).is_ok()
388 }
389}
390
391enum StateRef<'s> {
392 Owned(Box<SearchState>),
393 Shared(&'s SearchState),
394}
395
396pub struct FindRefs<'d, 'a, 's> {
398 dex: &'d Dex<'a>,
399 state: StateRef<'s>,
400}
401
402impl<'d, 'a, 's> FindRefs<'d, 'a, 's> {
403 pub fn new(dex: &'d Dex<'a>) -> Self {
405 FindRefs { dex, state: StateRef::Owned(Box::default()) }
406 }
407
408 pub fn with_state(dex: &'d Dex<'a>, state: &'s SearchState) -> Self {
410 FindRefs { dex, state: StateRef::Shared(state) }
411 }
412
413 fn state(&self) -> &SearchState {
414 match &self.state {
415 StateRef::Owned(s) => s,
416 StateRef::Shared(s) => s,
417 }
418 }
419
420 pub fn dex(&self) -> &'d Dex<'a> {
421 self.dex
422 }
423
424 pub fn string_locator(&self) -> Result<&StringLocator> {
425 let st = self.state();
426 built(&st.strings, &st.building[0], || StringLocator::build(self.dex))
427 }
428
429 pub fn type_locator(&self) -> Result<&TypeLocator> {
430 let st = self.state();
431 built(&st.types, &st.building[1], || TypeLocator::build(self.dex))
432 }
433
434 pub fn method_locator(&self) -> Result<&MethodLocator> {
435 let st = self.state();
436 built(&st.methods, &st.building[2], || MethodLocator::build(self.dex))
437 }
438
439 pub fn field_locator(&self) -> Result<&FieldLocator> {
440 let st = self.state();
441 built(&st.fields, &st.building[3], || FieldLocator::build(self.dex))
442 }
443
444 pub fn insn_locator(&self) -> Result<&InsnLocator> {
445 let st = self.state();
446 built(&st.insns, &st.building[4], || Ok(InsnLocator::build(self.dex)?))
447 }
448
449 pub fn call_graph(&self) -> Result<&CallGraph> {
451 let st = self.state();
452 built(&st.calls, &st.building[5], || Ok(CallGraph::build(self.dex.buf(), self.insn_locator()?)))
453 }
454
455 pub fn locate(&self, query: &Query) -> Result<BTreeSet<u32>> {
457 self.locate_compiled(&CompiledQuery::new(query)?)
458 }
459
460 pub fn locate_compiled(&self, query: &CompiledQuery) -> Result<BTreeSet<u32>> {
462 let buf = self.dex.buf();
463 match (query.kind(), query.parts()) {
464 (_, Parts::Pattern(None)) | (_, Parts::Literal(None)) => Ok(BTreeSet::new()),
465 (RefKind::String, Parts::Pattern(Some(re))) => Ok(self.string_locator()?.locate(buf, re)),
466 (RefKind::String, Parts::Literal(Some(f))) => Ok(self.string_locator()?.locate_literal(buf, f)),
467 (_, Parts::Relation(..)) => Ok(self.class_relation(query)?.unwrap_or_default().into_iter().collect()),
468 (RefKind::Type, Parts::Pattern(Some(re))) => Ok(self.type_locator()?.locate(buf, self.string_locator()?, re)),
469 (_, Parts::Call(c)) => self.locate_compiled(&c.target),
470 (_, Parts::Virtual(target, _)) => self.locate_compiled(target),
471 (RefKind::Method, Parts::Member(_) | Parts::MemberSet(_)) => {
472 self.method_locator()?.locate_compiled(self.dex, self.string_locator()?, self.type_locator()?, query)
473 }
474 (RefKind::Field, Parts::Member(_)) => {
475 self.field_locator()?.locate_compiled(self.dex, self.string_locator()?, self.type_locator()?, query)
476 }
477 _ => Ok(BTreeSet::new()),
479 }
480 }
481
482 pub fn cannot_match(&self, query: &CompiledQuery) -> Result<bool> {
491 if let Some(c) = query.compound() {
492 for part in &c.all {
494 if self.cannot_match(part)? {
495 return Ok(true);
496 }
497 }
498 return Ok(false);
499 }
500 match query.parts() {
501 Parts::Call(c) => self.cannot_match(&c.target),
502 Parts::Virtual(target, _) => self.cannot_match(target),
503 Parts::Relation(_, descriptor) => Ok(self.dex.find_type_idx(descriptor)?.is_none()),
505 Parts::Defined(Defined::Flags(_)) => Ok(false),
509 Parts::Defined(Defined::Annotation(a)) => Ok(self.dex.find_type_idx(&asc_dex::mutf8::encode(a))?.is_none()),
510 Parts::Defined(Defined::Members { classes, .. }) => {
511 for c in classes {
512 if self.dex.find_type_idx(&asc_dex::mutf8::encode(c))?.is_some() {
513 return Ok(false);
514 }
515 }
516 Ok(true)
517 }
518 Parts::Defined(Defined::Override(_) | Defined::In(_)) => Ok(false),
520 Parts::Member(query::CompiledMember { class: Some(query::CompiledClass::Precise(descriptor)), .. })
523 if self.dex.find_type_idx(descriptor)?.is_none() =>
524 {
525 Ok(true)
526 }
527 _ => Ok(self.locate_compiled(query)?.is_empty()),
528 }
529 }
530
531 pub fn find_ref(&self, query: &Query) -> Result<Vec<Hit>> {
533 self.find_ref_compiled(&CompiledQuery::new(query)?)
534 }
535
536 pub fn find_ref_compiled(&self, query: &CompiledQuery) -> Result<Vec<Hit>> {
543 if query.relation().is_some() || query.compound().is_some() || query.defined().is_some() {
544 return Ok(Vec::new());
545 }
546 let idxs = self.locate_compiled(query)?;
547 if idxs.is_empty() {
548 return Ok(Vec::new());
549 }
550 let mut hits = self.scan_idxs(query.kind(), &idxs)?;
551 if let Some(access) = query.query().field_access() {
552 hits.retain(|hit| hit.field_access == Some(access));
553 }
554 self.refine(query, hits)
555 }
556
557 fn refine(&self, query: &CompiledQuery, hits: Vec<Hit>) -> Result<Vec<Hit>> {
561 match query.parts() {
562 Parts::Call(c) => {
563 let hits = self.refine(&c.target, hits)?;
564 self.constrain(c, hits)
565 }
566 Parts::Virtual(_, via) => self.label_dispatch(via, hits),
567 _ => Ok(hits),
568 }
569 }
570
571 fn label_dispatch(&self, via: &BTreeMap<String, String>, mut hits: Vec<Hit>) -> Result<Vec<Hit>> {
574 if via.is_empty() {
575 return Ok(hits);
576 }
577 let mut labels: HashMap<u32, Option<&str>> = HashMap::new();
578 for hit in &mut hits {
579 let label = match labels.get(&hit.matched_idx) {
580 Some(&l) => l,
581 None => {
582 let m = self.dex.method_id(hit.matched_idx)?;
583 let name = self.dex.method_display(hit.matched_idx)?;
584 let proto = self.dex.proto_id(u32::from(m.proto_idx))?;
585 let mut sig = format!("{name}(");
586 for p in self.dex.type_list(proto.parameters_off)? {
587 sig.push_str(&self.dex.type_descriptor(u32::from(p))?);
588 }
589 sig.push(')');
590 sig.push_str(&self.dex.type_descriptor(proto.return_type_idx)?);
591 let l = via.get(&sig).or_else(|| via.get(&name)).map(String::as_str);
592 labels.insert(hit.matched_idx, l);
593 l
594 }
595 };
596 if let Some(label) = label {
597 hit.note = Some(format!("(dispatch {label})"));
598 }
599 }
600 Ok(hits)
601 }
602
603 pub fn find_many(&self, queries: &[&CompiledQuery]) -> Result<Vec<Vec<Hit>>> {
609 let mut out: Vec<Vec<Hit>> = (0..queries.len()).map(|_| Vec::new()).collect();
610 let mut located: Vec<Option<BTreeSet<u32>>> = Vec::with_capacity(queries.len());
611 for q in queries {
612 let skip = q.relation().is_some() || q.compound().is_some();
613 located.push(if skip { None } else { Some(self.locate_compiled(q)?).filter(|s| !s.is_empty()) });
614 }
615 for kind in [RefKind::String, RefKind::Type, RefKind::Method, RefKind::Field] {
616 let mine: Vec<usize> = (0..queries.len()).filter(|&i| queries[i].kind() == kind && located[i].is_some()).collect();
617 let union: BTreeSet<u32> = match mine.as_slice() {
618 [] => continue,
619 [only] => located[*only].clone().unwrap_or_default(),
620 _ => mine.iter().flat_map(|&i| located[i].iter().flatten().copied()).collect(),
621 };
622 for hit in self.scan_idxs(kind, &union)? {
623 for &i in &mine {
624 if located[i].as_ref().is_some_and(|s| s.contains(&hit.matched_idx)) {
625 let access = queries[i].query().field_access();
626 if access.is_none() || access == hit.field_access {
627 out[i].push(hit.clone());
628 }
629 }
630 }
631 }
632 }
633 for (q, hits) in queries.iter().zip(out.iter_mut()) {
634 *hits = self.refine(q, std::mem::take(hits))?;
635 }
636 Ok(out)
637 }
638
639 pub fn class_relation(&self, query: &CompiledQuery) -> Result<Option<Vec<u32>>> {
643 match query.parts() {
644 Parts::Relation(rel, descriptor) => Ok(Some(classdefs::related_classes(self.dex, *rel, descriptor)?)),
645 _ => Ok(None),
646 }
647 }
648
649 pub fn scan_idxs(&self, kind: RefKind, idxs: &BTreeSet<u32>) -> Result<Vec<Hit>> {
651 let buf = self.dex.buf();
652 let insn = self.insn_locator()?;
653 let scanner = CodeScanner::new(insn);
654 let max_idx = (buf.len() / 4) as u32;
659 let set = IdxSet::from_iter(idxs.iter().copied().filter(|&i| i < max_idx));
660 let matches = scanner.scan_instruction_candidates(buf, kind, &set);
665 let offsets: Vec<u32> = matches.iter().map(|m| m.0).collect();
666 let methods = insn.locate(buf, &offsets);
667 Ok(matches
668 .into_iter()
669 .zip(methods)
670 .map(|((offset, matched_idx), methods)| Hit {
671 offset,
672 matched_idx,
673 methods,
674 field_access: (kind == RefKind::Field).then(|| FieldAccess::from_opcode(buf[offset as usize])).flatten(),
675 note: None,
676 })
677 .collect())
678 }
679}