asc_findrefs/
query.rs

1//! Pre-compiled queries.
2//!
3//! A search over an APK runs the same [`Query`] against every DEX, so the
4//! regexes it needs are compiled once here and reused by every
5//! [`FindRefs`](crate::FindRefs) through `find_ref_compiled`.
6
7use crate::calls::CompiledCall;
8use crate::{ClassQuery, ClassRelation, Defined, FindError, MemberQuery, MemberSignature, Query, RefKind, Result, TypeMatch};
9use memchr::memmem::Finder;
10use regex::bytes::Regex;
11use std::collections::BTreeMap;
12
13/// Class constraint of a member query, ready to run.
14pub(crate) enum CompiledClass {
15    /// Exact descriptor (MUTF-8 encoded), looked up with a binary search over
16    /// the type table.
17    Precise(Vec<u8>),
18    /// Regex over the descriptor strings.
19    Fuzzy(Regex),
20}
21
22/// Member-name constraint of a member query, ready to run.
23pub(crate) enum CompiledName {
24    /// Substring match on the decoded name (precise class only).
25    Substring(String),
26    /// Regex over the raw string table.
27    Regex(Regex),
28    /// Whole-name equality; `bytes` is the MUTF-8 encoding of `name`.
29    Exact { name: String, bytes: Vec<u8> },
30}
31
32/// A [`TypeMatch`], MUTF-8 encoded like the descriptors it is compared with.
33pub(crate) struct CompiledType {
34    /// The whole descriptor, or `Name;` of a simple name.
35    bytes: Vec<u8>,
36    /// Array levels in front of a simple name; `None` for a descriptor.
37    simple_dims: Option<usize>,
38}
39
40impl CompiledType {
41    pub(crate) fn new(t: &TypeMatch) -> CompiledType {
42        match t {
43            TypeMatch::Descriptor(d) => CompiledType { bytes: asc_dex::mutf8::encode(d), simple_dims: None },
44            TypeMatch::Simple { name, dims } => {
45                CompiledType { bytes: asc_dex::mutf8::encode(&format!("{name};")), simple_dims: Some(*dims) }
46            }
47        }
48    }
49
50    /// Whether the raw descriptor `d` is of this type.
51    pub(crate) fn matches(&self, d: &[u8]) -> bool {
52        let Some(dims) = self.simple_dims else { return d == self.bytes };
53        // `[[Lpkg/Name;` for dims 2: the array levels, then `L`, any package
54        // ending in `/` (or none), then `Name;`.
55        let Some(class) = d.get(dims..).filter(|_| d[..dims].iter().all(|&b| b == b'[')) else { return false };
56        let Some(class) = class.strip_prefix(b"L") else { return false };
57        class.ends_with(&self.bytes) && (class.len() == self.bytes.len() || class[class.len() - self.bytes.len() - 1] == b'/')
58    }
59}
60
61/// A [`MemberSignature`] ready to run.
62pub(crate) enum CompiledSignature {
63    Method { params: Vec<CompiledType>, ret: Option<CompiledType> },
64    Field(CompiledType),
65}
66
67/// Member (method / field) query with its patterns compiled.
68pub(crate) struct CompiledMember {
69    pub class: Option<CompiledClass>,
70    pub name: Option<CompiledName>,
71    pub signature: Option<CompiledSignature>,
72}
73
74pub(crate) enum Parts {
75    /// String / type query; `None` when the pattern is empty (matches nothing).
76    Pattern(Option<Regex>),
77    /// Literal string query: the MUTF-8 encoded needle, `None` when empty.
78    Literal(Option<Finder<'static>>),
79    Member(CompiledMember),
80    /// Any of several method queries ([`Query::MethodSet`]).
81    MemberSet(Vec<CompiledMember>),
82    /// Class-definition query: the MUTF-8 encoded target descriptor.
83    Relation(ClassRelation, Vec<u8>),
84    /// Method-definition predicate ([`Query::Defined`]), resolved. It names no
85    /// constant pool index, so the engine answers it from the class
86    /// definitions rather than a code scan.
87    Defined(Defined),
88    Compound(Compound),
89    /// A method query whose call sites are judged by their values.
90    Call(CompiledCall),
91    /// A method set and the dispatch aliases among it ([`Query::Virtual`]).
92    Virtual(Box<CompiledQuery>, BTreeMap<String, String>),
93}
94
95/// The compiled parts of a [`Query::All`].
96#[non_exhaustive]
97pub struct Compound {
98    /// Every one must match in a method.
99    pub all: Vec<CompiledQuery>,
100    /// None may match in it.
101    pub none: Vec<CompiledQuery>,
102}
103
104/// A [`Query`] with every regular expression it needs compiled up front.
105///
106/// The string/type pattern, a fuzzy class pattern and a member name used as a
107/// regex are compiled here; a member name matched as a substring (precise
108/// class) or exactly (`exact_name`) is kept as the plain string. Empty patterns are kept as "no
109/// constraint" exactly as [`FindRefs::locate`](crate::FindRefs::locate)
110/// treats them.
111pub struct CompiledQuery {
112    query: Query,
113    kind: RefKind,
114    parts: Parts,
115}
116
117impl CompiledQuery {
118    /// Compile `q`.
119    ///
120    /// # Errors
121    ///
122    /// [`FindError::Pattern`] when one of its patterns is not a valid regex;
123    /// [`FindError::Invalid`] for a [`Query::Nested`] or an `override:` /
124    /// `in:` predicate the engine has not resolved yet, and for a compound
125    /// query with an empty `all`, a nested compound, or a class relation
126    /// among its parts.
127    pub fn new(q: &Query) -> Result<CompiledQuery> {
128        let parts = match q {
129            // A syntax-free ASCII word has the same language as a literal
130            // substring. Use the Finder path here (also used by rust_ASC);
131            // regex syntax and non-ASCII patterns retain regex semantics.
132            Query::String(p) if is_ascii_word(p) => Parts::Literal(Some(Finder::new(&asc_dex::mutf8::encode(p)).into_owned())),
133            Query::String(p) | Query::Type(p) => Parts::Pattern(compile_opt(p)?),
134            // MUTF-8, like the string data it is searched in, so a supplementary
135            // character (a surrogate pair there) is found too.
136            Query::StringLiteral(p) => Parts::Literal((!p.is_empty()).then(|| Finder::new(&asc_dex::mutf8::encode(p)).into_owned())),
137            Query::Method(m) => Parts::Member(compile_member(m, RefKind::Method)?),
138            Query::MethodSet(set) => Parts::MemberSet(set.iter().map(|m| compile_member(m, RefKind::Method)).collect::<Result<_>>()?),
139            Query::Nested(_) => {
140                return Err(FindError::Invalid(
141                    "method:{...} has to be resolved over every DEX first (by the engine), not compiled for one".into(),
142                ));
143            }
144            Query::Call { target, constraints } => Parts::Call(CompiledCall::new(target, constraints)?),
145            Query::Virtual { target, via, .. } => Parts::Virtual(Box::new(CompiledQuery::new(target)?), via.clone()),
146            Query::Defined(Defined::Override(_) | Defined::In(_)) => {
147                return Err(FindError::Invalid(
148                    "override: / in:{...} need the class hierarchy (every DEX); the engine resolves them before a search".into(),
149                ));
150            }
151            Query::Defined(d) => Parts::Defined(d.clone()),
152            Query::Field(m) | Query::FieldAccess { field: m, .. } => Parts::Member(compile_member(m, RefKind::Field)?),
153            Query::Extends(d) => Parts::Relation(ClassRelation::Extends, asc_dex::mutf8::encode(d)),
154            Query::Implements(d) => Parts::Relation(ClassRelation::Implements, asc_dex::mutf8::encode(d)),
155            Query::All { all, none } => {
156                if all.is_empty() {
157                    return Err(FindError::Invalid("a compound query needs at least one condition in `all`".into()));
158                }
159                let part = |q: &Query| match q {
160                    Query::All { .. } => Err(FindError::Invalid("compound queries do not nest".into())),
161                    Query::Extends(_) | Query::Implements(_) => {
162                        Err(FindError::Invalid("extends / impl name classes, not code: they cannot be combined".into()))
163                    }
164                    q => CompiledQuery::new(q),
165                };
166                Parts::Compound(Compound {
167                    all: all.iter().map(part).collect::<Result<_>>()?,
168                    none: none.iter().map(part).collect::<Result<_>>()?,
169                })
170            }
171        };
172        Ok(CompiledQuery { query: q.clone(), kind: q.kind(), parts })
173    }
174
175    /// The index table the query targets (see [`Query::kind`]).
176    #[inline]
177    pub fn kind(&self) -> RefKind {
178        self.kind
179    }
180
181    /// The query this was compiled from.
182    #[inline]
183    pub fn query(&self) -> &Query {
184        &self.query
185    }
186
187    /// The class relation an `Extends` / `Implements` query asks for.
188    #[inline]
189    pub fn relation(&self) -> Option<ClassRelation> {
190        self.query.relation()
191    }
192
193    /// The parts of a compound query ([`Query::All`]).
194    #[inline]
195    pub fn compound(&self) -> Option<&Compound> {
196        match &self.parts {
197            Parts::Compound(c) => Some(c),
198            _ => None,
199        }
200    }
201
202    /// The resolved method-definition predicate of a [`Query::Defined`].
203    #[inline]
204    pub fn defined(&self) -> Option<&Defined> {
205        match &self.parts {
206            Parts::Defined(d) => Some(d),
207            _ => None,
208        }
209    }
210
211    #[inline]
212    pub(crate) fn parts(&self) -> &Parts {
213        &self.parts
214    }
215}
216
217impl std::fmt::Debug for CompiledQuery {
218    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
219        f.debug_struct("CompiledQuery").field("query", &self.query).finish_non_exhaustive()
220    }
221}
222
223/// A nonempty ASCII letter/digit/underscore sequence contains no regex
224/// syntax and has no Unicode matching behavior, so literal search is exact.
225fn is_ascii_word(pattern: &str) -> bool {
226    !pattern.is_empty() && pattern.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_')
227}
228
229fn compile_opt(pattern: &str) -> Result<Option<Regex>> {
230    if pattern.is_empty() { Ok(None) } else { Ok(Some(Regex::new(pattern)?)) }
231}
232
233fn compile_member(m: &MemberQuery, kind: RefKind) -> Result<CompiledMember> {
234    let class = match m.class.as_ref().filter(|c| !c.pattern.is_empty()) {
235        None => None,
236        Some(ClassQuery { pattern, precise: true }) => Some(CompiledClass::Precise(asc_dex::mutf8::encode(pattern))),
237        Some(ClassQuery { pattern, precise: false }) => Some(CompiledClass::Fuzzy(Regex::new(pattern)?)),
238    };
239    let name = match m.name.as_deref().filter(|n| !n.is_empty()) {
240        None => None,
241        Some(n) if m.exact_name => Some(CompiledName::Exact { name: n.to_string(), bytes: asc_dex::mutf8::encode(n) }),
242        // With a precise class the member name is a substring match (as the
243        // Python implementation does); otherwise it is a regex.
244        Some(n) if matches!(class, Some(CompiledClass::Precise(_))) => Some(CompiledName::Substring(n.to_string())),
245        Some(n) => Some(CompiledName::Regex(Regex::new(n)?)),
246    };
247    let signature = match (&m.signature, kind) {
248        (None, _) => None,
249        (Some(MemberSignature::Method { params, ret }), RefKind::Method) => Some(CompiledSignature::Method {
250            params: params.iter().map(CompiledType::new).collect(),
251            ret: ret.as_ref().map(CompiledType::new),
252        }),
253        (Some(MemberSignature::Field(t)), RefKind::Field) => Some(CompiledSignature::Field(CompiledType::new(t))),
254        (Some(sig), _) => return Err(FindError::Invalid(format!("signature {sig} does not fit a {} query", kind.as_str()))),
255    };
256    Ok(CompiledMember { class, name, signature })
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    fn assert_send_sync<T: Send + Sync>() {}
264
265    #[test]
266    fn compiled_query_is_send_sync() {
267        assert_send_sync::<CompiledQuery>();
268    }
269
270    #[test]
271    fn compiles_only_what_is_used() {
272        let q = Query::Method(MemberQuery {
273            class: Some(ClassQuery { pattern: "Lexample/Test;".into(), precise: true }),
274            name: Some("(".into()), // invalid regex, but a substring here
275            exact_name: false,
276            signature: None,
277        });
278        let c = CompiledQuery::new(&q).unwrap();
279        assert_eq!(c.kind(), RefKind::Method);
280        assert_eq!(c.query(), &q);
281        let Parts::Member(m) = c.parts() else { panic!() };
282        assert!(matches!(m.class, Some(CompiledClass::Precise(_))));
283        assert!(matches!(m.name, Some(CompiledName::Substring(_))));
284
285        // An exact name is never compiled as a regex, whatever the class.
286        for class in [None, Some(ClassQuery { pattern: "Lexample/Test;".into(), precise: true })] {
287            let q = Query::Field(MemberQuery { class, name: Some("a(".into()), exact_name: true, signature: None });
288            let c = CompiledQuery::new(&q).unwrap();
289            let Parts::Member(m) = c.parts() else { panic!() };
290            assert!(matches!(&m.name, Some(CompiledName::Exact { name, bytes }) if name == "a(" && bytes == b"a("));
291        }
292        let q = Query::Method(MemberQuery { class: None, name: Some("\u{1F600}".into()), exact_name: true, signature: None });
293        let c = CompiledQuery::new(&q).unwrap();
294        let Parts::Member(m) = c.parts() else { panic!() };
295        assert!(matches!(&m.name, Some(CompiledName::Exact { bytes, .. }) if bytes == &[0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]));
296
297        let q = Query::Field(MemberQuery { class: None, name: Some("(".into()), exact_name: false, signature: None });
298        assert!(CompiledQuery::new(&q).is_err());
299        let q = Query::Field(MemberQuery {
300            class: Some(ClassQuery { pattern: "(".into(), precise: false }),
301            name: None,
302            exact_name: false,
303            signature: None,
304        });
305        assert!(CompiledQuery::new(&q).is_err());
306        let q = Query::String("[".into());
307        assert!(CompiledQuery::new(&q).is_err());
308
309        // A literal is never a regex; an empty one matches nothing.
310        let c = CompiledQuery::new(&Query::StringLiteral("[".into())).unwrap();
311        assert_eq!(c.kind(), RefKind::String);
312        assert!(matches!(c.parts(), Parts::Literal(Some(f)) if f.needle() == b"["));
313        let c = CompiledQuery::new(&Query::StringLiteral("\u{1F600}".into())).unwrap();
314        assert!(matches!(c.parts(), Parts::Literal(Some(f)) if f.needle() == [0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]));
315        assert!(matches!(CompiledQuery::new(&Query::StringLiteral(String::new())).unwrap().parts(), Parts::Literal(None)));
316
317        let c = CompiledQuery::new(&Query::String("token".into())).unwrap();
318        assert!(matches!(c.parts(), Parts::Literal(Some(f)) if f.needle() == b"token"));
319        let c = CompiledQuery::new(&Query::String("t.k".into())).unwrap();
320        assert!(matches!(c.parts(), Parts::Pattern(Some(_))));
321
322        let c = CompiledQuery::new(&Query::Implements("Lexample/Iface;".into())).unwrap();
323        assert_eq!(c.relation(), Some(ClassRelation::Implements));
324        assert!(matches!(c.parts(), Parts::Relation(ClassRelation::Implements, d) if d == b"Lexample/Iface;"));
325        assert_eq!(CompiledQuery::new(&Query::String("x".into())).unwrap().relation(), None);
326
327        let c = CompiledQuery::new(&Query::Type(String::new())).unwrap();
328        assert!(matches!(c.parts(), Parts::Pattern(None)));
329        let q = Query::Method(MemberQuery {
330            class: Some(ClassQuery { pattern: String::new(), precise: false }),
331            name: None,
332            exact_name: false,
333            signature: None,
334        });
335        let c = CompiledQuery::new(&q).unwrap();
336        let Parts::Member(m) = c.parts() else { panic!() };
337        assert!(m.class.is_none() && m.name.is_none());
338    }
339}
asc_findrefs/query.rs

asc_findrefs/
query.rs
