1use crate::{
26 CachedEntry, DexBytes, FindOptions, LogicalDex, Matched, MemoryPolicy, RefBatch, RefResult, RefSite, Shape, SharedMemory, TextRows,
27 available_memory,
28};
29use asc_cache::{ApkKey, DiskCache, LockExt, ProcessRegistry};
30use asc_dex::Dex;
31use asc_findrefs::{ClassRelation, Query, SearchState};
32use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
33use std::sync::{Arc, Mutex, OnceLock};
34use std::time::{Duration, Instant};
35
36pub(crate) const LOCATORS_KIND: &str = "locators";
38pub(crate) const L3_MAX_RESULT: usize = 4 << 20;
40const L3_MAX_TEXT: usize = 32 << 20;
44const RESULT_FORMAT: u32 = 6;
46const RESULT_END: u32 = u32::MAX;
49
50#[derive(Default)]
53pub(crate) struct Persisted {
54 tried: AtomicBool,
56 queued: AtomicU32,
58 cost_us: AtomicU64,
60}
61
62pub(crate) fn load_derived(disk: &DiskCache, key: &ApkKey, l: &LogicalDex, dex: &Dex<'_>) -> bool {
64 if l.persisted.tried.swap(true, Ordering::Relaxed) {
65 return false;
66 }
67 let Some(d) = disk.get_derived(key, &l.name, LOCATORS_KIND, asc_findrefs::PERSIST_FORMAT_VERSION) else { return false };
68 let cost = d.build_cost.as_micros() as u64;
69 let Some(state) = SearchState::from_persisted(dex, Arc::new(d)) else { return false };
70 l.persisted.queued.fetch_or(state.table_mask(), Ordering::Relaxed);
71 l.persisted.cost_us.fetch_add(cost, Ordering::Relaxed);
72 l.state.fill_from(state);
73 true
74}
75
76pub(crate) struct DerivedJob {
78 entry: String,
80 logical: usize,
81 name: String,
82 data: DexBytes,
83 state: Arc<SearchState>,
84 persisted: Arc<Persisted>,
85 claimed: u32,
87}
88
89impl DerivedJob {
90 pub(crate) fn held_bytes(&self) -> usize {
92 self.state.heap_bytes()
93 }
94
95 pub(crate) fn abandon(self) {
97 self.persisted.queued.fetch_and(!self.claimed, Ordering::Relaxed);
98 }
99}
100
101pub(crate) fn derived_job(disk: &DiskCache, entry: &CachedEntry, idx: usize, mask_before: u32, spent: Duration) -> Option<DerivedJob> {
105 let l = entry.logical().get(idx)?;
106 let mask = l.state.table_mask();
107 let p = &l.persisted;
108 if mask == mask_before && mask & !p.queued.load(Ordering::Relaxed) == 0 {
109 return None;
110 }
111 let cost = if mask != mask_before {
112 p.cost_us.fetch_add(spent.as_micros() as u64, Ordering::Relaxed) + spent.as_micros() as u64
113 } else {
114 p.cost_us.load(Ordering::Relaxed)
115 };
116 if !disk.admits_derived(Duration::from_micros(cost)) {
117 return None;
118 }
119 let claimed = mask & !p.queued.fetch_or(mask, Ordering::Relaxed);
120 if claimed == 0 {
121 return None;
122 }
123 Some(DerivedJob {
124 entry: entry.name.clone(),
125 logical: idx,
126 name: l.name.clone(),
127 data: Arc::clone(&l.data),
128 state: Arc::clone(&l.state),
129 persisted: Arc::clone(&l.persisted),
130 claimed,
131 })
132}
133
134pub(crate) fn write_derived(disk: &DiskCache, memory: Option<&SharedMemory>, key: &ApkKey, job: DerivedJob) {
137 let Ok(dex) = Dex::with_name(&job.data, &job.name) else { return };
138 let mask = job.state.table_mask();
139 let bytes = job.state.persist(&dex);
140 let cost = Duration::from_micros(job.persisted.cost_us.load(Ordering::Relaxed));
141 let version = asc_findrefs::PERSIST_FORMAT_VERSION;
142 if disk.put_derived(key, &job.name, LOCATORS_KIND, version, cost, &bytes).is_err() {
143 return;
144 }
145 drop(bytes);
146 let Some(memory) = memory else { return };
147 let Some(mapped) =
148 disk.get_derived(key, &job.name, LOCATORS_KIND, version).and_then(|d| SearchState::from_persisted(&dex, Arc::new(d)))
149 else {
150 return;
151 };
152 if mapped.table_mask() != mask || job.state.table_mask() != mask {
155 return;
156 }
157 let mut mem = crate::lock_memory(memory);
158 let Some(current) = mem.peek(key, &job.entry) else { return };
159 let still_ours = current.logical.get().and_then(|l| l.get(job.logical)).is_some_and(|l| Arc::ptr_eq(&l.state, &job.state));
160 if still_ours {
161 let new = Arc::new(current.with_state(job.logical, Arc::new(mapped)));
162 if mem.replace(key, &job.entry, ¤t, new) {
163 mem.recharge(key, &job.entry, crate::CachedEntry::charged_bytes);
164 }
165 }
166}
167
168fn tool_version() -> &'static str {
172 static VERSION: OnceLock<String> = OnceLock::new();
173 VERSION.get_or_init(|| {
174 let exe = std::env::current_exe().ok().and_then(|p| std::fs::metadata(p).ok());
175 let stamp = exe.map(|m| {
176 let mtime = m.modified().ok().and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map_or(0, |d| d.as_nanos());
177 format!("{}.{mtime}", m.len())
178 });
179 format!("{}+{}", env!("CARGO_PKG_VERSION"), stamp.unwrap_or_default())
180 })
181}
182
183pub(crate) fn findrefs_key(query: &Query, opts: FindOptions, shape: Shape) -> String {
185 let query = match shape {
186 Shape::Results => format!("{query:?}\0{opts:?}"),
187 Shape::Text => format!("{query:?}\0{opts:?}\0text"),
188 };
189 asc_cache::rendered_key("findrefs", tool_version(), &query)
190}
191
192pub(crate) fn result_limit(disk: &DiskCache, shape: Shape) -> usize {
194 match shape {
195 Shape::Results => L3_MAX_RESULT,
196 Shape::Text => usize::try_from(disk.capacity_for(0) / 8).unwrap_or(usize::MAX).min(L3_MAX_TEXT),
197 }
198}
199
200pub(crate) struct ResultRecorder {
205 out: Vec<u8>,
206 keep: bool,
207 limit: usize,
208}
209
210impl ResultRecorder {
211 pub(crate) fn new(limit: usize) -> ResultRecorder {
212 let mut out = Vec::new();
213 out.extend_from_slice(&RESULT_FORMAT.to_le_bytes());
214 ResultRecorder { out, keep: true, limit }
215 }
216
217 fn abandon(&mut self) {
219 self.keep = false;
220 self.out = Vec::new();
221 }
222
223 pub(crate) fn record(&mut self, b: &RefBatch) {
224 if !self.keep {
225 return;
226 }
227 let text_bytes = b.text.as_ref().map_or(0, |t| t.as_str().len() + 4 * t.len());
229 if b.error.is_some() || self.out.len().saturating_add(text_bytes) > self.limit {
230 self.abandon();
231 return;
232 }
233 put_str(&mut self.out, &b.entry_name);
234 put_u32(&mut self.out, b.results.len() as u32);
235 for r in &b.results {
236 put_str(&mut self.out, &r.dex_name);
237 put_u32(&mut self.out, r.method_idx);
238 put_str(&mut self.out, &r.method);
239 put_str(&mut self.out, &r.signature);
240 put_u32(&mut self.out, r.matched.len() as u32);
241 for m in &r.matched {
242 self.out.push(match m {
243 Matched::String(_) => 0,
244 Matched::Symbol(_) => 1,
245 Matched::Field { .. } => 2,
246 });
247 put_str(&mut self.out, m.as_str());
248 if let Matched::Field { access, .. } = m {
249 self.out.push(match access {
250 asc_findrefs::FieldAccess::Read => 1,
251 asc_findrefs::FieldAccess::Write => 2,
252 });
253 }
254 }
255 self.out.push(match r.relation {
256 None => 0,
257 Some(ClassRelation::Extends) => 1,
258 Some(ClassRelation::Implements) => 2,
259 });
260 put_u32(&mut self.out, r.sites.len() as u32);
261 for site in &r.sites {
262 put_u32(&mut self.out, site.offset);
263 put_u32(&mut self.out, site.matched as u32);
264 self.out.push(match site.field_access {
265 None => 0,
266 Some(asc_findrefs::FieldAccess::Read) => 1,
267 Some(asc_findrefs::FieldAccess::Write) => 2,
268 });
269 put_opt(&mut self.out, site.pc);
270 put_opt(&mut self.out, site.line);
271 }
272 match &r.context {
273 None => self.out.push(0),
274 Some(c) => {
275 self.out.push(1);
276 put_str(&mut self.out, c);
277 }
278 }
279 }
280 match &b.text {
281 None => self.out.push(0),
282 Some(t) => {
283 let Ok(ends) = t.ends().iter().map(|&e| u32::try_from(e)).collect::<Result<Vec<_>, _>>() else {
284 return self.abandon();
285 };
286 self.out.push(1);
287 put_str(&mut self.out, t.as_str());
288 put_u32(&mut self.out, ends.len() as u32);
289 for e in ends {
290 put_u32(&mut self.out, e);
291 }
292 }
293 }
294 if self.out.len() > self.limit {
295 self.abandon();
296 }
297 }
298
299 pub(crate) fn finish(mut self, notes: &[String]) -> Option<Vec<u8>> {
301 put_u32(&mut self.out, RESULT_END);
302 put_u32(&mut self.out, notes.len() as u32);
303 for n in notes {
304 put_str(&mut self.out, n);
305 }
306 self.keep.then_some(self.out)
307 }
308}
309
310fn put_u32(out: &mut Vec<u8>, v: u32) {
311 out.extend_from_slice(&v.to_le_bytes());
312}
313
314fn put_opt(out: &mut Vec<u8>, v: Option<u32>) {
315 match v {
316 None => out.push(0),
317 Some(v) => {
318 out.push(1);
319 put_u32(out, v);
320 }
321 }
322}
323
324fn put_str(out: &mut Vec<u8>, s: &str) {
325 put_u32(out, s.len() as u32);
326 out.extend_from_slice(s.as_bytes());
327}
328
329pub(crate) struct Recorded {
331 pub batches: Vec<RefBatch>,
332 pub notes: Vec<String>,
333}
334
335pub(crate) fn decode_batches(payload: impl AsRef<[u8]> + Send + Sync + 'static) -> Option<Recorded> {
338 let payload: Arc<dyn AsRef<[u8]> + Send + Sync> = Arc::new(payload);
339 let bytes = (*payload).as_ref();
340 let mut r = Cursor { bytes, at: 0 };
341 if r.u32()? != RESULT_FORMAT {
342 return None;
343 }
344 let mut out = Vec::new();
345 loop {
346 let len = r.u32()?;
347 if len == RESULT_END {
348 let notes = (0..r.u32()?).map(|_| r.str()).collect::<Option<Vec<_>>>()?;
349 return (r.at == bytes.len()).then_some(Recorded { batches: out, notes });
350 }
351 let entry_name = String::from_utf8(r.take(len as usize)?.to_vec()).ok()?;
352 let n = r.u32()? as usize;
353 let mut results = Vec::with_capacity(n.min(bytes.len()));
354 for _ in 0..n {
355 let dex_name = r.str()?;
356 let method_idx = r.u32()?;
357 let method = r.str()?;
358 let signature = r.str()?;
359 let m = r.u32()? as usize;
360 let matched = (0..m)
361 .map(|_| match r.u8()? {
362 0 => Some(Matched::String(r.str()?)),
363 1 => Some(Matched::Symbol(r.str()?)),
364 2 => {
365 let symbol = r.str()?;
366 let access = match r.u8()? {
367 1 => asc_findrefs::FieldAccess::Read,
368 2 => asc_findrefs::FieldAccess::Write,
369 _ => return None,
370 };
371 Some(Matched::Field { symbol, access })
372 }
373 _ => None,
374 })
375 .collect::<Option<Vec<_>>>()?;
376 let relation = match r.u8()? {
377 0 => None,
378 1 => Some(ClassRelation::Extends),
379 2 => Some(ClassRelation::Implements),
380 _ => return None,
381 };
382 let sites = (0..r.u32()?)
383 .map(|_| {
384 let offset = r.u32()?;
385 let matched = r.u32()? as usize;
386 let field_access = match r.u8()? {
387 0 => None,
388 1 => Some(asc_findrefs::FieldAccess::Read),
389 2 => Some(asc_findrefs::FieldAccess::Write),
390 _ => return None,
391 };
392 Some(RefSite { offset, matched, field_access, pc: r.opt()?, line: r.opt()? })
393 })
394 .collect::<Option<Vec<_>>>()?;
395 let context = match r.u8()? {
396 0 => None,
397 1 => Some(r.str()?),
398 _ => return None,
399 };
400 results.push(RefResult { dex_name, method_idx, method, signature, matched, relation, sites, context });
401 }
402 let text = match r.u8()? {
403 0 => None,
404 1 => {
405 let len = r.u32()? as usize;
406 let start = r.at;
407 r.take(len)?;
408 let ends = (0..r.u32()?).map(|_| r.u32().map(|e| e as usize)).collect::<Option<Vec<_>>>()?;
409 Some(TextRows::shared(Arc::clone(&payload), start..start + len, ends)?)
410 }
411 _ => return None,
412 };
413 out.push(RefBatch { entry_name, results, text, inflate_us: 0, search_us: 0, source: "result", error: None });
414 }
415}
416
417struct Cursor<'b> {
418 bytes: &'b [u8],
419 at: usize,
420}
421
422impl Cursor<'_> {
423 fn take(&mut self, n: usize) -> Option<&[u8]> {
424 let s = self.bytes.get(self.at..self.at.checked_add(n)?)?;
425 self.at += n;
426 Some(s)
427 }
428
429 fn u8(&mut self) -> Option<u8> {
430 Some(self.take(1)?[0])
431 }
432
433 fn u32(&mut self) -> Option<u32> {
434 Some(u32::from_le_bytes(self.take(4)?.try_into().ok()?))
435 }
436
437 fn opt(&mut self) -> Option<Option<u32>> {
439 match self.u8()? {
440 0 => Some(None),
441 1 => Some(Some(self.u32()?)),
442 _ => None,
443 }
444 }
445
446 fn str(&mut self) -> Option<String> {
447 let n = self.u32()? as usize;
448 String::from_utf8(self.take(n)?.to_vec()).ok()
449 }
450}
451
452const REBALANCE_EVERY: Duration = Duration::from_secs(1);
454
455struct Registration {
457 caches: usize,
458 registry: Option<ProcessRegistry>,
459 held: u64,
461 requests: usize,
463}
464
465impl Registration {
466 fn publish(&self) {
467 if let Some(registry) = &self.registry {
468 let _ = registry.update(self.held, self.requests > 0);
469 }
470 }
471}
472
473static REGISTRATION: Mutex<Registration> = Mutex::new(Registration { caches: 0, registry: None, held: 0, requests: 0 });
474
475#[must_use = "the process is busy only while the guard lives"]
480pub struct Busy(());
481
482impl Busy {
483 pub fn enter() -> Busy {
490 let mut reg = REGISTRATION.lock().expect("process registration mutex poisoned");
491 reg.requests += 1;
492 if reg.requests == 1 {
493 reg.publish();
494 }
495 Busy(())
496 }
497}
498
499impl Drop for Busy {
500 fn drop(&mut self) {
501 let mut reg = REGISTRATION.lock().expect("process registration mutex poisoned");
502 reg.requests -= 1;
503 if reg.requests == 0 {
504 reg.publish();
505 }
506 }
507}
508
509pub(crate) struct MemoryBudget {
513 auto: Option<MemoryPolicy>,
514 last: Mutex<Option<Instant>>,
515 published: AtomicU64,
516}
517
518impl MemoryBudget {
519 pub(crate) fn new(auto: Option<MemoryPolicy>) -> MemoryBudget {
520 let mut reg = REGISTRATION.lock().expect("process registration mutex poisoned");
521 reg.caches += 1;
522 if reg.registry.is_none() {
523 reg.registry = ProcessRegistry::default_dir().and_then(ProcessRegistry::open);
524 }
525 MemoryBudget { auto, last: Mutex::new(None), published: AtomicU64::new(0) }
526 }
527
528 pub(crate) fn maybe_rebalance(&self, memory: &SharedMemory) {
531 {
532 let mut last = self.last.lock_or_default();
533 if last.is_some_and(|t| t.elapsed() < REBALANCE_EVERY) {
534 return;
535 }
536 *last = Some(Instant::now());
537 }
538 let used = {
539 let m = crate::lock_memory(memory);
540 (m.used_bytes() + m.reserved_bytes()) as u64
541 };
542 let previous = self.published.swap(used, Ordering::Relaxed);
543 let (held, others) = {
544 let mut reg = REGISTRATION.lock().expect("process registration mutex poisoned");
545 reg.held = (reg.held + used).saturating_sub(previous);
546 reg.publish();
547 let Some(registry) = ®.registry else { return };
548 (reg.held, registry.others())
549 };
550 if let Some(policy) = self.auto {
551 let cap = policy.share(available_memory(), held, &others);
552 crate::lock_memory(memory).set_capacity(cap);
553 }
554 }
555}
556
557impl Drop for MemoryBudget {
558 fn drop(&mut self) {
559 let mut reg = REGISTRATION.lock().expect("process registration mutex poisoned");
560 reg.held = reg.held.saturating_sub(self.published.load(Ordering::Relaxed));
561 reg.caches -= 1;
562 if reg.caches == 0 {
563 reg.registry = None;
565 } else {
566 reg.publish();
567 }
568 }
569}
570
571#[cfg(test)]
572mod tests {
573 use super::*;
574
575 #[test]
576 fn results_roundtrip_and_errors_are_not_kept() {
577 let row = |rel| RefResult {
578 dex_name: "classes.dex".into(),
579 method_idx: 7,
580 method: "La;->b".into(),
581 signature: "La;->b()V".into(),
582 matched: vec![
583 Matched::String("x\ny".into()),
584 Matched::Symbol("La;->c:I".into()),
585 Matched::Field { symbol: "La;->f:I".into(), access: asc_findrefs::FieldAccess::Read },
586 ],
587 relation: rel,
588 sites: vec![
589 RefSite { offset: 9, matched: 1, field_access: None, pc: Some(3), line: None },
590 RefSite { offset: 11, matched: 2, field_access: Some(asc_findrefs::FieldAccess::Read), pc: Some(4), line: Some(8) },
591 ],
592 context: rel.map(|_| " const-string v0, \"x\"".to_string()),
593 };
594 let batch = |name: &str, results, error: Option<String>| RefBatch {
595 entry_name: name.into(),
596 results,
597 text: None,
598 inflate_us: 5,
599 search_us: 6,
600 source: "inflate",
601 error,
602 };
603 let mut rec = ResultRecorder::new(L3_MAX_RESULT);
604 rec.record(&batch("classes.dex", vec![row(None), row(Some(ClassRelation::Implements))], None));
605 rec.record(&batch("classes2.dex", vec![], None));
606 let bytes = rec.finish(&["a note".to_string()]).unwrap();
607 let Recorded { batches: back, notes } = decode_batches(bytes.clone()).unwrap();
608 assert_eq!(notes, ["a note"]);
609 assert_eq!(back.len(), 2);
610 assert_eq!(back[0].results, vec![row(None), row(Some(ClassRelation::Implements))]);
611 assert_eq!((back[1].entry_name.as_str(), back[1].source), ("classes2.dex", "result"));
612 assert!(decode_batches(bytes[..bytes.len() - 1].to_vec()).is_none());
613 let mut one = ResultRecorder::new(L3_MAX_RESULT);
615 one.record(&batch("classes2.dex", vec![], None));
616 let one = one.finish(&[]).unwrap();
617 assert!(decode_batches(one[..one.len() - 4].to_vec()).is_none());
618 assert_eq!(decode_batches(one).unwrap().batches.len(), 1);
619 assert!(decode_batches(Vec::new()).is_none());
620 let mut failed = ResultRecorder::new(L3_MAX_RESULT);
621 failed.record(&batch("classes.dex", vec![], Some("bad".into())));
622 assert!(failed.finish(&[]).is_none());
623 }
624
625 #[test]
626 fn text_rows_roundtrip_within_their_limit() {
627 let text = |lines: &[&str]| {
628 let bytes: Arc<dyn AsRef<[u8]> + Send + Sync> =
629 Arc::new(lines.iter().map(|l| format!("{l}\n")).collect::<String>().into_bytes());
630 let mut ends = Vec::new();
631 let mut at = 0;
632 for l in lines {
633 at += l.len();
634 ends.push(at);
635 at += 1;
636 }
637 let len = (*bytes).as_ref().len();
638 TextRows::shared(bytes, 0..len, ends).unwrap()
639 };
640 let batch = |name: &str, t: TextRows| RefBatch {
641 entry_name: name.into(),
642 results: Vec::new(),
643 text: Some(t),
644 inflate_us: 0,
645 search_us: 0,
646 source: "inflate",
647 error: None,
648 };
649 let a = text(&["classes.dex | La;->b()V | \"x\"", "classes.dex | La;->c()V | \"y\""]);
650 let mut rec = ResultRecorder::new(1 << 10);
651 rec.record(&batch("classes.dex", a.clone()));
652 rec.record(&batch("classes2.dex", text(&[])));
653 let Recorded { batches: back, notes } = decode_batches(rec.finish(&["n".to_string()]).unwrap()).unwrap();
654 assert_eq!(notes, ["n"]);
655 assert_eq!(back[0].text.as_ref(), Some(&a));
656 assert_eq!(back[0].row_count(), 2);
657 assert_eq!(back[1].text.as_ref().map(TextRows::len), Some(0));
658 let mut rec = ResultRecorder::new(64);
660 rec.record(&batch("classes.dex", a));
661 assert!(rec.out.is_empty() && rec.finish(&[]).is_none());
662 }
663}