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}
211
212impl DiscoveryQueue {
213    /// Record a discovery unless this exact key already has a terminal outcome.
214    pub fn insert(&mut self, discovery: Discovery) -> bool {
215        let key = discovery.key();
216        if matches!(
217            self.states.get(&key),
218            Some(
219                DiscoveryState::Lifted
220                    | DiscoveryState::Failed { .. }
221                    | DiscoveryState::Skipped { .. }
222            )
223        ) {
224            return false;
225        }
226        let inserted = !self.pending.contains_key(&key);
227        self.pending.entry(key.clone()).or_insert(discovery);
228        self.states.entry(key).or_insert(DiscoveryState::Pending);
229        inserted
230    }
231
232    pub fn drain(&mut self) -> Vec<Discovery> {
233        std::mem::take(&mut self.pending).into_values().collect()
234    }
235
236    pub fn iter(&self) -> impl Iterator<Item = &Discovery> + '_ {
237        self.pending.values()
238    }
239
240    pub fn is_empty(&self) -> bool {
241        self.pending.is_empty()
242    }
243
244    pub fn pending_len(&self) -> usize {
245        self.pending.len()
246    }
247
248    pub fn mark_lifted(&mut self, key: DiscoveryKey) {
249        self.states.insert(key, DiscoveryState::Lifted);
250    }
251
252    pub fn mark_failed(&mut self, key: DiscoveryKey, reason: impl Into<String>) {
253        self.states.insert(
254            key,
255            DiscoveryState::Failed {
256                reason: reason.into(),
257            },
258        );
259    }
260
261    pub fn mark_skipped(&mut self, key: DiscoveryKey, reason: impl Into<String>) {
262        self.states.insert(
263            key,
264            DiscoveryState::Skipped {
265                reason: reason.into(),
266            },
267        );
268    }
269
270    pub fn state(&self, key: &DiscoveryKey) -> Option<&DiscoveryState> {
271        self.states.get(key)
272    }
273
274    pub fn states(&self) -> impl Iterator<Item = (&DiscoveryKey, &DiscoveryState)> + '_ {
275        self.states.iter()
276    }
277
278    /// Every code address this queue lifted successfully, as a portable
279    /// [`CodeSeed`]. Exported from one analysis run and replayed into the next
280    /// (via [`Context::seed_code`](crate::context::Context::seed_code)) so the
281    /// lifter reaches jump-table targets in its first pass instead of waiting for
282    /// the analysis fixpoint to discover them round by round.
283    pub fn lifted_seeds(&self) -> Vec<CodeSeed> {
284        self.states
285            .iter()
286            .filter(|(_, state)| matches!(state, DiscoveryState::Lifted))
287            .map(|(key, _)| CodeSeed::from_key(key))
288            .collect()
289    }
290}
291
292/// A code address known to lift, captured from one analysis run to pre-seed the
293/// next. Mirrors a [`DiscoveryKey`] but drops the run-specific bits the lifter
294/// reconstructs on its own (the precise CFG edge is re-attached by the pass that
295/// originally found the target).
296#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
297pub enum CodeSeed {
298    /// A discovered function entry.
299    Function { target: Address },
300    /// A block at `target` inside the function entered at `function`.
301    Block {
302        function: Address,
303        target: Address,
304        edge_kind: EdgeKind,
305    },
306}
307
308impl CodeSeed {
309    fn from_key(key: &DiscoveryKey) -> Self {
310        match key.kind {
311            DiscoveryKind::Function { .. } => CodeSeed::Function { target: key.target },
312            DiscoveryKind::Block {
313                function,
314                edge_kind,
315            } => CodeSeed::Block {
316                function,
317                target: key.target,
318                edge_kind,
319            },
320        }
321    }
322
323    /// Rebuild the [`Discovery`] to enqueue, attributed to [`UserSeed`] so the
324    /// origin stays honest rather than impersonating the pass that first found it.
325    ///
326    /// [`UserSeed`]: DiscoveryProvenance::UserSeed
327    pub fn into_discovery(self) -> Discovery {
328        match self {
329            CodeSeed::Function { target } => {
330                Discovery::function(target).with_provenance(DiscoveryProvenance::UserSeed)
331            }
332            CodeSeed::Block {
333                function,
334                target,
335                edge_kind,
336            } => Discovery::block(target, function)
337                .with_edge_kind(edge_kind)
338                .with_provenance(DiscoveryProvenance::UserSeed),
339        }
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346
347    #[test]
348    fn dedups_exact_key_not_target() {
349        let mut q = DiscoveryQueue::default();
350        assert!(q.insert(Discovery::block(0x1000, 0x900)));
351        assert!(!q.insert(Discovery::block(0x1000, 0x900)));
352        assert!(q.insert(Discovery::block(0x1000, 0x900).from_block_addr(0x910)));
353        assert!(q.insert(Discovery::block(0x1000, 0x900).from_block_addr(0x920)));
354        assert!(q.insert(Discovery::block(0x1000, 0xa00)));
355        assert!(q.insert(Discovery::function(0x1000)));
356
357        let all: Vec<_> = q.iter().collect();
358        assert_eq!(all.len(), 5);
359    }
360
361    #[test]
362    fn terminal_state_blocks_requeue_of_same_key() {
363        let mut q = DiscoveryQueue::default();
364        let d = Discovery::function(0x10);
365        let key = d.key();
366        assert!(q.insert(d.clone()));
367        q.drain();
368        q.mark_failed(key, "decode");
369        assert!(!q.insert(d));
370        assert!(q.is_empty());
371    }
372
373    #[test]
374    fn terminal_target_from_one_source_does_not_suppress_another_source() {
375        let mut q = DiscoveryQueue::default();
376        let first = Discovery::block(0x1000, 0x900).from_block_addr(0x910);
377        let second = Discovery::block(0x1000, 0x900).from_block_addr(0x920);
378        let first_key = first.key();
379
380        assert!(q.insert(first));
381        q.drain();
382        q.mark_lifted(first_key);
383
384        assert!(q.insert(second));
385        assert_eq!(q.pending_len(), 1);
386    }
387
388    #[test]
389    fn drain_empties_pending_but_keeps_state() {
390        let mut q = DiscoveryQueue::default();
391        let d = Discovery::function(0x10);
392        let key = d.key();
393        q.insert(d);
394        let drained = q.drain();
395        assert_eq!(drained.len(), 1);
396        assert!(q.is_empty());
397        assert_eq!(q.state(&key), Some(&DiscoveryState::Pending));
398    }
399}
qcode/discovery.rs

qcode/
discovery.rs
