Skip to main content

asc_findrefs/
lib.rs

1//! Reference search over a single DEX without any global preprocessing.
2//!
3//! The pipeline mirrors the Python `findrefs` package:
4//!
5//! 1. A *locator* turns the user's query into a set of constant-pool indexes
6//!    (string, type, method or field ids). String matching runs a regex over
7//!    the raw string-data region; member lookups use small per-DEX maps.
8//! 2. The *code scanner* sweeps the raw code-item region once for opcodes that
9//!    carry an index of that kind and keeps the positions whose operand is in
10//!    the located set.
11//! 3. The *instruction locator* maps each surviving byte offset back to the
12//!    method that owns it through a dense 16-byte bucket table, and confirms
13//!    the offset really sits on an instruction boundary.
14//!
15//! String queries are a regex, or a plain substring found with SIMD `memmem`
16//! ([`Query::StringLiteral`]). `Extends` / `Implements` queries skip the
17//! pipeline: they compare `class_def_item`s with the target's type index
18//! ([`FindRefs::class_relation`]).
19//!
20//! Queries can be pre-compiled once per search ([`CompiledQuery`]) and run
21//! against many DEX files with [`FindRefs::find_ref_compiled`].
22
23#![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/// Errors from reference search.
50#[derive(Debug, thiserror::Error)]
51#[non_exhaustive]
52pub enum FindError {
53    /// The DEX is malformed where the search had to read it (an id table,
54    /// string data or `class_data_item` out of bounds or undecodable).
55    #[error(transparent)]
56    Dex(#[from] asc_dex::DexError),
57    /// A string, type, class or member-name pattern is not a valid regex.
58    #[error("invalid pattern: {0}")]
59    Pattern(#[from] regex::Error),
60    /// The query cannot be run on one DEX as it is: a predicate the engine
61    /// has to resolve first, or an ill-formed compound query.
62    #[error("invalid query: {0}")]
63    Invalid(String),
64}
65
66/// Result of reference search, failing with [`FindError`].
67pub type Result<T> = std::result::Result<T, FindError>;
68
69/// Class part of a member query.
70#[derive(Debug, Clone, PartialEq, Eq)]
71#[non_exhaustive]
72pub struct ClassQuery {
73    /// Exact Dalvik descriptor when `precise`, otherwise a fuzzy pattern.
74    pub pattern: String,
75    /// Match `pattern` as an exact descriptor rather than fuzzily.
76    pub precise: bool,
77}
78
79/// Method or field query: at least one of `class` / `name` must be set.
80#[derive(Debug, Clone, PartialEq, Eq, Default)]
81#[non_exhaustive]
82pub struct MemberQuery {
83    /// Declaring class of the member, any class when `None`.
84    pub class: Option<ClassQuery>,
85    /// Fuzzy member name (regex when `class` is absent or fuzzy, substring
86    /// when the class is precise), or the exact name when `exact_name` is set.
87    pub name: Option<String>,
88    /// Match `name` exactly (whatever the class constraint) instead of as a
89    /// regex / substring; used for "references to this member" lookups.
90    pub exact_name: bool,
91    /// Only the members of this signature: one overload of a method, or a
92    /// field of one type. Checked against the id tables, before any code is
93    /// scanned.
94    pub signature: Option<MemberSignature>,
95}
96
97impl ClassQuery {
98    /// The class of exactly this Dalvik descriptor (`Lcom/example/Main;`).
99    pub fn exact(descriptor: impl Into<String>) -> ClassQuery {
100        ClassQuery { pattern: descriptor.into(), precise: true }
101    }
102
103    /// The classes whose descriptor matches this fuzzy pattern.
104    pub fn fuzzy(pattern: impl Into<String>) -> ClassQuery {
105        ClassQuery { pattern: pattern.into(), precise: false }
106    }
107}
108
109impl MemberQuery {
110    /// This query with the class constraint `class`.
111    #[must_use]
112    pub fn with_class(self, class: Option<ClassQuery>) -> MemberQuery {
113        MemberQuery { class, ..self }
114    }
115
116    /// This query with the member name `name` (see [`MemberQuery::name`](#structfield.name)).
117    #[must_use]
118    pub fn with_name(self, name: Option<String>) -> MemberQuery {
119        MemberQuery { name, ..self }
120    }
121
122    /// This query matching the name exactly (`true`) or as a pattern.
123    #[must_use]
124    pub fn with_exact_name(self, exact_name: bool) -> MemberQuery {
125        MemberQuery { exact_name, ..self }
126    }
127
128    /// This query restricted to members of `signature`.
129    #[must_use]
130    pub fn with_signature(self, signature: Option<MemberSignature>) -> MemberQuery {
131        MemberQuery { signature, ..self }
132    }
133}
134
135/// The signature constraint of a [`MemberQuery`].
136#[derive(Debug, Clone, PartialEq, Eq)]
137pub enum MemberSignature {
138    /// A method's parameters, and its return type when one is given.
139    Method {
140        /// Parameter types, in order.
141        params: Vec<TypeMatch>,
142        /// Return type; any when `None`.
143        ret: Option<TypeMatch>,
144    },
145    /// A field's type.
146    Field(TypeMatch),
147}
148
149/// One type of a [`MemberSignature`].
150#[derive(Debug, Clone, PartialEq, Eq)]
151pub enum TypeMatch {
152    /// Exactly this descriptor: `I`, `Ljava/lang/String;`, `[B`.
153    Descriptor(String),
154    /// A class of this simple name in any package, `dims` levels of array
155    /// deep (`String[]` written without its package).
156    Simple {
157        /// Simple class name, without package or `L`/`;`.
158        name: String,
159        /// Array dimensions (`0` for the class itself).
160        dims: usize,
161    },
162}
163
164impl TypeMatch {
165    /// Whether the descriptor `descriptor` is of this type.
166    pub fn matches(&self, descriptor: &str) -> bool {
167        query::CompiledType::new(self).matches(&asc_dex::mutf8::encode(descriptor))
168    }
169}
170
171impl std::fmt::Display for TypeMatch {
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        match self {
174            TypeMatch::Descriptor(d) => f.write_str(d),
175            TypeMatch::Simple { name, dims } => write!(f, "{name}{}", "[]".repeat(*dims)),
176        }
177    }
178}
179
180impl std::fmt::Display for MemberSignature {
181    /// `(I, String)V`, `(I)` or `:I`, the way the query reads.
182    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183        match self {
184            MemberSignature::Method { params, ret } => {
185                f.write_str("(")?;
186                for (i, p) in params.iter().enumerate() {
187                    // Descriptors run together, as Dalvik writes them; a simple
188                    // name needs a separator to stay readable.
189                    let simple = |t: &TypeMatch| matches!(t, TypeMatch::Simple { .. });
190                    if i > 0 && (simple(p) || simple(&params[i - 1])) {
191                        f.write_str(", ")?;
192                    }
193                    write!(f, "{p}")?;
194                }
195                f.write_str(")")?;
196                ret.as_ref().map_or(Ok(()), |r| write!(f, "{r}"))
197            }
198            MemberSignature::Field(t) => write!(f, ":{t}"),
199        }
200    }
201}
202
203/// A reference search request.
204#[derive(Debug, Clone, PartialEq, Eq)]
205pub enum Query {
206    /// Strings matching a regex.
207    String(String),
208    /// Strings containing a plain substring (no regex syntax).
209    StringLiteral(String),
210    /// Type references whose descriptor matches a regex.
211    Type(String),
212    /// References to the methods a member query matches.
213    Method(MemberQuery),
214    /// References to any of these methods: a set read from a file, or the
215    /// methods a [`Query::Nested`] condition resolved to.
216    MethodSet(Vec<MemberQuery>),
217    /// References to the methods whose own code matches the inner query
218    /// (`method:{RULE}`). The inner query has to see every DEX, so whoever
219    /// holds them all (the engine) resolves this into a [`Query::MethodSet`]
220    /// first; compiling it as it is fails.
221    Nested(Box<Query>),
222    /// The calls that may dispatch to a method of `target` (a `MethodSet` or
223    /// `Nested`): those naming it, and those naming the same method on a
224    /// supertype of its class or on a subclass inheriting it
225    /// (`method:{RULE}` and `method:@FILE` by default). The engine resolves
226    /// it over the class hierarchy into `target` plus those methods, `via`
227    /// mapping each added one (`Lcls;->name(params)ret`, or `Lcls;->name`
228    /// for every overload) to the classes of `target` it stands for; a hit on
229    /// one carries the note `(dispatch LIST)`.
230    Virtual {
231        /// The methods whose possible dispatches are searched.
232        target: Box<Query>,
233        /// Each added method's signature (or bare `Lcls;->name`) mapped to
234        /// the `target` classes it stands for.
235        via: BTreeMap<String, String>,
236        /// Most classes of the program that may define a supertype's method
237        /// for the calls naming it to be followed (`None`: no limit). A
238        /// method every lambda or callback implements says nothing about
239        /// which one runs.
240        fan_out: Option<usize>,
241    },
242    /// References (reads and writes) to the fields a member query matches.
243    Field(MemberQuery),
244    /// Field references restricted to reads (`iget*` / `sget*`) or writes
245    /// (`iput*` / `sput*`). [`Query::Field`] keeps both.
246    FieldAccess {
247        /// The fields to look for.
248        field: MemberQuery,
249        /// Keep only reads or only writes.
250        access: FieldAccess,
251    },
252    /// The calls of a method query (`Method`, `MethodSet` or `Nested`) that
253    /// every constraint keeps (`method:X[arg1="id", ret=unused]`); see
254    /// [`CallConstraint`]. Its hits carry a note of the values judged.
255    Call {
256        /// The method query whose calls are judged.
257        target: Box<Query>,
258        /// Every constraint a call must satisfy.
259        constraints: Vec<CallConstraint>,
260    },
261    /// Classes whose superclass is exactly this descriptor (`Lcom/poc/Base;`).
262    /// Answered from `class_def_item`s, not by a code scan.
263    Extends(String),
264    /// Classes that directly implement exactly this interface descriptor.
265    Implements(String),
266    /// Methods matched by how they are declared, not by what their code
267    /// references (`flag:`, `annotation:`, `override:`, `in:{...}`). Answered
268    /// from the class definitions; the rows have no call site.
269    Defined(Defined),
270    /// Methods whose code matches every query of `all` and none of `none`
271    /// (string / type / method / field queries; at least one in `all`).
272    /// A row lists the constants each `all` query matched, in that order.
273    All {
274        /// Queries the method's code must all match.
275        all: Vec<Query>,
276        /// Queries the method's code must not match.
277        none: Vec<Query>,
278    },
279}
280
281impl Query {
282    /// The index table the query targets; class relations report
283    /// [`RefKind::Type`] (their results are classes).
284    pub fn kind(&self) -> RefKind {
285        match self {
286            Query::String(_) | Query::StringLiteral(_) => RefKind::String,
287            Query::Type(_) | Query::Extends(_) | Query::Implements(_) => RefKind::Type,
288            Query::Method(_) | Query::MethodSet(_) | Query::Nested(_) | Query::Virtual { .. } | Query::Call { .. } | Query::Defined(_) => {
289                RefKind::Method
290            }
291            Query::Field(_) | Query::FieldAccess { .. } => RefKind::Field,
292            // Rows of a compound query list constants of several kinds.
293            Query::All { all, .. } => all.first().map_or(RefKind::String, Query::kind),
294        }
295    }
296
297    /// The class relation of an `Extends` / `Implements` query.
298    pub fn relation(&self) -> Option<ClassRelation> {
299        match self {
300            Query::Extends(_) => Some(ClassRelation::Extends),
301            Query::Implements(_) => Some(ClassRelation::Implements),
302            _ => None,
303        }
304    }
305
306    /// The read/write restriction of a field query, if any.
307    pub fn field_access(&self) -> Option<FieldAccess> {
308        match self {
309            Query::FieldAccess { access, .. } => Some(*access),
310            _ => None,
311        }
312    }
313}
314
315/// A method-definition predicate ([`Query::Defined`]).
316///
317/// `Flags` and `Annotation` are matched against each method's own
318/// declaration. `Override` and `In` need the class hierarchy (which spans
319/// every DEX), so the engine resolves them into `Members` — a set of
320/// declaring classes, optionally with a name and proto — before a search;
321/// compiling them unresolved fails.
322#[derive(Debug, Clone, PartialEq, Eq)]
323pub enum Defined {
324    /// Carries all of these access flags (`flag:public,static`).
325    Flags(u32),
326    /// Annotated with this annotation type descriptor (`annotation:Lx/Y;`).
327    Annotation(String),
328    /// Declared in one of `classes` (exact descriptors), and, when set, with
329    /// this `name` and `proto` (`(params)ret`). `override:` and `in:{...}`
330    /// resolve to this.
331    Members {
332        /// Exact descriptors of the declaring classes.
333        classes: Vec<String>,
334        /// Method name, any when `None`.
335        name: Option<String>,
336        /// Method proto as `(params)ret`, any when `None`.
337        proto: Option<String>,
338    },
339    /// `override:TARGET`: methods overriding or implementing it. The engine
340    /// resolves it into `Members` over the class hierarchy.
341    Override(MemberQuery),
342    /// `in:{CLASS-RULE}`: methods declared in a matching class. The engine
343    /// resolves it into `Members`.
344    In(Box<Query>),
345}
346
347impl Defined {
348    /// A short label for the matched column of a definitional row.
349    pub fn label(&self) -> String {
350        match self {
351            Defined::Flags(_) => "flag".into(),
352            Defined::Annotation(a) => format!("@{a}"),
353            Defined::Members { name: Some(n), .. } => format!("override {n}"),
354            Defined::Members { .. } => "in".into(),
355            Defined::Override(m) => format!("override {}", m.name.as_deref().unwrap_or("")),
356            Defined::In(_) => "in".into(),
357        }
358    }
359}
360
361/// How a class relates to the target of an `Extends` / `Implements` query.
362#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
363pub enum ClassRelation {
364    /// The class's `superclass_idx` is the target.
365    Extends,
366    /// The target is in the class's `interfaces` list.
367    Implements,
368}
369
370impl ClassRelation {
371    /// The verb of the Python output line (`Lsub; extends Lbase;`).
372    pub fn verb(self) -> &'static str {
373        match self {
374            ClassRelation::Extends => "extends",
375            ClassRelation::Implements => "implements",
376        }
377    }
378}
379
380/// One instruction that references a located index.
381#[derive(Debug, Clone, PartialEq, Eq)]
382#[non_exhaustive]
383pub struct Hit {
384    /// Absolute byte offset of the instruction in the DEX.
385    pub offset: u32,
386    /// The constant-pool index the instruction references.
387    pub matched_idx: u32,
388    /// Owning method(s); `None` when the offset failed instruction verification.
389    pub methods: Option<MethodRef>,
390    /// Whether a field instruction reads or writes the field. `None` for
391    /// string, type and method references.
392    pub field_access: Option<FieldAccess>,
393    /// What the constraints of a [`Query::Call`] judged at this call site
394    /// (`[arg1="id", ret=unused]`, `?` first when one is undetermined),
395    /// and for a [`Query::Virtual`] alias, the methods it may dispatch to
396    /// (`(dispatch Lcls;)`).
397    pub note: Option<String>,
398}
399
400/// Lazily-built locator tables for one DEX buffer.
401///
402/// The tables hold offsets into the buffer they were built from (no
403/// references or pointers), so a `SearchState` must only ever be used
404/// together with that same buffer or lookups will silently return nonsense.
405/// The engine caches one per inflated DEX so repeated queries skip the builds.
406///
407/// Each table is built once: a second caller asking for a table under
408/// construction waits for it rather than building its own copy (a query
409/// arriving while a background prebuild is half way through a table gets
410/// that table sooner than by starting over).
411#[derive(Default)]
412pub struct SearchState {
413    strings: OnceLock<StringLocator>,
414    types: OnceLock<TypeLocator>,
415    methods: OnceLock<MethodLocator>,
416    fields: OnceLock<FieldLocator>,
417    insns: OnceLock<InsnLocator>,
418    calls: OnceLock<CallGraph>,
419    /// One lock per table, held while that table is built.
420    building: [Mutex<()>; 6],
421}
422
423/// `cell`'s value, built by `build` under `lock` the first time.
424fn built<'c, T>(cell: &'c OnceLock<T>, lock: &Mutex<()>, build: impl FnOnce() -> Result<T>) -> Result<&'c T> {
425    if let Some(v) = cell.get() {
426        return Ok(v);
427    }
428    // Guards no data, only the build: a build that panicked left nothing
429    // behind, so its poison means nothing.
430    let _building = lock.lock().unwrap_or_else(PoisonError::into_inner);
431    if let Some(v) = cell.get() {
432        return Ok(v);
433    }
434    let v = build()?;
435    Ok(cell.get_or_init(|| v))
436}
437
438impl SearchState {
439    /// A state with no table built yet.
440    pub fn new() -> Self {
441        Self::default()
442    }
443
444    /// Whether the instruction table has been built (a proxy for "warm").
445    pub fn is_warm(&self) -> bool {
446        self.insns.get().is_some()
447    }
448
449    /// Heap bytes of the tables built so far (0 for a fresh state). Grows as
450    /// queries of new kinds build their locators.
451    pub fn heap_bytes(&self) -> usize {
452        self.strings.get().map_or(0, StringLocator::heap_bytes)
453            + self.types.get().map_or(0, TypeLocator::heap_bytes)
454            + self.methods.get().map_or(0, MethodLocator::heap_bytes)
455            + self.fields.get().map_or(0, FieldLocator::heap_bytes)
456            + self.insns.get().map_or(0, InsnLocator::heap_bytes)
457            + self.calls.get().map_or(0, CallGraph::heap_bytes)
458    }
459
460    /// The call graph, if a query has built it.
461    pub fn call_graph(&self) -> Option<&CallGraph> {
462        self.calls.get()
463    }
464
465    /// Install a call graph loaded from elsewhere (e.g. an on-disk cache of
466    /// [`CallGraph::to_bytes`]); `false` if one is already built. It must
467    /// describe the same buffer this state is used with.
468    pub fn set_call_graph(&self, graph: CallGraph) -> bool {
469        self.calls.set(graph).is_ok()
470    }
471}
472
473enum StateRef<'s> {
474    Owned(Box<SearchState>),
475    Shared(&'s SearchState),
476}
477
478/// Reference search over one DEX.
479pub struct FindRefs<'d, 'a, 's> {
480    dex: &'d Dex<'a>,
481    state: StateRef<'s>,
482}
483
484impl<'d, 'a, 's> FindRefs<'d, 'a, 's> {
485    /// Search with private, throw-away locator tables.
486    pub fn new(dex: &'d Dex<'a>) -> Self {
487        FindRefs { dex, state: StateRef::Owned(Box::default()) }
488    }
489
490    /// Search reusing tables built earlier for the *same* buffer.
491    pub fn with_state(dex: &'d Dex<'a>, state: &'s SearchState) -> Self {
492        FindRefs { dex, state: StateRef::Shared(state) }
493    }
494
495    fn state(&self) -> &SearchState {
496        match &self.state {
497            StateRef::Owned(s) => s,
498            StateRef::Shared(s) => s,
499        }
500    }
501
502    /// The DEX being searched.
503    pub fn dex(&self) -> &'d Dex<'a> {
504        self.dex
505    }
506
507    /// The string table, built on first use.
508    ///
509    /// # Errors
510    ///
511    /// [`FindError::Dex`] when the `string_ids` or string data are out of
512    /// bounds. A failed build is not cached; the next call retries.
513    pub fn string_locator(&self) -> Result<&StringLocator> {
514        let st = self.state();
515        built(&st.strings, &st.building[0], || StringLocator::build(self.dex))
516    }
517
518    /// The type table, built on first use.
519    ///
520    /// # Errors
521    ///
522    /// [`FindError::Dex`] when the `type_ids` table is out of bounds.
523    pub fn type_locator(&self) -> Result<&TypeLocator> {
524        let st = self.state();
525        built(&st.types, &st.building[1], || TypeLocator::build(self.dex))
526    }
527
528    /// The method-id tables, built on first use.
529    ///
530    /// # Errors
531    ///
532    /// [`FindError::Dex`] when the `method_ids` table is out of bounds.
533    pub fn method_locator(&self) -> Result<&MethodLocator> {
534        let st = self.state();
535        built(&st.methods, &st.building[2], || MethodLocator::build(self.dex))
536    }
537
538    /// The field-id tables, built on first use.
539    ///
540    /// # Errors
541    ///
542    /// [`FindError::Dex`] when the `field_ids` table is out of bounds.
543    pub fn field_locator(&self) -> Result<&FieldLocator> {
544        let st = self.state();
545        built(&st.fields, &st.building[3], || FieldLocator::build(self.dex))
546    }
547
548    /// The instruction-to-method table, built on first use.
549    ///
550    /// # Errors
551    ///
552    /// [`FindError::Dex`] as for [`InsnLocator::build`].
553    pub fn insn_locator(&self) -> Result<&InsnLocator> {
554        let st = self.state();
555        built(&st.insns, &st.building[4], || Ok(InsnLocator::build(self.dex)?))
556    }
557
558    /// The DEX's call graph, built on first use (the instruction table too).
559    ///
560    /// # Errors
561    ///
562    /// The errors of [`FindRefs::insn_locator`].
563    pub fn call_graph(&self) -> Result<&CallGraph> {
564        let st = self.state();
565        built(&st.calls, &st.building[5], || Ok(CallGraph::build(self.dex.buf(), self.insn_locator()?)))
566    }
567
568    /// Resolve the query to constant-pool indexes (compiling its patterns).
569    ///
570    /// # Errors
571    ///
572    /// The errors of [`CompiledQuery::new`] and [`FindRefs::locate_compiled`].
573    pub fn locate(&self, query: &Query) -> Result<BTreeSet<u32>> {
574        self.locate_compiled(&CompiledQuery::new(query)?)
575    }
576
577    /// Resolve a pre-compiled query to constant-pool indexes.
578    ///
579    /// # Errors
580    ///
581    /// [`FindError::Dex`] when a locator table the query needs cannot be
582    /// built or an id it resolves through is out of bounds.
583    pub fn locate_compiled(&self, query: &CompiledQuery) -> Result<BTreeSet<u32>> {
584        let buf = self.dex.buf();
585        match (query.kind(), query.parts()) {
586            (_, Parts::Pattern(None)) | (_, Parts::Literal(None)) => Ok(BTreeSet::new()),
587            (RefKind::String, Parts::Pattern(Some(re))) => Ok(self.string_locator()?.locate(buf, re)),
588            (RefKind::String, Parts::Literal(Some(f))) => Ok(self.string_locator()?.locate_literal(buf, f)),
589            (_, Parts::Relation(..)) => Ok(self.class_relation(query)?.unwrap_or_default().into_iter().collect()),
590            (RefKind::Type, Parts::Pattern(Some(re))) => Ok(self.type_locator()?.locate(buf, self.string_locator()?, re)),
591            (_, Parts::Call(c)) => self.locate_compiled(&c.target),
592            (_, Parts::Virtual(target, _)) => self.locate_compiled(target),
593            (RefKind::Method, Parts::Member(_) | Parts::MemberSet(_)) => {
594                self.method_locator()?.locate_compiled(self.dex, self.string_locator()?, self.type_locator()?, query)
595            }
596            (RefKind::Field, Parts::Member(_)) => {
597                self.field_locator()?.locate_compiled(self.dex, self.string_locator()?, self.type_locator()?, query)
598            }
599            // `CompiledQuery::new` always pairs the kind with the matching parts.
600            _ => Ok(BTreeSet::new()),
601        }
602    }
603
604    /// Whether `query` certainly has no hit in this DEX, judged from the id
605    /// tables and the string data alone: no code is read, so this also works
606    /// on a DEX prefix that holds those (an error then means the prefix was
607    /// not enough, and the caller has to search the whole DEX).
608    ///
609    /// `true` exactly when [`FindRefs::find_ref_compiled`] (or, for a class
610    /// relation, [`FindRefs::class_relation`]; for a compound query, its rows)
611    /// could only come out empty.
612    ///
613    /// # Errors
614    ///
615    /// [`FindError::Dex`] when the type table or another table the judgement
616    /// reads is malformed or missing from the (possibly partial) buffer.
617    pub fn cannot_match(&self, query: &CompiledQuery) -> Result<bool> {
618        if let Some(c) = query.compound() {
619            // Rows need a hit of every `all` part in one method.
620            for part in &c.all {
621                if self.cannot_match(part)? {
622                    return Ok(true);
623                }
624            }
625            return Ok(false);
626        }
627        match query.parts() {
628            Parts::Call(c) => self.cannot_match(&c.target),
629            Parts::Virtual(target, _) => self.cannot_match(target),
630            // A class nobody references cannot be extended or implemented.
631            Parts::Relation(_, descriptor) => Ok(self.dex.find_type_idx(descriptor)?.is_none()),
632            // A definition predicate reads no code, so locating it yields
633            // nothing: rule it out only from the type table (flags could be
634            // on any method, so they never skip).
635            Parts::Defined(Defined::Flags(_)) => Ok(false),
636            Parts::Defined(Defined::Annotation(a)) => Ok(self.dex.find_type_idx(&asc_dex::mutf8::encode(a))?.is_none()),
637            Parts::Defined(Defined::Members { classes, .. }) => {
638                for c in classes {
639                    if self.dex.find_type_idx(&asc_dex::mutf8::encode(c))?.is_some() {
640                        return Ok(false);
641                    }
642                }
643                Ok(true)
644            }
645            // Unresolved predicates never reach a search.
646            Parts::Defined(Defined::Override(_) | Defined::In(_)) => Ok(false),
647            // Nor can its members be: one binary search, before the member
648            // tables are built (most DEX files do not reference the class).
649            Parts::Member(query::CompiledMember { class: Some(query::CompiledClass::Precise(descriptor)), .. })
650                if self.dex.find_type_idx(descriptor)?.is_none() =>
651            {
652                Ok(true)
653            }
654            _ => Ok(self.locate_compiled(query)?.is_empty()),
655        }
656    }
657
658    /// Full search: locate indexes, scan code, resolve owning methods.
659    ///
660    /// # Errors
661    ///
662    /// The errors of [`CompiledQuery::new`] and [`FindRefs::find_ref_compiled`].
663    pub fn find_ref(&self, query: &Query) -> Result<Vec<Hit>> {
664        self.find_ref_compiled(&CompiledQuery::new(query)?)
665    }
666
667    /// Full search with a query compiled once (typically shared by every DEX
668    /// of an APK).
669    ///
670    /// Class relations (`Extends` / `Implements`) reference no code and give
671    /// no hits here; [`FindRefs::class_relation`] answers them. Compound
672    /// queries give none either: run their parts with [`FindRefs::find_many`].
673    ///
674    /// # Errors
675    ///
676    /// [`FindError::Dex`] when the DEX is malformed where the search reads it
677    /// (locator tables, `class_data_item`s, or the ids a note is built from).
678    pub fn find_ref_compiled(&self, query: &CompiledQuery) -> Result<Vec<Hit>> {
679        if query.relation().is_some() || query.compound().is_some() || query.defined().is_some() {
680            return Ok(Vec::new());
681        }
682        let idxs = self.locate_compiled(query)?;
683        if idxs.is_empty() {
684            return Ok(Vec::new());
685        }
686        let mut hits = self.scan_idxs(query.kind(), &idxs)?;
687        if let Some(access) = query.query().field_access() {
688            hits.retain(|hit| hit.field_access == Some(access));
689        }
690        self.refine(query, hits)
691    }
692
693    /// The hits of `query`'s located indexes that it keeps, with their notes:
694    /// a [`Query::Virtual`] labels the calls it reached by dispatch, a
695    /// [`Query::Call`] judges the values of its target's calls.
696    fn refine(&self, query: &CompiledQuery, hits: Vec<Hit>) -> Result<Vec<Hit>> {
697        match query.parts() {
698            Parts::Call(c) => {
699                let hits = self.refine(&c.target, hits)?;
700                self.constrain(c, hits)
701            }
702            Parts::Virtual(_, via) => self.label_dispatch(via, hits),
703            _ => Ok(hits),
704        }
705    }
706
707    /// `hits` with a `(dispatch LIST)` note on each call [`Query::Virtual`]
708    /// added for dispatch.
709    fn label_dispatch(&self, via: &BTreeMap<String, String>, mut hits: Vec<Hit>) -> Result<Vec<Hit>> {
710        if via.is_empty() {
711            return Ok(hits);
712        }
713        let mut labels: HashMap<u32, Option<&str>> = HashMap::new();
714        for hit in &mut hits {
715            let label = match labels.get(&hit.matched_idx) {
716                Some(&l) => l,
717                None => {
718                    let m = self.dex.method_id(hit.matched_idx)?;
719                    let name = self.dex.method_display(hit.matched_idx)?;
720                    let proto = self.dex.proto_id(u32::from(m.proto_idx))?;
721                    let mut sig = format!("{name}(");
722                    for p in self.dex.type_list(proto.parameters_off)? {
723                        sig.push_str(&self.dex.type_descriptor(u32::from(p))?);
724                    }
725                    sig.push(')');
726                    sig.push_str(&self.dex.type_descriptor(proto.return_type_idx)?);
727                    let l = via.get(&sig).or_else(|| via.get(&name)).map(String::as_str);
728                    labels.insert(hit.matched_idx, l);
729                    l
730                }
731            };
732            if let Some(label) = label {
733                hit.note = Some(format!("(dispatch {label})"));
734            }
735        }
736        Ok(hits)
737    }
738
739    /// Hits of several queries with one code scan per index kind: the
740    /// located sets of the queries of a kind are merged, scanned once, and
741    /// each hit is handed to every query whose set holds its index. Gives,
742    /// per query, what [`FindRefs::find_ref_compiled`] would (class relations
743    /// and compound queries: nothing).
744    ///
745    /// # Errors
746    ///
747    /// The errors of [`FindRefs::find_ref_compiled`], for the first query
748    /// that fails.
749    pub fn find_many(&self, queries: &[&CompiledQuery]) -> Result<Vec<Vec<Hit>>> {
750        let mut out: Vec<Vec<Hit>> = (0..queries.len()).map(|_| Vec::new()).collect();
751        let mut located: Vec<Option<BTreeSet<u32>>> = Vec::with_capacity(queries.len());
752        for q in queries {
753            let skip = q.relation().is_some() || q.compound().is_some();
754            located.push(if skip { None } else { Some(self.locate_compiled(q)?).filter(|s| !s.is_empty()) });
755        }
756        for kind in [RefKind::String, RefKind::Type, RefKind::Method, RefKind::Field] {
757            let mine: Vec<usize> = (0..queries.len()).filter(|&i| queries[i].kind() == kind && located[i].is_some()).collect();
758            let union: BTreeSet<u32> = match mine.as_slice() {
759                [] => continue,
760                [only] => located[*only].clone().unwrap_or_default(),
761                _ => mine.iter().flat_map(|&i| located[i].iter().flatten().copied()).collect(),
762            };
763            for hit in self.scan_idxs(kind, &union)? {
764                for &i in &mine {
765                    if located[i].as_ref().is_some_and(|s| s.contains(&hit.matched_idx)) {
766                        let access = queries[i].query().field_access();
767                        if access.is_none() || access == hit.field_access {
768                            out[i].push(hit.clone());
769                        }
770                    }
771                }
772            }
773        }
774        for (q, hits) in queries.iter().zip(out.iter_mut()) {
775            *hits = self.refine(q, std::mem::take(hits))?;
776        }
777        Ok(out)
778    }
779
780    /// Answer an `Extends` / `Implements` query: the `class_idx` of every
781    /// class definition that extends / implements the target, in
782    /// `class_defs` order. `None` for any other query.
783    ///
784    /// # Errors
785    ///
786    /// [`FindError::Dex`] when a `class_def_item` or its interface list is
787    /// malformed.
788    pub fn class_relation(&self, query: &CompiledQuery) -> Result<Option<Vec<u32>>> {
789        match query.parts() {
790            Parts::Relation(rel, descriptor) => Ok(Some(classdefs::related_classes(self.dex, *rel, descriptor)?)),
791            _ => Ok(None),
792        }
793    }
794
795    /// Scan for already-located indexes of `kind`.
796    ///
797    /// # Errors
798    ///
799    /// The errors of [`FindRefs::insn_locator`].
800    pub fn scan_idxs(&self, kind: RefKind, idxs: &BTreeSet<u32>) -> Result<Vec<Hit>> {
801        let buf = self.dex.buf();
802        let insn = self.insn_locator()?;
803        let scanner = CodeScanner::new(insn);
804        // Every id table has at least 4-byte entries, so an index at or above
805        // `buf.len() / 4` cannot name anything in this DEX; dropping such
806        // targets also bounds the dense bitset for callers passing arbitrary
807        // indexes (`u32::MAX` would otherwise allocate 512 MiB).
808        let max_idx = (buf.len() / 4) as u32;
809        let set = IdxSet::from_iter(idxs.iter().copied().filter(|&i| i < max_idx));
810        // CodeScanner's public raw API retains unverified byte-pattern
811        // candidates. FindRefs asks its internal instruction-aware variant to
812        // omit odd offsets before allocating candidate storage; no DEX opcode
813        // can start there.
814        let matches = scanner.scan_instruction_candidates(buf, kind, &set);
815        let offsets: Vec<u32> = matches.iter().map(|m| m.0).collect();
816        let methods = insn.locate(buf, &offsets);
817        Ok(matches
818            .into_iter()
819            .zip(methods)
820            .map(|((offset, matched_idx), methods)| Hit {
821                offset,
822                matched_idx,
823                methods,
824                field_access: (kind == RefKind::Field).then(|| FieldAccess::from_opcode(buf[offset as usize])).flatten(),
825                note: None,
826            })
827            .collect())
828    }
829}