qcode/
discovery.rs

1//! Typed queue of code addresses discovered during lifting/analysis but not yet
2//! lifted into the IR.
3//!
4//! Durable discovery records use stable addresses, not context-local `FunctionId`
5//! or `BlockId` values. The queue is shared between a raw clean context and
6//! disposable optimized clones, so IDs from one context must not leak into the
7//! other.
8
9use std::collections::BTreeMap;
10
11pub type Address = u64;
12
13/// Why an address is believed to start a function.
14#[derive(
15    Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
16)]
17pub enum FunctionDiscoveryReason {
18    Entry,
19    CrtMain,
20    CallTarget,
21    TailCall,
22    UserSeed,
23    /// A constant code-pointer argument passed into a callee parameter that is
24    /// itself called indirectly (`call [@param]`). See the
25    /// `propagate_code_pointer_args` pass.
26    CodePointer,
27}
28
29/// The control-flow edge that exposed a block target.
30#[derive(
31    Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
32)]
33pub enum EdgeKind {
34    Entry,
35    DirectBranch,
36    ConditionalBranch,
37    Fallthrough,
38    CallTarget,
39    CallFallthrough,
40    JumpTableTarget,
41    TailCall,
42    Speculative,
43}
44
45/// Whether a discovered address starts a new function or extends an existing one.
46#[derive(
47    Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
48)]
49pub enum DiscoveryKind {
50    Function {
51        reason: FunctionDiscoveryReason,
52    },
53    Block {
54        /// Entry address of the owning function.
55        function: Address,
56        edge_kind: EdgeKind,
57    },
58}
59
60/// Stable identity of a discovery item.
61#[derive(
62    Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
63)]
64pub struct DiscoveryKey {
65    pub target: Address,
66    pub kind: DiscoveryKind,
67    /// Source block for edge-bearing discoveries. Two transfers to the same
68    /// target are distinct work items because the lifter must attach both
69    /// source edges after the target exists.
70    pub source_block: Option<Address>,
71}
72
73/// Why this discovery exists. Kept as debugging/UI metadata that records how a
74/// discovered address came to be queued.
75#[derive(Clone, Default, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
76pub enum DiscoveryProvenance {
77    LoaderEntry,
78
79    DirectLift {
80        source_addr: Address,
81    },
82
83    Optimization {
84        pass: String,
85        assumption: Option<String>,
86    },
87
88    UserSeed,
89
90    #[default]
91    Unknown,
92}
93/// A code address discovered but not yet lifted.
94#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
95pub struct Discovery {
96    pub target: Address,
97    pub kind: DiscoveryKind,
98    /// Symbol-name hint for the discovered code, when known.
99    pub name: Option<String>,
100    /// Stable source metadata. These are addresses so the item can move between
101    /// clean and optimized contexts.
102    pub source_function: Option<Address>,
103    pub source_block: Option<Address>,
104    pub source_addr: Option<Address>,
105    pub provenance: DiscoveryProvenance,
106}
107
108impl Discovery {
109    pub fn function(target: Address) -> Self {
110        Self {
111            target,
112            kind: DiscoveryKind::Function {
113                reason: FunctionDiscoveryReason::CallTarget,
114            },
115            name: None,
116            source_function: None,
117            source_block: None,
118            source_addr: None,
119            provenance: DiscoveryProvenance::Unknown,
120        }
121    }
122
123    pub fn entry(target: Address) -> Self {
124        Self::function(target)
125            .with_function_reason(FunctionDiscoveryReason::Entry)
126            .with_provenance(DiscoveryProvenance::LoaderEntry)
127    }
128
129    pub fn block(target: Address, function: Address) -> Self {
130        Self {
131            target,
132            kind: DiscoveryKind::Block {
133                function,
134                edge_kind: EdgeKind::DirectBranch,
135            },
136            name: None,
137            source_function: Some(function),
138            source_block: None,
139            source_addr: None,
140            provenance: DiscoveryProvenance::Unknown,
141        }
142    }
143
144    pub fn key(&self) -> DiscoveryKey {
145        DiscoveryKey {
146            target: self.target,
147            kind: self.kind.clone(),
148            source_block: self.source_block,
149        }
150    }
151
152    pub fn with_name(mut self, name: Option<String>) -> Self {
153        self.name = name;
154        self
155    }
156
157    pub fn from_addr(mut self, addr: Address) -> Self {
158        self.source_addr = Some(addr);
159        if matches!(self.provenance, DiscoveryProvenance::Unknown) {
160            self.provenance = DiscoveryProvenance::DirectLift { source_addr: addr };
161        }
162        self
163    }
164
165    pub fn from_block_addr(mut self, addr: Address) -> Self {
166        self.source_block = Some(addr);
167        self
168    }
169
170    pub fn with_edge_kind(mut self, edge_kind: EdgeKind) -> Self {
171        if let DiscoveryKind::Block {
172            edge_kind: existing,
173            ..
174        } = &mut self.kind
175        {
176            *existing = edge_kind;
177        }
178        self
179    }
180
181    pub fn with_function_reason(mut self, reason: FunctionDiscoveryReason) -> Self {
182        if let DiscoveryKind::Function {
183            reason: existing, ..
184        } = &mut self.kind
185        {
186            *existing = reason;
187        }
188        self
189    }
190
191    pub fn with_provenance(mut self, provenance: DiscoveryProvenance) -> Self {
192        self.provenance = provenance;
193        self
194    }
195}
196
197#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
198pub enum DiscoveryState {
199    Pending,
200    Lifted,
201    Failed { reason: String },
202    Skipped { reason: String },
203}
204
205/// The pending-discovery queue and durable outcome state.
206#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
207pub struct DiscoveryQueue {
208    pending: BTreeMap<DiscoveryKey, Discovery>,
209    states: BTreeMap<DiscoveryKey, DiscoveryState>,
210    /// The first [`Discovery`] recorded for each key, kept after it drains so
211    /// the provenance of a lifted (or failed) address stays inspectable.
212    records: BTreeMap<DiscoveryKey, Discovery>,
213}
214
215impl DiscoveryQueue {
216    /// Record a discovery unless this exact key already has a terminal outcome.
217    pub fn insert(&mut self, discovery: Discovery) -> bool {
218        let key = discovery.key();
219        if matches!(
220            self.states.get(&key),
221            Some(
222                DiscoveryState::Lifted
223                    | DiscoveryState::Failed { .. }
224                    | DiscoveryState::Skipped { .. }
225            )
226        ) {
227            return false;
228        }
229        let inserted = !self.pending.contains_key(&key);
230        self.records
231            .entry(key.clone())
232            .or_insert_with(|| discovery.clone());
233        self.pending.entry(key.clone()).or_insert(discovery);
234        self.states.entry(key).or_insert(DiscoveryState::Pending);
235        inserted
236    }
237
238    pub fn drain(&mut self) -> Vec<Discovery> {
239        std::mem::take(&mut self.pending).into_values().collect()
240    }
241
242    pub fn iter(&self) -> impl Iterator<Item = &Discovery> + '_ {
243        self.pending.values()
244    }
245
246    pub fn is_empty(&self) -> bool {
247        self.pending.is_empty()
248    }
249
250    pub fn pending_len(&self) -> usize {
251        self.pending.len()
252    }
253
254    pub fn mark_lifted(&mut self, key: DiscoveryKey) {
255        self.states.insert(key, DiscoveryState::Lifted);
256    }
257
258    pub fn mark_failed(&mut self, key: DiscoveryKey, reason: impl Into<String>) {
259        self.states.insert(
260            key,
261            DiscoveryState::Failed {
262                reason: reason.into(),
263            },
264        );
265    }
266
267    pub fn mark_skipped(&mut self, key: DiscoveryKey, reason: impl Into<String>) {
268        self.states.insert(
269            key,
270            DiscoveryState::Skipped {
271                reason: reason.into(),
272            },
273        );
274    }
275
276    pub fn state(&self, key: &DiscoveryKey) -> Option<&DiscoveryState> {
277        self.states.get(key)
278    }
279
280    pub fn states(&self) -> impl Iterator<Item = (&DiscoveryKey, &DiscoveryState)> + '_ {
281        self.states.iter()
282    }
283
284    /// Every discovery ever recorded, with its current outcome. Unlike
285    /// [`iter`](Self::iter) this includes drained (lifted, failed, skipped)
286    /// items, so it is the complete history of what was found and how.
287    pub fn records(
288        &self,
289    ) -> impl Iterator<Item = (&DiscoveryKey, &Discovery, &DiscoveryState)> + '_ {
290        self.states.iter().filter_map(|(key, state)| {
291            self.records
292                .get(key)
293                .map(|discovery| (key, discovery, state))
294        })
295    }
296
297    /// Every code address this queue lifted successfully, as a portable
298    /// [`CodeSeed`]. Exported from one analysis run and replayed into the next
299    /// (via [`Context::seed_code`](crate::context::Context::seed_code)) so the
300    /// lifter reaches jump-table targets in its first pass instead of waiting for
301    /// the analysis fixpoint to discover them round by round.
302    pub fn lifted_seeds(&self) -> Vec<CodeSeed> {
303        self.states
304            .iter()
305            .filter(|(_, state)| matches!(state, DiscoveryState::Lifted))
306            .map(|(key, _)| CodeSeed::from_key(key))
307            .collect()
308    }
309}
310
311/// A code address known to lift, captured from one analysis run to pre-seed the
312/// next. Mirrors a [`DiscoveryKey`] but drops the run-specific bits the lifter
313/// reconstructs on its own (the precise CFG edge is re-attached by the pass that
314/// originally found the target).
315#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
316pub enum CodeSeed {
317    /// A discovered function entry.
318    Function { target: Address },
319    /// A block at `target` inside the function entered at `function`.
320    Block {
321        function: Address,
322        target: Address,
323        edge_kind: EdgeKind,
324    },
325}
326
327impl CodeSeed {
328    fn from_key(key: &DiscoveryKey) -> Self {
329        match key.kind {
330            DiscoveryKind::Function { .. } => CodeSeed::Function { target: key.target },
331            DiscoveryKind::Block {
332                function,
333                edge_kind,
334            } => CodeSeed::Block {
335                function,
336                target: key.target,
337                edge_kind,
338            },
339        }
340    }
341
342    /// Rebuild the [`Discovery`] to enqueue, attributed to [`UserSeed`] so the
343    /// origin stays honest rather than impersonating the pass that first found it.
344    ///
345    /// [`UserSeed`]: DiscoveryProvenance::UserSeed
346    pub fn into_discovery(self) -> Discovery {
347        match self {
348            CodeSeed::Function { target } => {
349                Discovery::function(target).with_provenance(DiscoveryProvenance::UserSeed)
350            }
351            CodeSeed::Block {
352                function,
353                target,
354                edge_kind,
355            } => Discovery::block(target, function)
356                .with_edge_kind(edge_kind)
357                .with_provenance(DiscoveryProvenance::UserSeed),
358        }
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    #[test]
367    fn dedups_exact_key_not_target() {
368        let mut q = DiscoveryQueue::default();
369        assert!(q.insert(Discovery::block(0x1000, 0x900)));
370        assert!(!q.insert(Discovery::block(0x1000, 0x900)));
371        assert!(q.insert(Discovery::block(0x1000, 0x900).from_block_addr(0x910)));
372        assert!(q.insert(Discovery::block(0x1000, 0x900).from_block_addr(0x920)));
373        assert!(q.insert(Discovery::block(0x1000, 0xa00)));
374        assert!(q.insert(Discovery::function(0x1000)));
375
376        let all: Vec<_> = q.iter().collect();
377        assert_eq!(all.len(), 5);
378    }
379
380    #[test]
381    fn terminal_state_blocks_requeue_of_same_key() {
382        let mut q = DiscoveryQueue::default();
383        let d = Discovery::function(0x10);
384        let key = d.key();
385        assert!(q.insert(d.clone()));
386        q.drain();
387        q.mark_failed(key, "decode");
388        assert!(!q.insert(d));
389        assert!(q.is_empty());
390    }
391
392    #[test]
393    fn terminal_target_from_one_source_does_not_suppress_another_source() {
394        let mut q = DiscoveryQueue::default();
395        let first = Discovery::block(0x1000, 0x900).from_block_addr(0x910);
396        let second = Discovery::block(0x1000, 0x900).from_block_addr(0x920);
397        let first_key = first.key();
398
399        assert!(q.insert(first));
400        q.drain();
401        q.mark_lifted(first_key);
402
403        assert!(q.insert(second));
404        assert_eq!(q.pending_len(), 1);
405    }
406
407    #[test]
408    fn drain_empties_pending_but_keeps_state() {
409        let mut q = DiscoveryQueue::default();
410        let d = Discovery::function(0x10);
411        let key = d.key();
412        q.insert(d);
413        let drained = q.drain();
414        assert_eq!(drained.len(), 1);
415        assert!(q.is_empty());
416        assert_eq!(q.state(&key), Some(&DiscoveryState::Pending));
417    }
418}
qcode/discovery.rs

qcode/
discovery.rs
