Skip to main content

asc_engine/
lib.rs

1//! The asc-rs engine: stateless, zero-preprocessing queries against an APK.
2//!
3//! * [`Engine::get_class_dex`] inflates the DEX entries in parallel (smallest
4//!   first) and stops as soon as one defines the requested class.
5//! * [`Engine::find_refs`] runs a reference search over every DEX entry in
6//!   parallel and streams results back as each DEX completes.
7//! * [`Engine::get_class`] rebuilds a minimal DEX containing the target and
8//!   its `$` nested classes, then renders the target as smali.
9//! * [`Engine::list_classes`] lists the classes of every DEX entry in order.
10//! * [`Engine::manifest_xml`] decodes the binary `AndroidManifest.xml`.
11//!
12//! Optionally the engine keeps inflated DEX entries in an on-disk cache
13//! (shared across processes) and/or an in-memory LRU together with the locator
14//! tables built for them, so repeated queries skip inflation and table builds.
15//! The disk cache also keeps the locator tables (L2, mapped by later
16//! processes instead of rebuilt) and `find_refs` results (L3); see
17//! [`asc_cache`] for the tiers.
18//! Disk writes happen on a background thread so a query never waits on them;
19//! [`Engine::flush_cache_writes`] waits for them to land and [`Engine::close_cache_writes`]
20//! (run on drop) additionally stops queueing new ones.
21
22mod budget;
23mod callgraph;
24mod classes;
25mod compact;
26mod compose;
27pub mod condition;
28pub mod config;
29mod defs;
30mod hierarchy;
31pub mod hints;
32mod manifest;
33pub mod names;
34mod prebuild;
35mod resident;
36mod resources;
37mod rules;
38mod skip;
39mod sweep;
40pub mod text;
41mod tiers;
42
43use asc_apk::{Apk, Entry, Limits};
44use asc_cache::{ApkKey, DexData, DiskCache, LockExt, MemoryCache, Tier};
45use asc_dex::Dex;
46use asc_dex::container::{dex041_logical_offsets, logical_dex_buffers, logical_dex_iter};
47pub use asc_findrefs::CompiledQuery;
48use asc_findrefs::{FindRefs, Hit, Query, RefKind, SearchState};
49use asc_rebuild::Rebuilt;
50use std::borrow::Cow;
51use std::collections::{BTreeMap, HashSet};
52use std::path::{Path, PathBuf};
53use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
54use std::sync::{Arc, Condvar, Mutex, MutexGuard, OnceLock, Weak, mpsc};
55use std::time::{Duration, Instant, SystemTime};
56
57pub use asc_apk::Limits as InflateLimits;
58pub use asc_apk::archive_of;
59pub use asc_arsc;
60pub use asc_cache;
61pub use asc_cache::{MAPPED_CHARGE_DIVISOR, Threads, auto_threads, auto_threads_for};
62pub use asc_findrefs::{ClassQuery, ClassRelation, FieldAccess, MemberQuery};
63pub use asc_smali::Options as SmaliOptions;
64pub use budget::{Device, MemoryPolicy, available_memory, total_memory};
65pub use callgraph::{CallQuery, CallRoots, CallRow, CallTree, Direction};
66pub use classes::{ClassListing, DexClasses, EntryClasses, dex_classes};
67pub use compact::TextRows;
68pub use compose::{DEFAULT_FAN_OUT, MatchingMethods};
69pub use hierarchy::{HierarchyQuery, HierarchyRow, HierarchyTree, Lineage};
70pub use manifest::{Application, Component, DataSpec, IntentFilter, ManifestSummary, Permission};
71pub use prebuild::PrebuildReport;
72pub use resources::{ResourceField, ResourceRow};
73pub use rules::compile_rules;
74pub use sweep::{Claim, DexDedup};
75pub use tiers::Busy;
76
77/// Every error an [`Engine`] query can return.
78#[derive(Debug, thiserror::Error)]
79#[non_exhaustive]
80pub enum EngineError {
81    /// The APK (zip container) could not be opened or an entry not read.
82    #[error(transparent)]
83    Apk(#[from] asc_apk::ApkError),
84    /// A DEX entry is malformed.
85    #[error(transparent)]
86    Dex(#[from] asc_dex::DexError),
87    /// The reference search rejected the query or failed on a DEX.
88    #[error(transparent)]
89    Find(#[from] asc_findrefs::FindError),
90    /// The minimal DEX of a class could not be rebuilt.
91    #[error(transparent)]
92    Rebuild(#[from] asc_rebuild::RebuildError),
93    /// A class could not be disassembled to smali.
94    #[error(transparent)]
95    Smali(#[from] asc_smali::SmaliError),
96    /// A binary XML document (such as the manifest) could not be decoded.
97    #[error(transparent)]
98    Axml(#[from] asc_axml::AxmlError),
99    /// `resources.arsc` could not be decoded.
100    #[error(transparent)]
101    Arsc(#[from] asc_arsc::ArscError),
102    /// No DEX of the APK defines the class.
103    ///
104    /// `hint` (empty, or starting with `:` / `;`) says why and what to try.
105    #[error("class {class} not found in APK{hint}")]
106    ClassNotFound {
107        /// The class descriptor that was looked up.
108        class: String,
109        /// Appended to the message as is.
110        hint: String,
111    },
112    /// The query is invalid or names something the APK does not have; the
113    /// message is meant for the user.
114    #[error("{0}")]
115    Query(String),
116    /// A DEX entry is malformed in a way only noticed past its header.
117    #[error("{dex}: {reason}")]
118    BadDex {
119        /// The entry (or logical DEX) name.
120        dex: String,
121        /// What is wrong with it.
122        reason: String,
123    },
124    /// A worker thread panicked; the payload names what it was processing.
125    #[error("worker panicked while processing {0}")]
126    WorkerPanic(String),
127}
128
129/// The result type of the engine's fallible operations.
130pub type Result<T> = std::result::Result<T, EngineError>;
131
132/// Shared inflated bytes of a DEX (possibly one logical DEX of a container).
133pub type DexBytes = Arc<DexData>;
134
135/// A DEX (possibly one logical DEX of a 041 container) that defines a class.
136#[derive(Clone)]
137#[non_exhaustive]
138pub struct ClassHit {
139    /// The DEX entry name (`classes2.dex`, or a logical DEX of a container).
140    pub dex_name: String,
141    /// Its inflated bytes.
142    pub dex: DexBytes,
143    /// The descriptor of the class found: the one looked up, or the DEX
144    /// class a Java / Kotlin spelling of it stands for (`Lcom/a/Outer/Inner;`
145    /// → `Lcom/a/Outer$Inner;`) when only that exists.
146    pub descriptor: String,
147    /// In a program, the DEX files of other archives that define the class
148    /// too (the hit is the first archive's, in program order).
149    pub also_defined_in: Vec<String>,
150}
151
152/// Result of a full `getclass`.
153#[derive(Clone)]
154#[non_exhaustive]
155pub struct ClassSource {
156    /// The DEX entry that defines the class.
157    pub dex_name: String,
158    /// The class rendered ([`ClassHit::descriptor`]: differs from the one
159    /// asked for when that was a Java / Kotlin spelling of it).
160    pub descriptor: String,
161    /// The class rendered as smali.
162    pub source: String,
163    /// The minimal DEX containing the class and its `$` nested classes,
164    /// or `None` when the class could not be
165    /// rebuilt (see `warning`) and `source` was disassembled from the
166    /// original DEX instead.
167    pub rebuilt: Option<Rebuilt>,
168    /// Why the rebuild was skipped, or which nested classes it left out
169    /// (see [`Rebuilt::dropped`]), when either happened.
170    pub warning: Option<String>,
171    /// Microseconds spent locating and inflating the DEX.
172    pub scan_us: u128,
173    /// See [`ClassHit::also_defined_in`].
174    pub also_defined_in: Vec<String>,
175}
176
177/// One aggregated reference result: a method and what it references.
178///
179/// For a class relation query (`Extends` / `Implements`) the row is a class
180/// instead: `method_idx` is its type index, `method` and `signature` its
181/// descriptor, `matched` the target descriptor and `relation` is set.
182#[derive(Debug, Clone, PartialEq, Eq)]
183#[non_exhaustive]
184pub struct RefResult {
185    /// The DEX entry the row comes from.
186    pub dex_name: String,
187    /// The method's index in that DEX's method_ids table.
188    pub method_idx: u32,
189    /// `Lcls;->name`
190    pub method: String,
191    /// `Lcls;->name(params)ret`
192    pub signature: String,
193    /// The matched constants (strings, descriptors, members).
194    pub matched: Vec<Matched>,
195    /// Set for the rows of a class relation query.
196    pub relation: Option<ClassRelation>,
197    /// The referencing instructions, in offset order (none for class relations).
198    pub sites: Vec<RefSite>,
199    /// Smali around the sites, when asked for ([`RefContext`]).
200    pub context: Option<String>,
201}
202
203/// One instruction behind a [`RefResult`].
204#[derive(Debug, Clone, PartialEq, Eq)]
205#[non_exhaustive]
206pub struct RefSite {
207    /// Absolute byte offset of the instruction in the DEX.
208    pub offset: u32,
209    /// Index into the row's `matched` of the constant it references.
210    pub matched: usize,
211    /// Whether this site reads or writes a field. `None` for other queries.
212    pub field_access: Option<FieldAccess>,
213    /// Code units from the start of the method's code (filled with a context).
214    pub pc: Option<u32>,
215    /// Source line from the debug info (filled with a context, when known).
216    pub line: Option<u32>,
217}
218
219/// How much code to return with each `findrefs` row.
220#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
221pub enum RefContext {
222    /// The row alone (the Python output).
223    #[default]
224    None,
225    /// The referencing instructions with this many smali lines around each.
226    Lines(usize),
227    /// The whole referencing method.
228    Method,
229}
230
231/// Options of [`Engine::find_refs_opts`].
232#[derive(Debug, Clone, Copy)]
233#[non_exhaustive]
234pub struct FindOptions {
235    /// One row per method (with every matched constant) instead of one per instruction.
236    pub aggregate: bool,
237    /// How much code to return with each row.
238    pub context: RefContext,
239}
240
241impl FindOptions {
242    /// The options with [`aggregate`](#structfield.aggregate) set.
243    #[must_use]
244    pub fn aggregate(self, aggregate: bool) -> Self {
245        FindOptions { aggregate, ..self }
246    }
247
248    /// The options with [`context`](#structfield.context) set.
249    #[must_use]
250    pub fn context(self, context: RefContext) -> Self {
251        FindOptions { context, ..self }
252    }
253}
254
255impl Default for FindOptions {
256    fn default() -> Self {
257        FindOptions { aggregate: true, context: RefContext::None }
258    }
259}
260
261/// A constant a [`RefResult`] matched.
262#[derive(Debug, Clone, PartialEq, Eq)]
263#[non_exhaustive]
264pub enum Matched {
265    /// The value of a string constant.
266    String(String),
267    /// A type descriptor, `Lcls;->name(params)ret` or `Lcls;->name:type`.
268    Symbol(String),
269    /// A field symbol and how its instruction uses it.
270    Field {
271        /// `Lcls;->name:type`.
272        symbol: String,
273        /// Whether the instruction reads or writes it.
274        access: FieldAccess,
275    },
276}
277
278impl Matched {
279    /// The string value or the symbol, unquoted.
280    pub fn as_str(&self) -> &str {
281        match self {
282            Matched::String(s) | Matched::Symbol(s) | Matched::Field { symbol: s, .. } => s,
283        }
284    }
285
286    /// The field access represented by this match.
287    pub fn field_access(&self) -> Option<FieldAccess> {
288        match self {
289            Matched::Field { access, .. } => Some(*access),
290            _ => None,
291        }
292    }
293}
294
295/// Strings are quoted and escaped (one line, readable as JSON); symbols are
296/// written as they are, since descriptors never hold `, ` or a newline.
297impl std::fmt::Display for Matched {
298    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
299        match self {
300            Matched::String(s) => f.write_str(&text::quoted(s)),
301            Matched::Symbol(s) => f.write_str(s),
302            Matched::Field { symbol, access } => write!(f, "{symbol} [{}]", access.as_str()),
303        }
304    }
305}
306
307impl RefResult {
308    /// `dex | Lcls;->m(params)ret | "a", Lx;->f:I`, or `dex | Lsub; | extends
309    /// Lbase;` for a class relation: three `|` columns, one line.
310    pub fn line(&self) -> String {
311        let matched = self.matched.iter().map(Matched::to_string).collect::<Vec<_>>().join(", ");
312        match self.relation {
313            Some(rel) => format!("{} | {} | {} {matched}", self.dex_name, self.signature, rel.verb()),
314            None => format!("{} | {} | {matched}", self.dex_name, self.signature),
315        }
316    }
317
318    /// The row followed by its context lines (just the row without one).
319    pub fn text(&self) -> String {
320        let mut out = self.line();
321        if let Some(c) = &self.context {
322            out.push('\n');
323            out.push_str(c.trim_end_matches('\n'));
324        }
325        out
326    }
327}
328
329/// Results for one DEX, delivered as soon as its search completes.
330#[derive(Debug, Clone)]
331#[non_exhaustive]
332pub struct RefBatch {
333    /// The APK entry searched.
334    pub entry_name: String,
335    /// The rows found in it.
336    pub results: Vec<RefResult>,
337    /// The rows as text instead, for [`Engine::find_refs_text`] (`results`
338    /// is then empty).
339    pub text: Option<TextRows>,
340    /// Microseconds spent getting the DEX bytes (inflating or loading them).
341    pub inflate_us: u128,
342    /// Microseconds spent searching them.
343    pub search_us: u128,
344    /// Where the bytes came from: `inflate`, `memory`, `disk` or `stored`.
345    pub source: &'static str,
346    /// Why (part of) the entry could not be searched; `;`-separated for the
347    /// logical DEX files of a container.
348    pub error: Option<String>,
349}
350
351impl RefBatch {
352    /// The number of rows, in `results` or `text`.
353    pub fn row_count(&self) -> usize {
354        self.results.len() + self.text.as_ref().map_or(0, TextRows::len)
355    }
356
357    /// Record a failure of one logical DEX of this entry. A container holds
358    /// several, so failures accumulate instead of overwriting each other.
359    /// Consumers prefix `entry_name`; only a container's logical DEX (whose
360    /// name differs) adds its own.
361    fn note_error(&mut self, name: &str, e: &str) {
362        let msg = if name == self.entry_name { e.to_string() } else { format!("{name}: {e}") };
363        self.error = Some(match self.error.take() {
364            Some(prev) => format!("{prev}; {msg}"),
365            None => msg,
366        });
367    }
368}
369
370/// Summary of a whole `find_refs` run.
371#[derive(Debug, Clone, Default)]
372#[non_exhaustive]
373pub struct RefSummary {
374    /// DEX entries searched.
375    pub entries: usize,
376    /// Rows over all entries.
377    pub total_results: usize,
378    /// Wall-clock microseconds of the whole run.
379    pub total_us: u128,
380    /// What the search left out on purpose, e.g. a supertype method a
381    /// dispatch set did not follow; one line each, for `#` lines or hints.
382    pub notes: Vec<String>,
383}
384
385/// Cache configuration for an [`Engine`].
386#[non_exhaustive]
387pub struct CacheConfig {
388    /// In-memory LRU capacity for inflated entries (0 disables it).
389    pub memory_bytes: usize,
390    /// On-disk cache shared across processes.
391    pub disk: Option<DiskCache>,
392    /// Treat `memory_bytes` as a starting point and keep resizing the LRU to
393    /// this process's share of one budget for every resident host the user
394    /// runs ([`MemoryPolicy::share`]), instead of a fixed capacity.
395    pub memory_auto: Option<MemoryPolicy>,
396    /// Heap queued disk writes may hold before new ones are deferred
397    /// (an L1 entry is inflated again when its turn comes, L2 tables are
398    /// dropped); a memory cache's budget lowers it further.
399    pub write_queue_bytes: usize,
400    /// Disk writer threads (`None`: [`CacheConfig::auto_writers`]).
401    pub writers: Option<usize>,
402}
403
404impl CacheConfig {
405    /// The configuration with [`memory_bytes`](#structfield.memory_bytes) set.
406    #[must_use]
407    pub fn memory_bytes(self, memory_bytes: usize) -> Self {
408        CacheConfig { memory_bytes, ..self }
409    }
410
411    /// The configuration with [`disk`](#structfield.disk) set.
412    #[must_use]
413    pub fn disk(self, disk: Option<DiskCache>) -> Self {
414        CacheConfig { disk, ..self }
415    }
416
417    /// The configuration with [`memory_auto`](#structfield.memory_auto) set.
418    #[must_use]
419    pub fn memory_auto(self, memory_auto: Option<MemoryPolicy>) -> Self {
420        CacheConfig { memory_auto, ..self }
421    }
422
423    /// The configuration with [`write_queue_bytes`](#structfield.write_queue_bytes) set.
424    #[must_use]
425    pub fn write_queue_bytes(self, write_queue_bytes: usize) -> Self {
426        CacheConfig { write_queue_bytes, ..self }
427    }
428
429    /// The configuration with [`writers`](#structfield.writers) set.
430    #[must_use]
431    pub fn writers(self, writers: Option<usize>) -> Self {
432        CacheConfig { writers, ..self }
433    }
434}
435
436impl CacheConfig {
437    /// Default of [`CacheConfig::write_queue_bytes`]. A cold multi-APK sweep
438    /// inflates far faster than the writers persist, and each queued job
439    /// keeps its DEX bytes or tables alive until written; without a bound the
440    /// heap grows with the backlog (tens of GB over a phone's APKs).
441    pub const WRITE_QUEUE: usize = 512 << 20;
442    /// Most automatic writer threads: writing is I/O, but a deferred L1
443    /// write inflates its entry again first.
444    pub const MAX_AUTO_WRITERS: usize = 4;
445
446    /// Disk cache in the default directory, no in-memory cache (CLI default).
447    pub fn disk_default() -> CacheConfig {
448        CacheConfig { disk: DiskCache::open_default(), ..CacheConfig::default() }
449    }
450
451    /// Disk cache plus an in-memory LRU of a fixed capacity.
452    pub fn resident(memory_bytes: usize) -> CacheConfig {
453        CacheConfig { memory_bytes, disk: DiskCache::open_default(), ..CacheConfig::default() }
454    }
455
456    /// One writer per 8 CPUs, 1..=[`CacheConfig::MAX_AUTO_WRITERS`].
457    pub fn auto_writers(cpus: usize) -> usize {
458        (cpus / 8).clamp(1, CacheConfig::MAX_AUTO_WRITERS)
459    }
460}
461
462impl Default for CacheConfig {
463    fn default() -> CacheConfig {
464        CacheConfig { memory_bytes: 0, disk: None, memory_auto: None, write_queue_bytes: CacheConfig::WRITE_QUEUE, writers: None }
465    }
466}
467
468/// A cache several [`Engine`]s share, so a process serving many APKs keeps
469/// one memory budget (the LRU is keyed by APK) and one disk-writer thread.
470/// Engines attached with [`Engine::with_shared_cache`] leave queued disk
471/// writes alone when dropped; dropping the `SharedCache` flushes them and
472/// stops queueing (as dropping a private engine does).
473pub struct SharedCache(Option<Arc<EngineCache>>);
474
475impl SharedCache {
476    /// Build the cache [`Engine::with_cache`] would; warnings (an unusable
477    /// disk cache directory) go to `logger`.
478    pub fn new(config: CacheConfig, logger: impl Fn(&str) + Send + Sync + 'static) -> SharedCache {
479        let mut holder = Engine::new("").with_logger(logger).with_cache(config);
480        SharedCache(holder.cache.take())
481    }
482
483    /// Bytes currently held by the shared in-memory cache.
484    pub fn memory_bytes(&self) -> usize {
485        self.0.as_ref().and_then(|c| c.memory.as_ref()).map(|m| lock_memory(m).used_bytes()).unwrap_or(0)
486    }
487
488    /// Budget of the shared in-memory cache (0 when it is disabled).
489    pub fn memory_capacity(&self) -> usize {
490        self.0.as_ref().and_then(|c| c.memory.as_ref()).map(|m| lock_memory(m).capacity_bytes()).unwrap_or(0)
491    }
492
493    /// Bytes counted against the memory budget besides the cache's entries
494    /// ([`SharedCache::set_reserved`]).
495    pub fn memory_reserved(&self) -> usize {
496        self.0.as_ref().and_then(|c| c.memory.as_ref()).map(|m| lock_memory(m).reserved_bytes()).unwrap_or(0)
497    }
498
499    /// Count `bytes` the host holds outside the cache (the class graphs and
500    /// resource tables of its engines, [`Engine::resident_bytes`]) against
501    /// the memory budget: the cache evicts entries to make room for them.
502    pub fn set_reserved(&self, bytes: usize) {
503        if let Some(m) = self.0.as_ref().and_then(|c| c.memory.as_ref()) {
504            lock_memory(m).set_reserved(bytes);
505        }
506    }
507
508    /// Block until every queued disk-cache write has landed.
509    pub fn flush_writes(&self) {
510        if let Some(c) = &self.0 {
511            c.flush(false);
512        }
513    }
514}
515
516impl Drop for SharedCache {
517    fn drop(&mut self) {
518        if let Some(c) = &self.0 {
519            c.flush(true);
520        }
521    }
522}
523
524/// One logical DEX of a cached entry with its lazily built locator tables.
525/// The tables are shared with the entry that replaces this one once its
526/// bytes are remapped from the disk cache (same bytes, so still valid).
527#[non_exhaustive]
528pub struct LogicalDex {
529    /// The entry name, or `entry!classesN.dex` for a container's logical DEX.
530    pub name: String,
531    /// Its bytes (the entry's own when it is not a container).
532    pub data: DexBytes,
533    /// Locator tables built for it on first search.
534    pub state: Arc<SearchState>,
535    persisted: Arc<tiers::Persisted>,
536}
537
538/// A cached inflated entry.
539#[non_exhaustive]
540pub struct CachedEntry {
541    /// The APK entry name.
542    pub name: String,
543    /// The inflated entry.
544    pub data: DexBytes,
545    logical: OnceLock<Vec<LogicalDex>>,
546}
547
548impl CachedEntry {
549    fn new(name: &str, data: DexBytes) -> CachedEntry {
550        CachedEntry { name: name.to_string(), data, logical: OnceLock::new() }
551    }
552
553    /// Logical DEX files of this entry (one, unless it is a 041 container).
554    pub fn logical(&self) -> &[LogicalDex] {
555        self.logical.get_or_init(|| {
556            let bufs = logical_dex_buffers(&self.name, &self.data);
557            bufs.into_iter()
558                .map(|(name, buf)| {
559                    let data = match buf {
560                        Cow::Borrowed(_) => Arc::clone(&self.data),
561                        Cow::Owned(v) => Arc::new(DexData::Owned(v)),
562                    };
563                    LogicalDex { name, data, state: Arc::new(SearchState::new()), persisted: Arc::default() }
564                })
565                .collect()
566        })
567    }
568
569    /// Bytes this entry keeps resident once searched: the inflated entry plus,
570    /// for a 041 container, one full-container copy per further logical DEX
571    /// (what `logical()` materialises). Computed from the headers alone so
572    /// accounting does not itself trigger those copies.
573    pub fn resident_bytes(&self) -> usize {
574        self.data.len() * dex041_logical_offsets(&self.data).len()
575    }
576
577    /// Bytes of the locator tables built (on the heap) or mapped (from L2)
578    /// so far for this entry's logical DEX files (0 before the first search).
579    pub fn state_bytes(&self) -> usize {
580        self.logical.get().map_or(0, |l| l.iter().map(|d| d.state.heap_bytes() + d.state.mapped_bytes()).sum())
581    }
582
583    /// What the in-memory cache charges for this entry: heap bytes in full
584    /// (the inflated entry when not mapped, a container's logical copies,
585    /// tables built on the heap) and file-backed bytes (the entry mapped from
586    /// L1 or the APK, tables mapped from L2) at 1/[`MAPPED_CHARGE_DIVISOR`].
587    pub fn charged_bytes(&self) -> usize {
588        let copies = dex041_logical_offsets(&self.data).len();
589        let len = self.data.len();
590        let own =
591            if self.is_mapped() { (len / MAPPED_CHARGE_DIVISOR) * copies.min(1) + len * copies.saturating_sub(1) } else { len * copies };
592        let tables = self
593            .logical
594            .get()
595            .map_or(0, |l| l.iter().map(|d| d.state.heap_bytes() + d.state.mapped_bytes() / MAPPED_CHARGE_DIVISOR).sum::<usize>());
596        own + tables
597    }
598
599    /// The same entry over `data` (identical bytes from another source, e.g.
600    /// the disk cache mapping), sharing the locator tables (built so far and
601    /// later) of the logical DEX files listed so far.
602    /// A container's normalised logical copies are owned and are kept as is.
603    fn with_data(&self, data: DexBytes) -> CachedEntry {
604        let logical = OnceLock::new();
605        if let Some(old) = self.logical.get() {
606            let moved = old
607                .iter()
608                .map(|l| LogicalDex {
609                    name: l.name.clone(),
610                    data: if Arc::ptr_eq(&l.data, &self.data) { Arc::clone(&data) } else { Arc::clone(&l.data) },
611                    state: Arc::clone(&l.state),
612                    persisted: Arc::clone(&l.persisted),
613                })
614                .collect();
615            let _ = logical.set(moved);
616        }
617        CachedEntry { name: self.name.clone(), data, logical }
618    }
619
620    /// The same entry with the locator tables of logical DEX `idx` replaced
621    /// by `state` (the same tables viewed from L2).
622    fn with_state(&self, idx: usize, state: Arc<SearchState>) -> CachedEntry {
623        let logical = OnceLock::new();
624        if let Some(old) = self.logical.get() {
625            let moved = old
626                .iter()
627                .enumerate()
628                .map(|(i, l)| LogicalDex {
629                    name: l.name.clone(),
630                    data: Arc::clone(&l.data),
631                    state: if i == idx { Arc::clone(&state) } else { Arc::clone(&l.state) },
632                    persisted: Arc::clone(&l.persisted),
633                })
634                .collect();
635            let _ = logical.set(moved);
636        }
637        CachedEntry { name: self.name.clone(), data: Arc::clone(&self.data), logical }
638    }
639
640    /// Whether the entry's bytes are mapped from the disk cache.
641    pub fn is_mapped(&self) -> bool {
642        matches!(*self.data, DexData::Mapped(_))
643    }
644}
645
646/// A queued disk write for the APK of `key` at `apk`.
647struct WriteJob {
648    key: ApkKey,
649    apk: Arc<Path>,
650    kind: Job,
651}
652
653/// How often a call waiting on its workers checks for cancellation.
654const CANCEL_POLL: Duration = Duration::from_millis(20);
655
656/// The flag of a call nobody cancels.
657pub(crate) static NOT_CANCELLED: AtomicBool = AtomicBool::new(false);
658
659/// What a cancelled call returns.
660pub(crate) fn cancelled() -> EngineError {
661    asc_apk::ApkError::Cancelled.into()
662}
663
664impl WriteJob {
665    /// Heap the job keeps alive until written. L3 payloads are small and
666    /// not counted.
667    fn held_bytes(&self) -> usize {
668        match &self.kind {
669            Job::Dex { data, .. } => data.len(),
670            Job::Derived(d) => d.held_bytes(),
671            Job::Graph { graph, .. } => graph.heap_bytes(),
672            Job::Rendered { .. } | Job::Reinflate { .. } => 0,
673        }
674    }
675
676    /// The job without the bytes it holds, for a full queue. An L1 entry is
677    /// inflated again from the APK when its turn comes, so a cold sweep still
678    /// leaves L1 complete; L2 tables and the class graph cannot be rebuilt
679    /// by a writer and are released (`None`) for a later query to queue again.
680    fn defer(self) -> Option<WriteJob> {
681        let WriteJob { key, apk, kind } = self;
682        let kind = match kind {
683            Job::Dex { entry, limits, remap, .. } => Job::Reinflate { entry, limits, remap },
684            Job::Derived(d) => {
685                d.abandon();
686                return None;
687            }
688            Job::Graph { .. } => return None,
689            kind @ (Job::Rendered { .. } | Job::Reinflate { .. }) => kind,
690        };
691        Some(WriteJob { key, apk, kind })
692    }
693
694    /// Dropped unwritten: an L1 entry is simply written by a later query;
695    /// L2 tables must be released for one to queue them again.
696    fn abandon(self) {
697        if let Job::Derived(d) = self.kind {
698            d.abandon();
699        }
700    }
701}
702
703enum Job {
704    /// An inflated entry for L1: the bytes are shared with the query that
705    /// produced them, so queueing costs nothing.
706    Dex {
707        entry: Entry,
708        /// The APK's read limits, should the write be deferred.
709        limits: Limits,
710        data: DexBytes,
711        /// The in-memory cache entry holding `data`, if any: once the file is
712        /// in place its heap bytes are swapped for a mapping of that file.
713        remap: Option<Weak<CachedEntry>>,
714    },
715    /// An L1 write deferred by a full queue: the writer inflates `entry`
716    /// again (if the APK is unchanged and the file still missing).
717    Reinflate { entry: Entry, limits: Limits, remap: Option<Weak<CachedEntry>> },
718    /// Locator tables for L2.
719    Derived(tiers::DerivedJob),
720    /// A query result for L3.
721    Rendered { query: String, payload: Vec<u8> },
722    /// The class graph for L2, serialised by the writer; `cost` is what
723    /// building it took.
724    Graph { graph: Arc<hierarchy::ClassGraph>, cost: Duration },
725}
726
727struct Writer {
728    /// Each job with its [`WriteJob::held_bytes`] when queued.
729    tx: mpsc::Sender<(WriteJob, usize)>,
730    handles: Vec<std::thread::JoinHandle<()>>,
731}
732
733type SharedMemory = Arc<Mutex<MemoryCache<CachedEntry>>>;
734
735/// Lock the memory cache, emptying it if a panic left it poisoned (its byte
736/// accounting may be half-updated); see [`LockExt`].
737pub(crate) fn lock_memory(memory: &Mutex<MemoryCache<CachedEntry>>) -> MutexGuard<'_, MemoryCache<CachedEntry>> {
738    memory.lock_or_reset(MemoryCache::clear)
739}
740
741struct EngineCache {
742    memory: Option<SharedMemory>,
743    /// Publishes (and, when automatic, adjusts) the memory cache's size.
744    budget: Option<tiers::MemoryBudget>,
745    disk: Option<Arc<DiskCache>>,
746    writer: Mutex<Option<Writer>>,
747    /// Set by `flush`: later writes are dropped rather than handed to a new
748    /// writer thread the exiting process would kill mid-write.
749    closed: AtomicBool,
750    /// Entries currently being inflated, so concurrent queries share one
751    /// inflate instead of racing (and replacing each other's locator tables).
752    inflight: Mutex<HashSet<String>>,
753    inflight_done: Condvar,
754    /// L2 writes queued and not yet written: their tables stay on the heap
755    /// until written and mapped, so a prebuild waits while this is high.
756    derived_pending: Arc<AtomicUsize>,
757    /// Heap held by queued writes, bounded by [`EngineCache::queue_cap`].
758    pending_bytes: Arc<AtomicUsize>,
759    /// [`CacheConfig::write_queue_bytes`].
760    queue_cap: usize,
761    /// [`CacheConfig::writers`].
762    writers: Option<usize>,
763}
764
765impl EngineCache {
766    /// Heap the queued writes may hold: [`CacheConfig::write_queue_bytes`], or the
767    /// memory cache's budget when smaller. The budget is this process's share
768    /// of the user's (`--memory-mb`, or split over the registered processes):
769    /// sixteen servers sweeping at once would otherwise each queue 512 MiB
770    /// whatever their caches may hold.
771    fn queue_cap(&self) -> usize {
772        let budget = self.memory.as_ref().map_or(usize::MAX, |m| lock_memory(m).capacity_bytes());
773        self.queue_cap.min(budget)
774    }
775
776    /// Queue a disk write; falls back to a synchronous write only when no
777    /// writer thread can be started. A write that would take the queued heap
778    /// past [`EngineCache::queue_cap`] is queued without its bytes instead
779    /// ([`WriteJob::defer`]): the query never waits for the disk, the heap
780    /// stays bounded, and L1 still ends up complete.
781    fn put_async(&self, job: WriteJob) {
782        let Some(disk) = &self.disk else { return };
783        // Before the writer lock: never hold it while waiting for the memory cache's.
784        let cap = self.queue_cap();
785        let mut writer = self.writer.lock_or_default();
786        if self.closed.load(Ordering::Relaxed) {
787            return;
788        }
789        let pending = self.pending_bytes.load(Ordering::Relaxed);
790        let job = if pending > 0 && pending + job.held_bytes() > cap {
791            match job.defer() {
792                Some(job) => job,
793                None => return,
794            }
795        } else {
796            job
797        };
798        let held = job.held_bytes();
799        if writer.is_none() {
800            match self.spawn_writers(disk) {
801                Some(w) => *writer = Some(w),
802                None => {
803                    run_job(disk, self.memory.as_ref(), job, &mut WriterState::default());
804                    disk.evict_lazily();
805                    return;
806                }
807            }
808        }
809        if let Some(w) = writer.as_ref() {
810            let derived = matches!(job.kind, Job::Derived(_));
811            if derived {
812                self.derived_pending.fetch_add(1, Ordering::Relaxed);
813            }
814            self.pending_bytes.fetch_add(held, Ordering::Relaxed);
815            if let Err(mpsc::SendError((job, _))) = w.tx.send((job, held)) {
816                if derived {
817                    self.derived_pending.fetch_sub(1, Ordering::Relaxed);
818                }
819                self.pending_bytes.fetch_sub(held, Ordering::Relaxed);
820                job.abandon();
821                *writer = None;
822            }
823        }
824    }
825
826    /// Wait for every queued disk write to land. With `close`, writes queued
827    /// afterwards are dropped (the next query re-inflates and re-queues them).
828    fn flush(&self, close: bool) {
829        let taken = {
830            let mut writer = self.writer.lock_or_default();
831            // Under the lock, so a concurrent `put_async` either queued its job
832            // before the take or sees the flag.
833            if close {
834                self.closed.store(true, Ordering::Relaxed);
835            }
836            writer.take()
837        };
838        if let Some(Writer { tx, handles }) = taken {
839            drop(tx);
840            for handle in handles {
841                let _ = handle.join();
842            }
843        }
844    }
845
846    /// Start the writer threads, sharing one queue; `None` when not even one
847    /// can be started.
848    fn spawn_writers(&self, disk: &Arc<DiskCache>) -> Option<Writer> {
849        let (tx, rx) = mpsc::channel::<(WriteJob, usize)>();
850        let rx = Arc::new(Mutex::new(rx));
851        let count = self.writers.unwrap_or_else(|| CacheConfig::auto_writers(asc_cache::cpus())).max(1);
852        let handles: Vec<_> = (0..count)
853            .map_while(|_| {
854                let (rx, disk, memory) = (Arc::clone(&rx), Arc::clone(disk), self.memory.clone());
855                let (derived_pending, pending_bytes) = (Arc::clone(&self.derived_pending), Arc::clone(&self.pending_bytes));
856                std::thread::Builder::new()
857                    .name("asc-cache-writer".into())
858                    .spawn(move || {
859                        // Writes (and re-inflating deferred entries) must not take CPU from queries.
860                        prebuild::lower_priority();
861                        let mut state = WriterState::default();
862                        loop {
863                            // Only the receive holds the lock, so the writers run their jobs in parallel.
864                            let Ok((job, held)) = rx.lock().expect("writer queue mutex poisoned").recv() else { break };
865                            let derived = matches!(job.kind, Job::Derived(_));
866                            run_job(&disk, memory.as_ref(), job, &mut state);
867                            pending_bytes.fetch_sub(held, Ordering::Relaxed);
868                            if derived {
869                                derived_pending.fetch_sub(1, Ordering::Relaxed);
870                            }
871                            // Throttled by bytes written, so a long burst still keeps
872                            // the directory within budget without a walk per entry.
873                            disk.evict_lazily();
874                        }
875                    })
876                    .ok()
877            })
878            .collect();
879        (!handles.is_empty()).then_some(Writer { tx, handles })
880    }
881
882    /// Account for bytes just charged to the memory cache.
883    fn charged(&self) {
884        if let (Some(budget), Some(memory)) = (&self.budget, &self.memory) {
885            budget.maybe_rebalance(memory);
886        }
887    }
888
889    fn begin_inflate(&self, id: String) -> InflightGuard<'_> {
890        let mut set = self.inflight.lock().expect("inflight mutex poisoned");
891        while set.contains(&id) {
892            set = self.inflight_done.wait(set).expect("inflight mutex poisoned");
893        }
894        set.insert(id.clone());
895        InflightGuard { cache: self, id }
896    }
897}
898
899/// What one writer thread keeps between jobs.
900#[derive(Default)]
901struct WriterState {
902    /// APKs whose location the cache already recorded.
903    noted: HashSet<ApkKey>,
904    /// The archive deferred L1 writes last inflated from: a sweep defers
905    /// runs of entries of the same APK.
906    opened: Option<(Arc<Path>, ApkKey, Apk)>,
907}
908
909impl WriterState {
910    /// `path` opened with `limits`, if it is still the archive `key` names.
911    fn apk(&mut self, path: &Arc<Path>, key: &ApkKey, limits: Limits) -> Option<&Apk> {
912        if !matches!(&self.opened, Some((p, k, _)) if p == path && k == key) {
913            self.opened = None;
914            let apk = Apk::open(path).ok()?.with_limits(limits);
915            // Rewritten since the job was queued: its entries are not the ones to store.
916            if ApkKey::for_apk(&apk).ok()? != *key {
917                return None;
918            }
919            self.opened = Some((Arc::clone(path), key.clone(), apk));
920        }
921        self.opened.as_ref().map(|(_, _, apk)| apk)
922    }
923}
924
925/// Run a queued write on a writer thread. The first write for an APK also
926/// records where it lives.
927fn run_job(disk: &DiskCache, memory: Option<&SharedMemory>, job: WriteJob, state: &mut WriterState) {
928    let WriteJob { key, apk, kind } = job;
929    if state.noted.insert(key.clone()) {
930        let _ = disk.note_source(&key, &apk);
931    }
932    match kind {
933        Job::Dex { entry, data, remap, .. } => {
934            if disk.put(&key, &entry.name, &data).is_ok() {
935                remap_entry(disk, memory, key, entry.name, data, remap);
936            }
937        }
938        Job::Reinflate { entry, limits, remap } => {
939            // Written meanwhile, by a later query of the entry or another process.
940            if disk.get(&key, &entry.name, entry.uncomp_size).is_some() {
941                return;
942            }
943            let Some(archive) = state.apk(&apk, &key, limits) else { return };
944            let Ok(Cow::Owned(bytes)) = archive.read_entry(&entry) else { return };
945            let data: DexBytes = Arc::new(DexData::Owned(bytes));
946            if disk.put(&key, &entry.name, &data).is_ok() {
947                remap_entry(disk, memory, key, entry.name, data, remap);
948            }
949        }
950        Job::Derived(d) => tiers::write_derived(disk, memory, &key, d),
951        Job::Rendered { query, payload } => {
952            let _ = disk.put_rendered(&key, &query, &payload);
953        }
954        Job::Graph { graph, cost } => {
955            let _ = disk.put_derived(&key, hierarchy::GRAPH_NAME, hierarchy::GRAPH_KIND, hierarchy::GRAPH_FORMAT, cost, graph.bytes());
956        }
957    }
958}
959
960/// Swap a freshly written entry's heap bytes for a mapping of its cache file:
961/// the pages become file-backed (reclaimable by the OS and shared with other
962/// processes) and the heap copy is freed once the queries still holding the
963/// old entry drop it. The mapping is compared byte for byte first, so a file
964/// changed underneath (another writer, a cleared cache) is never served.
965fn remap_entry(
966    disk: &DiskCache,
967    memory: Option<&SharedMemory>,
968    key: ApkKey,
969    name: String,
970    data: DexBytes,
971    remap: Option<Weak<CachedEntry>>,
972) {
973    // The job's own handle on the bytes goes first, so the heap copy is freed
974    // as soon as the superseded entry is.
975    drop(data);
976    let (Some(memory), Some(entry)) = (memory, remap.and_then(|w| w.upgrade())) else { return };
977    if entry.is_mapped() {
978        return;
979    }
980    let Some(mapped) = disk.get(&key, &name, entry.data.len() as u64) else { return };
981    if *mapped != **entry.data {
982        return;
983    }
984    // A query may still be about to search the entry: list its logical DEX
985    // files now (or wait for the query listing them) so the swapped entry
986    // shares the locator tables that query builds, instead of starting
987    // without any while the query fills the superseded entry's.
988    entry.logical();
989    let new = Arc::new(entry.with_data(Arc::new(mapped)));
990    let mut memory = lock_memory(memory);
991    if memory.replace(&key, &name, &entry, new) {
992        // Now file-backed: charged less. Measured under the lock, so a query
993        // recharging concurrently charges the tables it built either way.
994        memory.recharge(&key, &name, CachedEntry::charged_bytes);
995    }
996}
997
998/// Make freed DEX-sized blocks go straight back to the OS. glibc serves
999/// blocks above its mmap threshold with `mmap`, but raises that threshold to
1000/// the size of every such block freed (up to 32 MiB), after which inflated
1001/// DEX copies come from the heap arenas and stay resident after an eviction
1002/// or a remap, fragmented between longer-lived locator tables. Fixing the
1003/// threshold disables that adjustment. Process-wide; done once, by the first
1004/// engine that keeps entries in memory.
1005fn tune_allocator() {
1006    #[cfg(all(target_os = "linux", target_env = "gnu"))]
1007    {
1008        static ONCE: std::sync::Once = std::sync::Once::new();
1009        ONCE.call_once(|| {
1010            unsafe extern "C" {
1011                fn mallopt(param: std::os::raw::c_int, value: std::os::raw::c_int) -> std::os::raw::c_int;
1012            }
1013            const M_MMAP_THRESHOLD: std::os::raw::c_int = -3;
1014            // SAFETY: mallopt only adjusts glibc allocator parameters.
1015            unsafe {
1016                mallopt(M_MMAP_THRESHOLD, 1 << 20);
1017            }
1018        });
1019    }
1020}
1021
1022struct InflightGuard<'c> {
1023    cache: &'c EngineCache,
1024    id: String,
1025}
1026
1027impl Drop for InflightGuard<'_> {
1028    fn drop(&mut self) {
1029        self.cache.inflight.lock().expect("inflight mutex poisoned").remove(&self.id);
1030        self.cache.inflight_done.notify_all();
1031    }
1032}
1033
1034/// Cache activity counters.
1035#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1036#[non_exhaustive]
1037pub struct CacheStats {
1038    /// Entries served from the in-memory LRU.
1039    pub memory_hits: u64,
1040    /// Entries mapped from the disk cache (L1) instead of inflated.
1041    pub disk_hits: u64,
1042    /// Derived indexes mapped from L2: the locator tables of a logical DEX,
1043    /// and the class graph.
1044    pub derived_hits: u64,
1045    /// Entries inflated from the APK.
1046    pub inflates: u64,
1047    /// Deflate entries a search left out: a streamed prefix showed it could
1048    /// not match them, so they were never inflated whole.
1049    pub skipped: u64,
1050}
1051
1052type Logger = Arc<dyn Fn(&str) + Send + Sync>;
1053
1054/// Where an entry's bytes came from.
1055#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1056enum Source {
1057    Memory,
1058    Disk,
1059    Stored,
1060    Inflate,
1061    /// Not loaded: a prefix showed the search cannot match it.
1062    Skipped,
1063}
1064
1065impl Source {
1066    fn as_str(self) -> &'static str {
1067        match self {
1068            Source::Memory => "memory",
1069            Source::Disk => "disk",
1070            Source::Stored => "stored",
1071            Source::Inflate => "inflate",
1072            Source::Skipped => "skipped",
1073        }
1074    }
1075}
1076
1077enum Loaded<'a> {
1078    Borrowed(Cow<'a, [u8]>),
1079    Cached(Arc<CachedEntry>),
1080}
1081
1082impl<'a> Loaded<'a> {
1083    fn bytes(&self) -> &[u8] {
1084        match self {
1085            Loaded::Borrowed(c) => c,
1086            Loaded::Cached(e) => &e.data,
1087        }
1088    }
1089
1090    /// Logical DEX buffers, moving the inflated bytes instead of copying them
1091    /// when the entry is a plain (non-container) DEX. A container's logical
1092    /// DEXes are copied one at a time, as the caller reaches them.
1093    fn into_logical(self, entry_name: &str) -> Box<dyn Iterator<Item = (String, DexBytes)> + 'a> {
1094        match self {
1095            Loaded::Cached(e) => {
1096                let logical: Vec<_> = e.logical().iter().map(|l| (l.name.clone(), Arc::clone(&l.data))).collect();
1097                Box::new(logical.into_iter())
1098            }
1099            Loaded::Borrowed(c) => {
1100                let offsets = dex041_logical_offsets(&c);
1101                if offsets.len() <= 1 {
1102                    return Box::new(std::iter::once((entry_name.to_string(), Arc::new(DexData::Owned(c.into_owned())))));
1103                }
1104                let name = entry_name.to_string();
1105                Box::new(offsets.into_iter().enumerate().map(move |(i, off)| {
1106                    let data = asc_dex::container::normalize_dex041_logical(&c, off).into_owned();
1107                    (format!("{name}!classes{}.dex", i + 1), Arc::new(DexData::Owned(data)))
1108                }))
1109            }
1110        }
1111    }
1112}
1113
1114/// Engine over one APK.
1115pub struct Engine {
1116    apk_path: PathBuf,
1117    threads: Threads,
1118    debug: bool,
1119    logger: Logger,
1120    cache: Option<Arc<EngineCache>>,
1121    /// `cache` belongs to a [`SharedCache`], which closes it instead of us.
1122    shared_cache: bool,
1123    stats: Arc<Stats>,
1124    /// `(size, mtime, key)` of the APK the key was last computed for.
1125    key_cache: Mutex<Option<(u64, Option<SystemTime>, ApkKey)>>,
1126    /// The APK last opened, handed out again while its files are unchanged
1127    /// (a program of many jars costs a mapping and a directory parse each).
1128    opened: Mutex<Option<Apk>>,
1129    limits: Limits,
1130    /// Parsed `resources.arsc` for the APK version it was read from.
1131    resources: Mutex<Option<resources::CachedTable>>,
1132    /// Background table builds ([`Engine::prebuild`]) and the searches they
1133    /// give way to.
1134    prebuild: prebuild::Prebuild,
1135    /// The class hierarchy, for the APK version it was built from.
1136    class_graph: Mutex<Option<hierarchy::CachedGraph>>,
1137    /// Held while the class graph is built, outside `class_graph`'s lock.
1138    graph_build: Mutex<()>,
1139    /// The names decompiled source shows ([`Engine::with_name_mappings`]).
1140    name_mappings: Arc<names::NameMappings>,
1141}
1142
1143/// An APK opened for a search whose DEX entries are scheduled by the caller,
1144/// e.g. across many APKs at once ([`Engine::search_target`]).
1145#[non_exhaustive]
1146pub struct SearchTarget {
1147    apk: Apk,
1148    key: Option<ApkKey>,
1149    /// Root DEX entries, smallest compressed first.
1150    pub entries: Vec<Entry>,
1151}
1152
1153#[derive(Default)]
1154struct Stats {
1155    memory_hits: AtomicUsize,
1156    disk_hits: AtomicUsize,
1157    derived_hits: AtomicUsize,
1158    inflates: AtomicUsize,
1159    skipped: AtomicUsize,
1160}
1161
1162impl Engine {
1163    /// An engine over `apk_path`: an APK, a jar, a bare DEX file, or a
1164    /// program of several read as one (a directory, or paths joined with
1165    /// `,`; see [`Apk::open`]).
1166    pub fn new(apk_path: impl AsRef<Path>) -> Engine {
1167        Engine {
1168            apk_path: apk_path.as_ref().to_path_buf(),
1169            threads: Threads::Auto,
1170            debug: false,
1171            logger: Arc::new(|m| println!("{m}")),
1172            cache: None,
1173            shared_cache: false,
1174            stats: Arc::new(Stats::default()),
1175            key_cache: Mutex::new(None),
1176            opened: Mutex::new(None),
1177            limits: Limits::default(),
1178            resources: Mutex::new(None),
1179            prebuild: prebuild::Prebuild::default(),
1180            class_graph: Mutex::new(None),
1181            graph_build: Mutex::new(()),
1182            name_mappings: Arc::default(),
1183        }
1184    }
1185
1186    /// Refuse entries that inflate past `limits` (zip bombs).
1187    pub fn with_limits(mut self, limits: Limits) -> Engine {
1188        self.limits = limits;
1189        self
1190    }
1191
1192    /// Use `threads` worker threads per search (or pick them automatically).
1193    pub fn with_threads(mut self, threads: impl Into<Threads>) -> Engine {
1194        self.threads = threads.into();
1195        self
1196    }
1197
1198    /// Log timings and cache decisions of each query through the logger.
1199    pub fn with_debug(mut self, debug: bool) -> Engine {
1200        self.debug = debug;
1201        self
1202    }
1203
1204    /// Route debug output somewhere other than stdout.
1205    pub fn with_logger(mut self, logger: impl Fn(&str) + Send + Sync + 'static) -> Engine {
1206        self.logger = Arc::new(logger);
1207        self
1208    }
1209
1210    /// Enable caching of inflated entries. A disk cache whose directory holds
1211    /// foreign files is dropped with a warning through the logger (set it
1212    /// first), instead of every write failing silently in the background.
1213    pub fn with_cache(mut self, mut config: CacheConfig) -> Engine {
1214        if let Some(Err(e)) = config.disk.as_ref().map(DiskCache::check_usable) {
1215            (self.logger)(&format!("Warning: {e}; disk cache disabled"));
1216            config.disk = None;
1217        }
1218        if config.memory_bytes > 0 {
1219            tune_allocator();
1220        }
1221        if config.memory_bytes == 0 && config.disk.is_none() {
1222            self.cache = None;
1223        } else {
1224            let cache = EngineCache {
1225                memory: (config.memory_bytes > 0).then(|| Arc::new(Mutex::new(MemoryCache::new(config.memory_bytes)))),
1226                budget: (config.memory_bytes > 0).then(|| tiers::MemoryBudget::new(config.memory_auto)),
1227                disk: config.disk.map(Arc::new),
1228                writer: Mutex::new(None),
1229                closed: AtomicBool::new(false),
1230                inflight: Mutex::new(HashSet::new()),
1231                inflight_done: Condvar::new(),
1232                derived_pending: Arc::new(AtomicUsize::new(0)),
1233                pending_bytes: Arc::new(AtomicUsize::new(0)),
1234                queue_cap: config.write_queue_bytes,
1235                writers: config.writers,
1236            };
1237            // Registers this process and takes its share right away.
1238            cache.charged();
1239            self.cache = Some(Arc::new(cache));
1240        }
1241        self
1242    }
1243
1244    /// Use `cache` (shared with other engines) instead of a private one.
1245    pub fn with_shared_cache(mut self, cache: &SharedCache) -> Engine {
1246        self.cache = cache.0.clone();
1247        self.shared_cache = true;
1248        self
1249    }
1250
1251    /// The APK (or program) this engine queries.
1252    pub fn apk_path(&self) -> &Path {
1253        &self.apk_path
1254    }
1255
1256    /// The thread count for a search starting now.
1257    pub fn threads(&self) -> usize {
1258        self.threads.count()
1259    }
1260
1261    /// Cache activity of this engine so far.
1262    pub fn cache_stats(&self) -> CacheStats {
1263        CacheStats {
1264            memory_hits: self.stats.memory_hits.load(Ordering::Relaxed) as u64,
1265            disk_hits: self.stats.disk_hits.load(Ordering::Relaxed) as u64,
1266            derived_hits: self.stats.derived_hits.load(Ordering::Relaxed) as u64,
1267            inflates: self.stats.inflates.load(Ordering::Relaxed) as u64,
1268            skipped: self.stats.skipped.load(Ordering::Relaxed) as u64,
1269        }
1270    }
1271
1272    /// Heap this engine holds outside the memory cache, roughly: the class
1273    /// graph and the parsed resource table it keeps for later calls. A host
1274    /// keeping many engines counts it against its memory budget
1275    /// ([`SharedCache::set_reserved`]).
1276    pub fn resident_bytes(&self) -> usize {
1277        let graph = self.class_graph.lock_or_default().as_ref().map_or(0, |(_, g)| g.heap_bytes());
1278        let table = self.resources.lock_or_default().as_ref().map_or(0, resources::CachedTable::heap_bytes);
1279        graph + table
1280    }
1281
1282    /// Bytes currently charged to the in-memory cache: entry bytes plus the
1283    /// locator tables built for them.
1284    pub fn memory_cache_bytes(&self) -> usize {
1285        self.cache.as_ref().and_then(|c| c.memory.as_ref()).map(|m| lock_memory(m).used_bytes()).unwrap_or(0)
1286    }
1287
1288    /// Drop this APK's entries from the in-memory cache (shared or not), and
1289    /// its parsed resource table, and return the DEX bytes released; queries
1290    /// still holding them keep them alive.
1291    pub fn release_memory_cache(&self) -> usize {
1292        self.release_resources();
1293        self.class_graph.lock_or_default().take();
1294        let Some(memory) = self.cache.as_ref().and_then(|c| c.memory.as_ref()) else { return 0 };
1295        let key = self.key_cache.lock_or_default().as_ref().map(|(_, _, key)| key.clone());
1296        let Some(key) = key.or_else(|| self.open().ok().and_then(|apk| self.apk_key(&apk))) else { return 0 };
1297        lock_memory(memory).remove_apk(&key)
1298    }
1299
1300    /// Entry bytes of the in-memory cache that are mapped from the disk cache
1301    /// (file-backed pages the OS can reclaim) rather than held on the heap.
1302    pub fn memory_cache_mapped_bytes(&self) -> usize {
1303        let Some(m) = self.cache.as_ref().and_then(|c| c.memory.as_ref()) else { return 0 };
1304        lock_memory(m).values().filter(|e| e.is_mapped()).map(|e| e.data.len()).sum()
1305    }
1306
1307    /// Bytes of the in-memory cache's locator tables that are viewed in L2
1308    /// files (file-backed pages shared with other processes) rather than
1309    /// held on the heap.
1310    pub fn memory_cache_mapped_table_bytes(&self) -> usize {
1311        let Some(m) = self.cache.as_ref().and_then(|c| c.memory.as_ref()) else { return 0 };
1312        let m = lock_memory(m);
1313        m.values().filter_map(|e| e.logical.get()).flatten().map(|l| l.state.mapped_bytes()).sum()
1314    }
1315
1316    /// Block until every queued disk-cache write has been renamed into place.
1317    pub fn flush_cache_writes(&self) {
1318        if let Some(c) = &self.cache {
1319            c.flush(false);
1320        }
1321    }
1322
1323    /// [`Engine::flush_cache_writes`] for process exit: also stops queueing,
1324    /// so a worker still holding the engine cannot start a writer thread that
1325    /// exiting would kill mid-write. Run on drop.
1326    pub fn close_cache_writes(&self) {
1327        if let Some(c) = &self.cache {
1328            c.flush(true);
1329        }
1330    }
1331
1332    fn log(&self, msg: impl AsRef<str>) {
1333        if self.debug {
1334            (self.logger)(msg.as_ref());
1335        }
1336    }
1337
1338    /// Open the APK, or reuse the one opened last while [`Apk::version_of`]
1339    /// finds the same files, sizes and times.
1340    ///
1341    /// # Errors
1342    ///
1343    /// [`EngineError::Apk`] when the APK cannot be opened or its central
1344    /// directory read.
1345    pub fn open(&self) -> Result<Apk> {
1346        let mut opened = self.opened.lock_or_default();
1347        // Nothing to compare with on the first open: opening reads the
1348        // version anyway, so do not walk a program's directories twice.
1349        if let Some(apk) = opened.as_ref()
1350            && *apk.version() == Apk::version_of(&self.apk_path)?
1351        {
1352            return Ok(apk.clone().with_limits(self.limits));
1353        }
1354        let apk = Apk::open(&self.apk_path)?.with_limits(self.limits);
1355        *opened = Some(apk.clone());
1356        Ok(apk)
1357    }
1358
1359    /// `apk_path` (what [`Engine::new`] takes) with relative paths resolved
1360    /// against `base`, each path of a program on its own.
1361    pub fn resolve_path(apk_path: &Path, base: &Path) -> PathBuf {
1362        Apk::resolve(apk_path, base)
1363    }
1364
1365    /// Whether `apk_path` opens: a readable file, or a program with code.
1366    pub fn can_open(apk_path: &Path) -> bool {
1367        Apk::open(apk_path).is_ok()
1368    }
1369
1370    /// One line per archive of a program left out of its code (no root DEX,
1371    /// e.g. a jar deodexed into `oat/`, or not a zip); none for one file.
1372    pub fn program_notes(&self) -> Vec<String> {
1373        if self.apk_path.is_file() {
1374            return Vec::new();
1375        }
1376        self.open().map(|apk| apk.notes().to_vec()).unwrap_or_default()
1377    }
1378
1379    /// Open the APK and list its DEX entries for
1380    /// [`Engine::search_target_entry`].
1381    ///
1382    /// # Errors
1383    ///
1384    /// [`EngineError::Apk`] when the APK cannot be opened or its central
1385    /// directory read.
1386    pub fn search_target(&self) -> Result<SearchTarget> {
1387        let apk = self.open()?;
1388        let entries = apk.dex_entries()?;
1389        let key = self.apk_key(&apk);
1390        Ok(SearchTarget { apk, key, entries })
1391    }
1392
1393    /// Search one entry of `target` (what [`Engine::find_refs_with`] runs per
1394    /// entry), through this engine's caches.
1395    pub fn search_target_entry(
1396        &self,
1397        target: &SearchTarget,
1398        entry: &Entry,
1399        query: &CompiledQuery,
1400        aggregate: bool,
1401        cancel: &AtomicBool,
1402    ) -> RefBatch {
1403        search_entry(
1404            &target.apk,
1405            target.key.as_ref(),
1406            entry,
1407            query,
1408            FindOptions { aggregate, context: RefContext::None },
1409            Shape::Results,
1410            None,
1411            cancel,
1412            self.cache.as_deref(),
1413            &self.stats,
1414        )
1415    }
1416
1417    /// Root-level DEX entries, smallest compressed first.
1418    ///
1419    /// # Errors
1420    ///
1421    /// [`EngineError::Apk`] when the APK cannot be opened or its central
1422    /// directory read.
1423    pub fn dex_entries(&self) -> Result<Vec<Entry>> {
1424        Ok(self.open()?.dex_entries()?)
1425    }
1426
1427    /// Cache key for `apk` when caching is enabled. Hashing the central
1428    /// directory is done once per APK version; later calls only `stat`.
1429    pub fn apk_key(&self, apk: &Apk) -> Option<ApkKey> {
1430        self.cache.as_ref()?;
1431        let ident = apk.stamp().ok()?;
1432        let mut cached = self.key_cache.lock_or_default();
1433        if let Some((len, mtime, key)) = cached.as_ref()
1434            && (*len, *mtime) == ident
1435        {
1436            return Some(key.clone());
1437        }
1438        let key = ApkKey::for_apk(apk).ok()?;
1439        *cached = Some((ident.0, ident.1, key.clone()));
1440        Some(key)
1441    }
1442
1443    /// Inflated bytes of `entry`, through the caches when enabled. Stored
1444    /// entries are copied only when an in-memory cache wants to keep them.
1445    ///
1446    /// # Errors
1447    ///
1448    /// [`EngineError::Apk`] when the entry cannot be read or inflated, or
1449    /// inflates past the [`InflateLimits`].
1450    pub fn entry_bytes(&self, apk: &Apk, key: Option<&ApkKey>, entry: &Entry) -> Result<DexBytes> {
1451        Ok(match load_entry(apk, key, entry, self.cache.as_deref(), &self.stats, None)?.0 {
1452            Loaded::Borrowed(c) => Arc::new(DexData::Owned(c.into_owned())),
1453            Loaded::Cached(e) => Arc::clone(&e.data),
1454        })
1455    }
1456
1457    /// Find the DEX that defines `descriptor` (`Lcom/poc/Main;`).
1458    ///
1459    /// DEX entries are inflated in parallel, smallest first. Entries that are
1460    /// not cached are first streamed only as far as their id tables and string
1461    /// data to answer the probe cheaply; only a hit is inflated fully.
1462    ///
1463    /// A class several DEX files define resolves as a class loader would:
1464    /// the first archive of a program, then `classes.dex`, `classes2.dex`,
1465    /// and so on. Once a DEX defines it, DEX files after it are no
1466    /// longer read; workers still inflating one at that moment finish on
1467    /// their own in the background. A program reads every archive, to name
1468    /// the others that define the class too ([`ClassHit::also_defined_in`]).
1469    ///
1470    /// A class no DEX defines is `Ok(None)`.
1471    ///
1472    /// # Errors
1473    ///
1474    /// [`EngineError::Apk`] when the APK cannot be opened, and the error of a
1475    /// DEX that failed to read or parse when no other DEX defines the class.
1476    pub fn get_class_dex(&self, descriptor: &str) -> Result<Option<ClassHit>> {
1477        self.get_class_dex_cancellable(descriptor, &NOT_CANCELLED)
1478    }
1479
1480    /// [`Engine::get_class_dex`] that stops, with `ApkError::Cancelled`, once
1481    /// `cancel` is raised.
1482    ///
1483    /// # Errors
1484    ///
1485    /// As [`Engine::get_class_dex`], plus the cancellation.
1486    pub fn get_class_dex_cancellable(&self, descriptor: &str, cancel: &AtomicBool) -> Result<Option<ClassHit>> {
1487        if let Some(hit) = self.get_class_dex_exact(descriptor, cancel)? {
1488            return Ok(Some(hit));
1489        }
1490        self.get_class_dex_denormalized(descriptor, cancel)
1491    }
1492
1493    /// The class `descriptor` stands for when it is written as decompiled
1494    /// Java / Kotlin source spells it (`com.a.Outer.Inner`, `Outer.Anonymous1`,
1495    /// `Cdo`, `` `in` ``): the first of [`names::denormalize_class_candidates`]
1496    /// the APK defines, else a nested class of the named outer class whose
1497    /// `dalvik.annotation.InnerClass` name is the one written (an obfuscated
1498    /// `Outer$a` the source shows as `Outer.RealName`). The APK's class list
1499    /// is read once (the miss it follows costs as much for its hint).
1500    fn get_class_dex_denormalized(&self, descriptor: &str, cancel: &AtomicBool) -> Result<Option<ClassHit>> {
1501        let candidates = names::denormalize_class_candidates(descriptor);
1502        if candidates.is_empty() {
1503            return Ok(None);
1504        }
1505        let defined: HashSet<String> =
1506            self.scan_classes(None, |_, _| {})?.into_iter().filter_map(|e| e.result.ok()).flatten().flat_map(|d| d.descriptors).collect();
1507        let found = |hit: Option<ClassHit>, real: &str| {
1508            hit.map(|mut h| {
1509                self.log(format!("[APK] getclass {descriptor} read as {real}"));
1510                h.descriptor = real.to_string();
1511                h
1512            })
1513        };
1514        if let Some(real) = candidates.iter().find(|c| defined.contains(c.as_str())) {
1515            return Ok(found(self.get_class_dex_exact(real, cancel)?, real));
1516        }
1517        // Several candidates name the same outer class: read each once.
1518        let mut outers_seen = HashSet::new();
1519        for candidate in std::iter::once(descriptor).chain(candidates.iter().map(String::as_str)) {
1520            let Some((outer_stem, inner_name)) = candidate.strip_suffix(';').and_then(|s| s.rsplit_once('$')) else { continue };
1521            let outer = format!("{outer_stem};");
1522            if !defined.contains(&outer) || !outers_seen.insert(outer.clone()) {
1523                continue;
1524            }
1525            let Some(hit) = self.get_class_dex_exact(&outer, cancel)? else { continue };
1526            if let Some(real) = classes::inner_class_named(&hit.dex, &outer, inner_name) {
1527                return Ok(found(self.get_class_dex_exact(&real, cancel)?, &real));
1528            }
1529        }
1530        Ok(None)
1531    }
1532
1533    /// [`Engine::get_class_dex_cancellable`] of exactly `descriptor`.
1534    fn get_class_dex_exact(&self, descriptor: &str, cancel: &AtomicBool) -> Result<Option<ClassHit>> {
1535        if cancel.load(Ordering::Relaxed) {
1536            return Err(cancelled());
1537        }
1538        let t0 = Instant::now();
1539        let apk = Arc::new(self.open()?);
1540        let entries = Arc::new(apk.dex_entries()?);
1541        let key = Arc::new(self.apk_key(&apk));
1542        self.log(format!("[APK] scan setup={} us entries={}", t0.elapsed().as_micros(), entries.len()));
1543        if entries.is_empty() {
1544            return Ok(None);
1545        }
1546        // DEX descriptors are MUTF-8: a supplementary character is a surrogate pair,
1547        // not the 4-byte UTF-8 sequence `as_bytes` would compare against.
1548        let target: Arc<Vec<u8>> = Arc::new(asc_dex::mutf8::encode(descriptor));
1549        let next = Arc::new(AtomicUsize::new(0));
1550        let stop = Arc::new(AtomicBool::new(false));
1551        // Rank of the best DEX defining the class so far: later ones need not be read.
1552        let best = Arc::new(AtomicUsize::new(usize::MAX));
1553        let ranks: Arc<Vec<usize>> = Arc::new(load_ranks(&entries));
1554        let exhaustive = apk.archive_count() > 1;
1555        let workers = self.threads().min(entries.len()).max(1);
1556        let (tx, rx) = mpsc::channel::<(usize, std::result::Result<Option<ClassHit>, EngineError>)>();
1557        let debug = self.debug;
1558        for w in 0..workers {
1559            let (apk, entries, key, target, next, stop, tx, logger, cache, stats, best, ranks) = (
1560                Arc::clone(&apk),
1561                Arc::clone(&entries),
1562                Arc::clone(&key),
1563                Arc::clone(&target),
1564                Arc::clone(&next),
1565                Arc::clone(&stop),
1566                tx.clone(),
1567                Arc::clone(&self.logger),
1568                self.cache.clone(),
1569                Arc::clone(&self.stats),
1570                Arc::clone(&best),
1571                Arc::clone(&ranks),
1572            );
1573            std::thread::spawn(move || {
1574                while !stop.load(Ordering::Relaxed) {
1575                    let i = next.fetch_add(1, Ordering::Relaxed);
1576                    let Some(entry) = entries.get(i) else { break };
1577                    let rank = ranks[i];
1578                    // A DEX the class loader would read after the best hit cannot win.
1579                    if !exhaustive && rank > best.load(Ordering::Relaxed) {
1580                        let _ = tx.send((i, Ok(None)));
1581                        continue;
1582                    }
1583                    let t0 = Instant::now();
1584                    // Cached bytes answer directly; otherwise probe a prefix so
1585                    // most entries never get fully inflated.
1586                    let cached = cache.as_deref().and_then(|c| lookup_cached(c, key.as_ref().as_ref(), entry, &stats));
1587                    let (loaded, source) = match cached {
1588                        Some((e, source)) => (Loaded::Cached(e), source),
1589                        None => {
1590                            match probe_entry(&apk, entry, &target) {
1591                                Ok(Some(false)) => {
1592                                    if debug {
1593                                        logger(&format!(
1594                                            "[APK] [T{w:04x}] '{}' prefix-probe={} us hit=false",
1595                                            entry.name,
1596                                            t0.elapsed().as_micros()
1597                                        ));
1598                                    }
1599                                    let _ = tx.send((i, Ok(None)));
1600                                    continue;
1601                                }
1602                                Ok(Some(true)) | Ok(None) => {}
1603                                Err(e) => {
1604                                    let _ = tx.send((i, Err(e)));
1605                                    continue;
1606                                }
1607                            }
1608                            match load_entry(&apk, key.as_ref().as_ref(), entry, cache.as_deref(), &stats, Some(&stop)) {
1609                                Ok(l) => l,
1610                                Err(EngineError::Apk(asc_apk::ApkError::Cancelled)) => break,
1611                                Err(e) => {
1612                                    let _ = tx.send((i, Err(e)));
1613                                    continue;
1614                                }
1615                            }
1616                        }
1617                    };
1618                    let t1 = Instant::now();
1619                    let mut found = None;
1620                    // Exactly one message per entry: the caller counts them to
1621                    // know when every entry has answered.
1622                    let mut entry_err: Option<EngineError> = None;
1623                    for (name, data) in loaded.into_logical(&entry.name) {
1624                        if stop.load(Ordering::Relaxed) {
1625                            return;
1626                        }
1627                        match asc_dex::quick::dex_defines_class(&data, &target) {
1628                            Ok(true) => {
1629                                found =
1630                                    Some(ClassHit { dex_name: name, dex: data, descriptor: String::new(), also_defined_in: Vec::new() });
1631                                break;
1632                            }
1633                            Ok(false) => {}
1634                            Err(e) => {
1635                                entry_err.get_or_insert(e.into());
1636                            }
1637                        }
1638                    }
1639                    let is_hit = found.is_some();
1640                    if debug {
1641                        logger(&format!(
1642                            "[APK] [T{w:04x}] '{}' {}={} us lookup={} us hit={}",
1643                            entry.name,
1644                            source.as_str(),
1645                            (t1 - t0).as_micros(),
1646                            t1.elapsed().as_micros(),
1647                            is_hit
1648                        ));
1649                    }
1650                    if let Some(hit) = found {
1651                        best.fetch_min(rank, Ordering::Relaxed);
1652                        let _ = tx.send((i, Ok(Some(hit))));
1653                        continue;
1654                    }
1655                    let _ = tx.send((i, entry_err.map_or(Ok(None), Err)));
1656                }
1657            });
1658        }
1659        drop(tx);
1660        let mut hits: Vec<(usize, ClassHit)> = Vec::new();
1661        let mut last_err = None;
1662        let mut answered = vec![false; entries.len()];
1663        let mut seen = 0usize;
1664        loop {
1665            let (i, msg) = match rx.recv_timeout(CANCEL_POLL) {
1666                Ok(msg) => msg,
1667                Err(mpsc::RecvTimeoutError::Timeout) if cancel.load(Ordering::Relaxed) => {
1668                    stop.store(true, Ordering::Relaxed);
1669                    return Err(cancelled());
1670                }
1671                Err(mpsc::RecvTimeoutError::Timeout) => continue,
1672                Err(mpsc::RecvTimeoutError::Disconnected) => break,
1673            };
1674            seen += 1;
1675            answered[i] = true;
1676            match msg {
1677                Ok(Some(h)) => hits.push((ranks[i], h)),
1678                Ok(None) => {}
1679                Err(e) => last_err = Some(e),
1680            }
1681            if seen == entries.len() {
1682                break;
1683            }
1684            // The best hit wins once every DEX the class loader reads before it has answered.
1685            let best = hits.iter().map(|(rank, _)| *rank).min();
1686            if let Some(best) = best.filter(|_| !exhaustive)
1687                && ranks.iter().zip(&answered).all(|(&rank, &done)| done || rank >= best)
1688            {
1689                break;
1690            }
1691        }
1692        stop.store(true, Ordering::Relaxed);
1693        self.log(format!("[APK] getclass total={} us", t0.elapsed().as_micros()));
1694        hits.sort_by_key(|(rank, _)| *rank);
1695        let mut hits = hits.into_iter().map(|(_, h)| h);
1696        match (hits.next(), last_err) {
1697            (Some(mut h), _) => {
1698                h.descriptor = descriptor.to_string();
1699                let archive = archive_of(&h.dex_name).map(str::to_string);
1700                h.also_defined_in = hits.map(|o| o.dex_name).filter(|n| archive_of(n).map(str::to_string) != archive).collect();
1701                Ok(Some(h))
1702            }
1703            (None, Some(e)) => Err(e),
1704            (None, None) => Ok(None),
1705        }
1706    }
1707
1708    /// Locate, rebuild and disassemble a class.
1709    ///
1710    /// # Errors
1711    ///
1712    /// [`EngineError::ClassNotFound`] when no DEX defines it, the errors of
1713    /// [`Engine::get_class_dex`], and [`EngineError::Rebuild`] /
1714    /// [`EngineError::Smali`] when the class cannot be rendered at all (a class
1715    /// the rebuild cannot represent is disassembled from the original DEX
1716    /// instead, with a [`ClassSource::warning`]).
1717    pub fn get_class(&self, descriptor: &str) -> Result<ClassSource> {
1718        self.get_class_with(descriptor, SmaliOptions::default())
1719    }
1720
1721    /// [`Engine::get_class`] with smali rendering options.
1722    ///
1723    /// # Errors
1724    ///
1725    /// As [`Engine::get_class`].
1726    pub fn get_class_with(&self, descriptor: &str, opts: SmaliOptions) -> Result<ClassSource> {
1727        self.get_class_method(descriptor, opts, None)
1728    }
1729
1730    /// [`Engine::get_class_with`] that renders only the method whose
1731    /// `name(params)ret` is `method` (no fields) when one is given. A malformed
1732    /// signature is rejected before the APK is read; a class without such a
1733    /// method is an error. The rebuilt DEX still holds the whole class.
1734    ///
1735    /// # Errors
1736    ///
1737    /// As [`Engine::get_class`], plus [`EngineError::Smali`] for a malformed
1738    /// `method` or one the class does not have.
1739    pub fn get_class_method(&self, descriptor: &str, opts: SmaliOptions, method: Option<&str>) -> Result<ClassSource> {
1740        self.get_class_method_cancellable(descriptor, opts, method, &NOT_CANCELLED)
1741    }
1742
1743    /// [`Engine::get_class_method`] that stops looking for the class, with
1744    /// `ApkError::Cancelled`, once `cancel` is raised.
1745    ///
1746    /// # Errors
1747    ///
1748    /// As [`Engine::get_class_method`], plus the cancellation.
1749    pub fn get_class_method_cancellable(
1750        &self,
1751        descriptor: &str,
1752        opts: SmaliOptions,
1753        method: Option<&str>,
1754        cancel: &AtomicBool,
1755    ) -> Result<ClassSource> {
1756        if let Some(m) = method {
1757            asc_smali::normalize_method_signature(m)?;
1758        }
1759        let t0 = Instant::now();
1760        let hit = self
1761            .get_class_dex_cancellable(descriptor, cancel)?
1762            .ok_or_else(|| EngineError::ClassNotFound { class: descriptor.to_string(), hint: self.class_not_found_hint(descriptor) })?;
1763        let scan_us = t0.elapsed().as_micros();
1764        let descriptor = hit.descriptor.as_str();
1765        let (rebuilt, source, warning) = match decompile_dex_class_method(&hit.dex, descriptor, opts, method) {
1766            Ok((r, s)) => {
1767                let warning = (!r.dropped.is_empty()).then(|| {
1768                    let dropped = r.dropped.iter().map(ToString::to_string).collect::<Vec<_>>().join("; ");
1769                    format!("nested classes left out of the rebuilt DEX: {dropped}")
1770                });
1771                (Some(r), s, warning)
1772            }
1773            // The class is there but cannot be rebuilt faithfully (undefined
1774            // opcode, invoke-custom, too many members for 16-bit indexes...):
1775            // still show it, disassembled from the original DEX, and say why
1776            // there is no DEX to dump.
1777            Err(EngineError::Rebuild(e)) if e.is_input_limitation() => {
1778                let dex = Dex::with_name(&hit.dex, &hit.dex_name)?;
1779                let source = asc_smali::disassemble_class_method(&dex, descriptor, opts, method)?;
1780                let warning = format!("class not rebuilt ({e}); disassembled from the original DEX instead");
1781                self.log(format!("[DEBUG] {warning}"));
1782                (None, source, Some(warning))
1783            }
1784            Err(e) => return Err(e),
1785        };
1786        if self.debug
1787            && let Some(r) = &rebuilt
1788        {
1789            let t = r.timings;
1790            self.log(format!(
1791                    "[DEBUG] Hollower Time: {} us\n[DEBUG] Collect Time: {} us\n[DEBUG] Remapper Time: {} us\n[DEBUG] Builder Time: {} us\n[DEBUG] Total Extraction Time: {} us",
1792                    t.hollow_us,
1793                    t.collect_us,
1794                    t.map_us,
1795                    t.build_us,
1796                    t.total_us()
1797                ));
1798        }
1799        Ok(ClassSource {
1800            descriptor: descriptor.to_string(),
1801            dex_name: hit.dex_name,
1802            source,
1803            rebuilt,
1804            warning,
1805            scan_us,
1806            also_defined_in: hit.also_defined_in,
1807        })
1808    }
1809
1810    /// Search every DEX entry for references. `sink` receives one batch per
1811    /// entry in completion order. Returns a summary once all entries finish.
1812    ///
1813    /// A query that finds no row as written is searched again with its names
1814    /// read as decompiled Java / Kotlin source spells DEX names
1815    /// ([`names::denormalize_query`]); [`RefSummary::notes`] says so when that
1816    /// finds rows.
1817    ///
1818    /// # Errors
1819    ///
1820    /// [`EngineError::Query`] / [`EngineError::Find`] for an invalid query,
1821    /// and [`EngineError::Apk`] when the APK cannot be opened. A DEX entry that
1822    /// fails to load or search does not fail the run: its batch carries the
1823    /// error in [`RefBatch::error`].
1824    pub fn find_refs(&self, query: &Query, mut sink: impl FnMut(RefBatch)) -> Result<RefSummary> {
1825        self.find_refs_with(query, true, &AtomicBool::new(false), &mut sink)
1826    }
1827
1828    /// `find_refs` with control over aggregation and cancellation. An invalid
1829    /// pattern is reported once, up front, instead of once per DEX.
1830    ///
1831    /// # Errors
1832    ///
1833    /// As [`Engine::find_refs`].
1834    pub fn find_refs_with(
1835        &self,
1836        query: &Query,
1837        aggregate: bool,
1838        cancel: &AtomicBool,
1839        sink: &mut dyn FnMut(RefBatch),
1840    ) -> Result<RefSummary> {
1841        self.find_refs_opts(query, FindOptions { aggregate, context: RefContext::None }, cancel, sink)
1842    }
1843
1844    /// `find_refs` with every option: aggregation, context, cancellation.
1845    ///
1846    /// # Errors
1847    ///
1848    /// As [`Engine::find_refs`].
1849    pub fn find_refs_opts(
1850        &self,
1851        query: &Query,
1852        opts: FindOptions,
1853        cancel: &AtomicBool,
1854        sink: &mut dyn FnMut(RefBatch),
1855    ) -> Result<RefSummary> {
1856        self.find_refs_shaped(query, opts, Shape::Results, cancel, sink)
1857    }
1858
1859    /// The rows of [`Engine::find_refs`] (one per method, no context) as
1860    /// text: each batch holds them in [`RefBatch::text`], exactly the
1861    /// [`RefResult::line`]s `find_refs` would give, with no structured rows.
1862    /// A broad query renders and replays much faster this way.
1863    ///
1864    /// # Errors
1865    ///
1866    /// As [`Engine::find_refs`].
1867    pub fn find_refs_text(&self, query: &Query, cancel: &AtomicBool, sink: &mut dyn FnMut(RefBatch)) -> Result<RefSummary> {
1868        self.find_refs_shaped(query, FindOptions::default(), Shape::Text, cancel, sink)
1869    }
1870
1871    fn find_refs_shaped(
1872        &self,
1873        query: &Query,
1874        opts: FindOptions,
1875        shape: Shape,
1876        cancel: &AtomicBool,
1877        sink: &mut dyn FnMut(RefBatch),
1878    ) -> Result<RefSummary> {
1879        let t0 = Instant::now();
1880        compose::check_patterns(query)?;
1881        if let Some(engine) = self.transient(compose::needs_program(query))? {
1882            return engine.find_refs_shaped(query, opts, shape, cancel, sink);
1883        }
1884        let apk = self.open()?;
1885        let mut entries = apk.dex_entries()?;
1886        // Every entry has to be searched, so run the largest first: the biggest
1887        // inflate is the critical path and must not wait behind the small ones.
1888        entries.reverse();
1889        let key = self.apk_key(&apk);
1890        let mut summary = RefSummary { entries: entries.len(), ..Default::default() };
1891        if entries.is_empty() {
1892            return Ok(summary);
1893        }
1894        // L3: the rows of an earlier run of this query on this APK.
1895        let disk = self.cache.as_ref().and_then(|c| c.disk.as_ref()).filter(|d| d.uses(Tier::Rendered));
1896        let rendered = disk.zip(key.as_ref()).map(|(d, k)| (d, k, tiers::findrefs_key(query, opts, shape)));
1897        if let Some(recorded) = rendered.as_ref().and_then(|(d, k, q)| d.get_rendered_view(k, q)).and_then(tiers::decode_batches) {
1898            for batch in recorded.batches {
1899                summary.total_results += batch.row_count();
1900                sink(batch);
1901            }
1902            summary.notes = recorded.notes;
1903            summary.total_us = t0.elapsed().as_micros();
1904            self.log(format!("[APK] for_each_findrefs total={} us count={} source=result", summary.total_us, entries.len()));
1905            return Ok(summary);
1906        }
1907        // `method:{RULE}`: the methods RULE matches first, then the calls of
1908        // them. L3 keys on the query as given, so a replay skips this too.
1909        let resolved = self.resolve(std::slice::from_ref(query), cancel)?;
1910        summary.notes = resolved.notes.into_iter().next().unwrap_or_default();
1911        let _query = self.prebuild.query();
1912        let active = resolved.queries.as_ref().map_or(query, |r| &r[0]);
1913        let compiled = CompiledQuery::new(active)?;
1914        let mut recorder = rendered.as_ref().map(|(d, _, _)| tiers::ResultRecorder::new(tiers::result_limit(d, shape)));
1915        let workers = self.threads().min(entries.len()).max(1);
1916        let cache = self.cache.as_deref();
1917        let stats = &self.stats;
1918        // One search of every entry, each batch handed to `on_batch` as it completes.
1919        let run = |compiled: &CompiledQuery, on_batch: &mut dyn FnMut(RefBatch)| {
1920            let next = AtomicUsize::new(0);
1921            let (tx, rx) = mpsc::channel::<RefBatch>();
1922            let probe = skip::Probe::new(vec![compiled]);
1923            let probe = &probe;
1924            std::thread::scope(|s| {
1925                for _ in 0..workers {
1926                    let tx = tx.clone();
1927                    let (apk, entries, next, key) = (&apk, &entries, &next, key.as_ref());
1928                    s.spawn(move || {
1929                        loop {
1930                            if cancel.load(Ordering::Relaxed) {
1931                                break;
1932                            }
1933                            let i = next.fetch_add(1, Ordering::Relaxed);
1934                            let Some(entry) = entries.get(i) else { break };
1935                            let batch = search_entry(apk, key, entry, compiled, opts, shape, Some(probe), cancel, cache, stats);
1936                            if tx.send(batch).is_err() {
1937                                break;
1938                            }
1939                        }
1940                    });
1941                }
1942                drop(tx);
1943                for batch in rx {
1944                    if self.debug {
1945                        self.log(format!(
1946                            "[APK] '{}' {}={} us process={} us results={}{}",
1947                            batch.entry_name,
1948                            batch.source,
1949                            batch.inflate_us,
1950                            batch.search_us,
1951                            batch.row_count(),
1952                            batch.error.as_ref().map(|e| format!(" error={e}")).unwrap_or_default()
1953                        ));
1954                    }
1955                    on_batch(batch);
1956                }
1957            });
1958        };
1959        let mut emit = |batch: RefBatch| {
1960            summary.total_results += batch.row_count();
1961            if let Some(r) = recorder.as_mut() {
1962                r.record(&batch);
1963            }
1964            sink(batch);
1965        };
1966        // Names written as decompiled Java / Kotlin source spells them
1967        // (`com.a.Outer.Inner`, `Cdo`, `a$2`) get a second search when the
1968        // query as written finds nothing. Until a row shows up the (empty)
1969        // batches are held back, so a query that matches streams as before.
1970        let fallback = names::denormalize_query(active).and_then(|q| CompiledQuery::new(&q).ok());
1971        let mut held: Option<Vec<RefBatch>> = fallback.as_ref().map(|_| Vec::new());
1972        run(&compiled, &mut |batch| match held.as_mut() {
1973            Some(h) if batch.row_count() == 0 => h.push(batch),
1974            Some(h) => {
1975                h.drain(..).for_each(&mut emit);
1976                held = None;
1977                emit(batch);
1978            }
1979            None => emit(batch),
1980        });
1981        let mut read_as_source = false;
1982        if let (Some(held), Some(fallback)) = (held, fallback) {
1983            if !cancel.load(Ordering::Relaxed) && held.iter().all(|b| b.error.is_none()) {
1984                let mut batches = Vec::with_capacity(held.len());
1985                run(&fallback, &mut |batch| batches.push(batch));
1986                read_as_source = batches.iter().any(|b| b.row_count() > 0);
1987                batches.into_iter().for_each(&mut emit);
1988            } else {
1989                held.into_iter().for_each(&mut emit);
1990            }
1991        }
1992        if read_as_source {
1993            summary.notes.push(
1994                "nothing matched the names as written; these rows match them read as Java/Kotlin source spells DEX names \
1995                 (Outer.Inner -> Outer$Inner, AnonymousN -> $N, Cdo -> do, a$2 -> a, `in` -> in)"
1996                    .to_string(),
1997            );
1998        }
1999        if let (Some((_, k, q)), Some(payload)) = (rendered, recorder.and_then(|r| r.finish(&summary.notes))) {
2000            // A cancelled run may be missing entries.
2001            if !cancel.load(Ordering::Relaxed) {
2002                let job = WriteJob { key: k.clone(), apk: Arc::from(apk.path()), kind: Job::Rendered { query: q, payload } };
2003                self.cache.as_ref().expect("rendered implies a cache").put_async(job);
2004            }
2005        }
2006        summary.total_us = t0.elapsed().as_micros();
2007        self.log(format!("[APK] for_each_findrefs total={} us count={} workers={}", summary.total_us, entries.len(), workers));
2008        Ok(summary)
2009    }
2010
2011    /// Fill the context of rows an earlier [`Engine::find_refs`] returned
2012    /// (typically one page of them): only their DEX files are loaded, from
2013    /// the cache when it holds them. Returns the entries that failed to load.
2014    ///
2015    /// # Errors
2016    ///
2017    /// [`EngineError::Apk`] when the APK cannot be opened or its central
2018    /// directory read.
2019    pub fn fill_context(&self, rows: &mut [RefResult], ctx: RefContext) -> Result<Vec<String>> {
2020        if ctx == RefContext::None || rows.is_empty() {
2021            return Ok(Vec::new());
2022        }
2023        let _query = self.prebuild.query();
2024        let wanted: HashSet<String> = rows.iter().map(|r| entry_of(&r.dex_name).to_string()).collect();
2025        let (views, errors) = self.load_views(Some(&wanted), &NOT_CANCELLED)?;
2026        let mut errors = errors;
2027        for v in &views {
2028            let r = Dex::with_name(&v.data, &v.name).map_err(EngineError::from).and_then(|dex| {
2029                attach_context(&FindRefs::with_state(&dex, &v.state), rows.iter_mut().filter(|r| r.dex_name == v.name), ctx)
2030            });
2031            if let Err(e) = r {
2032                errors.push(format!("{}: {e}", v.name));
2033            }
2034        }
2035        self.recharge_views(&views);
2036        Ok(errors)
2037    }
2038
2039    /// Load DEX entries (every one, or those named in `only`) in parallel,
2040    /// as logical DEX files in entry order with their locator tables (the
2041    /// cached ones when the memory cache holds the entry). Entries that fail
2042    /// to load are reported as `name: error` strings.
2043    /// Stops with `ApkError::Cancelled` once `cancel` is raised.
2044    fn load_views(&self, only: Option<&HashSet<String>>, cancel: &AtomicBool) -> Result<(Vec<DexView>, Vec<String>)> {
2045        let apk = self.open()?;
2046        let mut entries = apk.dex_entries()?;
2047        if let Some(only) = only {
2048            entries.retain(|e| only.contains(&e.name));
2049        }
2050        let key = self.apk_key(&apk);
2051        let next = AtomicUsize::new(0);
2052        let workers = self.threads().min(entries.len()).max(1);
2053        let cache = self.cache.as_deref();
2054        let l2 = cache.and_then(|c| c.disk.as_deref()).filter(|d| d.uses(Tier::Derived)).zip(key.as_ref());
2055        let mut slots: Vec<Option<std::result::Result<Vec<DexView>, String>>> = (0..entries.len()).map(|_| None).collect();
2056        let done = Mutex::new(&mut slots);
2057        std::thread::scope(|s| {
2058            for _ in 0..workers {
2059                s.spawn(|| {
2060                    loop {
2061                        if cancel.load(Ordering::Relaxed) {
2062                            break;
2063                        }
2064                        let i = next.fetch_add(1, Ordering::Relaxed);
2065                        let Some(entry) = entries.get(i) else { break };
2066                        let r = load_entry(&apk, key.as_ref(), entry, cache, &self.stats, Some(cancel))
2067                            .map(|(loaded, _)| match loaded {
2068                                Loaded::Cached(e) => e
2069                                    .logical()
2070                                    .iter()
2071                                    .enumerate()
2072                                    .map(|(idx, l)| {
2073                                        if let (Some((disk, k)), Ok(dex)) = (l2, Dex::with_name(&l.data, &l.name))
2074                                            && tiers::load_derived(disk, k, l, &dex)
2075                                        {
2076                                            self.stats.derived_hits.fetch_add(1, Ordering::Relaxed);
2077                                        }
2078                                        DexView {
2079                                            name: l.name.clone(),
2080                                            data: Arc::clone(&l.data),
2081                                            state: Arc::clone(&l.state),
2082                                            entry: Some(Arc::clone(&e)),
2083                                            idx,
2084                                            mask: l.state.table_mask(),
2085                                            loaded: Instant::now(),
2086                                        }
2087                                    })
2088                                    .collect(),
2089                                borrowed => borrowed
2090                                    .into_logical(&entry.name)
2091                                    .map(|(name, data)| DexView {
2092                                        name,
2093                                        data,
2094                                        state: Arc::new(SearchState::new()),
2095                                        entry: None,
2096                                        idx: 0,
2097                                        mask: 0,
2098                                        loaded: Instant::now(),
2099                                    })
2100                                    .collect(),
2101                            })
2102                            .map_err(|e| format!("{}: {e}", entry.name));
2103                        done.lock().expect("result slots mutex poisoned")[i] = Some(r);
2104                    }
2105                });
2106            }
2107        });
2108        if cancel.load(Ordering::Relaxed) {
2109            return Err(cancelled());
2110        }
2111        let mut views = Vec::new();
2112        let mut errors = Vec::new();
2113        for slot in slots.into_iter().flatten() {
2114            match slot {
2115                Ok(mut v) => views.append(&mut v),
2116                Err(e) => errors.push(e),
2117            }
2118        }
2119        Ok((views, errors))
2120    }
2121
2122    /// Charge the memory cache for the tables a pass over `views` built,
2123    /// and queue the new tables for L2.
2124    fn recharge_views(&self, views: &[DexView]) {
2125        let Some(cache) = self.cache.as_deref() else { return };
2126        let Some(key) = self.key_cache.lock_or_default().as_ref().map(|k| k.2.clone()) else { return };
2127        let disk = cache.disk.as_deref().filter(|d| d.uses(Tier::Derived));
2128        let mut seen = HashSet::new();
2129        for v in views {
2130            let Some(e) = &v.entry else { continue };
2131            if let Some(job) = disk.and_then(|d| tiers::derived_job(d, e, v.idx, v.mask, v.loaded.elapsed())) {
2132                cache.put_async(WriteJob { key: key.clone(), apk: Arc::from(self.apk_path.as_path()), kind: Job::Derived(job) });
2133            }
2134            if seen.insert(Arc::as_ptr(e)) {
2135                recharge(Some(cache), Some(&key), e);
2136            }
2137        }
2138    }
2139
2140    /// Decode `AndroidManifest.xml`.
2141    ///
2142    /// # Errors
2143    ///
2144    /// [`EngineError::Apk`] when the APK or the manifest entry cannot be read,
2145    /// [`EngineError::Axml`] when the manifest is not valid binary XML.
2146    pub fn manifest_xml(&self) -> Result<String> {
2147        let apk = self.open()?;
2148        let data = apk.read("AndroidManifest.xml")?;
2149        Ok(asc_axml::decode(&data)?)
2150    }
2151}
2152
2153impl Drop for Engine {
2154    fn drop(&mut self) {
2155        self.prebuild.cancel();
2156        // A shared cache outlives this engine; its `SharedCache` owner closes it.
2157        if !self.shared_cache {
2158            self.close_cache_writes();
2159        }
2160    }
2161}
2162
2163/// One logical DEX loaded for a pass over the whole APK.
2164struct DexView {
2165    name: String,
2166    data: DexBytes,
2167    state: Arc<SearchState>,
2168    /// The memory cache entry holding it, recharged after the pass.
2169    entry: Option<Arc<CachedEntry>>,
2170    /// Its index among the entry's logical DEX files.
2171    idx: usize,
2172    /// Its tables (`SearchState::table_mask`) when loaded, and when that was:
2173    /// tables built during the pass are queued for L2 with the pass's time.
2174    mask: u32,
2175    loaded: Instant,
2176}
2177
2178/// The APK entry a (logical) DEX name belongs to: `classes2.dex!classes1.dex`
2179/// is the first logical DEX of the `classes2.dex` container.
2180fn entry_of(dex_name: &str) -> &str {
2181    dex_name.split('!').next().unwrap_or(dex_name)
2182}
2183
2184/// Memory or disk cache lookup (no inflation).
2185/// Each entry's place in the order a class loader reads DEX files: archive
2186/// by archive in program order, then `classes.dex`, `classes2.dex`, ... by
2187/// number, other names last.
2188fn load_ranks(entries: &[Entry]) -> Vec<usize> {
2189    let key = |e: &Entry| {
2190        let inner = e.name.rsplit('/').next().unwrap_or(&e.name);
2191        let number = match inner.strip_prefix("classes").and_then(|rest| rest.strip_suffix(".dex")) {
2192            Some("") => 1,
2193            Some(n) if !n.starts_with('0') => n.parse().unwrap_or(u64::MAX),
2194            _ => u64::MAX,
2195        };
2196        (e.archive, number, inner.to_string())
2197    };
2198    let mut order: Vec<usize> = (0..entries.len()).collect();
2199    order.sort_by_cached_key(|&i| key(&entries[i]));
2200    let mut ranks = vec![0; entries.len()];
2201    for (rank, i) in order.into_iter().enumerate() {
2202        ranks[i] = rank;
2203    }
2204    ranks
2205}
2206
2207fn lookup_cached(cache: &EngineCache, key: Option<&ApkKey>, entry: &Entry, stats: &Stats) -> Option<(Arc<CachedEntry>, Source)> {
2208    let key = key?;
2209    if let Some(mem) = &cache.memory
2210        && let Some(e) = lock_memory(mem).get(key, &entry.name)
2211    {
2212        stats.memory_hits.fetch_add(1, Ordering::Relaxed);
2213        return Some((e, Source::Memory));
2214    }
2215    if entry.method == asc_apk::METHOD_DEFLATE
2216        && let Some(disk) = &cache.disk
2217        && let Some(data) = disk.get(key, &entry.name, entry.uncomp_size)
2218    {
2219        stats.disk_hits.fetch_add(1, Ordering::Relaxed);
2220        let e = Arc::new(CachedEntry::new(&entry.name, Arc::new(data)));
2221        if let Some(mem) = &cache.memory {
2222            // The declared size may be 0 (undeclared) and omits a
2223            // container's logical copies; charge what will be resident.
2224            let size = e.charged_bytes();
2225            lock_memory(mem).insert(key, &entry.name, Arc::clone(&e), size);
2226            cache.charged();
2227        }
2228        return Some((e, Source::Disk));
2229    }
2230    None
2231}
2232
2233/// Bytes of `entry`: cache → (stored: borrow) → inflate (+ fill caches).
2234fn load_entry<'a>(
2235    apk: &'a Apk,
2236    key: Option<&ApkKey>,
2237    entry: &Entry,
2238    cache: Option<&EngineCache>,
2239    stats: &Stats,
2240    cancel: Option<&AtomicBool>,
2241) -> Result<(Loaded<'a>, Source)> {
2242    if let (Some(c), Some(k)) = (cache, key) {
2243        if let Some((e, source)) = lookup_cached(c, Some(k), entry, stats) {
2244            return Ok((Loaded::Cached(e), source));
2245        }
2246        // One inflate per entry at a time; whoever waited re-checks the cache.
2247        let _guard = c.begin_inflate(format!("{}/{}", k.as_str(), entry.name));
2248        if let Some((e, source)) = lookup_cached(c, Some(k), entry, stats) {
2249            return Ok((Loaded::Cached(e), source));
2250        }
2251        let bytes = apk.read_entry_cancellable(entry, cancel)?;
2252        let (data, keep_in_memory, source) = match bytes {
2253            // Stored entries are already free to read; keep them only for the
2254            // locator tables that ride along with a cached entry (in memory,
2255            // or through L2). Mapped from the APK file rather than copied, so
2256            // they stay file pages the OS can share and reclaim.
2257            Cow::Borrowed(b) => {
2258                if c.memory.is_none() && !c.disk.as_ref().is_some_and(|d| d.uses(Tier::Derived)) {
2259                    return Ok((Loaded::Borrowed(Cow::Borrowed(b)), Source::Stored));
2260                }
2261                let offset = b.as_ptr() as usize - apk.archive_data(entry).as_ptr() as usize;
2262                let data =
2263                    DexData::map_region(apk.archive_path(entry), offset as u64, b.len()).unwrap_or_else(|_| DexData::Owned(b.to_vec()));
2264                (data, c.memory.is_some(), Source::Stored)
2265            }
2266            Cow::Owned(v) => {
2267                stats.inflates.fetch_add(1, Ordering::Relaxed);
2268                (DexData::Owned(v), c.memory.is_some(), Source::Inflate)
2269            }
2270        };
2271        let e = Arc::new(CachedEntry::new(&entry.name, Arc::new(data)));
2272        let mut remap = None;
2273        if keep_in_memory && let Some(mem) = &c.memory {
2274            // Account for the container copies too, so the LRU cap holds.
2275            let size = e.charged_bytes();
2276            lock_memory(mem).insert(k, &entry.name, Arc::clone(&e), size);
2277            c.charged();
2278            remap = Some(Arc::downgrade(&e));
2279        }
2280        // Small entries are cheaper to inflate again than to write.
2281        if source == Source::Inflate && c.disk.as_ref().is_some_and(|d| d.admits_dex(e.data.len() as u64)) {
2282            let job = WriteJob {
2283                key: k.clone(),
2284                apk: Arc::from(apk.path()),
2285                kind: Job::Dex { entry: entry.clone(), limits: apk.limits(), data: Arc::clone(&e.data), remap },
2286            };
2287            c.put_async(job);
2288        }
2289        return Ok((Loaded::Cached(e), source));
2290    }
2291    let bytes = apk.read_entry_cancellable(entry, cancel)?;
2292    let source = if matches!(bytes, Cow::Owned(_)) {
2293        stats.inflates.fetch_add(1, Ordering::Relaxed);
2294        Source::Inflate
2295    } else {
2296        Source::Stored
2297    };
2298    Ok((Loaded::Borrowed(bytes), source))
2299}
2300
2301/// Decide from a streamed prefix whether `entry` defines the class.
2302/// `Ok(Some(x))` is a definitive answer; `Ok(None)` means "inflate it fully".
2303fn probe_entry(apk: &Apk, entry: &Entry, target: &[u8]) -> Result<Option<bool>> {
2304    /// Entries below this are cheaper to inflate whole.
2305    const MIN_GAIN: u64 = 256 * 1024;
2306    /// Room after the last string offset for the string's own bytes.
2307    const SLACK: usize = 64 * 1024;
2308    /// Only worth a streaming (slower) decoder below this share of the entry.
2309    const WORTH_BELOW: f64 = 0.70;
2310    if entry.method != asc_apk::METHOD_DEFLATE || entry.uncomp_size <= MIN_GAIN {
2311        return Ok(None);
2312    }
2313    let size = entry.uncomp_size as usize;
2314    let prefix = apk.inflate_until(entry, |out| {
2315        if out.len() >= 8 && &out[..8] == asc_dex::container::DEX041_MAGIC {
2316            return asc_apk::PrefixStep::Abort;
2317        }
2318        match asc_dex::quick::probe_needed_len(out) {
2319            Some(needed) if (needed as f64) < size as f64 * WORTH_BELOW => asc_apk::PrefixStep::Continue(needed + SLACK),
2320            Some(_) => asc_apk::PrefixStep::Abort,
2321            None => asc_apk::PrefixStep::Continue(0),
2322        }
2323    })?;
2324    let bytes = match prefix {
2325        asc_apk::Prefix::Complete(data) | asc_apk::Prefix::Partial(data) => data,
2326        asc_apk::Prefix::Aborted(_) => return Ok(None),
2327    };
2328    // A truncated string read means the answer is not provable from the prefix.
2329    Ok(asc_dex::quick::dex_defines_class(&bytes, target).ok())
2330}
2331
2332#[allow(clippy::too_many_arguments)]
2333fn search_entry(
2334    apk: &Apk,
2335    key: Option<&ApkKey>,
2336    entry: &Entry,
2337    query: &CompiledQuery,
2338    opts: FindOptions,
2339    shape: Shape,
2340    probe: Option<&skip::Probe<'_>>,
2341    cancel: &AtomicBool,
2342    cache: Option<&EngineCache>,
2343    stats: &Stats,
2344) -> RefBatch {
2345    let run = |fr: &FindRefs<'_, '_, '_>| {
2346        Ok(vec![match shape {
2347            Shape::Results => Rows::Results(rows_compiled(fr, query, opts)?),
2348            Shape::Text => Rows::Text(compact::text_rows(fr, query)?),
2349        }])
2350    };
2351    search_entry_with(apk, key, entry, 1, &run, probe, cancel, cache, stats).pop().expect("one batch per query")
2352}
2353
2354/// Rows of `n` queries over one logical DEX, one per query.
2355pub(crate) type RunRows<'q> = dyn Fn(&FindRefs<'_, '_, '_>) -> Result<Vec<Rows>> + Sync + 'q;
2356
2357/// The rows of one query over one logical DEX, in the form asked for.
2358pub(crate) enum Rows {
2359    Results(Vec<RefResult>),
2360    Text(TextRows),
2361}
2362
2363/// The form [`Engine::find_refs_opts`] (`Results`) and
2364/// [`Engine::find_refs_text`] (`Text`) return rows in.
2365#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2366pub(crate) enum Shape {
2367    Results,
2368    Text,
2369}
2370
2371/// Load one entry (through the caches) and run `run` over each of its
2372/// logical DEX files: one batch per query, all sharing the load. With a
2373/// `probe`, an entry no cache keeps may first be shown not to match from a
2374/// prefix, and then is not loaded at all.
2375#[allow(clippy::too_many_arguments)]
2376pub(crate) fn search_entry_with(
2377    apk: &Apk,
2378    key: Option<&ApkKey>,
2379    entry: &Entry,
2380    n: usize,
2381    run: &RunRows<'_>,
2382    probe: Option<&skip::Probe<'_>>,
2383    cancel: &AtomicBool,
2384    cache: Option<&EngineCache>,
2385    stats: &Stats,
2386) -> Vec<RefBatch> {
2387    let t0 = Instant::now();
2388    let empty = RefBatch {
2389        entry_name: entry.name.clone(),
2390        results: Vec::new(),
2391        text: None,
2392        inflate_us: 0,
2393        search_us: 0,
2394        source: "inflate",
2395        error: None,
2396    };
2397    let mut batches = vec![empty; n];
2398    if probe.is_some_and(|p| p.skips(apk, entry, cache, stats)) {
2399        let inflate_us = t0.elapsed().as_micros();
2400        for b in &mut batches {
2401            (b.source, b.inflate_us) = (Source::Skipped.as_str(), inflate_us);
2402        }
2403        return batches;
2404    }
2405    let (loaded, source) = match load_entry(apk, key, entry, cache, stats, Some(cancel)) {
2406        Ok(l) => l,
2407        Err(e) => {
2408            let msg = e.to_string();
2409            for b in &mut batches {
2410                b.error = Some(msg.clone());
2411            }
2412            return batches;
2413        }
2414    };
2415    let t1 = Instant::now();
2416    let mut add = |name: &str, r: Result<Vec<Rows>>| match r {
2417        Ok(rows) => {
2418            for (b, rows) in batches.iter_mut().zip(rows) {
2419                match rows {
2420                    Rows::Results(mut rows) => b.results.append(&mut rows),
2421                    Rows::Text(t) => match &mut b.text {
2422                        Some(text) => text.append(&t),
2423                        None => b.text = Some(t),
2424                    },
2425                }
2426            }
2427        }
2428        Err(e) => {
2429            let msg = e.to_string();
2430            for b in batches.iter_mut() {
2431                b.note_error(name, &msg);
2432            }
2433        }
2434    };
2435    match &loaded {
2436        Loaded::Cached(e) => {
2437            let l2 = cache.and_then(|c| c.disk.as_deref().filter(|d| d.uses(Tier::Derived)).map(|d| (c, d))).zip(key);
2438            for (i, l) in e.logical().iter().enumerate() {
2439                let r = Dex::with_name(&l.data, &l.name).map_err(EngineError::from).and_then(|dex| {
2440                    if let Some(((_, disk), k)) = l2
2441                        && tiers::load_derived(disk, k, l, &dex)
2442                    {
2443                        stats.derived_hits.fetch_add(1, Ordering::Relaxed);
2444                    }
2445                    let (mask, t) = (l.state.table_mask(), Instant::now());
2446                    let rows = run(&FindRefs::with_state(&dex, &l.state));
2447                    if let Some(((c, disk), k)) = l2
2448                        && let Some(job) = tiers::derived_job(disk, e, i, mask, t.elapsed())
2449                    {
2450                        c.put_async(WriteJob { key: k.clone(), apk: Arc::from(apk.path()), kind: Job::Derived(job) });
2451                    }
2452                    rows
2453                });
2454                add(&l.name, r);
2455            }
2456            recharge(cache, key, e);
2457        }
2458        Loaded::Borrowed(_) => {
2459            for (name, buf) in logical_dex_iter(&entry.name, loaded.bytes()) {
2460                let r = Dex::with_name(&buf, &name).map_err(EngineError::from).and_then(|dex| run(&FindRefs::new(&dex)));
2461                add(&name, r);
2462            }
2463        }
2464    }
2465    let (inflate_us, search_us) = ((t1 - t0).as_micros(), t1.elapsed().as_micros());
2466    for b in &mut batches {
2467        (b.source, b.inflate_us, b.search_us) = (source.as_str(), inflate_us, search_us);
2468    }
2469    batches
2470}
2471
2472/// Charge the in-memory cache for the locator tables a search just built for
2473/// `e`, evicting other entries if that exceeds the capacity. A writer may
2474/// have swapped `e` for a remapped copy sharing those tables meanwhile, so
2475/// whatever is cached under its name now is what gets measured.
2476fn recharge(cache: Option<&EngineCache>, key: Option<&ApkKey>, e: &CachedEntry) {
2477    let (Some(c), Some(key)) = (cache, key) else { return };
2478    let Some(mem) = &c.memory else { return };
2479    if lock_memory(mem).recharge(key, &e.name, CachedEntry::charged_bytes) {
2480        c.charged();
2481    }
2482}
2483
2484/// The constant a hit matched, with the note of a call-site constraint
2485/// (`Lc;->m(I)V [arg0=1]`) when it has one.
2486fn hit_matched(dex: &Dex<'_>, kind: RefKind, idx: u32, note: Option<&str>, access: Option<FieldAccess>) -> Result<Matched> {
2487    let m = matched_name(dex, kind, idx)?;
2488    if let (RefKind::Field, Some(access)) = (kind, access) {
2489        return Ok(Matched::Field { symbol: m.as_str().to_string(), access });
2490    }
2491    Ok(match note {
2492        Some(note) => Matched::Symbol(format!("{} {note}", m.as_str())),
2493        None => m,
2494    })
2495}
2496
2497fn matched_name(dex: &Dex<'_>, kind: RefKind, idx: u32) -> Result<Matched> {
2498    Ok(match kind {
2499        RefKind::String => Matched::String(dex.string(idx)?),
2500        RefKind::Type => Matched::Symbol(dex.type_descriptor(idx)?),
2501        RefKind::Method => Matched::Symbol(asc_smali::method_signature(dex, idx)?),
2502        RefKind::Field => Matched::Symbol(asc_smali::field_signature(dex, idx)?),
2503    })
2504}
2505
2506impl RefResult {
2507    /// The row of method `m` of `dex`, with no context yet.
2508    fn method_row(dex: &Dex<'_>, m: u32, matched: Vec<Matched>, sites: Vec<RefSite>) -> Result<RefResult> {
2509        Ok(RefResult {
2510            dex_name: dex.name().to_string(),
2511            method_idx: m,
2512            method: dex.method_display(m)?,
2513            signature: asc_smali::method_signature(dex, m)?,
2514            matched,
2515            relation: None,
2516            sites,
2517            context: None,
2518        })
2519    }
2520}
2521
2522/// Run a reference search over one DEX buffer and return structured results.
2523///
2524/// # Errors
2525///
2526/// [`EngineError::Dex`] when `buf` is not a valid DEX,
2527/// [`EngineError::Find`] for an invalid query.
2528pub fn find_refs_in_dex(dex_name: &str, buf: &[u8], query: &Query, aggregate: bool) -> Result<Vec<RefResult>> {
2529    find_refs_in_dex_opts(dex_name, buf, query, FindOptions { aggregate, context: RefContext::None })
2530}
2531
2532/// [`find_refs_in_dex`] with a context.
2533///
2534/// # Errors
2535///
2536/// As [`find_refs_in_dex`].
2537pub fn find_refs_in_dex_opts(dex_name: &str, buf: &[u8], query: &Query, opts: FindOptions) -> Result<Vec<RefResult>> {
2538    let rows = find_refs_in_dex_compiled(dex_name, buf, &CompiledQuery::new(query)?, opts)?;
2539    if !rows.is_empty() {
2540        return Ok(rows);
2541    }
2542    // As [`Engine::find_refs`]: names spelled as decompiled source get a second try.
2543    match names::denormalize_query(query).and_then(|q| CompiledQuery::new(&q).ok()) {
2544        Some(fallback) => find_refs_in_dex_compiled(dex_name, buf, &fallback, opts),
2545        None => Ok(rows),
2546    }
2547}
2548
2549fn find_refs_in_dex_compiled(dex_name: &str, buf: &[u8], query: &CompiledQuery, opts: FindOptions) -> Result<Vec<RefResult>> {
2550    let dex = Dex::with_name(buf, dex_name)?;
2551    rows_compiled(&FindRefs::new(&dex), query, opts)
2552}
2553
2554/// Result rows of `query` over the DEX `fr` searches: class relation rows
2555/// straight from the class definitions, otherwise the code scan's hits.
2556fn rows_compiled(fr: &FindRefs<'_, '_, '_>, query: &CompiledQuery, opts: FindOptions) -> Result<Vec<RefResult>> {
2557    let dex = fr.dex();
2558    if let Some(c) = query.compound() {
2559        return rules::compound_rows(fr, c, opts);
2560    }
2561    if let Some(d) = query.defined() {
2562        let mut rows = defined_rows(dex, d)?;
2563        if opts.context != RefContext::None {
2564            attach_context(fr, &mut rows, opts.context)?;
2565        }
2566        return Ok(rows);
2567    }
2568    let Some(classes) = fr.class_relation(query)? else {
2569        let hits = fr.find_ref_compiled(query)?;
2570        let mut rows = results_from_hits(dex, query.kind(), &hits, opts.aggregate)?;
2571        if opts.context != RefContext::None {
2572            attach_context(fr, &mut rows, opts.context)?;
2573        }
2574        return Ok(rows);
2575    };
2576    let target = match query.query() {
2577        Query::Extends(t) | Query::Implements(t) => t.as_str(),
2578        _ => "",
2579    };
2580    let dex_name = dex.name().to_string();
2581    classes
2582        .into_iter()
2583        .map(|c| {
2584            let class = dex.type_descriptor(c)?;
2585            Ok(RefResult {
2586                dex_name: dex_name.clone(),
2587                method_idx: c,
2588                method: class.clone(),
2589                signature: class,
2590                matched: vec![Matched::Symbol(target.to_string())],
2591                relation: query.relation(),
2592                sites: Vec::new(),
2593                context: None,
2594            })
2595        })
2596        .collect()
2597}
2598
2599/// Rows of a method-definition predicate: one per matching method, with no
2600/// call site (the `matched` column names the predicate).
2601pub(crate) fn defined_rows(dex: &Dex<'_>, d: &asc_findrefs::Defined) -> Result<Vec<RefResult>> {
2602    let label = d.label();
2603    defs::defined_methods(dex, d)?
2604        .into_iter()
2605        .map(|m| RefResult::method_row(dex, m, vec![Matched::Symbol(label.clone())], Vec::new()))
2606        .collect()
2607}
2608
2609/// Turn raw hits into display results (grouped per method when `aggregate`).
2610pub(crate) fn results_from_hits(dex: &Dex<'_>, kind: RefKind, hits: &[Hit], aggregate: bool) -> Result<Vec<RefResult>> {
2611    if !aggregate {
2612        let mut out = Vec::new();
2613        for hit in hits {
2614            let Some(methods) = &hit.methods else { continue };
2615            let matched = hit_matched(dex, kind, hit.matched_idx, hit.note.as_deref(), hit.field_access)?;
2616            for &m in methods.ids() {
2617                let site = RefSite { offset: hit.offset, matched: 0, field_access: hit.field_access, pc: None, line: None };
2618                out.push(RefResult::method_row(dex, m, vec![matched.clone()], vec![site])?);
2619            }
2620        }
2621        return Ok(out);
2622    }
2623    // method -> (matched (index, note, access)s, (offset, matched key) per hit)
2624    type Key<'h> = (u32, Option<&'h str>, Option<FieldAccess>);
2625    type Group<'h> = (std::collections::BTreeSet<Key<'h>>, Vec<(u32, Key<'h>)>);
2626    let mut grouped: BTreeMap<u32, Group<'_>> = BTreeMap::new();
2627    for hit in hits {
2628        let Some(methods) = &hit.methods else { continue };
2629        let key = (hit.matched_idx, hit.note.as_deref(), hit.field_access);
2630        for &m in methods.ids() {
2631            let g = grouped.entry(m).or_default();
2632            g.0.insert(key);
2633            g.1.push((hit.offset, key));
2634        }
2635    }
2636    let mut out = Vec::with_capacity(grouped.len());
2637    for (m, (idxs, offsets)) in grouped {
2638        let idxs: Vec<Key<'_>> = idxs.into_iter().collect();
2639        let matched = idxs.iter().map(|&(i, note, access)| hit_matched(dex, kind, i, note, access)).collect::<Result<Vec<_>>>()?;
2640        let sites = offsets
2641            .into_iter()
2642            .map(|(offset, idx)| RefSite {
2643                offset,
2644                matched: idxs.binary_search(&idx).unwrap_or(0),
2645                field_access: idx.2,
2646                pc: None,
2647                line: None,
2648            })
2649            .collect();
2650        out.push(RefResult::method_row(dex, m, matched, sites)?);
2651    }
2652    Ok(out)
2653}
2654
2655/// Fill the sites' `pc` / `line` and the `context` of rows found in `fr`'s DEX.
2656/// A method that cannot be rendered gets a `# context unavailable` comment.
2657fn attach_context<'r>(fr: &FindRefs<'_, '_, '_>, rows: impl IntoIterator<Item = &'r mut RefResult>, ctx: RefContext) -> Result<()> {
2658    let dex = fr.dex();
2659    let bounds = fr.insn_locator()?.method_bounds();
2660    for row in rows {
2661        let start = bounds.get(row.method_idx as usize).copied().unwrap_or(u32::MAX);
2662        let listing = match asc_smali::disassemble_method(dex, row.method_idx, SmaliOptions::default()) {
2663            Ok(l) => l,
2664            Err(e) => {
2665                row.context = Some(format!("  # context unavailable: {e}\n"));
2666                continue;
2667            }
2668        };
2669        for site in &mut row.sites {
2670            if let Some(pc) = site.offset.checked_sub(start).filter(|_| start != u32::MAX).map(|b| b / 2) {
2671                site.pc = Some(pc);
2672                site.line = listing.source_line(pc);
2673            }
2674        }
2675        row.context = Some(render_context(&listing, &row.sites, ctx));
2676    }
2677    Ok(())
2678}
2679
2680/// Most smali lines one row's context shows; beyond it the rest is summarised.
2681const CONTEXT_MAX_LINES: usize = 150;
2682/// A whole method longer than [`CONTEXT_MAX_LINES`] shows this many lines
2683/// around each reference instead.
2684const METHOD_FALLBACK_LINES: usize = 10;
2685
2686/// Context block of one row: an `@` line listing the sites, then the smali
2687/// (`>` marks the referencing instructions; windows are separated by `--`).
2688fn render_context(listing: &asc_smali::MethodListing, sites: &[RefSite], ctx: RefContext) -> String {
2689    let mut out = String::from("  @");
2690    for (i, s) in sites.iter().enumerate() {
2691        out.push_str(if i == 0 { " " } else { ", " });
2692        match (s.pc, s.line) {
2693            (Some(pc), Some(line)) => out.push_str(&format!("pc 0x{pc:x} line {line}")),
2694            (Some(pc), None) => out.push_str(&format!("pc 0x{pc:x}")),
2695            _ => out.push_str(&format!("offset 0x{:x}", s.offset)),
2696        }
2697        if let Some(access) = s.field_access {
2698            out.push_str(&format!(" {}", access.as_str()));
2699        }
2700    }
2701    out.push('\n');
2702    let lines: Vec<&str> = listing.text.lines().collect();
2703    let mut hits: Vec<usize> = sites.iter().filter_map(|s| s.pc).filter_map(|pc| listing.line_of(pc)).map(|l| l as usize).collect();
2704    hits.sort_unstable();
2705    hits.dedup();
2706    let around = |n: usize| {
2707        let mut w: Vec<(usize, usize)> = Vec::new();
2708        for &h in &hits {
2709            let (a, b) = (h.saturating_sub(n), (h + n + 1).min(lines.len()));
2710            match w.last_mut() {
2711                Some(last) if a <= last.1 => last.1 = last.1.max(b),
2712                _ => w.push((a, b)),
2713            }
2714        }
2715        w
2716    };
2717    let windows = match ctx {
2718        RefContext::None => return out,
2719        RefContext::Method if lines.len() <= CONTEXT_MAX_LINES => vec![(0, lines.len())],
2720        RefContext::Method => {
2721            out.push_str(&format!("  # method has {} lines; showing {METHOD_FALLBACK_LINES} around each reference\n", lines.len()));
2722            around(METHOD_FALLBACK_LINES)
2723        }
2724        RefContext::Lines(n) => around(n),
2725    };
2726    let mut shown = 0;
2727    for (i, &(a, b)) in windows.iter().enumerate() {
2728        if shown + (b - a) > CONTEXT_MAX_LINES {
2729            let rest: usize = windows[i..].iter().map(|w| w.1 - w.0).sum::<usize>();
2730            out.push_str(&format!("  # cut at {CONTEXT_MAX_LINES} lines ({rest} more); read the method with getclass\n"));
2731            break;
2732        }
2733        if i > 0 {
2734            out.push_str("  --\n");
2735        }
2736        for (j, line) in lines[a..b].iter().enumerate() {
2737            out.push_str(if hits.binary_search(&(a + j)).is_ok() { "  > " } else { "    " });
2738            out.push_str(line);
2739            out.push('\n');
2740        }
2741        shown += b - a;
2742    }
2743    out
2744}
2745
2746/// Python `AscHandler.findrefs` equivalent: formatted lines for one DEX.
2747///
2748/// # Errors
2749///
2750/// As [`find_refs_in_dex`].
2751pub fn find_refs_lines(dex_name: &str, buf: &[u8], query: &Query, aggregate: bool) -> Result<Vec<String>> {
2752    Ok(find_refs_in_dex(dex_name, buf, query, aggregate)?.iter().map(RefResult::line).collect())
2753}
2754
2755/// Rebuild `descriptor` out of `buf` and disassemble it.
2756///
2757/// # Errors
2758///
2759/// [`EngineError::Dex`] when `buf` is not a valid DEX,
2760/// [`EngineError::Rebuild`] when the class is missing or cannot be rebuilt,
2761/// [`EngineError::Smali`] when it cannot be disassembled.
2762pub fn decompile_dex_class(buf: &[u8], descriptor: &str, opts: SmaliOptions) -> Result<(Rebuilt, String)> {
2763    decompile_dex_class_method(buf, descriptor, opts, None)
2764}
2765
2766/// [`decompile_dex_class`] rendering only `method` (`name(params)ret`) when given.
2767///
2768/// # Errors
2769///
2770/// As [`decompile_dex_class`], plus [`EngineError::Smali`] for a malformed
2771/// `method` or one the class does not have.
2772pub fn decompile_dex_class_method(buf: &[u8], descriptor: &str, opts: SmaliOptions, method: Option<&str>) -> Result<(Rebuilt, String)> {
2773    let dex = Dex::parse(buf)?;
2774    let rebuilt = asc_rebuild::rebuild_class(&dex, descriptor)?;
2775    let new_dex = Dex::parse(&rebuilt.bytes)?;
2776    let source = asc_smali::disassemble_class_method(&new_dex, descriptor, opts, method)?;
2777    Ok((rebuilt, source))
2778}
2779
2780#[cfg(test)]
2781mod tests {
2782    use super::*;
2783
2784    /// An APK of `n` deflated copies of the fixture DEX, an L1-only cache
2785    /// (admitting entries of any size), and the entry names.
2786    fn multi_dex_apk(dir: &Path, n: usize) -> (PathBuf, DiskCache, Vec<String>) {
2787        let dex = asc_fixture::make_dex();
2788        let names: Vec<String> = (1..=n).map(|i| if i == 1 { "classes.dex".into() } else { format!("classes{i}.dex") }).collect();
2789        let entries: Vec<(&str, &[u8])> = names.iter().map(|name| (name.as_str(), dex.as_slice())).collect();
2790        let apk = dir.join("multi.apk");
2791        std::fs::write(&apk, asc_fixture::make_apk(&entries, true)).unwrap();
2792        (apk, l1_cache(&dir.join("cache")), names)
2793    }
2794
2795    fn l1_cache(dir: &Path) -> DiskCache {
2796        DiskCache::open(dir).with_capacity(1 << 30).with_min_dex_bytes(0).with_tiers(&[Tier::Dex])
2797    }
2798
2799    #[test]
2800    fn a_full_write_queue_defers_l1_writes_instead_of_dropping_them() {
2801        let dir = tempfile::tempdir().unwrap();
2802        let (apk, disk, names) = multi_dex_apk(dir.path(), 8);
2803        let mut engine = Engine::new(&apk).with_threads(4).with_cache(CacheConfig { disk: Some(disk), ..CacheConfig::default() });
2804        // Room for one entry's bytes: every other write is queued without them.
2805        Arc::get_mut(engine.cache.as_mut().unwrap()).unwrap().queue_cap = 1;
2806        engine.find_refs(&Query::String("token".into()), |batch| assert!(batch.error.is_none())).unwrap();
2807        engine.flush_cache_writes();
2808        let cache = engine.cache.as_deref().unwrap();
2809        assert_eq!(cache.pending_bytes.load(Ordering::Relaxed), 0);
2810        let key = engine.apk_key(&Apk::open(&apk).unwrap()).unwrap();
2811        let len = asc_fixture::make_dex().len() as u64;
2812        for name in &names {
2813            assert!(cache.disk.as_ref().unwrap().get(&key, name, len).is_some(), "{name} missing from L1");
2814        }
2815    }
2816
2817    #[test]
2818    fn the_write_queue_stays_within_the_memory_budget() {
2819        let dir = tempfile::tempdir().unwrap();
2820        let disk = || Some(DiskCache::open(dir.path().join("cache")));
2821        let cap = |engine: &Engine| engine.cache.as_deref().unwrap().queue_cap();
2822        assert_eq!(cap(&Engine::new("x.apk").with_cache(CacheConfig { disk: disk(), ..CacheConfig::default() })), CacheConfig::WRITE_QUEUE);
2823        let small = Engine::new("x.apk").with_cache(CacheConfig { memory_bytes: 64 << 20, disk: disk(), ..CacheConfig::default() });
2824        assert_eq!(cap(&small), 64 << 20);
2825        let large = Engine::new("x.apk").with_cache(CacheConfig { memory_bytes: 4 << 30, disk: disk(), ..CacheConfig::default() });
2826        assert_eq!(cap(&large), CacheConfig::WRITE_QUEUE);
2827        let set = Engine::new("x.apk").with_cache(CacheConfig {
2828            memory_bytes: 4 << 30,
2829            disk: disk(),
2830            write_queue_bytes: 1 << 20,
2831            ..CacheConfig::default()
2832        });
2833        assert_eq!(cap(&set), 1 << 20);
2834        assert_eq!((CacheConfig::auto_writers(4), CacheConfig::auto_writers(16), CacheConfig::auto_writers(64)), (1, 2, 4));
2835    }
2836
2837    #[test]
2838    fn a_deferred_write_inflates_its_entry_again_unless_the_apk_changed() {
2839        let dir = tempfile::tempdir().unwrap();
2840        let (apk_path, disk, names) = multi_dex_apk(dir.path(), 2);
2841        let apk = Apk::open(&apk_path).unwrap();
2842        let key = ApkKey::for_apk(&apk).unwrap();
2843        let entries = apk.dex_entries().unwrap();
2844        let limits = apk.limits();
2845        let job = |entry: &Entry| WriteJob {
2846            key: key.clone(),
2847            apk: Arc::from(apk_path.as_path()),
2848            kind: Job::Dex { entry: entry.clone(), limits, data: Arc::new(DexData::Owned(vec![0; 64])), remap: None },
2849        };
2850        // Deferring lets go of the bytes; the writer inflates the entry itself.
2851        let deferred = job(&entries[0]).defer().unwrap();
2852        assert!(matches!(deferred.kind, Job::Reinflate { .. }));
2853        assert_eq!(deferred.held_bytes(), 0);
2854        run_job(&disk, None, deferred, &mut WriterState::default());
2855        let written = disk.get(&key, &entries[0].name, entries[0].uncomp_size).expect("inflated again and written");
2856        assert_eq!(&*written, &*apk.read_entry(&entries[0]).unwrap());
2857        // Replaced since the job was queued: nothing is written for the old key.
2858        drop(apk);
2859        let replacement = dir.path().join("replacement.apk");
2860        std::fs::write(&replacement, asc_fixture::make_apk(&[(names[0].as_str(), &asc_fixture::make_dex()[..])], false)).unwrap();
2861        std::fs::rename(&replacement, &apk_path).unwrap();
2862        let other = l1_cache(&dir.path().join("other"));
2863        run_job(&other, None, job(&entries[1]).defer().unwrap(), &mut WriterState::default());
2864        assert!(other.get(&key, &entries[1].name, entries[1].uncomp_size).is_none());
2865    }
2866
2867    /// Whichever of a query's search and the writer's remap of the entry it
2868    /// just inflated runs first, the cached entry ends up mapped, keeping the
2869    /// tables the search built, and charged exactly for what it holds.
2870    #[test]
2871    fn a_remap_racing_a_search_keeps_and_charges_its_tables() {
2872        #[derive(Debug, Clone, Copy)]
2873        enum Order {
2874            RemapFirst,
2875            SearchFirst,
2876        }
2877        let dir = tempfile::tempdir().unwrap();
2878        let (apk_path, disk, _) = multi_dex_apk(dir.path(), 1);
2879        let key = ApkKey::for_apk(&Apk::open(&apk_path).unwrap()).unwrap();
2880        let engine = Engine::new(&apk_path).with_cache(CacheConfig { memory_bytes: 64 << 20, disk: Some(disk), ..CacheConfig::default() });
2881        let cache = engine.cache.as_deref().unwrap();
2882        let (memory, disk) = (cache.memory.as_ref().unwrap(), cache.disk.as_deref().unwrap());
2883        let search = |e: &Arc<CachedEntry>| {
2884            for l in e.logical() {
2885                let dex = Dex::with_name(&l.data, &l.name).unwrap();
2886                let fr = FindRefs::with_state(&dex, &l.state);
2887                prebuild::TABLES.iter().for_each(|build| build(&fr).unwrap());
2888            }
2889            recharge(Some(cache), Some(&key), e);
2890        };
2891        for (dex, order) in [asc_fixture::make_dex(), asc_fixture::make_dex041_container()]
2892            .into_iter()
2893            .flat_map(|dex| [Order::RemapFirst, Order::SearchFirst].map(|order| (dex.clone(), order)))
2894        {
2895            lock_memory(memory).clear();
2896            let name = "classes.dex";
2897            let e = Arc::new(CachedEntry::new(name, Arc::new(DexData::Owned(dex))));
2898            lock_memory(memory).insert(&key, name, Arc::clone(&e), e.charged_bytes());
2899            disk.put(&key, name, &e.data).unwrap();
2900            let remap = || remap_entry(disk, Some(memory), key.clone(), name.into(), Arc::clone(&e.data), Some(Arc::downgrade(&e)));
2901            match order {
2902                Order::RemapFirst => {
2903                    remap();
2904                    search(&e);
2905                }
2906                Order::SearchFirst => {
2907                    search(&e);
2908                    remap();
2909                }
2910            }
2911            let memory = lock_memory(memory);
2912            let current = memory.peek(&key, name).unwrap();
2913            assert!(current.is_mapped(), "{order:?}: not remapped");
2914            assert_eq!(current.state_bytes(), e.state_bytes(), "{order:?}: the remapped entry lost the search's tables");
2915            assert!(current.state_bytes() > 0);
2916            assert_eq!(memory.used_bytes(), current.charged_bytes(), "{order:?}: charged for something else than what is cached");
2917        }
2918    }
2919}