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`].
96pub struct Compound {
97    /// Every one must match in a method.
98    pub all: Vec<CompiledQuery>,
99    /// None may match in it.
100    pub none: Vec<CompiledQuery>,
101}
102
103/// A [`Query`] with every regular expression it needs compiled up front.
104///
105/// The string/type pattern, a fuzzy class pattern and a member name used as a
106/// regex are compiled here; a member name matched as a substring (precise
107/// class) or exactly (`exact_name`) is kept as the plain string. Empty patterns are kept as "no
108/// constraint" exactly as [`FindRefs::locate`](crate::FindRefs::locate)
109/// treats them.
110pub struct CompiledQuery {
111    query: Query,
112    kind: RefKind,
113    parts: Parts,
114}
115
116impl CompiledQuery {
117    /// Compile `q`. Fails only when one of its patterns is not a valid regex.
118    pub fn new(q: &Query) -> Result<CompiledQuery> {
119        let parts = match q {
120            // A syntax-free ASCII word has the same language as a literal
121            // substring. Use the Finder path here (also used by rust_ASC);
122            // regex syntax and non-ASCII patterns retain regex semantics.
123            Query::String(p) if is_ascii_word(p) => Parts::Literal(Some(Finder::new(&asc_dex::mutf8::encode(p)).into_owned())),
124            Query::String(p) | Query::Type(p) => Parts::Pattern(compile_opt(p)?),
125            // MUTF-8, like the string data it is searched in, so a supplementary
126            // character (a surrogate pair there) is found too.
127            Query::StringLiteral(p) => Parts::Literal((!p.is_empty()).then(|| Finder::new(&asc_dex::mutf8::encode(p)).into_owned())),
128            Query::Method(m) => Parts::Member(compile_member(m, RefKind::Method)?),
129            Query::MethodSet(set) => Parts::MemberSet(set.iter().map(|m| compile_member(m, RefKind::Method)).collect::<Result<_>>()?),
130            Query::Nested(_) => {
131                return Err(FindError::Invalid(
132                    "method:{...} has to be resolved over every DEX first (by the engine), not compiled for one".into(),
133                ));
134            }
135            Query::Call { target, constraints } => Parts::Call(CompiledCall::new(target, constraints)?),
136            Query::Virtual { target, via, .. } => Parts::Virtual(Box::new(CompiledQuery::new(target)?), via.clone()),
137            Query::Defined(Defined::Override(_) | Defined::In(_)) => {
138                return Err(FindError::Invalid(
139                    "override: / in:{...} need the class hierarchy (every DEX); the engine resolves them before a search".into(),
140                ));
141            }
142            Query::Defined(d) => Parts::Defined(d.clone()),
143            Query::Field(m) | Query::FieldAccess { field: m, .. } => Parts::Member(compile_member(m, RefKind::Field)?),
144            Query::Extends(d) => Parts::Relation(ClassRelation::Extends, asc_dex::mutf8::encode(d)),
145            Query::Implements(d) => Parts::Relation(ClassRelation::Implements, asc_dex::mutf8::encode(d)),
146            Query::All { all, none } => {
147                if all.is_empty() {
148                    return Err(FindError::Invalid("a compound query needs at least one condition in `all`".into()));
149                }
150                let part = |q: &Query| match q {
151                    Query::All { .. } => Err(FindError::Invalid("compound queries do not nest".into())),
152                    Query::Extends(_) | Query::Implements(_) => {
153                        Err(FindError::Invalid("extends / impl name classes, not code: they cannot be combined".into()))
154                    }
155                    q => CompiledQuery::new(q),
156                };
157                Parts::Compound(Compound {
158                    all: all.iter().map(part).collect::<Result<_>>()?,
159                    none: none.iter().map(part).collect::<Result<_>>()?,
160                })
161            }
162        };
163        Ok(CompiledQuery { query: q.clone(), kind: q.kind(), parts })
164    }
165
166    #[inline]
167    pub fn kind(&self) -> RefKind {
168        self.kind
169    }
170
171    #[inline]
172    pub fn query(&self) -> &Query {
173        &self.query
174    }
175
176    /// The class relation an `Extends` / `Implements` query asks for.
177    #[inline]
178    pub fn relation(&self) -> Option<ClassRelation> {
179        self.query.relation()
180    }
181
182    /// The parts of a compound query ([`Query::All`]).
183    #[inline]
184    pub fn compound(&self) -> Option<&Compound> {
185        match &self.parts {
186            Parts::Compound(c) => Some(c),
187            _ => None,
188        }
189    }
190
191    /// The resolved method-definition predicate of a [`Query::Defined`].
192    #[inline]
193    pub fn defined(&self) -> Option<&Defined> {
194        match &self.parts {
195            Parts::Defined(d) => Some(d),
196            _ => None,
197        }
198    }
199
200    #[inline]
201    pub(crate) fn parts(&self) -> &Parts {
202        &self.parts
203    }
204}
205
206impl std::fmt::Debug for CompiledQuery {
207    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
208        f.debug_struct("CompiledQuery").field("query", &self.query).finish_non_exhaustive()
209    }
210}
211
212/// A nonempty ASCII letter/digit/underscore sequence contains no regex
213/// syntax and has no Unicode matching behavior, so literal search is exact.
214fn is_ascii_word(pattern: &str) -> bool {
215    !pattern.is_empty() && pattern.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_')
216}
217
218fn compile_opt(pattern: &str) -> Result<Option<Regex>> {
219    if pattern.is_empty() { Ok(None) } else { Ok(Some(Regex::new(pattern)?)) }
220}
221
222fn compile_member(m: &MemberQuery, kind: RefKind) -> Result<CompiledMember> {
223    let class = match m.class.as_ref().filter(|c| !c.pattern.is_empty()) {
224        None => None,
225        Some(ClassQuery { pattern, precise: true }) => Some(CompiledClass::Precise(asc_dex::mutf8::encode(pattern))),
226        Some(ClassQuery { pattern, precise: false }) => Some(CompiledClass::Fuzzy(Regex::new(pattern)?)),
227    };
228    let name = match m.name.as_deref().filter(|n| !n.is_empty()) {
229        None => None,
230        Some(n) if m.exact_name => Some(CompiledName::Exact { name: n.to_string(), bytes: asc_dex::mutf8::encode(n) }),
231        // With a precise class the member name is a substring match (as the
232        // Python implementation does); otherwise it is a regex.
233        Some(n) if matches!(class, Some(CompiledClass::Precise(_))) => Some(CompiledName::Substring(n.to_string())),
234        Some(n) => Some(CompiledName::Regex(Regex::new(n)?)),
235    };
236    let signature = match (&m.signature, kind) {
237        (None, _) => None,
238        (Some(MemberSignature::Method { params, ret }), RefKind::Method) => Some(CompiledSignature::Method {
239            params: params.iter().map(CompiledType::new).collect(),
240            ret: ret.as_ref().map(CompiledType::new),
241        }),
242        (Some(MemberSignature::Field(t)), RefKind::Field) => Some(CompiledSignature::Field(CompiledType::new(t))),
243        (Some(sig), _) => return Err(FindError::Invalid(format!("signature {sig} does not fit a {} query", kind.as_str()))),
244    };
245    Ok(CompiledMember { class, name, signature })
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    fn assert_send_sync<T: Send + Sync>() {}
253
254    #[test]
255    fn compiled_query_is_send_sync() {
256        assert_send_sync::<CompiledQuery>();
257    }
258
259    #[test]
260    fn compiles_only_what_is_used() {
261        let q = Query::Method(MemberQuery {
262            class: Some(ClassQuery { pattern: "Lexample/Test;".into(), precise: true }),
263            name: Some("(".into()), // invalid regex, but a substring here
264            exact_name: false,
265            signature: None,
266        });
267        let c = CompiledQuery::new(&q).unwrap();
268        assert_eq!(c.kind(), RefKind::Method);
269        assert_eq!(c.query(), &q);
270        let Parts::Member(m) = c.parts() else { panic!() };
271        assert!(matches!(m.class, Some(CompiledClass::Precise(_))));
272        assert!(matches!(m.name, Some(CompiledName::Substring(_))));
273
274        // An exact name is never compiled as a regex, whatever the class.
275        for class in [None, Some(ClassQuery { pattern: "Lexample/Test;".into(), precise: true })] {
276            let q = Query::Field(MemberQuery { class, name: Some("a(".into()), exact_name: true, signature: None });
277            let c = CompiledQuery::new(&q).unwrap();
278            let Parts::Member(m) = c.parts() else { panic!() };
279            assert!(matches!(&m.name, Some(CompiledName::Exact { name, bytes }) if name == "a(" && bytes == b"a("));
280        }
281        let q = Query::Method(MemberQuery { class: None, name: Some("\u{1F600}".into()), exact_name: true, signature: None });
282        let c = CompiledQuery::new(&q).unwrap();
283        let Parts::Member(m) = c.parts() else { panic!() };
284        assert!(matches!(&m.name, Some(CompiledName::Exact { bytes, .. }) if bytes == &[0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]));
285
286        let q = Query::Field(MemberQuery { class: None, name: Some("(".into()), exact_name: false, signature: None });
287        assert!(CompiledQuery::new(&q).is_err());
288        let q = Query::Field(MemberQuery {
289            class: Some(ClassQuery { pattern: "(".into(), precise: false }),
290            name: None,
291            exact_name: false,
292            signature: None,
293        });
294        assert!(CompiledQuery::new(&q).is_err());
295        let q = Query::String("[".into());
296        assert!(CompiledQuery::new(&q).is_err());
297
298        // A literal is never a regex; an empty one matches nothing.
299        let c = CompiledQuery::new(&Query::StringLiteral("[".into())).unwrap();
300        assert_eq!(c.kind(), RefKind::String);
301        assert!(matches!(c.parts(), Parts::Literal(Some(f)) if f.needle() == b"["));
302        let c = CompiledQuery::new(&Query::StringLiteral("\u{1F600}".into())).unwrap();
303        assert!(matches!(c.parts(), Parts::Literal(Some(f)) if f.needle() == [0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]));
304        assert!(matches!(CompiledQuery::new(&Query::StringLiteral(String::new())).unwrap().parts(), Parts::Literal(None)));
305
306        let c = CompiledQuery::new(&Query::String("token".into())).unwrap();
307        assert!(matches!(c.parts(), Parts::Literal(Some(f)) if f.needle() == b"token"));
308        let c = CompiledQuery::new(&Query::String("t.k".into())).unwrap();
309        assert!(matches!(c.parts(), Parts::Pattern(Some(_))));
310
311        let c = CompiledQuery::new(&Query::Implements("Lexample/Iface;".into())).unwrap();
312        assert_eq!(c.relation(), Some(ClassRelation::Implements));
313        assert!(matches!(c.parts(), Parts::Relation(ClassRelation::Implements, d) if d == b"Lexample/Iface;"));
314        assert_eq!(CompiledQuery::new(&Query::String("x".into())).unwrap().relation(), None);
315
316        let c = CompiledQuery::new(&Query::Type(String::new())).unwrap();
317        assert!(matches!(c.parts(), Parts::Pattern(None)));
318        let q = Query::Method(MemberQuery {
319            class: Some(ClassQuery { pattern: String::new(), precise: false }),
320            name: None,
321            exact_name: false,
322            signature: None,
323        });
324        let c = CompiledQuery::new(&q).unwrap();
325        let Parts::Member(m) = c.parts() else { panic!() };
326        assert!(m.class.is_none() && m.name.is_none());
327    }
328}
asc_findrefs/query.rs

asc_findrefs/
query.rs
