qcode/
address_index.rs

1//! Disposable address lookup over an immutable qcode module snapshot.
2//!
3//! [`AddressIndex`] is derived state: consumers build it for a [`Context`], keep
4//! it only while that context remains structurally unchanged, and rebuild it
5//! after mutation. It is intentionally not stored in or serialized with the IR.
6
7use rustc_hash::FxHashMap;
8
9use crate::{
10    context::Context,
11    error::{Error, ErrorTy, Result},
12    value::{BlockId, FunctionBody, FunctionId, ValueId},
13};
14
15/// A live module entity selected by a machine address.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum AddressTarget {
18    Function(FunctionId),
19    Block(BlockId),
20}
21
22/// An immutable, disposable address-to-entity snapshot.
23#[derive(Debug, Clone, Default, PartialEq, Eq)]
24pub struct AddressIndex {
25    targets: FxHashMap<u64, AddressTarget>,
26    /// Addresses known to start a block because something branches to them.
27    ///
28    /// Learned, not derived: an address becomes a boundary the first time a
29    /// block has to be split at it. Remembering that is what stops a run from
30    /// being folded across the same address again on the next pass, which for a
31    /// loop header — the branch target that is discovered *after* the run
32    /// through it — would otherwise repeat forever.
33    boundaries: rustc_hash::FxHashSet<u64>,
34}
35
36impl AddressIndex {
37    /// Computes an index for the current live shape of `ctx`.
38    ///
39    /// Primary and extra block addresses are indexed. Functions are installed
40    /// last so a function wins the intentional collision with its entry block.
41    pub fn analyze(ctx: &Context<'_>) -> Self {
42        let mut targets = FxHashMap::default();
43
44        for block_id in ctx.block_ids() {
45            let block = ctx.block(block_id);
46            if let Some(address) = block.address {
47                targets
48                    .entry(address)
49                    .or_insert(AddressTarget::Block(block_id));
50            }
51            for &address in &block.extra_addresses {
52                targets
53                    .entry(address)
54                    .or_insert(AddressTarget::Block(block_id));
55            }
56        }
57
58        for function in ctx.functions() {
59            if let Some(address) = function.address() {
60                targets.insert(address, AddressTarget::Function(function.id));
61            }
62        }
63
64        Self {
65            targets,
66            boundaries: rustc_hash::FxHashSet::default(),
67        }
68    }
69
70    /// Recomputes this index after a structural mutation that changes several
71    /// addresses at once (for example function splitting or block rehoming).
72    pub fn refresh(&mut self, ctx: &Context<'_>) {
73        *self = Self::analyze(ctx);
74    }
75
76    /// Re-point `addr` from a relocated block `old` to its clone `new`, in place.
77    ///
78    /// The incremental analogue of a [`refresh`](Self::refresh) after a block
79    /// rehome: the caller already knows exactly which address moved and where, so
80    /// there is no need to re-scan the whole module. A no-op unless `addr` is
81    /// currently indexed to `old` — this preserves [`analyze`](Self::analyze)'s
82    /// function-over-block precedence (a function entry that deliberately shadows
83    /// its root block, or another block that already owns the address, is left
84    /// untouched).
85    pub fn rehome_block(&mut self, addr: u64, old: BlockId, new: BlockId) {
86        if self.targets.get(&addr) == Some(&AddressTarget::Block(old)) {
87            self.targets.insert(addr, AddressTarget::Block(new));
88        }
89    }
90
91    /// Drops `address` from the index, so it resolves to nothing until it is
92    /// registered again. Used when a block stops covering an address.
93    pub fn forget(&mut self, address: u64) {
94        self.targets.remove(&address);
95    }
96
97    /// Points `address` at `block`, whatever it pointed at before.
98    ///
99    /// For a caller that has just made `block` cover an address another block
100    /// used to — absorbing that block, typically, which leaves the index
101    /// naming something deleted.
102    pub fn set_block(&mut self, address: u64, block: BlockId) {
103        self.targets.insert(address, AddressTarget::Block(block));
104    }
105
106    /// Records that `address` starts a block, and must keep starting one.
107    pub fn mark_boundary(&mut self, address: u64) {
108        self.boundaries.insert(address);
109    }
110
111    /// Whether `address` is known to start a block.
112    pub fn is_boundary(&self, address: u64) -> bool {
113        self.boundaries.contains(&address)
114    }
115
116    /// Registers one address-bearing entity during module construction.
117    ///
118    /// A function and one of its own blocks may intentionally share an entry
119    /// address; the function remains the indexed target and the block becomes
120    /// its root. Every other collision is rejected.
121    pub fn register(
122        &mut self,
123        ctx: &mut Context<'_>,
124        address: u64,
125        target: AddressTarget,
126    ) -> Result<()> {
127        let Some(existing) = self.targets.get(&address).copied() else {
128            self.targets.insert(address, target);
129            return Ok(());
130        };
131        if existing == target {
132            return Ok(());
133        }
134
135        match (existing, target) {
136            (AddressTarget::Function(function), AddressTarget::Block(block))
137            | (AddressTarget::Block(block), AddressTarget::Function(function)) => {
138                // A split first registers a rootless function over a block that
139                // is still stored in the old function. Root only once storage
140                // ownership matches; rehoming refreshes the index afterwards.
141                if block.func == function {
142                    FunctionBody::from_id_mut(ctx, function).ensure_root(block)?;
143                }
144                self.targets
145                    .insert(address, AddressTarget::Function(function));
146                Ok(())
147            }
148            _ => Err(Error::spanless(ErrorTy::DuplicateAddress(
149                address,
150                existing.into(),
151            ))),
152        }
153    }
154
155    /// Returns the live target registered at `address` in this snapshot.
156    pub fn get(&self, address: u64) -> Option<AddressTarget> {
157        self.targets.get(&address).copied()
158    }
159
160    /// Returns the function registered at `address`, if that is the target kind.
161    pub fn function_at(&self, address: u64) -> Option<FunctionId> {
162        match self.get(address) {
163            Some(AddressTarget::Function(id)) => Some(id),
164            _ => None,
165        }
166    }
167
168    /// Returns the block registered at `address`, if that is the target kind.
169    pub fn block_at(&self, address: u64) -> Option<BlockId> {
170        match self.get(address) {
171            Some(AddressTarget::Block(id)) => Some(id),
172            _ => None,
173        }
174    }
175
176    /// Returns the number of distinct indexed addresses.
177    pub fn len(&self) -> usize {
178        self.targets.len()
179    }
180
181    /// Returns whether the snapshot contains no addresses.
182    pub fn is_empty(&self) -> bool {
183        self.targets.is_empty()
184    }
185}
186
187impl From<AddressTarget> for ValueId {
188    fn from(target: AddressTarget) -> Self {
189        match target {
190            AddressTarget::Function(id) => id.into(),
191            AddressTarget::Block(id) => id.into(),
192        }
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use crate::value::QCodeMut;
200    use crate::value::{BasicBlock, FunctionBody};
201
202    #[test]
203    fn function_wins_its_entry_block_collision() {
204        let mut ctx = Context::new();
205        let mut index = AddressIndex::analyze(&ctx);
206        let function = FunctionBody::make_at_addr_indexed(&mut ctx, &mut index, 0x1000, None).id;
207        let root = BasicBlock::make(&mut ctx, function)
208            .with_address_indexed(&mut index, 0x1000)
209            .id;
210
211        assert_eq!(index.get(0x1000), Some(AddressTarget::Function(function)));
212        assert_eq!(index.function_at(0x1000), Some(function));
213        assert_eq!(index.block_at(0x1000), None);
214        assert_eq!(ctx.function(function).root_id(), Some(root.local));
215    }
216
217    #[test]
218    fn indexes_primary_and_extra_block_addresses() {
219        let mut ctx = Context::new();
220        let mut index = AddressIndex::analyze(&ctx);
221        let function = ctx.anon_function();
222        let block = BasicBlock::make(&mut ctx, function)
223            .with_address_indexed(&mut index, 0x2000)
224            .id;
225        ctx.block_mut(block)
226            .extra_addresses
227            .extend([0x2001, 0x2002]);
228
229        index.refresh(&ctx);
230        assert_eq!(index.block_at(0x2000), Some(block));
231        assert_eq!(index.block_at(0x2001), Some(block));
232        assert_eq!(index.block_at(0x2002), Some(block));
233        assert_eq!(index.len(), 3);
234    }
235
236    #[test]
237    fn excludes_deleted_body_shape() {
238        let mut ctx = Context::new();
239        let mut index = AddressIndex::analyze(&ctx);
240        let function = ctx.anon_function();
241        let block = BasicBlock::make(&mut ctx, function)
242            .with_address_indexed(&mut index, 0x3000)
243            .id;
244        ctx.delete_block(block);
245
246        index.refresh(&ctx);
247        assert_eq!(index.get(0x3000), None);
248        assert!(index.is_empty());
249    }
250}
qcode/address_index.rs

qcode/
address_index.rs
