Skip to content

perf: optimize APK query pipeline #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MG1937 merged 1 commit into from
Oct 20, 2026
Merged

perf: optimize APK query pipeline #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
220 changes: 220 additions & 0 deletions src/asc_client/apk/dex_reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,226 @@ fn _decode_name(raw_name: &[u8]) -> String {
}
}

/// Smallest DEX prefix that contains every class descriptor string.
///
/// R8/D8 emit `class_defs` before the data section, so the id tables pin down
/// each descriptor's `string_data_off`; only those strings (plus a NUL) need to
/// be decompressed to list class names.
/// Smallest DEX prefix that can answer "does this DEX define `target`?".
///
/// The answer needs only the id tables plus the descriptor string; code items,
/// debug info and class_data are irrelevant. A class lookup therefore inflates
/// ~10-30% of each DEX it rejects instead of the whole file.
pub(crate) fn class_defines_prefix_len(buf: &[u8], target: &[u8]) -> Result<usize, String> {
const SLACK: usize = 1024;
if buf.len() < 0x70 {
return Ok(0x70);
}
if buf[..3] != *b"dex" {
// Not a DEX prefix we can reason about; ask for everything.
return Ok(buf.len().saturating_add(SLACK));
}
let string_ids_size = _STRUCT_I::unpack_from(buf, 0x38) as usize;
let string_ids_off = _STRUCT_I::unpack_from(buf, 0x3C) as usize;
let type_ids_size = _STRUCT_I::unpack_from(buf, 0x40) as usize;
let type_ids_off = _STRUCT_I::unpack_from(buf, 0x44) as usize;
let class_defs_size = _STRUCT_I::unpack_from(buf, 0x60) as usize;
let class_defs_off = _STRUCT_I::unpack_from(buf, 0x64) as usize;

let tables_end = string_ids_off
.saturating_add(string_ids_size * 4)
.max(type_ids_off.saturating_add(type_ids_size * 4))
.max(class_defs_off.saturating_add(class_defs_size * 32));
if buf.len() < tables_end {
return Ok(tables_end);
}

let Ok(type_idx) = _find_type_idx(buf, target) else {
return Ok(buf.len().saturating_add(SLACK));
};
if type_idx < 0 {
// Descriptor absent: the id tables alone already give the answer.
return Ok(tables_end);
}
let string_idx = _STRUCT_I::unpack_from(buf, type_ids_off + (type_idx as usize) * 4) as usize;
if string_idx >= string_ids_size {
return Ok(tables_end);
}
let string_off = _STRUCT_I::unpack_from(buf, string_ids_off + (string_idx << 2)) as usize;
match _read_string_data_bytes(buf, string_off) {
Ok(raw) => {
let start = raw.as_ptr() as usize - buf.as_ptr() as usize;
Ok((start + raw.len() + 1).max(tables_end))
}
Err(_) => Ok(buf.len().saturating_add(SLACK)),
}
}

pub(crate) fn class_names_prefix_len(buf: &[u8]) -> Result<usize, String> {
const SLACK: usize = 4096;
if buf.len() < 0x70 {
return Ok(0x70);
}
if buf[..3] != *b"dex" {
return Err("invalid DEX header".to_owned());
}
let string_ids_size = _STRUCT_I::unpack_from(buf, 0x38) as usize;
let string_ids_off = _STRUCT_I::unpack_from(buf, 0x3C) as usize;
let type_ids_size = _STRUCT_I::unpack_from(buf, 0x40) as usize;
let type_ids_off = _STRUCT_I::unpack_from(buf, 0x44) as usize;
let class_defs_size = _STRUCT_I::unpack_from(buf, 0x60) as usize;
let class_defs_off = _STRUCT_I::unpack_from(buf, 0x64) as usize;

// Probe one string_data offset before requesting the complete id tables.
// A late string pool is a hint that prefix decoding will be expensive. Ask
// for the full entry early so the loader's existing fallback can restart
// after only a header-sized probe, rather than redoing all of the id tables.
// Physical string order is not guaranteed by DEX: this is only a choice of
// decoder, never an assumption used to omit classes or validate an offset.
if string_ids_size != 0 {
let first_id_end = string_ids_off
.checked_add(4)
.ok_or_else(|| "bad string_ids range".to_owned())?;
if buf.len() < first_id_end {
return Ok(first_id_end);
}
let first_data_off = _STRUCT_I::unpack_from(buf, string_ids_off) as u64;
let file_size = _STRUCT_I::unpack_from(buf, 0x20) as u64;
if first_data_off < file_size && first_data_off * 5 > file_size * 3 {
return Ok(usize::MAX);
}
}

let Some(id_end) = class_defs_off.checked_add(class_defs_size * 32) else {
return Ok(buf.len());
};
let tables_end = id_end
.max(string_ids_off.saturating_add(string_ids_size * 4))
.max(type_ids_off.saturating_add(type_ids_size * 4));
if buf.len() < tables_end {
return Ok(tables_end);
}

let descriptor_off = |class_index: usize| -> Option<usize> {
let class_off = class_defs_off + class_index * 32;
let type_idx = _STRUCT_I::unpack_from(buf, class_off) as usize;
if type_idx >= type_ids_size {
return None;
}
let string_idx = _STRUCT_I::unpack_from(buf, type_ids_off + (type_idx << 2)) as usize;
if string_idx >= string_ids_size {
return None;
}
Some(_STRUCT_I::unpack_from(buf, string_ids_off + (string_idx << 2)) as usize)
};

let mut max_start = tables_end;
for class_index in 0..class_defs_size {
if let Some(offset) = descriptor_off(class_index) {
max_start = max_start.max(offset);
}
}
if buf.len() < max_start.saturating_add(SLACK) {
return Ok(max_start.saturating_add(SLACK));
}

let mut max_end = max_start.saturating_add(1);
for class_index in 0..class_defs_size {
let Some(offset) = descriptor_off(class_index) else {
continue;
};
match _read_string_data_bytes(buf, offset) {
Ok(raw) => {
let start = raw.as_ptr() as usize - buf.as_ptr() as usize;
max_end = max_end.max(start + raw.len() + 1);
}
// The terminator is past the current prefix; ask for more.
Err(_) => return Ok(buf.len().saturating_add(SLACK)),
}
}
Ok(max_end)
}

pub(crate) fn write_dex_class_names(
buf: &[u8],
prefix: Option<&str>,
prefix_bytes: Option<&[u8]>,
out: &mut Vec<u8>,
) -> Result<usize, String> {
if buf.len() < 0x70 || buf[..3] != *b"dex" {
return Err("invalid DEX header".to_owned());
}

let string_ids_size = _STRUCT_I::unpack_from(buf, 0x38) as usize;
let string_ids_off = _STRUCT_I::unpack_from(buf, 0x3C) as usize;
let type_ids_size = _STRUCT_I::unpack_from(buf, 0x40) as usize;
let type_ids_off = _STRUCT_I::unpack_from(buf, 0x44) as usize;
let class_defs_size = _STRUCT_I::unpack_from(buf, 0x60) as usize;
let class_defs_off = _STRUCT_I::unpack_from(buf, 0x64) as usize;

if class_defs_size == 0 {
return Ok(0);
}
if string_ids_off + string_ids_size * 4 > buf.len() {
return Err("bad string_ids range".to_owned());
}
if type_ids_off + type_ids_size * 4 > buf.len() {
return Err("bad type_ids range".to_owned());
}
if class_defs_off + class_defs_size * 32 > buf.len() {
return Err("bad class_defs range".to_owned());
}

let start_len = out.len();
let mut count = 0usize;
let class_defs_end = class_defs_off + class_defs_size * 32;
for class_def_off in (class_defs_off..class_defs_end).step_by(32) {
let type_idx = _STRUCT_I::unpack_from(buf, class_def_off) as usize;
if type_idx >= type_ids_size {
return Err("bad class_def->type_idx".to_owned());
}
let string_idx = _STRUCT_I::unpack_from(buf, type_ids_off + (type_idx << 2)) as usize;
if string_idx >= string_ids_size {
return Err("bad type_id->string_idx".to_owned());
}
let string_off = _STRUCT_I::unpack_from(buf, string_ids_off + (string_idx << 2)) as usize;
if string_off >= buf.len() {
return Err("bad string_data_off".to_owned());
}
let raw_name = _read_string_data_bytes(buf, string_off)
.map_err(|_| "bad string_data_off".to_owned())?;

// ASCII prefixes can be rejected before decoding every class name.
if let Some(prefix_bytes) = prefix_bytes {
if !raw_name.starts_with(prefix_bytes) {
continue;
}
}
let name = std::str::from_utf8(raw_name);
match name {
Ok(name) => {
if prefix.is_some() && prefix_bytes.is_none() && !name.starts_with(prefix.unwrap()) {
continue;
}
out.reserve(name.len() + 1);
out.extend_from_slice(name.as_bytes());
}
Err(_) => {
let name = _decode_name(raw_name);
if prefix.is_some_and(|prefix| !name.starts_with(prefix)) {
continue;
}
out.reserve(name.len() + 1);
out.extend_from_slice(name.as_bytes());
}
}
out.push(b'\n');
count += 1;
}
let _ = start_len;
Ok(count)
}

pub(crate) fn list_dex_classes(
buf: &[u8],
prefix: Option<&str>,
Expand Down
120 changes: 120 additions & 0 deletions src/asc_client/apk/loader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,126 @@ fn _validate_inflated_size(entry: &DexEntry, data: Vec<u8>) -> Result<Vec<u8>, S
Ok(data)
}

/// Inflate only as much of an entry as `need` asks for.
///
/// `listclass` only reads the DEX id tables plus the class descriptor strings,
/// which live in the first ~20-35% of a D8/R8 DEX; the remaining code/debug
/// sections never need to leave the deflate stream. `need` is re-evaluated on
/// each round and returns the prefix length required by the current data.
pub(crate) fn _inflate_dex_prefix_into<F>(
buf: &[u8],
entry: &DexEntry,
out: &mut Vec<u8>,
decompressor: &mut Decompressor,
need: F,
) -> Result<(), String>
where
F: Fn(&[u8]) -> Result<usize, String>,
{
let comp_view = _compressed_data(buf, entry)?;
out.clear();
match entry.comp_method {
0 => {
out.extend_from_slice(comp_view);
return Ok(());
}
8 => {}
method => return Err(format!("unsupported compression method: {method}")),
}

let total = entry.uncomp_size as usize;
// Low-level streaming API: `flate2`'s Read adapter copies through an
// intermediate buffer, while `decompress` writes straight into `out`.
let mut decoder = flate2::Decompress::new(false);
let mut produced = 0usize;
let mut step = 1usize << 20;
let mut round = 0usize;
loop {
let target = need(&out[..produced])?.min(total);
if produced >= target {
out.truncate(produced);
return Ok(());
}
// Some APKs keep class descriptors near the end of the DEX. Streaming
// (flate2) decodes slower than the one-shot libdeflate path, so switch
// back when the saving would be small.
if round > 0 && target * 5 > total * 3 {
return _inflate_dex_into_with(buf, entry, out, decompressor);
}
round += 1;

let want = (target - produced).min(step);
let consumed = decoder.total_in() as usize;
out.resize(produced + want, 0);
decoder
.decompress(
&comp_view[consumed..],
&mut out[produced..produced + want],
flate2::FlushDecompress::None,
)
.map_err(|error| format!("flate2 cannot decompress {}: {error}", entry.name))?;
produced = decoder.total_out() as usize;
if produced == 0 || (produced < target && consumed == decoder.total_in() as usize) {
out.truncate(produced);
return Err(format!(
"size mismatch: expect {}, got {}",
entry.uncomp_size,
produced
));
}
step = (step * 2).min(8 << 20);
}
}

/// Inflate one entry into a caller-owned buffer so a worker can reuse the same
/// pages across entries. Fresh multi-MB allocations must be zero-filled on
/// first touch, which roughly halves libdeflate's throughput; reusing a warm
/// buffer measured ~915 MB/s against ~460 MB/s for a fresh `Vec`.
pub(crate) fn _inflate_dex_into(
buf: &[u8],
entry: &DexEntry,
out: &mut Vec<u8>,
) -> Result<(), String> {
_inflate_dex_into_with(buf, entry, out, &mut Decompressor::new())
}

/// Same as [`_inflate_dex_into`] with a reusable decompressor: libdeflate
/// builds its decode tables per instance, which measured ~7% of inflate time.
pub(crate) fn _inflate_dex_into_with(
buf: &[u8],
entry: &DexEntry,
out: &mut Vec<u8>,
decompressor: &mut Decompressor,
) -> Result<(), String> {
let comp_view = _compressed_data(buf, entry)?;
out.clear();
match entry.comp_method {
0 => out.extend_from_slice(comp_view),
8 => {
let size = entry.uncomp_size as usize;
out.reserve(size);
// SAFETY: `reserve` guarantees capacity for `size` bytes and every
// byte is written by the decompressor before it is observed.
unsafe { out.set_len(size) };
let actual_size = decompressor
.deflate_decompress(comp_view, out)
.map_err(|error| {
format!("libdeflate cannot decompress {}: {error}", entry.name)
})?;
out.truncate(actual_size);
}
method => return Err(format!("unsupported compression method: {method}")),
}
if entry.uncomp_size != 0 && out.len() != entry.uncomp_size as usize {
return Err(format!(
"size mismatch: expect {}, got {}",
entry.uncomp_size,
out.len()
));
}
Ok(())
}

pub(crate) fn _inflate_dex(
buf: &[u8],
entry: &DexEntry,
Expand Down
Loading
Loading