1use crate::value::QCodeMut;
4use std::{borrow::Cow, fmt::Display};
5
6use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
7
8use crate::{
9 assumption::{Certainty, KnownContradiction, PassName, Proposition, Truth, Violation},
10 error::{Error, ErrorTy, Result},
11 pass_scope,
12 space::{LocalMemorySpaceId, MemorySpaceId, Space, SpaceId, SpaceStore},
13 types::TypeManager,
14 value::{
15 BasicBlock, BlockParamRef, FunctionBody, FunctionId, FunctionRef, Instruction, ModuleView,
16 QCodeView, TempId, TempSpaceId, ValueId,
17 block::{BlockId, BlockRef, EdgeData, EdgeId},
18 block_param::{BlockParam, BlockParamId},
19 insn::{InstructionId, InstructionRef, Mnemonic, PCodeOpId},
20 literal::{LiteralId, LiteralRef},
21 registry::ValueRegistry,
22 varnode::{Varnode, VarnodeId, VarnodeRef, register::RegisterId},
23 },
24};
25use jstd::registry::{self, Identified, Registry};
26
27#[derive(Default, Clone, serde::Serialize)]
55pub struct Context<'str> {
56 pub shared: Shared<'str>,
62
63 #[serde(default)]
68 pub interfaces: Registry<FunctionId, crate::value::function::FunctionInterface<'str>>,
69
70 pub bodies: Registry<FunctionId, FunctionBody<'str>>,
76}
77
78#[derive(serde::Deserialize)]
79struct ContextWire<'str> {
80 shared: Shared<'str>,
81 #[serde(default)]
82 interfaces: Registry<FunctionId, crate::value::function::FunctionInterface<'str>>,
83 bodies: Registry<FunctionId, FunctionBody<'str>>,
84}
85
86impl<'de, 'str> serde::Deserialize<'de> for Context<'str> {
87 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
88 where
89 D: serde::Deserializer<'de>,
90 {
91 let ContextWire {
92 shared,
93 interfaces,
94 mut bodies,
95 } = ContextWire::deserialize(deserializer)?;
96 if interfaces.len() != bodies.len() {
97 return Err(serde::de::Error::custom(
98 "function body/interface registries drifted",
99 ));
100 }
101 for mut body in bodies.iter_mut() {
102 let id = body.id;
103 body.rehydrate_id(id);
104 }
105 Ok(Self {
106 shared,
107 interfaces,
108 bodies,
109 })
110 }
111}
112
113#[derive(Default, Clone, serde::Serialize, serde::Deserialize)]
121pub struct Shared<'str> {
122 pub default_space: SpaceId,
123
124 pub(crate) spaces: Registry<SpaceId, Space>,
126
127 pub pcode_ops: Registry<PCodeOpId, Box<str>>,
129
130 pub named_spaces: HashMap<Box<str>, SpaceId>,
132
133 pub(crate) name_map: NameTable<'str>,
139
140 pub registers: HashMap<RegisterId, VarnodeId>,
142
143 pub values: ValueRegistry<'str>,
145
146 pub types: TypeManager,
148
149 #[serde(default)]
156 pub(crate) protections_known: bool,
157
158 #[serde(default)]
162 pub(crate) primary_entrypoint: Option<u64>,
163
164 #[serde(default)]
171 pub(crate) discoveries: crate::discovery::DiscoveryQueue,
172
173 #[serde(default)]
178 pub(crate) target_os: TargetOs,
179
180 #[serde(default)]
186 pub(crate) linked_libraries: Vec<String>,
187
188 #[serde(default)]
194 pub(crate) ignored_functions: HashSet<u64>,
195
196 #[serde(skip)]
204 pub(crate) assumed_call_convention: Option<crate::assumption::AssumedCallEffect>,
205}
206
207impl SpaceStore for Shared<'_> {
210 fn spaces(&self) -> &Registry<SpaceId, Space> {
211 &self.spaces
212 }
213}
214
215impl SpaceStore for Context<'_> {
217 fn spaces(&self) -> &Registry<SpaceId, Space> {
218 &self.shared.spaces
219 }
220}
221
222impl<'str> Shared<'str> {
223 pub fn get_named(&self, name: &str) -> Option<ValueId> {
226 self.name_map.get(name)
227 }
228
229 pub fn varnode(&self, id: VarnodeId) -> &crate::value::Varnode<'str> {
231 &self.values.varnodes[id]
232 }
233
234 pub fn space(&self, id: SpaceId) -> &Space {
236 &self.spaces[id]
237 }
238
239 pub fn spaces(&self) -> impl Iterator<Item = Identified<SpaceId, &Space>> + '_ {
242 self.spaces.iter()
243 }
244
245 pub fn get_const(&self, value: u64, size: usize) -> ValueId {
249 let type_id = self.types.get_or_make_int(size);
250 ValueId::Literal(self.values.get_or_make_typed_literal(value, type_id, size))
251 }
252
253 pub fn pcode_op(&mut self, name: &str) -> PCodeOpId {
259 if let Some(op) = self.pcode_ops.iter().find(|op| op.as_ref() == name) {
260 return op.id;
261 }
262 self.pcode_ops.push(Box::from(name))
263 }
264
265 pub fn vm_interrupt_op(&mut self) -> PCodeOpId {
268 self.pcode_op(crate::value::insn::VM_INTERRUPT)
269 }
270
271 pub fn get_bool_const(&self, value: bool) -> ValueId {
274 let type_id = self.types.get_or_make_bool();
275 ValueId::Literal(
276 self.values
277 .get_or_make_typed_literal(u64::from(value), type_id, 1),
278 )
279 }
280
281 pub fn get_typed_const(&self, value: u64, type_id: crate::types::TypeId) -> ValueId {
284 let size = self.types.size_of(type_id);
285 ValueId::Literal(self.values.get_or_make_typed_literal(value, type_id, size))
286 }
287
288 pub fn get_bytes(&self, data: Vec<u8>) -> ValueId {
291 let i8_ty = self.types.get_or_make_int(1);
292 let type_id = self.types.get_or_make_array(i8_ty, data.len());
293 ValueId::Bytes(
294 self.values
295 .bytes
296 .push(crate::value::Bytes { data, type_id }),
297 )
298 }
299
300 pub fn bytes_display(&self, id: crate::value::BytesId) -> crate::value::BytesDisplay {
305 self.values
306 .bytes_display
307 .get(&id)
308 .copied()
309 .unwrap_or_default()
310 }
311
312 pub fn get_typed_bytes(&self, data: Vec<u8>, type_id: crate::types::TypeId) -> ValueId {
316 ValueId::Bytes(
317 self.values
318 .bytes
319 .push(crate::value::Bytes { data, type_id }),
320 )
321 }
322
323 pub fn truth(&self, prop: Proposition) -> Option<Truth> {
327 self.values.truths.get(&prop).copied()
328 }
329
330 pub fn assumed_call_convention(&self) -> Option<&crate::assumption::AssumedCallEffect> {
336 self.assumed_call_convention.as_ref()
337 }
338
339 pub fn varnodes(&self) -> impl Iterator<Item = crate::value::VarnodeRef<'str, '_>> + '_ {
342 self.values
343 .varnodes
344 .iter()
345 .map(move |v| crate::value::Varnode::from_id(self, v.id))
346 }
347
348 pub fn varnode_count(&self) -> usize {
351 self.values.varnodes.len()
352 }
353
354 pub fn stored_type_of(&self, id: ValueId) -> Option<crate::types::TypeId> {
361 match id {
362 ValueId::Literal(lid) => Some(self.values.literals[lid].type_id),
363 ValueId::Bytes(bid) => Some(self.values.bytes[bid].type_id),
364 ValueId::Varnode(vid) => self.values.varnode_types.get(&vid).copied(),
365 ValueId::Poison(pid) => Some(self.values.poisons[pid].type_id),
366 ValueId::Instruction(_)
367 | ValueId::BlockParam(_)
368 | ValueId::BasicBlock(_)
369 | ValueId::Temp(_)
370 | ValueId::Function(_) => None,
371 }
372 }
373}
374
375pub use wazabin_binary::TargetOs;
379
380impl<'str> Context<'str> {
381 pub fn new() -> Self {
387 let mut ctx = Self::default();
388 ctx.shared.spaces.push(Space::new(Some("const"), 1, 8));
390 let default_space = Space::new(Some("ram"), 1, 8);
392 ctx.shared.default_space = ctx.shared.spaces.push(default_space);
393 ctx
394 }
395
396 pub fn try_get_space(&self, name: &str) -> Option<SpaceId> {
399 self.shared.named_spaces.get(name).copied()
400 }
401
402 pub fn get_or_make_named_space(&mut self, name: &str) -> SpaceId {
407 if let Some(id) = self.try_get_space(name) {
408 return id;
409 }
410 let default_id = self.shared.default_space;
411 if self.shared.spaces[default_id].name.as_deref() == Some(name) {
412 return default_id;
413 }
414 let default = &self.shared.spaces[self.shared.default_space];
415 let space = Space::new(Some(name), default.word_size, default.addr_size);
416 self.add_space(space)
417 }
418
419 pub fn add_space(&mut self, space: Space) -> SpaceId {
421 let name_key: Option<Box<str>> = space.name.clone();
422 let id = self.shared.spaces.push(space);
423 if let Some(name) = name_key {
424 self.shared.named_spaces.insert(name, id);
425 }
426 id
427 }
428
429 pub fn space_count(&self) -> usize {
431 self.shared.spaces.len()
432 }
433
434 pub fn spaces(&self) -> impl Iterator<Item = Identified<SpaceId, &Space>> + '_ {
437 self.shared.spaces()
438 }
439
440 pub fn set_primary_entrypoint(&mut self, entrypoint: Option<u64>) {
441 self.shared.primary_entrypoint = entrypoint;
442 }
443
444 pub fn primary_entrypoint(&self) -> Option<u64> {
445 self.shared.primary_entrypoint
446 }
447
448 pub fn set_ignored_functions(&mut self, addrs: HashSet<u64>) {
452 self.shared.ignored_functions = addrs;
453 }
454
455 pub fn ignored_functions(&self) -> &HashSet<u64> {
457 &self.shared.ignored_functions
458 }
459
460 pub fn is_function_ignored(&self, addr: Option<u64>) -> bool {
463 addr.is_some_and(|a| self.shared.ignored_functions.contains(&a))
464 }
465
466 pub fn set_target_os(&mut self, os: TargetOs) {
469 self.shared.target_os = os;
470 }
471
472 pub fn target_os(&self) -> TargetOs {
474 self.shared.target_os
475 }
476
477 pub fn set_linked_libraries(&mut self, libs: Vec<String>) {
480 self.shared.linked_libraries = libs;
481 }
482
483 pub fn linked_libraries(&self) -> &[String] {
486 &self.shared.linked_libraries
487 }
488
489 pub fn load_spaces(&mut self, spaces: registry::Registry<SpaceId, Space>) {
491 self.shared.spaces = spaces;
492 }
493
494 pub fn mark_protections_known(&mut self) {
498 self.shared.protections_known = true;
499 }
500
501 pub fn protections_known(&self) -> bool {
503 self.shared.protections_known
504 }
505
506 pub fn assume_executable(
525 &mut self,
526 binary: &dyn wazabin_binary::BinaryFormat,
527 addr: u64,
528 ) -> bool {
529 let bounds = binary.segment_bounds(addr);
530 if let Some((start, end)) = bounds
531 && let Some(known) = self.known(Proposition::ExecutableMemory { start, end })
532 {
533 return known;
534 }
535 if !self.shared.protections_known || binary.is_executable(addr) {
536 return true;
537 }
538 if let Some((start, end)) = bounds {
539 self.set_known(Proposition::ExecutableMemory { start, end }, false);
540 }
541 false
542 }
543
544 pub fn discover(&mut self, discovery: crate::discovery::Discovery) -> bool {
547 self.shared.discoveries.insert(discovery)
548 }
549
550 pub fn discover_code(&mut self, func_entry: u64, source_block: u64, target: u64) {
559 self.shared.discoveries.insert(
560 crate::discovery::Discovery::block(target, func_entry)
561 .with_edge_kind(crate::discovery::EdgeKind::JumpTableTarget)
562 .from_block_addr(source_block)
563 .with_provenance(crate::discovery::DiscoveryProvenance::Optimization {
564 pass: "handle_jump_tables".to_string(),
565 assumption: None,
566 }),
567 );
568 }
569
570 pub fn drain_discoveries(&mut self) -> Vec<crate::discovery::Discovery> {
572 self.shared.discoveries.drain()
573 }
574
575 pub fn discoveries(&self) -> impl Iterator<Item = &crate::discovery::Discovery> + '_ {
577 self.shared.discoveries.iter()
578 }
579
580 pub fn has_no_discoveries(&self) -> bool {
582 self.shared.discoveries.is_empty()
583 }
584
585 pub fn lifted_code_seeds(&self) -> Vec<crate::discovery::CodeSeed> {
590 self.shared.discoveries.lifted_seeds()
591 }
592
593 pub fn seed_code(&mut self, seeds: impl IntoIterator<Item = crate::discovery::CodeSeed>) {
599 for seed in seeds {
600 self.shared.discoveries.insert(seed.into_discovery());
601 }
602 }
603
604 pub fn mark_discovery_lifted(&mut self, key: crate::discovery::DiscoveryKey) {
605 self.shared.discoveries.mark_lifted(key);
606 }
607
608 pub fn mark_discovery_failed(
609 &mut self,
610 key: crate::discovery::DiscoveryKey,
611 reason: impl Into<String>,
612 ) {
613 self.shared.discoveries.mark_failed(key, reason);
614 }
615
616 pub fn mark_discovery_skipped(
617 &mut self,
618 key: crate::discovery::DiscoveryKey,
619 reason: impl Into<String>,
620 ) {
621 self.shared.discoveries.mark_skipped(key, reason);
622 }
623
624 pub fn get_or_make_block(&mut self, addr: u64, func: FunctionId) -> BlockId {
629 let mut addresses = crate::address_index::AddressIndex::analyze(self);
630 self.get_or_make_block_indexed(&mut addresses, addr, func)
631 }
632
633 #[track_caller]
637 pub fn get_or_make_block_indexed(
638 &mut self,
639 addresses: &mut crate::address_index::AddressIndex,
640 addr: u64,
641 func: FunctionId,
642 ) -> BlockId {
643 use crate::address_index::AddressTarget;
644
645 if let Some(AddressTarget::Function(owner)) = addresses.get(addr) {
646 assert_eq!(
647 owner, func,
648 "cannot create a block at an address owned by another function"
649 );
650 }
651 let existing = match addresses.get(addr) {
652 Some(AddressTarget::Block(block)) => Some(block),
653 Some(AddressTarget::Function(function)) => FunctionBody::from_id(self, function)
654 .root()
655 .map(|root| root.id),
656 None => None,
657 };
658 match existing {
659 Some(block) => {
660 if self.block(block).address != Some(addr)
670 && block.func == func
671 && self.block(block).extra_addresses.contains(&addr)
672 {
673 return self.split_block_at_address(addresses, block, addr);
674 }
675 if block.func != func {
676 let stored = FunctionBody::from_id(self, block.func);
677 let requested = FunctionBody::from_id(self, func);
678 let parent = Some(block.func);
682 let caller = std::panic::Location::caller();
683 let detail = format!(
684 "cannot reuse a block stored in another function arena: block={block:?} address=0x{addr:x}; stored={:?} name={:?} entry={:?} parent={parent:?}; requested={:?} name={:?} entry={:?}; caller={caller}",
685 block.func,
686 stored.name(),
687 stored.address(),
688 func,
689 requested.name(),
690 requested.address(),
691 );
692 log::error!(
693 target: "qcode::arena",
694 "{detail}\nbacktrace:\n{}",
695 std::backtrace::Backtrace::force_capture()
696 );
697 panic!("{detail}");
698 }
699 block
700 }
701 None => {
702 BasicBlock::make(self, func)
703 .with_address_indexed(addresses, addr)
704 .id
705 }
706 }
707 }
708
709 pub fn split_block_at_address(
729 &mut self,
730 addresses: &mut crate::address_index::AddressIndex,
731 block: BlockId,
732 addr: u64,
733 ) -> BlockId {
734 addresses.mark_boundary(addr);
737 let tail = BasicBlock::make(self, block.func).id;
738
739 self.bodies[block.func].clear_block_instructions(block);
742
743 let absorbed = std::mem::take(&mut self.block_mut(block).extra_addresses);
747 for absorbed_addr in absorbed {
748 addresses.forget(absorbed_addr);
749 }
750 BasicBlock::from_id_mut(self, tail)
751 .in_function(block.func)
752 .with_address_indexed(addresses, addr);
753 tail
754 }
755
756 pub fn split_block_before(&mut self, block: BlockId, insn: InstructionId) -> BlockId {
760 self.bodies[block.func].split_block_before(block, insn)
761 }
762
763 pub fn builder(&mut self, block: BlockId) -> crate::builder::Builder<'str, '_> {
766 let body = &mut self.bodies[block.func];
767 crate::builder::Builder::new(body, &self.shared, &self.interfaces, block)
768 }
769
770 pub fn builder_at(&mut self, address: u64) -> crate::builder::Builder<'str, '_> {
773 use crate::address_index::AddressTarget;
774
775 let mut addresses = crate::address_index::AddressIndex::analyze(self);
776 let block = match addresses.get(address) {
777 Some(AddressTarget::Function(function)) => self.bodies[function]
778 .root_id()
779 .map(|local| BlockId::new(function, local))
780 .unwrap_or_else(|| {
781 self.get_or_make_block_indexed(&mut addresses, address, function)
782 }),
783 Some(AddressTarget::Block(block)) => block,
784 None => {
785 let function = FunctionBody::make(self, Cow::Owned(format!("blk_{address:x}")))
786 .expect("anonymous host function")
787 .id;
788 self.get_or_make_block_indexed(&mut addresses, address, function)
789 }
790 };
791 let mut builder = self.builder(block);
792 builder.set_address(address);
793 builder
794 }
795
796 pub fn bytes_display(&self, id: crate::value::BytesId) -> crate::value::BytesDisplay {
799 self.shared
800 .values
801 .bytes_display
802 .get(&id)
803 .copied()
804 .unwrap_or_default()
805 }
806
807 pub fn set_bytes_display(
811 &mut self,
812 id: crate::value::BytesId,
813 mode: crate::value::BytesDisplay,
814 ) {
815 if mode == crate::value::BytesDisplay::Auto {
816 self.shared.values.bytes_display.remove(&id);
817 } else {
818 self.shared.values.bytes_display.insert(id, mode);
819 }
820 }
821
822 pub fn block_ids(&self) -> Vec<BlockId> {
824 self.functions().flat_map(|f| f.block_ids()).collect()
825 }
826
827 pub fn instruction_ids(&self) -> Vec<InstructionId> {
831 let mut ids: Vec<_> = self.functions().flat_map(|f| f.instruction_ids()).collect();
832 ids.sort_unstable();
833 ids
834 }
835
836 pub fn function_ids(&self) -> Vec<FunctionId> {
838 self.interfaces.iter().map(|i| i.id).collect()
839 }
840
841 pub fn anon_function(&mut self) -> FunctionId {
847 let name = self.get_unique_name(std::borrow::Cow::Borrowed("anon"));
848 crate::value::FunctionBody::make(self, name)
849 .expect("unique anon function name")
850 .id
851 }
852
853 pub fn instruction_arena_stats(&self) -> (usize, usize) {
855 let mut total = 0;
856 let mut dead = 0;
857 for f in self.bodies.iter() {
858 total += f.insns.issued_len();
859 dead += f.insns.issued_len() - f.insns.len();
860 }
861 (total, dead)
862 }
863
864 pub fn body_arena_stats(&self) -> crate::value::BodyArenaStats {
870 let mut total = crate::value::BodyArenaStats::default();
871 for body in self.bodies.iter() {
872 total.add_assign(body.arena_stats());
873 }
874 total
875 }
876
877 pub fn shrink_bodies_to_fit(&mut self) {
884 for mut body in self.bodies.iter_mut() {
885 body.shrink_to_fit();
886 }
887 }
888
889 pub fn instructions(&self) -> impl Iterator<Item = InstructionRef<'str, '_>> + '_ {
891 self.instruction_ids()
892 .into_iter()
893 .map(move |id| Instruction::from_id(self, id))
894 }
895
896 pub fn blocks(&self) -> impl Iterator<Item = BlockRef<'str, '_>> + '_ {
898 self.block_ids()
899 .into_iter()
900 .map(move |id| BlockRef::from_id(self, id))
901 }
902
903 pub fn functions(&self) -> FunctionIter<'str, '_> {
905 FunctionIter {
906 ctx: self,
907 inner: self.bodies.iter(),
908 }
909 }
910
911 pub fn iter(&self) -> FunctionIter<'str, '_> {
914 self.functions()
915 }
916
917 pub fn varnodes(&self) -> impl Iterator<Item = VarnodeRef<'str, '_>> + '_ {
918 self.shared.varnodes()
919 }
920
921 pub fn varnode_count(&self) -> usize {
926 self.shared.varnode_count()
927 }
928
929 pub fn remove_cfg_edge(&mut self, func: FunctionId, edge_id: EdgeId) {
933 self.bodies[func].remove_cfg_edge(edge_id);
934 }
935
936 pub fn rehome_owned_blocks(
956 &mut self,
957 addresses: &mut crate::address_index::AddressIndex,
958 target: FunctionId,
959 olds: &[BlockId],
960 ) -> HashMap<BlockId, BlockId> {
961 let mut needed_temps: HashSet<TempId> = HashSet::default();
965 let mut needed_temp_spaces: HashSet<TempSpaceId> = HashSet::default();
966 for &old in olds {
967 for ¶m_local in &self.block(old).params {
968 let param = self.block_param(BlockParamId::new(old.func, param_local));
969 if let Some(crate::value::LocalValueId::Temp(temp)) = param.origin {
970 needed_temps.insert(TempId::new(old.func, temp));
971 }
972 if let Some(MemorySpaceId::Temp(space)) = self.shared.types.space_of(param.type_id)
973 {
974 needed_temp_spaces.insert(space);
975 }
976 }
977 for &insn_local in &self.block(old).instructions {
978 let insn = self.instruction(InstructionId::new(old.func, insn_local));
979 for arg in insn.mnemonic().args() {
980 if let crate::value::LocalValueId::Temp(temp) = arg {
981 needed_temps.insert(TempId::new(old.func, temp));
982 }
983 }
984 let explicit_space = match insn.mnemonic() {
985 Mnemonic::Load(load) => Some(load.space),
986 Mnemonic::Store(store) => Some(store.space),
987 _ => None,
988 };
989 if let Some(LocalMemorySpaceId::Temp(space)) = explicit_space {
990 needed_temp_spaces.insert(TempSpaceId::new(old.func, space));
991 }
992 if let Some(MemorySpaceId::Temp(space)) = self.shared.types.space_of(insn.type_id) {
993 needed_temp_spaces.insert(space);
994 }
995 }
996 }
997 for &temp in &needed_temps {
998 let data = &self.bodies[temp.func].temps[temp.local];
999 needed_temp_spaces.insert(TempSpaceId::new(temp.func, data.space));
1000 }
1001
1002 let mut needed_temp_spaces: Vec<_> = needed_temp_spaces.into_iter().collect();
1003 needed_temp_spaces.sort_unstable();
1004 let mut temp_space_map: HashMap<TempSpaceId, TempSpaceId> = HashMap::default();
1005 for old in needed_temp_spaces {
1006 if old.func == target {
1007 continue;
1008 }
1009 let space = self.bodies[old.func].temp_spaces[old.local].clone();
1010 let new = self.bodies[target].push_temp_space(space);
1011 temp_space_map.insert(old, new);
1012 }
1013
1014 let mut needed_temps: Vec<_> = needed_temps.into_iter().collect();
1015 needed_temps.sort_unstable();
1016 let mut value_map: HashMap<ValueId, ValueId> = HashMap::default();
1017 for old in needed_temps {
1018 if old.func == target {
1019 continue;
1020 }
1021 let mut temp = self.bodies[old.func].temps[old.local].clone();
1022 temp.space = temp_space_map[&TempSpaceId::new(old.func, temp.space)].local;
1023 if let Some(name) = temp.name.take() {
1024 temp.name = Some(self.bodies[target].names.unique(name));
1025 }
1026 let new = self.bodies[target].push_temp(temp);
1027 value_map.insert(ValueId::Temp(old), ValueId::Temp(new));
1028 }
1029
1030 let mut block_map: HashMap<BlockId, BlockId> = HashMap::default();
1033 for &old in olds {
1034 let new = BasicBlock::clone_block_into(self, old, target, &mut value_map);
1035 block_map.insert(old, new);
1036 }
1037
1038 for (&old, &new) in &block_map {
1043 let old_params = self.block(old).params.clone();
1044 let new_params = self.block(new).params.clone();
1045 for (old_local, new_local) in old_params.into_iter().zip(new_params) {
1046 let old_param = BlockParamId::new(old.func, old_local);
1047 let new_param = BlockParamId::new(new.func, new_local);
1048 let type_id = remap_rehomed_type(
1049 self,
1050 self.block_param(new_param).type_id,
1051 target,
1052 &temp_space_map,
1053 );
1054 self.block_param_mut(new_param).type_id = type_id;
1055 let Some(origin) = self.block_param(new_param).origin else {
1056 continue;
1057 };
1058 let qualified = origin.qualify(old.func);
1059 let remapped = value_map.get(&qualified).copied().unwrap_or(qualified);
1060 debug_assert!(
1061 remapped.owning_function().is_none_or(|f| f == target),
1062 "rehome: relocated block param {old_param:?} has an origin in another \
1063 function ({qualified:?}); the relocated set is not closed",
1064 );
1065 self.block_param_mut(new_param).origin = Some(remapped.localize(new.func));
1066 }
1067
1068 let insns = self.block(new).instructions.clone();
1069 for insn_local in insns {
1070 let insn_id = InstructionId::new(new.func, insn_local);
1071 let type_id = remap_rehomed_type(
1072 self,
1073 self.instruction(insn_id).type_id,
1074 target,
1075 &temp_space_map,
1076 );
1077 self.instruction_mut(insn_id).type_id = type_id;
1078 let mut mnemonic = self.instruction(insn_id).mnemonic().clone();
1079 let mut pairs = Vec::new();
1080 for arg in mnemonic.args() {
1081 let qualified = arg.qualify(old.func);
1085 if let Some(&new_val) = value_map.get(&qualified) {
1086 pairs.push((arg, new_val.localize(new.func)));
1087 } else if let Some(new_lit) =
1088 remap_symbolic_block_literal(&self.shared.values.literals, arg, &block_map)
1089 {
1090 pairs.push((arg, new_lit));
1091 } else {
1092 debug_assert!(
1097 qualified.owning_function().is_none_or(|f| f == target),
1098 "rehome: relocated block references a value in another \
1099 function ({qualified:?}); the relocated set is not closed",
1100 );
1101 }
1102 }
1103 crate::value::block::substitute_operands(&mut mnemonic, &pairs);
1104 remap_rehomed_memory_space(&mut mnemonic, old.func, target, &temp_space_map);
1105 remap_block_targets(&mut mnemonic, old.func, new.func, &block_map);
1106 *self.instruction_mut(insn_id).mnemonic_mut() = mnemonic;
1107 }
1108 }
1109
1110 let mut incident: HashSet<(FunctionId, EdgeId)> = HashSet::default();
1117 for &old in olds {
1118 incident.extend(self.block(old).edges.iter().map(|&e| (old.func, e)));
1119 }
1120 let mut incident: Vec<_> = incident.into_iter().collect();
1121 incident.sort_unstable();
1122 for (edge_func, edge) in incident {
1123 let EdgeData { from, to } = *self.edge(edge_func, edge);
1124 let from = BlockId::new(edge_func, from);
1125 let to = BlockId::new(edge_func, to);
1126 let new_from = block_map.get(&from).copied().unwrap_or(from);
1127 let new_to = block_map.get(&to).copied().unwrap_or(to);
1128 self.add_cfg_edge(new_from, new_to);
1129 }
1130
1131 for &old in olds {
1137 let Some(addr) = self.block(old).address else {
1138 continue;
1139 };
1140 let new = block_map[&old];
1141 let extra = self.block(old).extra_addresses.clone();
1142 addresses.rehome_block(addr, old, new);
1143 for &e in &extra {
1144 addresses.rehome_block(e, old, new);
1145 }
1146 self.block_mut(new).extra_addresses = extra;
1147 self.block_mut(new).address = Some(addr);
1148 }
1149
1150 for &old in olds {
1153 BasicBlock::from_id_mut(self, old).delete();
1154 }
1155
1156 self.rebuild_users(target);
1159 block_map
1160 }
1161
1162 fn function_registered_at_block(
1166 &self,
1167 addresses: &crate::address_index::AddressIndex,
1168 block: BlockId,
1169 ) -> Option<FunctionId> {
1170 self.block(block)
1171 .address
1172 .and_then(|addr| addresses.function_at(addr))
1173 }
1174
1175 fn split_tail(
1184 &self,
1185 addresses: &crate::address_index::AddressIndex,
1186 block: BlockId,
1187 g: FunctionId,
1188 ) -> Vec<BlockId> {
1189 let mut seen: HashSet<BlockId> = HashSet::default();
1190 seen.insert(block);
1191 let mut queue = vec![block];
1192 while let Some(b) = queue.pop() {
1193 let succs: Vec<BlockId> = BasicBlock::from_id(self, b)
1194 .successors()
1195 .map(|(_, s)| s)
1196 .collect();
1197 for s in succs {
1198 if seen.contains(&s) {
1199 continue;
1200 }
1201 if let Some(entry_func) = self.function_registered_at_block(addresses, s)
1205 && entry_func != g
1206 {
1207 continue;
1208 }
1209 seen.insert(s);
1210 queue.push(s);
1211 }
1212 }
1213 let mut tail: Vec<BlockId> = seen.into_iter().collect();
1214 tail.sort_unstable_by_key(|&b| (self.block(b).address, b.local, b.func));
1215 tail
1216 }
1217
1218 pub fn split_function_at(&mut self, block: BlockId) -> FunctionId {
1242 let mut addresses = crate::address_index::AddressIndex::analyze(self);
1243 self.split_function_at_indexed(&mut addresses, block)
1244 }
1245
1246 pub fn split_function_at_indexed(
1249 &mut self,
1250 addresses: &mut crate::address_index::AddressIndex,
1251 block: BlockId,
1252 ) -> FunctionId {
1253 use crate::value::insn::{Branch, CBranch, Callee, TailCall};
1254
1255 let addr = self
1256 .block(block)
1257 .address
1258 .expect("split_function_at: block has no machine address");
1259
1260 let g = match addresses.function_at(addr) {
1265 Some(existing) => existing,
1266 None => FunctionBody::make_at_addr_indexed(self, addresses, addr, None).id,
1267 };
1268
1269 loop {
1282 let tail_set: HashSet<BlockId> =
1283 self.split_tail(addresses, block, g).into_iter().collect();
1284 let mut promote: Option<BlockId> = None;
1285 'scan: for b in self.block_ids() {
1286 if tail_set.contains(&b) {
1287 continue;
1289 }
1290 let Some(mnemonic) = BasicBlock::from_id(self, b)
1291 .instructions()
1292 .last()
1293 .map(|t| t.mnemonic().clone())
1294 else {
1295 continue;
1296 };
1297 let targets = match &mnemonic {
1298 Mnemonic::Branch(Branch { target, .. }) => vec![*target],
1299 Mnemonic::CBranch(CBranch {
1300 success_block,
1301 failure_block,
1302 ..
1303 }) => vec![*success_block, *failure_block],
1304 _ => vec![],
1305 };
1306 for t in targets {
1307 let tid = BlockId::new(b.func, t);
1308 if tid == block || !tail_set.contains(&tid) {
1312 continue;
1313 }
1314 promote = Some(tid);
1323 break 'scan;
1324 }
1325 }
1326 match promote {
1327 Some(tid) => {
1328 self.split_function_at_indexed(addresses, tid);
1329 }
1330 None => break,
1331 }
1332 }
1333
1334 let mut tail = self.split_tail(addresses, block, g);
1337 let mut tail_set: HashSet<BlockId> = tail.iter().copied().collect();
1338
1339 let mut prev_owners: HashSet<FunctionId> = HashSet::default();
1342 for &b in &tail {
1343 prev_owners.insert(b.func);
1345 }
1346
1347 let effective_owner = |_ctx: &Context, candidate: BlockId| {
1351 if tail_set.contains(&candidate) {
1352 Some(g)
1353 } else {
1354 Some(candidate.func)
1356 }
1357 };
1358
1359 let foreign_entry =
1362 |ctx: &Context, target: BlockId, owner: FunctionId| -> Option<FunctionId> {
1363 let callee = if target == block {
1364 g
1365 } else {
1366 ctx.function_registered_at_block(addresses, target)?
1367 };
1368 (callee != owner).then_some(callee)
1369 };
1370
1371 let mut tail_calls: Vec<(InstructionId, FunctionId)> = Vec::new();
1377 let mut cond_calls: Vec<(
1384 InstructionId,
1385 BlockId,
1386 FunctionId,
1387 crate::value::LocalBlockId,
1388 )> = Vec::new();
1389 let relevant: Vec<BlockId> = self.block_ids();
1390 for b in relevant {
1391 let Some(owner) = effective_owner(self, b) else {
1392 continue;
1393 };
1394 let Some((term_id, mnemonic)) = BasicBlock::from_id(self, b)
1395 .instructions()
1396 .last()
1397 .map(|t| (t.id, t.mnemonic().clone()))
1398 else {
1399 continue;
1400 };
1401 match mnemonic {
1404 Mnemonic::Branch(Branch { target, .. }) => {
1405 if let Some(callee) = foreign_entry(self, BlockId::new(b.func, target), owner) {
1406 tail_calls.push((term_id, callee));
1407 }
1408 }
1409 Mnemonic::CBranch(CBranch {
1410 success_block,
1411 failure_block,
1412 ..
1413 }) => {
1414 if let Some(callee) =
1415 foreign_entry(self, BlockId::new(b.func, success_block), owner)
1416 {
1417 cond_calls.push((term_id, b, callee, success_block));
1418 }
1419 if let Some(callee) =
1420 foreign_entry(self, BlockId::new(b.func, failure_block), owner)
1421 {
1422 cond_calls.push((term_id, b, callee, failure_block));
1423 }
1424 }
1425 _ => {}
1426 }
1427 }
1428
1429 for (insn, callee) in tail_calls {
1430 self.replace_instruction_mnemonic(
1431 insn,
1432 Mnemonic::TailCall(TailCall {
1433 target: Callee::Real(callee),
1434 args: vec![],
1435 }),
1436 );
1437 }
1438 for (insn, owner_block, callee, arm_target) in cond_calls {
1439 let tramp = BasicBlock::make(self, owner_block.func).id;
1441 (self).builder(tramp).push_tail_call(callee);
1442 self.add_cfg_edge(owner_block, tramp);
1443
1444 if tail_set.contains(&owner_block) {
1452 tail.push(tramp);
1453 tail_set.insert(tramp);
1454 }
1455
1456 let Mnemonic::CBranch(mut cb) = self.instruction(insn).mnemonic().clone() else {
1457 continue;
1458 };
1459 let tramp_local = tramp.localize(insn.func);
1464 if cb.success_block == arm_target {
1465 cb.success_block = tramp_local;
1466 }
1467 if cb.failure_block == arm_target {
1468 cb.failure_block = tramp_local;
1469 }
1470 self.replace_instruction_mnemonic(insn, Mnemonic::CBranch(cb));
1471 }
1472
1473 let moved_owner = |candidate: BlockId| {
1482 if tail_set.contains(&candidate) {
1483 g
1484 } else {
1485 candidate.func
1486 }
1487 };
1488 let mut stale: HashSet<(FunctionId, EdgeId)> = HashSet::default();
1489 for &b in &tail {
1490 for edge in self.block(b).edges.iter().copied() {
1491 let &EdgeData { from, to } = self.edge(b.func, edge);
1492 let from = BlockId::new(b.func, from);
1493 let to = BlockId::new(b.func, to);
1494 let cross = moved_owner(from) != moved_owner(to);
1495 let touches_tail = tail_set.contains(&from) || tail_set.contains(&to);
1496 if cross && touches_tail {
1497 stale.insert((b.func, edge));
1498 }
1499 }
1500 }
1501 let mut stale: Vec<_> = stale.into_iter().collect();
1502 stale.sort_unstable();
1503 for (func, edge) in stale {
1504 self.remove_cfg_edge(func, edge);
1505 }
1506
1507 let moved = self.rehome_owned_blocks(addresses, g, &tail);
1511 self.bodies[g].set_root_id(Some(moved[&block].local));
1512
1513 self.recompute_instruction_addrs(g);
1515 for owner in prev_owners {
1516 if owner != g {
1517 self.recompute_instruction_addrs(owner);
1518 }
1519 }
1520
1521 g
1522 }
1523
1524 fn recompute_instruction_addrs(&mut self, func: FunctionId) {
1527 let blocks = FunctionBody::from_id(self, func).block_ids();
1528 let mut addrs = std::collections::BTreeSet::new();
1529 for b in blocks {
1530 for insn in BasicBlock::from_id(self, b).instructions() {
1531 if let Some(a) = insn.address() {
1532 addrs.insert(a);
1533 }
1534 }
1535 }
1536 self.bodies[func].instruction_addrs = addrs;
1537 }
1538
1539 fn rebuild_users(&mut self, func: FunctionId) {
1543 let live: Vec<InstructionId> = FunctionBody::from_id(self, func).instruction_ids();
1544 let users = &mut self.bodies[func].users;
1545 users.clear();
1546 for id in live {
1547 let args = self.bodies[func].insns[id.local].mnemonic().args();
1548 let users = &mut self.bodies[func].users;
1549 for arg in args {
1550 users.entry(arg).or_default().push(id.localize(func));
1551 }
1552 }
1553 }
1554
1555 pub fn assume_true(&mut self, prop: Proposition) -> bool {
1560 self.assume(prop, true)
1561 }
1562
1563 pub fn assume_false(&mut self, prop: Proposition) -> bool {
1565 self.assume(prop, false)
1566 }
1567
1568 fn assume(&mut self, prop: Proposition, value: bool) -> bool {
1569 match self.shared.values.truths.get(&prop) {
1570 Some(t) => t.value == value,
1571 None => {
1572 self.shared.values.truths.insert(
1573 prop,
1574 Truth {
1575 value,
1576 certainty: Certainty::Assumed,
1577 pass: PassName(pass_scope::current_pass()),
1578 },
1579 );
1580 true
1581 }
1582 }
1583 }
1584
1585 pub fn set_known(&mut self, prop: Proposition, value: bool) -> bool {
1593 let pass = PassName(pass_scope::current_pass());
1594 let novel = match self.shared.values.truths.get(&prop) {
1595 Some(prior) => {
1596 if prior.certainty == Certainty::Known && prior.value != value {
1601 self.shared
1602 .values
1603 .known_contradictions
1604 .push(KnownContradiction {
1605 prop,
1606 known: prior.value,
1607 proven: value,
1608 known_pass: prior.pass,
1609 proven_pass: pass,
1610 });
1611 return false;
1612 }
1613 if prior.certainty == Certainty::Assumed && prior.value != value {
1614 self.shared.values.violations.push(Violation {
1615 prop,
1616 assumed: prior.value,
1617 assuming_pass: prior.pass,
1618 asserting_pass: pass,
1619 });
1620 true
1621 } else {
1622 false
1623 }
1624 }
1625 None => true,
1626 };
1627 self.shared.values.truths.insert(
1628 prop,
1629 Truth {
1630 value,
1631 certainty: Certainty::Known,
1632 pass,
1633 },
1634 );
1635 novel
1636 }
1637
1638 pub fn seed_known(&mut self, prop: Proposition, value: bool, pass: PassName) {
1648 let prior = self.shared.values.truths.insert(
1649 prop,
1650 Truth {
1651 value,
1652 certainty: Certainty::Known,
1653 pass,
1654 },
1655 );
1656 debug_assert!(prior.is_none(), "seeding {prop:?} over an existing truth");
1657 }
1658
1659 pub fn truth(&self, prop: Proposition) -> Option<Truth> {
1661 self.shared.values.truths.get(&prop).copied()
1662 }
1663
1664 pub fn set_assumed_call_convention(
1670 &mut self,
1671 effect: Option<crate::assumption::AssumedCallEffect>,
1672 ) {
1673 self.shared.assumed_call_convention = effect;
1674 }
1675
1676 pub fn assumed_call_convention(&self) -> Option<&crate::assumption::AssumedCallEffect> {
1680 self.shared.assumed_call_convention.as_ref()
1681 }
1682
1683 pub fn known(&self, prop: Proposition) -> Option<bool> {
1685 self.truth(prop)
1686 .filter(|t| t.certainty == Certainty::Known)
1687 .map(|t| t.value)
1688 }
1689
1690 pub fn truths(&self) -> impl Iterator<Item = (Proposition, Truth)> + '_ {
1692 self.shared.values.truths.iter().map(|(&p, &t)| (p, t))
1693 }
1694
1695 pub fn known_facts(&self) -> impl Iterator<Item = (Proposition, bool, PassName)> + '_ {
1698 self.truths()
1699 .filter(|(_, t)| t.certainty == Certainty::Known)
1700 .map(|(p, t)| (p, t.value, t.pass))
1701 }
1702
1703 pub fn violations(&self) -> &[Violation] {
1706 &self.shared.values.violations
1707 }
1708
1709 pub fn known_contradictions(&self) -> &[KnownContradiction] {
1713 &self.shared.values.known_contradictions
1714 }
1715
1716 pub fn get_literal_value(&self, id: LiteralId) -> u64 {
1718 self.shared.values.literals[id].value
1719 }
1720
1721 pub fn get_insn(&self, id: InstructionId) -> InstructionRef<'str, '_> {
1723 InstructionRef::from_id(self, id)
1724 }
1725
1726 pub fn body(&self, fid: FunctionId) -> &crate::value::FunctionBody<'str> {
1736 &self.bodies[fid]
1737 }
1738
1739 pub fn body_mut(&mut self, fid: FunctionId) -> &mut crate::value::FunctionBody<'str> {
1741 &mut self.bodies[fid]
1742 }
1743
1744 pub fn push_insn(&mut self, func: FunctionId, insn: Instruction<'str>) -> InstructionId {
1759 let args = insn.mnemonic().args();
1760 let local = self.bodies[func].insns.push(insn);
1761 let id = InstructionId::new(func, local);
1762 for arg in args {
1763 self.bodies[func]
1764 .users
1765 .entry(arg)
1766 .or_default()
1767 .push(id.localize(func));
1768 }
1769 id
1770 }
1771
1772 pub fn instruction(&self, id: InstructionId) -> &Instruction<'str> {
1774 &self.bodies[id.func].insns[id.local]
1775 }
1776
1777 pub fn instruction_mut(&mut self, id: InstructionId) -> &mut Instruction<'str> {
1779 &mut self.bodies[id.func].insns[id.local]
1780 }
1781
1782 pub fn contains_instruction(&self, id: InstructionId) -> bool {
1784 Into::<usize>::into(id.func) < self.bodies.len()
1785 && self.bodies[id.func].insns.contains(id.local)
1786 }
1787
1788 pub fn block(&self, id: BlockId) -> &BasicBlock<'str> {
1790 &self.bodies[id.func].blocks[id.local]
1791 }
1792
1793 pub fn block_mut(&mut self, id: BlockId) -> &mut BasicBlock<'str> {
1795 &mut self.bodies[id.func].blocks[id.local]
1796 }
1797
1798 pub fn contains_block(&self, id: BlockId) -> bool {
1800 Into::<usize>::into(id.func) < self.bodies.len()
1801 && self.bodies[id.func].blocks.contains(id.local)
1802 }
1803
1804 pub fn block_param(&self, id: BlockParamId) -> &BlockParam<'str> {
1806 &self.bodies[id.func].params[id.local]
1807 }
1808
1809 pub fn block_param_mut(&mut self, id: BlockParamId) -> &mut BlockParam<'str> {
1811 &mut self.bodies[id.func].params[id.local]
1812 }
1813
1814 pub fn contains_block_param(&self, id: BlockParamId) -> bool {
1816 Into::<usize>::into(id.func) < self.bodies.len()
1817 && self.bodies[id.func].params.contains(id.local)
1818 }
1819
1820 pub fn edge(&self, func: FunctionId, id: EdgeId) -> &EdgeData {
1822 &self.bodies[func].edges[id]
1823 }
1824
1825 pub fn edge_mut(&mut self, func: FunctionId, id: EdgeId) -> &mut EdgeData {
1827 &mut self.bodies[func].edges[id]
1828 }
1829
1830 pub fn users_of(&self, value: ValueId) -> Vec<InstructionId> {
1835 match value.owning_function() {
1836 Some(func) => self.bodies[func].users_of(value),
1837 None => Vec::new(),
1838 }
1839 }
1840
1841 pub fn has_users(&self, value: ValueId) -> bool {
1843 match value.owning_function() {
1844 Some(func) => self.bodies[func].has_users(value),
1845 None => false,
1846 }
1847 }
1848
1849 pub fn push_block(&mut self, func: FunctionId, block: BasicBlock<'str>) -> BlockId {
1850 let local = self.bodies[func].blocks.push(block);
1851 let id = BlockId::new(func, local);
1852 self.bodies[func].roster.push(local);
1854 id
1855 }
1856
1857 pub fn push_block_param(&mut self, func: FunctionId, param: BlockParam<'str>) -> BlockParamId {
1858 let local = self.bodies[func].params.push(param);
1859 BlockParamId::new(func, local)
1860 }
1861
1862 pub fn push_edge(&mut self, func: FunctionId, edge: EdgeData) -> EdgeId {
1863 self.bodies[func].edges.push(edge)
1864 }
1865
1866 pub fn push_function(
1869 &mut self,
1870 interface: crate::value::function::FunctionInterface<'str>,
1871 body: FunctionBody<'str>,
1872 ) -> FunctionId {
1873 let expected = FunctionId::from(self.bodies.len());
1874 assert_eq!(
1875 body.id(),
1876 expected,
1877 "function body id does not match its registry slot"
1878 );
1879 let id = self.bodies.push(body);
1880 let iid = self.interfaces.push(interface);
1881 debug_assert_eq!(
1882 Into::<usize>::into(id),
1883 Into::<usize>::into(iid),
1884 "function body/interface registries drifted"
1885 );
1886 id
1887 }
1888
1889 pub fn get_register(&self, id: RegisterId) -> VarnodeRef<'str, '_> {
1892 Varnode::from_id(self, self.shared.registers[&id])
1893 }
1894
1895 pub fn get_const(&self, value: u64, size: usize) -> LiteralRef<'str, '_> {
1897 let type_id = self.shared.types.get_or_make_int(size);
1898 let id = self
1899 .shared
1900 .values
1901 .get_or_make_typed_literal(value, type_id, size);
1902 LiteralRef::from_id(self, id)
1903 }
1904
1905 pub fn get_bool_const(&self, value: bool) -> LiteralRef<'str, '_> {
1908 let type_id = self.shared.types.get_or_make_bool();
1909 let id = self
1910 .shared
1911 .values
1912 .get_or_make_typed_literal(u64::from(value), type_id, 1);
1913 LiteralRef::from_id(self, id)
1914 }
1915
1916 pub fn get_poison(&self, type_id: crate::types::TypeId) -> ValueId {
1920 ValueId::Poison(self.shared.values.push_poison(type_id))
1921 }
1922
1923 pub fn get_typed_const(
1929 &self,
1930 value: u64,
1931 type_id: crate::types::TypeId,
1932 ) -> LiteralRef<'str, '_> {
1933 let size = self.shared.types.size_of(type_id);
1934 let id = self
1935 .shared
1936 .values
1937 .get_or_make_typed_literal(value, type_id, size);
1938 LiteralRef::from_id(self, id)
1939 }
1940
1941 pub fn get_bytes(&self, data: Vec<u8>) -> crate::value::BytesRef<'str, '_> {
1949 let i8_ty = self.shared.types.get_or_make_int(1);
1950 let type_id = self.shared.types.get_or_make_array(i8_ty, data.len());
1951 self.get_typed_bytes(data, type_id)
1952 }
1953
1954 pub fn get_typed_bytes(
1960 &self,
1961 data: Vec<u8>,
1962 type_id: crate::types::TypeId,
1963 ) -> crate::value::BytesRef<'str, '_> {
1964 let id = self
1965 .shared
1966 .values
1967 .bytes
1968 .push(crate::value::Bytes { data, type_id });
1969 crate::value::BytesRef::from_id(self, id)
1970 }
1971
1972 pub fn type_of(&self, id: ValueId) -> crate::types::TypeId {
1977 match id {
1978 ValueId::Literal(lid) => self.shared.values.literals[lid].type_id,
1979 ValueId::Bytes(bid) => self.shared.values.bytes[bid].type_id,
1980 ValueId::Instruction(iid) => self.instruction(iid).type_id,
1981 ValueId::BlockParam(pid) => self.block_param(pid).type_id,
1982 ValueId::Varnode(vid) => {
1983 if let Some(&ty) = self.shared.values.varnode_types.get(&vid) {
1984 return ty;
1985 }
1986 let size = self.shared.values.varnodes[vid].size_bytes();
1987 self.shared.types.get_or_make_int(size)
1988 }
1989 ValueId::Temp(id) => self
1990 .shared
1991 .types
1992 .get_or_make_int(self.bodies[id.func].temps[id.local].size),
1993 ValueId::Poison(pid) => self.shared.values.poisons[pid].type_id,
1994 ValueId::BasicBlock(_) | ValueId::Function(_) => self.shared.types.get_or_make_int(0),
1997 }
1998 }
1999
2000 pub fn stored_type_of(&self, id: ValueId) -> Option<crate::types::TypeId> {
2006 match id {
2007 ValueId::Literal(lid) => Some(self.shared.values.literals[lid].type_id),
2008 ValueId::Bytes(bid) => Some(self.shared.values.bytes[bid].type_id),
2009 ValueId::Instruction(iid) => Some(self.instruction(iid).type_id),
2010 ValueId::BlockParam(pid) => Some(self.block_param(pid).type_id),
2011 ValueId::Varnode(vid) => self.shared.values.varnode_types.get(&vid).copied(),
2012 ValueId::Poison(pid) => Some(self.shared.values.poisons[pid].type_id),
2013 ValueId::Temp(_) => None,
2014 ValueId::BasicBlock(_) | ValueId::Function(_) => None,
2015 }
2016 }
2017
2018 pub fn set_varnode_type(&mut self, varnode: VarnodeId, type_id: crate::types::TypeId) {
2023 self.shared.values.varnode_types.insert(varnode, type_id);
2024 }
2025
2026 pub fn users(&self, value: impl Into<ValueId>) -> Vec<InstructionId> {
2034 self.users_of(value.into())
2035 }
2036
2037 pub fn users_across_functions(&self, value: impl Into<ValueId>) -> Vec<InstructionId> {
2042 let value = value.into();
2043 if value.owning_function().is_some() {
2044 self.users_of(value)
2045 } else {
2046 self.functions().flat_map(|f| f.users_of(value)).collect()
2047 }
2048 }
2049
2050 pub fn view(&self) -> ModuleView<'_, 'str> {
2059 ModuleView::new(self)
2060 }
2061 pub fn shr(&self) -> &Shared<'str> {
2066 &self.shared
2067 }
2068 pub fn function(&self, f: FunctionId) -> &FunctionBody<'str> {
2070 &self.bodies[f]
2071 }
2072 pub fn function_mut(&mut self, f: FunctionId) -> &mut FunctionBody<'str> {
2074 &mut self.bodies[f]
2075 }
2076
2077 pub fn block_ref(&self, id: BlockId) -> BlockRef<'str, '_, ModuleView<'_, 'str>> {
2079 self.view().block_ref(id)
2080 }
2081 pub fn insn_ref(&self, id: InstructionId) -> InstructionRef<'str, '_, ModuleView<'_, 'str>> {
2083 self.view().insn_ref(id)
2084 }
2085 pub fn param_ref(&self, id: BlockParamId) -> BlockParamRef<'str, '_, ModuleView<'_, 'str>> {
2087 self.view().param_ref(id)
2088 }
2089 pub fn function_ref(&self, id: FunctionId) -> FunctionRef<'str, '_, ModuleView<'_, 'str>> {
2091 self.view().function_ref(id)
2092 }
2093
2094 pub fn push_mnemonic(
2096 &mut self,
2097 func: FunctionId,
2098 mnemonic: Mnemonic,
2099 size: usize,
2100 ) -> InstructionId {
2101 let type_id = self.shared.types.get_or_make_int(size);
2102 self.push_insn(func, Instruction::new(type_id, mnemonic))
2103 }
2104
2105 pub fn push_mnemonic_with_type(
2108 &mut self,
2109 func: FunctionId,
2110 mnemonic: Mnemonic,
2111 type_id: crate::types::TypeId,
2112 ) -> InstructionId {
2113 self.push_insn(func, Instruction::new(type_id, mnemonic))
2114 }
2115
2116 pub fn make_block(&mut self, func: FunctionId) -> BlockId {
2119 self.push_block(func, BasicBlock::detached())
2120 }
2121
2122 pub fn register_local_name(
2125 &mut self,
2126 id: ValueId,
2127 name: Cow<'str, str>,
2128 old_name: Option<&str>,
2129 ) -> Result<()> {
2130 let existing = match id.name_scope_function() {
2131 Some(func) => self
2132 .function(func)
2133 .names
2134 .get(&name)
2135 .map(|id| id.qualify(func)),
2136 None => self.get_named(&name),
2137 };
2138 if let Some(existing) = existing {
2139 return if existing == id {
2140 Ok(())
2141 } else {
2142 Err(Error::spanless(ErrorTy::DuplicateName(name.to_string())))
2143 };
2144 }
2145 match id.name_scope_function() {
2146 Some(func) => self
2147 .function_mut(func)
2148 .names
2149 .register(name, id.localize(func), old_name),
2150 None => self.update_name(name, id, old_name),
2151 }
2152 }
2153
2154 pub(crate) fn set_address_indexed(
2156 &mut self,
2157 addresses: &mut crate::address_index::AddressIndex,
2158 addr: u64,
2159 id: ValueId,
2160 ) -> crate::error::Result<()> {
2161 let target = match id {
2162 ValueId::Function(id) => crate::address_index::AddressTarget::Function(id),
2163 ValueId::BasicBlock(id) => crate::address_index::AddressTarget::Block(id),
2164 _ => unreachable!("only functions and blocks have module addresses"),
2165 };
2166 addresses.register(self, addr, target)
2167 }
2168
2169 pub fn update_name(
2172 &mut self,
2173 name: Cow<'str, str>,
2174 id: ValueId,
2175 old_name: Option<&str>,
2176 ) -> Result<()> {
2177 match id.name_scope_function() {
2178 Some(func) => self.bodies[func]
2179 .names
2180 .register(name, id.localize(func), old_name),
2181 None => self.shared.name_map.register(name, id, old_name),
2182 }
2183 }
2184
2185 pub fn get_named_in_scope(&self, id: ValueId, name: &str) -> Option<ValueId> {
2190 match id.name_scope_function() {
2191 Some(func) => self.bodies[func].names.get(name).map(|id| id.qualify(func)),
2192 None => self.shared.name_map.get(name),
2193 }
2194 }
2195
2196 pub fn get_named(&self, name: &str) -> Option<ValueId> {
2205 self.shared.name_map.get(name)
2206 }
2207
2208 pub fn get_unique_name(&mut self, name: Cow<'str, str>) -> Cow<'str, str> {
2213 self.shared.name_map.unique(name)
2214 }
2215
2216 pub fn get_unique_name_in(&mut self, func: FunctionId, name: Cow<'str, str>) -> Cow<'str, str> {
2220 self.bodies[func].names.unique(name)
2221 }
2222}
2223
2224#[derive(Clone, serde::Serialize, serde::Deserialize)]
2235pub struct NameTable<'str, Id = ValueId> {
2236 map: HashMap<Cow<'str, str>, Id>,
2238 #[serde(skip)]
2242 suffix_hint: HashMap<String, u32>,
2243}
2244
2245impl<Id> Default for NameTable<'_, Id> {
2246 fn default() -> Self {
2247 Self {
2248 map: HashMap::default(),
2249 suffix_hint: HashMap::default(),
2250 }
2251 }
2252}
2253
2254impl<'str, Id: Copy + Eq> NameTable<'str, Id> {
2255 pub(crate) fn entries(&self) -> impl Iterator<Item = (&str, Id)> + '_ {
2256 self.map.iter().map(|(name, &value)| (name.as_ref(), value))
2257 }
2258
2259 pub fn get(&self, name: &str) -> Option<Id> {
2261 self.map.get(name).copied()
2262 }
2263
2264 pub fn contains(&self, name: &str) -> bool {
2266 self.map.contains_key(name)
2267 }
2268
2269 pub fn register(&mut self, name: Cow<'str, str>, id: Id, old_name: Option<&str>) -> Result<()> {
2273 if let Some(old_name) = old_name {
2274 self.forget(old_name);
2275 }
2276 match self.map.insert(name.clone(), id) {
2277 Some(_) => Err(Error::spanless(ErrorTy::DuplicateName(name.to_string()))),
2278 None => Ok(()),
2279 }
2280 }
2281
2282 pub fn forget(&mut self, name: &str) {
2286 self.map.remove(name);
2287 if let Some((base, suffix)) = split_generated_suffix(name)
2288 && let Some(hint) = self.suffix_hint.get_mut(base)
2289 {
2290 *hint = (*hint).min(suffix);
2291 }
2292 }
2293
2294 pub fn unique(&mut self, name: Cow<'str, str>) -> Cow<'str, str> {
2299 use std::fmt::Write as _;
2300
2301 if !self.map.contains_key(&name) {
2302 return name;
2303 }
2304 let base: &str = &name;
2305 let mut suffix = self.suffix_hint.get(base).copied().unwrap_or(1).max(1);
2306 let mut unique_name = format!("{base}_{suffix}");
2307 while self.map.contains_key(unique_name.as_str()) {
2308 suffix += 1;
2309 unique_name.clear();
2310 let _ = write!(unique_name, "{base}_{suffix}");
2311 }
2312 self.suffix_hint.insert(base.to_string(), suffix);
2313 Cow::Owned(unique_name)
2314 }
2315}
2316
2317fn split_generated_suffix(name: &str) -> Option<(&str, u32)> {
2322 let (base, digits) = name.rsplit_once('_')?;
2323 if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) {
2324 return None;
2325 }
2326 Some((base, digits.parse().ok()?))
2327}
2328
2329fn remap_rehomed_type(
2332 ctx: &Context<'_>,
2333 type_id: crate::types::TypeId,
2334 target: FunctionId,
2335 temp_space_map: &HashMap<TempSpaceId, TempSpaceId>,
2336) -> crate::types::TypeId {
2337 let Some(MemorySpaceId::Temp(old_space)) = ctx.shared.types.space_of(type_id) else {
2338 return type_id;
2339 };
2340 let Some(&new_space) = temp_space_map.get(&old_space) else {
2341 debug_assert_eq!(
2342 old_space.func, target,
2343 "rehome: result type references unmapped foreign temporary space {old_space:?}"
2344 );
2345 return type_id;
2346 };
2347 ctx.shared.types.get_or_make_space_address(
2348 ctx.shared.types.size_of(type_id),
2349 MemorySpaceId::Temp(new_space),
2350 )
2351}
2352
2353fn remap_rehomed_memory_space(
2356 mnemonic: &mut Mnemonic,
2357 old_func: FunctionId,
2358 target: FunctionId,
2359 temp_space_map: &HashMap<TempSpaceId, TempSpaceId>,
2360) {
2361 let remap = |space: &mut LocalMemorySpaceId| {
2362 let LocalMemorySpaceId::Temp(old_local) = *space else {
2363 return;
2364 };
2365 let old = TempSpaceId::new(old_func, old_local);
2366 if let Some(&new) = temp_space_map.get(&old) {
2367 *space = LocalMemorySpaceId::Temp(new.local);
2368 } else {
2369 debug_assert_eq!(
2370 old_func, target,
2371 "rehome: mnemonic references unmapped foreign temporary space {old:?}"
2372 );
2373 }
2374 };
2375 match mnemonic {
2376 Mnemonic::Load(load) => remap(&mut load.space),
2377 Mnemonic::Store(store) => remap(&mut store.space),
2378 _ => {}
2379 }
2380}
2381
2382fn remap_symbolic_block_literal(
2408 literals: &crate::value::interner::LiteralInterner,
2409 arg: crate::value::LocalValueId,
2410 block_map: &HashMap<BlockId, BlockId>,
2411) -> Option<crate::value::LocalValueId> {
2412 use crate::value::literal::SymbolicRef;
2413
2414 let crate::value::LocalValueId::Literal(lid) = arg else {
2415 return None;
2416 };
2417 let literal = literals[lid].clone();
2418 let Some(SymbolicRef::Block(old_block)) = literal.symbolic else {
2419 return None;
2420 };
2421 let &new_block = block_map.get(&old_block)?;
2422 let new_lit = literals.push_literal(crate::value::literal::Literal {
2423 symbolic: Some(SymbolicRef::Block(new_block)),
2424 ..literal
2425 });
2426 Some(crate::value::LocalValueId::Literal(new_lit))
2427}
2428
2429fn remap_block_targets(
2430 mnemonic: &mut Mnemonic,
2431 old_func: FunctionId,
2432 new_func: FunctionId,
2433 block_map: &HashMap<BlockId, BlockId>,
2434) {
2435 let remap = |b: &mut crate::value::LocalBlockId| {
2436 if let Some(&new) = block_map.get(&BlockId::new(old_func, *b)) {
2437 *b = new.localize(new_func);
2438 }
2439 };
2440 match mnemonic {
2441 Mnemonic::Branch(branch) => remap(&mut branch.target),
2442 Mnemonic::CBranch(cbranch) => {
2443 remap(&mut cbranch.success_block);
2444 remap(&mut cbranch.failure_block);
2445 }
2446 _ => {}
2447 }
2448}
2449
2450impl Display for Context<'_> {
2451 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2452 self.functions().try_for_each(|fun| fun.fmt(f))?;
2453
2454 self.blocks()
2455 .filter(|block| block.parent().is_none())
2456 .try_for_each(|block| block.fmt(f))
2457 }
2458}
2459
2460pub struct FunctionIter<'str, 'ctx> {
2461 ctx: &'ctx Context<'str>,
2462 inner: registry::Iter<'ctx, FunctionId, FunctionBody<'str>>,
2463}
2464
2465impl<'str, 'ctx> Iterator for FunctionIter<'str, 'ctx> {
2466 type Item = FunctionRef<'str, 'ctx>;
2467
2468 fn next(&mut self) -> Option<Self::Item> {
2469 let ctx = self.ctx;
2470 self.inner.next().map(|f| FunctionRef::from_id(ctx, f.id))
2471 }
2472}
2473
2474impl<'str, 'ctx> IntoIterator for &'ctx Context<'str> {
2475 type Item = FunctionRef<'str, 'ctx>;
2476 type IntoIter = FunctionIter<'str, 'ctx>;
2477
2478 fn into_iter(self) -> Self::IntoIter {
2479 self.iter()
2480 }
2481}
2482
2483#[cfg(test)]
2484mod tests {
2485 use super::*;
2486 use crate::value::{
2487 BasicBlock, FunctionBody, ValueId,
2488 insn::{Binary, Binop, Call, Callee, IntBinop, Load, Mnemonic},
2489 };
2490 use wazabin_qcode_macro::qcode;
2491
2492 fn make_fn_with_blocks(ctx: &mut Context<'static>, name: &'static str, n: usize) -> FunctionId {
2493 let f = FunctionBody::make(ctx, name.into()).unwrap().id;
2495 for _ in 0..n {
2496 BasicBlock::make(ctx, f);
2497 }
2498 f
2499 }
2500
2501 #[test]
2502 #[should_panic(expected = "cannot reuse a block stored in another function arena")]
2503 fn get_or_make_block_rejects_foreign_storage_at_address() {
2504 let mut ctx = Context::new();
2505 let a = FunctionBody::make(&mut ctx, "address_owner".into())
2506 .unwrap()
2507 .id;
2508 let b = FunctionBody::make(&mut ctx, "address_requester".into())
2509 .unwrap()
2510 .id;
2511 BasicBlock::make(&mut ctx, a).with_address(0x1000);
2512
2513 ctx.get_or_make_block(0x1000, b);
2514 }
2515
2516 #[test]
2517 #[should_panic(expected = "cannot create a block at an address owned by another function")]
2518 fn get_or_make_block_rejects_foreign_function_address_without_root() {
2519 let mut ctx = Context::new();
2520 FunctionBody::make_at_addr(&mut ctx, 0x1000, None);
2521 let requester = FunctionBody::make(&mut ctx, "address_requester".into())
2522 .unwrap()
2523 .id;
2524
2525 ctx.get_or_make_block(0x1000, requester);
2526 }
2527
2528 #[test]
2529 fn functions_iter_yields_all_functions() {
2530 let mut ctx = Context::new();
2531 let alpha = make_fn_with_blocks(&mut ctx, "alpha", 1);
2532 let beta = make_fn_with_blocks(&mut ctx, "beta", 1);
2533
2534 let names: Vec<_> = ctx.functions().map(|f| f.name().to_string()).collect();
2535 assert!(names.contains(&"alpha".to_string()));
2536 assert!(names.contains(&"beta".to_string()));
2537 assert_eq!(names.len(), 2);
2538 assert_eq!(ctx.function_ids(), vec![alpha, beta]);
2539 assert_eq!(ctx.function_ids().len(), ctx.interfaces.len());
2540 }
2541
2542 #[test]
2543 fn body_view_reads_match_module_reads() {
2544 use crate::value::{BodyView, FunctionId, FunctionRef, ModuleView, QCodeView};
2545
2546 let mut ctx = Context::new();
2547 qcode!(
2548 ctx,
2549 "
2550 fn foo:
2551 <bb1>
2552 if i8 1 goto <bb2> else goto <bb3>;
2553 <bb2>
2554 goto <bb3>;
2555 <bb3>
2556 return at 0;
2557 "
2558 );
2559 let fid = FunctionBody::from_name(&ctx, "foo").unwrap().id();
2560 let fid = ValueId::as_function(fid).unwrap();
2561
2562 type Snap = (String, Vec<(String, Vec<String>, Vec<String>, usize)>);
2567 fn snapshot<'a, 'str: 'a>(view: impl QCodeView<'a, 'str>, fid: FunctionId) -> Snap {
2568 let f = FunctionRef::new(view, fid);
2569 let blocks = f
2570 .blocks()
2571 .map(|b| {
2572 let name = b.name().unwrap_or("?").to_string();
2573 let mut succ: Vec<String> = b
2574 .successors()
2575 .map(|(_, s)| BlockRef::new(view, s).name().unwrap_or("?").to_string())
2576 .collect();
2577 succ.sort();
2578 let ops: Vec<String> =
2579 b.instructions().map(|i| i.opcode().to_string()).collect();
2580 (name, succ, ops, b.num_params())
2581 })
2582 .collect();
2583 (f.name().to_string(), blocks)
2584 }
2585
2586 let module_snap = snapshot(ModuleView::new(&ctx), fid);
2587 assert!(!module_snap.1.is_empty(), "sanity: foo has blocks");
2588
2589 let checked = BodyView::new(&ctx.bodies[fid], &ctx.shared, &ctx.interfaces);
2592 let checked_snap = snapshot(checked, fid);
2593 assert_eq!(
2594 module_snap, checked_snap,
2595 "reads through BodyView must match the module reads"
2596 );
2597 }
2598
2599 #[test]
2600 fn body_mut_mut_matches_module_mut() {
2601 use crate::value::{
2602 BlockParam, FunctionId, FunctionRef, InstructionId, Renameable,
2603 block::BlockId,
2604 block_param::BlockParamId,
2605 util::{base_ref::BaseRef, body_mut::BodyMut},
2606 };
2607
2608 fn build(mut ctx: &mut Context<'static>) -> (FunctionId, BlockId, BlockId, InstructionId) {
2609 qcode!(
2610 ctx,
2611 "
2612 varnode i64 x;
2613 fn foo:
2614 <entry>
2615 %a = load(x:8, &x);
2616 %b = load(x:8, &x);
2617 goto <bb1>;
2618 <bb1>
2619 return at %a;
2620 "
2621 );
2622 let fid = foo;
2623 let entry = FunctionRef::from_id(ctx, fid).root().unwrap().id;
2624 let bb1 = FunctionRef::from_id(ctx, fid)
2625 .blocks()
2626 .map(|b| b.id)
2627 .find(|&b| b != entry)
2628 .unwrap();
2629 let insns = BasicBlock::from_id(ctx, entry).instruction_ids();
2630 (fid, entry, bb1, insns[0])
2631 }
2632
2633 fn add_param(ctx: &mut Context<'static>, bb1: BlockId) -> BlockParamId {
2635 BasicBlock::from_id_mut(ctx, bb1).push_param(8).id
2636 }
2637
2638 type MSnap = Vec<(String, Option<String>, Vec<usize>, Vec<String>, Vec<String>)>;
2641 fn snap(ctx: &Context, fid: FunctionId) -> MSnap {
2642 FunctionRef::from_id(ctx, fid)
2643 .blocks()
2644 .map(|b| {
2645 let name = b.name().unwrap_or("?").to_string();
2646 let comment = b.comment().map(str::to_string);
2647 let params: Vec<usize> = b.params().map(|p| p.size()).collect();
2648 let ops: Vec<String> =
2649 b.instructions().map(|i| i.opcode().to_string()).collect();
2650 let mut succ: Vec<String> = b
2651 .successors()
2652 .map(|(_, s)| {
2653 BasicBlock::from_id(ctx, s)
2654 .name()
2655 .unwrap_or("?")
2656 .to_string()
2657 })
2658 .collect();
2659 succ.sort();
2660 (name, comment, params, ops, succ)
2661 })
2662 .collect()
2663 }
2664
2665 let mut ctx_a = Context::new();
2667 let (fid, entry, bb1, a) = build(&mut ctx_a);
2668 let param = add_param(&mut ctx_a, bb1);
2669 let b = BasicBlock::from_id(&ctx_a, entry).instruction_ids()[1];
2670 BasicBlock::from_id_mut(&mut ctx_a, entry).set_comment(Some("c".into()));
2671 BasicBlock::from_id_mut(&mut ctx_a, entry)
2672 .rename("start".into())
2673 .unwrap();
2674 let e = ctx_a.add_cfg_edge(entry, bb1);
2675 ctx_a.remove_cfg_edge(entry.func, e);
2676 ctx_a.replace_instruction(a, ValueId::Instruction(b));
2677 BlockParam::from_id_mut(&mut ctx_a, param).set_size(4);
2678 let snap_a = snap(&ctx_a, fid);
2679
2680 let mut ctx_b = Context::new();
2682 let (fid_b, entry_b, bb1_b, a_b) = build(&mut ctx_b);
2683 let param_b = add_param(&mut ctx_b, bb1_b);
2684 let b_b = BasicBlock::from_id(&ctx_b, entry_b).instruction_ids()[1];
2685
2686 {
2687 let mut host = BodyMut::new(&mut ctx_b.bodies[fid_b], &ctx_b.shared, &ctx_b.interfaces);
2688 let mut r = BaseRef::new(host.reborrow(), entry_b);
2689 r.set_comment(Some("c".into()));
2690 let mut r = BaseRef::new(host.reborrow(), entry_b);
2691 r.rename("start".into()).unwrap();
2692 let e = host.add_cfg_edge(entry_b, bb1_b);
2693 host.remove_cfg_edge(e);
2694 host.replace_instruction(a_b, ValueId::Instruction(b_b));
2695 let mut r = BaseRef::new(host.reborrow(), param_b);
2696 r.set_size(4);
2697 }
2698 let snap_b = snap(&ctx_b, fid_b);
2699
2700 assert_eq!(
2701 snap_a, snap_b,
2702 "mutations through a pass-scoped host must match the module-path mutations"
2703 );
2704 }
2705
2706 #[test]
2707 fn into_iterator_for_context_matches_functions() {
2708 let mut ctx = Context::new();
2709 make_fn_with_blocks(&mut ctx, "f1", 1);
2710 make_fn_with_blocks(&mut ctx, "f2", 1);
2711
2712 let via_method: Vec<_> = ctx.functions().map(|f| f.id()).collect();
2713 let via_into: Vec<_> = (&ctx).into_iter().map(|f| f.id()).collect();
2714 assert_eq!(via_method, via_into);
2715 }
2716
2717 #[test]
2718 fn blocks_iter_yields_all_blocks() {
2719 let mut ctx = Context::new();
2720 make_fn_with_blocks(&mut ctx, "g", 3);
2721
2722 let count = ctx.blocks().count();
2723 assert_eq!(count, 3);
2724 }
2725
2726 #[test]
2727 fn instructions_iter_yields_all_instructions() {
2728 let mut ctx = Context::new();
2729
2730 qcode!(
2731 ctx,
2732 "
2733 varnode i64 ptr;
2734
2735 <block>
2736 store(ptr:8, &ptr <- i64 0x1234);
2737 return at ptr;
2738 "
2739 );
2740
2741 let count = ctx.instructions().count();
2742 assert!(count >= 1, "expected at least one instruction, got {count}");
2743 }
2744
2745 #[test]
2746 fn move_insn_before_preserves_id_and_supports_arbitrary_anchors() {
2747 let mut ctx = Context::new();
2748 qcode!(
2749 ctx,
2750 "
2751 fn f:
2752 <source>
2753 %a = i64 0x1 + i64 0x2;
2754 %free = i64 0x5 + i64 0x6;
2755 goto <target>;
2756 <target>
2757 %b = i64 0x3 + i64 0x4;
2758 %consumer = %a + %b;
2759 return %consumer;
2760 "
2761 );
2762
2763 assert!(ctx.users(a).contains(&consumer));
2764 ctx.move_insn_before(a, b);
2765
2766 assert!(ctx.contains_instruction(a), "moving keeps the ID live");
2767 assert_eq!(ctx.get_insn(a).parent().map(|block| block.id), Some(target));
2768 assert!(
2769 !BasicBlock::from_id(&ctx, source)
2770 .instruction_ids()
2771 .contains(&a)
2772 );
2773 assert_eq!(
2774 BasicBlock::from_id(&ctx, target).instruction_ids()[..3],
2775 [a, b, consumer]
2776 );
2777 assert!(
2778 ctx.users(a).contains(&consumer),
2779 "moving preserves use-map entries"
2780 );
2781
2782 ctx.move_insn_before(b, a);
2784 assert_eq!(
2785 BasicBlock::from_id(&ctx, target).instruction_ids()[..3],
2786 [b, a, consumer]
2787 );
2788
2789 let return_id = *BasicBlock::from_id(&ctx, target)
2791 .instruction_ids()
2792 .last()
2793 .unwrap();
2794 ctx.move_insn_before(free, return_id);
2795 assert_eq!(
2796 BasicBlock::from_id(&ctx, target).instruction_ids()[..4],
2797 [b, a, consumer, free]
2798 );
2799 }
2800
2801 #[test]
2802 fn remove_instruction_removes_from_block() {
2803 let mut ctx = Context::new();
2804 qcode!(
2805 ctx,
2806 "
2807 varnode i64 x;
2808 <block>
2809 %a = load(x:8, &x);
2810 %b = load(x:8, &x);
2811 return at %a;
2812 "
2813 );
2814 let block_ref = BasicBlock::from_id(&ctx, block);
2815 let ids = block_ref.instruction_ids();
2816 let load_a = ids[0];
2817 let original_len = ids.len();
2818
2819 ctx.remove_instruction(load_a);
2820
2821 let remaining = BasicBlock::from_id(&ctx, block).instruction_ids();
2822 assert_eq!(remaining.len(), original_len - 1);
2823 assert!(!remaining.contains(&load_a));
2824 }
2825
2826 #[test]
2827 fn remove_instruction_drops_payload() {
2828 let mut ctx = Context::new();
2829 qcode!(
2830 ctx,
2831 "
2832 varnode i64 x;
2833 <block>
2834 %a = load(x:8, &x);
2835 return at %a;
2836 "
2837 );
2838 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2839
2840 ctx.remove_instruction(load_id);
2841
2842 assert!(!ctx.contains_instruction(load_id));
2843 }
2844
2845 #[test]
2846 fn remove_instruction_frees_name() {
2847 let mut ctx = Context::new();
2848 qcode!(
2849 ctx,
2850 "
2851 varnode i64 x;
2852 <block>
2853 %a = load(x:8, &x);
2854 return at %a;
2855 "
2856 );
2857 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2858 assert!(
2860 ctx.get_named_in_scope(load_id.into(), "a").is_some(),
2861 "name should be in map before removal"
2862 );
2863
2864 ctx.remove_instruction(load_id);
2865
2866 assert!(
2867 ctx.get_named_in_scope(load_id.into(), "a").is_none(),
2868 "name should be gone after removal"
2869 );
2870 assert!(!ctx.contains_instruction(load_id));
2871 }
2872
2873 #[test]
2874 fn remove_instruction_frees_name_for_reuse() {
2875 let mut ctx = Context::new();
2876 qcode!(
2877 ctx,
2878 "
2879 varnode i64 x;
2880 <block>
2881 %a = load(x:8, &x);
2882 return at %a;
2883 "
2884 );
2885 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2886
2887 ctx.remove_instruction(load_id);
2888
2889 qcode!(
2891 ctx,
2892 "
2893 varnode i64 y;
2894 <block2>
2895 %a = load(y:8, &y);
2896 return at %a;
2897 "
2898 );
2899 let a2 = BasicBlock::from_id(&ctx, block2).instruction_ids()[0];
2900 assert!(
2901 ctx.get_named_in_scope(a2.into(), "a").is_some(),
2902 "name should be reusable after removal"
2903 );
2904 }
2905
2906 #[test]
2907 fn remove_instruction_updates_users_map() {
2908 let mut ctx = Context::new();
2909 qcode!(
2910 ctx,
2911 "
2912 varnode i64 x;
2913 <block>
2914 %a = load(x:8, &x);
2915 %b = %a + i64 1;
2916 return at %b;
2917 "
2918 );
2919 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
2920 let load_id = ids[0];
2921 let add_id = ids[1];
2922
2923 assert!(
2924 ctx.users(load_id).contains(&add_id),
2925 "add should be a user of load before removal"
2926 );
2927
2928 ctx.remove_instruction(add_id);
2929
2930 assert!(
2931 ctx.users(load_id).is_empty(),
2932 "load should have no users after add is removed"
2933 );
2934 }
2935
2936 #[test]
2937 fn removed_instruction_is_absent_and_not_iterated() {
2938 let mut ctx = Context::new();
2943 qcode!(
2944 ctx,
2945 "
2946 varnode i64 x;
2947 <block>
2948 %a = load(x:8, &x);
2949 %dead = %a + i64 1;
2950 return at i64 0;
2951 "
2952 );
2953 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
2954 let dead_id = ids[1]; assert!(
2957 ctx.instructions().any(|i| i.id == dead_id),
2958 "the instruction is iterated while live"
2959 );
2960
2961 ctx.remove_instruction(dead_id);
2962
2963 assert!(!ctx.contains_instruction(dead_id));
2964 assert!(
2965 !ctx.instructions().any(|i| i.id == dead_id),
2966 "a deleted instruction must not be yielded by ctx.instructions()"
2967 );
2968 }
2969
2970 #[test]
2971 fn replace_instruction_mnemonic_rewrites_callind_users() {
2972 let mut ctx = Context::new();
2973 qcode!(
2974 ctx,
2975 "
2976 varnode i64 ptr;
2977 <block>
2978 call [ptr];
2979 "
2980 );
2981 let call_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2982 let ptr = match ctx.get_insn(call_id).mnemonic() {
2983 Mnemonic::CallInd(call) => call.ptr.qualify(call_id.func),
2984 other => panic!("expected CallInd, got {other:?}"),
2985 };
2986 assert_eq!(ctx.users_across_functions(ptr), vec![call_id]);
2988
2989 let target = FunctionBody::make(&mut ctx, "target".into()).unwrap().id;
2990 ctx.replace_instruction_mnemonic(
2991 call_id,
2992 Mnemonic::Call(Call {
2993 target: Callee::Real(target),
2994 args: vec![],
2995 clobbers: vec![],
2996 tag: Default::default(),
2997 }),
2998 );
2999
3000 assert!(
3001 ctx.users_across_functions(ptr).is_empty(),
3002 "old indirect pointer should no longer list the rewritten call"
3003 );
3004 assert!(matches!(
3005 ctx.get_insn(call_id).mnemonic(),
3006 Mnemonic::Call(Call {
3007 target: actual,
3008 args,
3009 ..
3010 }) if *actual == Callee::Real(target) && args.is_empty()
3011 ));
3012 }
3013
3014 #[test]
3015 fn users_across_functions_keeps_ssa_users_in_the_owning_function() {
3016 let mut ctx = Context::new();
3017 qcode!(
3018 ctx,
3019 "
3020 fn f:
3021 <f_entry>
3022 %fx = i64 1 + i64 2;
3023 %fuse = %fx + i64 3;
3024 return at %fuse;
3025 fn g:
3026 <g_entry>
3027 %gx = i64 4 + i64 5;
3028 %guse = %gx + i64 6;
3029 return at %guse;
3030 "
3031 );
3032 let f_ids = BasicBlock::from_id(&ctx, f_entry).instruction_ids();
3033 let g_ids = BasicBlock::from_id(&ctx, g_entry).instruction_ids();
3034 assert_eq!(
3035 f_ids[0].local, g_ids[0].local,
3036 "precondition: arena-local ids collide"
3037 );
3038 assert_eq!(
3039 ctx.users_across_functions(ValueId::Instruction(f_ids[0])),
3040 vec![f_ids[1]],
3041 "an SSA query must not pick up the same local key from another function"
3042 );
3043 }
3044
3045 #[test]
3046 fn replace_instruction_mnemonic_moves_operand_users() {
3047 let mut ctx = Context::new();
3048 qcode!(
3049 ctx,
3050 "
3051 varnode i64 x;
3052 varnode i64 y;
3053 <block>
3054 %a = load(x:8, x);
3055 return at %a;
3056 "
3057 );
3058 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3059 let old_ptr = ValueId::Varnode(x);
3060 let new_ptr = ValueId::Varnode(y);
3061 assert_eq!(ctx.users_across_functions(old_ptr), vec![load_id]);
3063 assert!(ctx.users_across_functions(new_ptr).is_empty());
3064
3065 ctx.replace_instruction_mnemonic(
3066 load_id,
3067 Mnemonic::Load(Load {
3068 space: ctx.shared.default_space.into(),
3069 ptr: new_ptr.localize(load_id.func),
3070 size: 8,
3071 }),
3072 );
3073
3074 assert!(ctx.users_across_functions(old_ptr).is_empty());
3075 assert_eq!(ctx.users_across_functions(new_ptr), vec![load_id]);
3076 }
3077
3078 #[test]
3079 fn replace_instruction_mnemonic_tracks_repeated_operands() {
3080 let mut ctx = Context::new();
3081 qcode!(
3082 ctx,
3083 "
3084 varnode i64 x;
3085 varnode i64 y;
3086 <block>
3087 %a = load(x:8, x);
3088 return at %a;
3089 "
3090 );
3091 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3092 let old_ptr = ValueId::Varnode(x);
3093 let new_arg = ValueId::Varnode(y);
3094
3095 ctx.replace_instruction_mnemonic(
3096 load_id,
3097 Mnemonic::Binop(Binary {
3098 op: Binop::Int(IntBinop::Add),
3099 lhs: new_arg.localize(load_id.func),
3100 rhs: new_arg.localize(load_id.func),
3101 }),
3102 );
3103
3104 assert!(ctx.users_across_functions(old_ptr).is_empty());
3105 assert_eq!(
3106 ctx.users_across_functions(new_arg),
3107 vec![load_id, load_id],
3108 "a mnemonic using the same operand twice should record both uses"
3109 );
3110 }
3111
3112 #[test]
3113 fn remove_instruction_unparented_noop() {
3114 let mut ctx = Context::new();
3115 qcode!(
3116 ctx,
3117 "
3118 varnode i64 x;
3119 <block>
3120 %a = load(x:8, &x);
3121 return at %a;
3122 "
3123 );
3124 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3125
3126 ctx.instruction_mut(load_id).parent = None;
3129
3130 ctx.remove_instruction(load_id);
3132
3133 assert!(ctx.get_named("a").is_none());
3134 }
3135
3136 #[test]
3137 fn add_cfg_edge_returns_id_and_remove_unlinks_both_blocks() {
3138 let mut ctx = Context::new();
3139 let f = ctx.anon_function();
3141 let a = BasicBlock::make(&mut ctx, f).id;
3142 let b = BasicBlock::make(&mut ctx, f).id;
3143 let c = BasicBlock::make(&mut ctx, f).id;
3144
3145 let edge = ctx.add_cfg_edge(a, b);
3146 let surviving_edge = ctx.add_cfg_edge(b, c);
3147 assert_eq!(
3148 BasicBlock::from_id(&ctx, a)
3149 .successors()
3150 .collect::<Vec<_>>(),
3151 vec![(edge, b)]
3152 );
3153 assert_eq!(
3154 BasicBlock::from_id(&ctx, b)
3155 .predecessors()
3156 .collect::<Vec<_>>(),
3157 vec![(edge, a)]
3158 );
3159
3160 ctx.remove_cfg_edge(a.func, edge);
3161 assert!(BasicBlock::from_id(&ctx, a).successors().next().is_none());
3162 assert!(BasicBlock::from_id(&ctx, b).predecessors().next().is_none());
3163 assert!(!ctx.bodies[a.func].edges.contains(edge));
3164 let surviving = ctx.edge(a.func, surviving_edge);
3165 assert_eq!(
3166 surviving.from, b.local,
3167 "swap removal must preserve the source"
3168 );
3169 assert_eq!(
3170 surviving.to, c.local,
3171 "swap removal must preserve the target"
3172 );
3173 assert_eq!(ctx.bodies[a.func].edges.len(), 1);
3174
3175 let self_edge = ctx.add_cfg_edge(a, a);
3176 ctx.remove_cfg_edge(a.func, self_edge);
3177 assert!(!ctx.bodies[a.func].edges.contains(self_edge));
3178 assert!(ctx.block(a).edges.is_empty());
3179
3180 let parallel_a = ctx.add_cfg_edge(a, b);
3181 let parallel_b = ctx.add_cfg_edge(a, b);
3182 ctx.remove_cfg_edge(a.func, parallel_a);
3183 assert!(!ctx.bodies[a.func].edges.contains(parallel_a));
3184 assert!(ctx.bodies[a.func].edges.contains(parallel_b));
3185 assert_eq!(
3186 BasicBlock::from_id(&ctx, a)
3187 .successors()
3188 .collect::<Vec<_>>(),
3189 vec![(parallel_b, b)],
3190 );
3191 }
3192
3193 #[test]
3194 fn truth_map_tracks_four_states_and_conflicts() {
3195 let mut ctx = Context::new();
3196 let callee = FunctionBody::make(&mut ctx, "callee".into()).unwrap().id;
3197 let prop = Proposition::FunctionReturns(callee);
3198
3199 assert!(ctx.assume_true(prop));
3201 assert!(ctx.assume_true(prop));
3202 assert!(!ctx.assume_false(prop));
3203 assert_eq!(ctx.known(prop), None, "assumed is not known");
3204
3205 let snapshot = ctx.clone();
3207
3208 let scope = pass_scope::enter("verifier");
3211 assert!(ctx.set_known(prop, false), "overturning is novel");
3212 drop(scope);
3213 assert_eq!(ctx.known(prop), Some(false));
3214 let [v] = ctx.violations() else {
3215 panic!("expected one violation")
3216 };
3217 assert_eq!(v.prop, prop);
3218 assert!(v.assumed);
3219 assert_eq!(v.asserting_pass, "verifier");
3220
3221 assert!(!ctx.set_known(prop, false));
3223
3224 assert!(snapshot.violations().is_empty());
3226 assert_eq!(snapshot.known(prop), None);
3227
3228 assert!(!ctx.assume_true(prop));
3230 assert!(ctx.assume_false(prop));
3231 }
3232
3233 #[test]
3234 fn seeded_facts_are_not_novel() {
3235 let mut ctx = Context::new();
3236 let callee = FunctionBody::make(&mut ctx, "exit".into()).unwrap().id;
3237 let prop = Proposition::FunctionReturns(callee);
3238
3239 ctx.seed_known(prop, false, PassName("seed"));
3240 assert_eq!(ctx.known(prop), Some(false));
3241 assert!(!ctx.assume_true(prop), "seeded fact blocks opposite assume");
3242 assert!(
3243 !ctx.set_known(prop, false),
3244 "re-proving a seed is not novel"
3245 );
3246 assert!(ctx.violations().is_empty());
3247 }
3248
3249 #[test]
3250 fn discovered_code_records_and_survives_round_trip() {
3251 let mut ctx = Context::new();
3252 ctx.discover_code(0x1000, 0x10f0, 0x1100);
3253 ctx.discover_code(0x1000, 0x10f0, 0x1200);
3254 ctx.discover_code(0x1000, 0x10f0, 0x1100); let targets: Vec<u64> = ctx.discoveries().map(|d| d.target).collect();
3257 assert_eq!(targets, vec![0x1100, 0x1200]);
3258
3259 let config = bincode::config::standard();
3260 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3261 let (restored, _): (Context<'static>, usize) =
3262 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3263 assert_eq!(
3264 restored.discoveries().map(|d| d.target).collect::<Vec<_>>(),
3265 targets
3266 );
3267 }
3268
3269 #[test]
3270 fn assume_executable_narrows_once_protections_known() {
3271 let mut ctx = Context::new();
3272 let mut image = crate::memory_image::MemoryImage::default();
3273 image.add_segment(0x1000, vec![0u8; 4], true, false); image.add_segment(0x2000, vec![0u8; 4], false, true); let binary: &dyn wazabin_binary::BinaryFormat = ℑ
3276
3277 assert!(ctx.assume_executable(binary, 0x1000));
3280 assert!(ctx.assume_executable(binary, 0x2000));
3281 assert!(ctx.assume_executable(binary, 0x9999));
3282
3283 ctx.mark_protections_known();
3284 assert!(
3285 ctx.assume_executable(binary, 0x1000),
3286 "code region stays liftable"
3287 );
3288 assert!(
3289 !ctx.assume_executable(binary, 0x2000),
3290 "data region is skipped once protections are known"
3291 );
3292 assert!(
3293 !ctx.assume_executable(binary, 0x9999),
3294 "unmapped is skipped once known"
3295 );
3296 assert_eq!(
3298 ctx.known(Proposition::ExecutableMemory {
3299 start: 0x2000,
3300 end: 0x2004,
3301 }),
3302 Some(false),
3303 );
3304 }
3305
3306 #[test]
3307 fn assume_executable_honors_region_override() {
3308 let mut ctx = Context::new();
3309 let mut image = crate::memory_image::MemoryImage::default();
3310 image.add_segment(0x1000, vec![0u8; 4], true, false); image.add_segment(0x2000, vec![0u8; 4], false, true); let binary: &dyn wazabin_binary::BinaryFormat = ℑ
3313 ctx.mark_protections_known();
3314
3315 ctx.seed_known(
3317 Proposition::ExecutableMemory {
3318 start: 0x2000,
3319 end: 0x2004,
3320 },
3321 true,
3322 PassName("override"),
3323 );
3324 ctx.seed_known(
3325 Proposition::ExecutableMemory {
3326 start: 0x1000,
3327 end: 0x1004,
3328 },
3329 false,
3330 PassName("override"),
3331 );
3332
3333 assert!(
3334 ctx.assume_executable(binary, 0x2000),
3335 "override wins over the non-executable segment flag"
3336 );
3337 assert!(
3338 !ctx.assume_executable(binary, 0x1000),
3339 "override wins over the executable segment flag"
3340 );
3341 }
3342
3343 #[test]
3344 fn context_survives_bincode_round_trip() {
3345 let mut ctx = Context::new();
3346 qcode!(
3347 ctx,
3348 "
3349 varnode i64 ptr;
3350 <block>
3351 %a = load(ptr:8, &ptr);
3352 %b = %a + i64 0x10;
3353 store(ptr:8, &ptr <- i64 0x1234);
3354 return at %b;
3355 "
3356 );
3357
3358 let some_space = ctx.get_or_make_named_space("scratch");
3360 let sa = ctx.shared.types.get_or_make_space_address(8, some_space);
3361 let sa_size = ctx.shared.types.size_of(sa);
3362
3363 let blocks_before = ctx.block_ids().len();
3364 let insns_before = ctx.instruction_ids().len();
3365 let funcs_before = ctx.function_ids().len();
3366
3367 let config = bincode::config::standard();
3368 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3369 let (restored, _): (Context<'static>, usize) =
3370 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3371
3372 assert_eq!(restored.block_ids().len(), blocks_before);
3373 assert_eq!(restored.instruction_ids().len(), insns_before);
3374 assert_eq!(restored.function_ids().len(), funcs_before);
3375 for function_id in restored.function_ids() {
3376 assert_eq!(restored.bodies[function_id].id(), function_id);
3377 }
3378 assert_eq!(restored.shared.types.size_of(sa), sa_size);
3380 assert_eq!(
3381 restored.shared.types.space_of(sa),
3382 Some(crate::space::MemorySpaceId::Shared(some_space))
3383 );
3384 }
3385
3386 #[test]
3387 fn compact_edge_arena_preserves_ids_across_round_trip() {
3388 let mut ctx = Context::new();
3389 let function = ctx.anon_function();
3390 let a = BasicBlock::make(&mut ctx, function).id;
3391 let b = BasicBlock::make(&mut ctx, function).id;
3392 let c = BasicBlock::make(&mut ctx, function).id;
3393 let d = BasicBlock::make(&mut ctx, function).id;
3394 let first = ctx.add_cfg_edge(a, b);
3395 let removed = ctx.add_cfg_edge(b, c);
3396 let last = ctx.add_cfg_edge(c, d);
3397 ctx.remove_cfg_edge(function, removed);
3398
3399 let physical_order: Vec<_> = ctx.bodies[function]
3400 .edges
3401 .iter()
3402 .map(|edge| edge.id)
3403 .collect();
3404 assert_eq!(physical_order, vec![first, last]);
3405
3406 let config = bincode::config::standard();
3407 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3408 let (mut restored, _): (Context<'static>, usize) =
3409 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3410
3411 assert!(!restored.bodies[function].edges.contains(removed));
3412 assert_eq!(
3413 restored.bodies[function]
3414 .edges
3415 .iter()
3416 .map(|edge| edge.id)
3417 .collect::<Vec<_>>(),
3418 physical_order,
3419 );
3420 assert_eq!(restored.edge(function, first).to, b.local);
3421 assert_eq!(restored.edge(function, last).from, c.local);
3422
3423 let fresh = restored.add_cfg_edge(a, d);
3424 assert!(fresh > last);
3425 assert_ne!(fresh, removed, "removed edge IDs must never be reused");
3426 }
3427
3428 #[test]
3429 fn compact_instruction_arena_preserves_ids_across_round_trip() {
3430 let mut ctx = Context::new();
3431 qcode!(
3432 ctx,
3433 "
3434 <block>
3435 %first = i64 1 + i64 2;
3436 %removed = i64 3 + i64 4;
3437 return at %first;
3438 "
3439 );
3440 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
3441 let first = ids[0];
3442 let removed = ids[1];
3443 let last = ids[2];
3444 ctx.remove_instruction(removed);
3445
3446 let physical_order: Vec<_> = ctx.bodies[first.func]
3447 .insns
3448 .iter()
3449 .map(|insn| insn.id)
3450 .collect();
3451 assert_eq!(physical_order, vec![first.local, last.local]);
3452
3453 let config = bincode::config::standard();
3454 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3455 let (mut restored, _): (Context<'static>, usize) =
3456 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3457
3458 assert!(!restored.contains_instruction(removed));
3459 assert_eq!(
3460 restored.bodies[first.func]
3461 .insns
3462 .iter()
3463 .map(|insn| insn.id)
3464 .collect::<Vec<_>>(),
3465 physical_order,
3466 );
3467 assert!(restored.contains_instruction(first));
3468 assert!(restored.contains_instruction(last));
3469
3470 let template = restored.instruction(last).clone();
3471 let fresh = restored.push_insn(first.func, template);
3472 assert!(fresh.local > last.local);
3473 assert_ne!(
3474 fresh, removed,
3475 "removed instruction IDs must never be reused"
3476 );
3477 }
3478
3479 #[test]
3480 fn compact_param_arena_preserves_ids_across_round_trip() {
3481 let mut ctx = Context::new();
3482 let function = ctx.anon_function();
3483 let block = BasicBlock::make(&mut ctx, function).id;
3484 let first = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3485 let removed = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3486 let last = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3487
3488 ctx.block_mut(block).params.remove(1);
3489 ctx.block_param_mut(last).index = 1;
3490 ctx.remove_block_param(removed);
3491
3492 let physical_order: Vec<_> = ctx.bodies[function]
3493 .params
3494 .iter()
3495 .map(|param| param.id)
3496 .collect();
3497 assert_eq!(physical_order, vec![first.local, last.local]);
3498 assert_eq!(ctx.block_param(first).index, 0);
3499 assert_eq!(ctx.block_param(last).index, 1);
3500
3501 let config = bincode::config::standard();
3502 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3503 let (mut restored, _): (Context<'static>, usize) =
3504 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3505
3506 assert!(!restored.contains_block_param(removed));
3507 assert_eq!(
3508 restored.bodies[function]
3509 .params
3510 .iter()
3511 .map(|param| param.id)
3512 .collect::<Vec<_>>(),
3513 physical_order,
3514 );
3515 assert!(restored.contains_block_param(first));
3516 assert!(restored.contains_block_param(last));
3517
3518 let fresh = BasicBlock::from_id_mut(&mut restored, block)
3519 .push_param(8)
3520 .id;
3521 assert!(fresh.local > last.local);
3522 assert_ne!(fresh, removed, "removed parameter IDs must never be reused");
3523 }
3524
3525 #[test]
3526 fn compact_block_arena_preserves_ids_across_round_trip() {
3527 let mut ctx = Context::new();
3528 let function = ctx.anon_function();
3529 let first = BasicBlock::make(&mut ctx, function).id;
3530 let removed = BasicBlock::make(&mut ctx, function).id;
3531 let last = BasicBlock::make(&mut ctx, function).id;
3532 FunctionBody::from_id_mut(&mut ctx, function)
3533 .set_root(first)
3534 .expect("set root");
3535
3536 ctx.delete_block(removed);
3537
3538 let physical_order: Vec<_> = ctx.bodies[function]
3539 .blocks
3540 .iter()
3541 .map(|block| block.id)
3542 .collect();
3543 assert_eq!(physical_order, vec![first.local, last.local]);
3544 assert_eq!(ctx.block_ids(), vec![first, last]);
3545
3546 let config = bincode::config::standard();
3547 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3548 let (mut restored, _): (Context<'static>, usize) =
3549 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3550
3551 assert!(!restored.contains_block(removed));
3552 assert_eq!(
3553 restored.bodies[function]
3554 .blocks
3555 .iter()
3556 .map(|block| block.id)
3557 .collect::<Vec<_>>(),
3558 physical_order,
3559 );
3560 assert!(restored.contains_block(first));
3561 assert!(restored.contains_block(last));
3562 assert_eq!(
3563 FunctionBody::from_id(&restored, function)
3564 .root()
3565 .map(|block| block.id),
3566 Some(first),
3567 );
3568
3569 let fresh = BasicBlock::make(&mut restored, function).id;
3570 assert!(fresh.local > last.local);
3571 assert_ne!(fresh, removed, "removed block IDs must never be reused");
3572 }
3573
3574 #[test]
3575 fn deleting_root_clears_function_root() {
3576 let mut ctx = Context::new();
3577 let function = ctx.anon_function();
3578 let root = BasicBlock::make(&mut ctx, function).id;
3579 FunctionBody::from_id_mut(&mut ctx, function)
3580 .set_root(root)
3581 .expect("set root");
3582
3583 ctx.delete_block(root);
3584
3585 assert!(!ctx.contains_block(root));
3586 assert!(FunctionBody::from_id(&ctx, function).root().is_none());
3587 assert!(ctx.block_ids().is_empty());
3588 }
3589
3590 #[test]
3591 fn get_unique_name_resumes_probe_and_reuses_freed_suffixes() {
3592 use crate::value::VarnodeId;
3593
3594 let mut ctx = Context::new();
3595 let id = ValueId::Varnode(VarnodeId::from(0usize));
3596
3597 fn take(ctx: &mut Context<'static>, id: ValueId, base: &str) -> String {
3599 let name = ctx
3600 .get_unique_name(Cow::Owned(base.to_string()))
3601 .to_string();
3602 ctx.update_name(Cow::Owned(name.clone()), id, None).unwrap();
3603 name
3604 }
3605
3606 assert_eq!(take(&mut ctx, id, "tmp"), "tmp");
3608 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_1");
3609 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_2");
3610 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_3");
3611
3612 assert_eq!(take(&mut ctx, id, "x"), "x");
3614 assert_eq!(take(&mut ctx, id, "x"), "x_1");
3615
3616 ctx.update_name(Cow::Borrowed("relocated"), id, Some("tmp_1"))
3619 .unwrap();
3620 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_1");
3621 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_4");
3623 }
3624
3625 mod split_function_at {
3628 use super::*;
3629
3630 use crate::value::insn::{Callee, Mnemonic, TailCall};
3631 use crate::value::{BasicBlock, FunctionBody, Instruction, Value};
3632 use std::borrow::Cow;
3633
3634 fn block_at(ctx: &mut Context<'static>, func: FunctionId, addr: u64) -> BlockId {
3635 BasicBlock::make(ctx, func).with_address(addr).id
3636 }
3637
3638 fn branch_at(ctx: &mut Context<'static>, block: BlockId, target: BlockId, addr: u64) {
3639 let id = (ctx).builder(block).push_branch(target).id;
3640 Instruction::from_id_mut(ctx, id).set_address(addr);
3641 }
3642
3643 fn cbranch_at(
3644 ctx: &mut Context<'static>,
3645 block: BlockId,
3646 success: BlockId,
3647 failure: BlockId,
3648 addr: u64,
3649 ) {
3650 let cond = ctx.get_const(1, 1).id();
3651 let id = (ctx).builder(block).push_cbranch(cond, success, failure).id;
3652 Instruction::from_id_mut(ctx, id).set_address(addr);
3653 }
3654
3655 fn return_at(ctx: &mut Context<'static>, block: BlockId, addr: u64) {
3656 let zero = ctx.get_const(0, 8).id();
3657 let id = (ctx).builder(block).push_return(zero).id;
3658 Instruction::from_id_mut(ctx, id).set_address(addr);
3659 }
3660
3661 fn block_at_addr(ctx: &Context, func: FunctionId, addr: u64) -> BlockId {
3662 FunctionBody::from_id(ctx, func)
3663 .block_ids()
3664 .into_iter()
3665 .find(|b| ctx.block(*b).address == Some(addr))
3666 .unwrap_or_else(|| panic!("{func:?} has no block at {addr:#x}"))
3667 }
3668
3669 fn addrs(ctx: &Context, func: FunctionId) -> Vec<u64> {
3670 let mut got: Vec<u64> = FunctionBody::from_id(ctx, func)
3671 .block_ids()
3672 .into_iter()
3673 .filter_map(|b| ctx.block(b).address)
3674 .collect();
3675 got.sort_unstable();
3676 got
3677 }
3678
3679 #[test]
3684 fn splits_absorbed_body_reusing_the_stub() {
3685 let mut ctx = Context::new();
3686 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("thunk"))).id;
3687 let b0 = block_at(&mut ctx, f, 0x1000);
3688 let b1 = block_at(&mut ctx, f, 0x2000);
3689 let b2 = block_at(&mut ctx, f, 0x2005);
3690 branch_at(&mut ctx, b0, b1, 0x1000);
3691 branch_at(&mut ctx, b1, b2, 0x2000);
3692 return_at(&mut ctx, b2, 0x2005);
3693 {
3694 let mut func = FunctionBody::from_id_mut(&mut ctx, f);
3695 func.set_root(b0).unwrap();
3696 }
3697 let g = FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("real"))).id;
3699
3700 let split_g = ctx.split_function_at(b1);
3701 assert_eq!(
3702 split_g, g,
3703 "the split must reuse the existing stub at 0x2000"
3704 );
3705
3706 assert_eq!(addrs(&ctx, f), vec![0x1000]);
3707 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2005]);
3708 let g_entry = block_at_addr(&ctx, g, 0x2000);
3709 assert_eq!(ctx.bodies[g].root_id(), Some(g_entry.local));
3710
3711 for b in FunctionBody::from_id(&ctx, g).block_ids() {
3713 assert_eq!(b.func, g);
3714 }
3715
3716 let f_entry = block_at_addr(&ctx, f, 0x1000);
3718 assert_eq!(BasicBlock::from_id(&ctx, f_entry).successors().count(), 0);
3719 let term = BasicBlock::from_id(&ctx, f_entry)
3720 .instructions()
3721 .last()
3722 .map(|i| i.mnemonic().clone());
3723 assert!(
3724 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
3725 "thunk branch must become TailCall(G), got {term:?}",
3726 );
3727 }
3728
3729 #[test]
3730 fn split_rehomes_temporary_values_spaces_and_pointer_types() {
3731 let mut ctx = Context::new();
3732 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3733 let entry = block_at(&mut ctx, f, 0x1000);
3734 let tail = block_at(&mut ctx, f, 0x2000);
3735 branch_at(&mut ctx, entry, tail, 0x1000);
3736 FunctionBody::from_id_mut(&mut ctx, f)
3737 .set_root(entry)
3738 .unwrap();
3739
3740 let temp = ctx
3741 .builder(tail)
3742 .make_named_temp(Cow::Borrowed("scratch"), 8);
3743 ctx.builder(entry)
3744 .make_named_temp(Cow::Borrowed("unused"), 4);
3745 let temp_space = ctx.bodies[f].temps[temp.local].space;
3746 let load = {
3747 let mut builder = ctx.builder(tail);
3748 let ValueId::Instruction(load) = builder
3749 .push_load::<false>(
3750 ValueId::Temp(temp),
3751 8,
3752 LocalMemorySpaceId::Temp(temp_space),
3753 )
3754 .id()
3755 else {
3756 unreachable!()
3757 };
3758 builder.push_return(ValueId::Instruction(load));
3759 load
3760 };
3761 let pointer_type = ctx
3762 .shared
3763 .types
3764 .get_or_make_space_address(8, MemorySpaceId::Temp(TempSpaceId::new(f, temp_space)));
3765 ctx.instruction_mut(load).type_id = pointer_type;
3766
3767 let g =
3768 FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("discovered"))).id;
3769 assert_eq!(ctx.split_function_at(tail), g);
3770
3771 let diagnostics = crate::verify_body_arena_integrity(&ctx);
3772 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
3773 assert_eq!(ctx.bodies[g].temp_spaces.len(), 1);
3774 assert_eq!(ctx.bodies[g].temps.len(), 1);
3775 assert_eq!(ctx.bodies[f].temps.len(), 2, "source arenas remain intact");
3776
3777 let moved_load = FunctionBody::from_id(&ctx, g)
3778 .blocks()
3779 .flat_map(|block| block.instructions())
3780 .find(|insn| matches!(insn.mnemonic(), Mnemonic::Load(_)))
3781 .expect("load moved with the split");
3782 let Mnemonic::Load(moved) = moved_load.mnemonic() else {
3783 unreachable!()
3784 };
3785 let LocalMemorySpaceId::Temp(moved_space) = moved.space else {
3786 panic!("load lost temporary-space provenance")
3787 };
3788 assert!(matches!(moved.ptr, crate::value::LocalValueId::Temp(_)));
3789 assert!(usize::from(moved_space) < ctx.bodies[g].temp_spaces.len());
3790 assert_eq!(
3791 ctx.shared.types.space_of(moved_load.type_id()),
3792 Some(MemorySpaceId::Temp(TempSpaceId::new(g, moved_space)))
3793 );
3794
3795 let rendered = FunctionBody::from_id(&ctx, g).to_string();
3797 assert!(rendered.contains("scratch"));
3798 }
3799
3800 #[test]
3801 fn split_stops_at_a_foreign_rootless_stub_address() {
3802 let mut ctx = Context::new();
3803 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3804 let entry = block_at(&mut ctx, f, 0x1000);
3805 let split = block_at(&mut ctx, f, 0x2000);
3806 let foreign_entry = block_at(&mut ctx, f, 0x3000);
3807 let foreign_body = block_at(&mut ctx, f, 0x3005);
3808 branch_at(&mut ctx, entry, split, 0x1000);
3809 branch_at(&mut ctx, split, foreign_entry, 0x2000);
3810 branch_at(&mut ctx, foreign_entry, foreign_body, 0x3000);
3811 return_at(&mut ctx, foreign_body, 0x3005);
3812 FunctionBody::from_id_mut(&mut ctx, f)
3813 .set_root(entry)
3814 .unwrap();
3815
3816 let g = FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("g"))).id;
3817 let h = FunctionBody::make_at_addr(&mut ctx, 0x3000, Some(Cow::Borrowed("h"))).id;
3818 assert!(FunctionBody::from_id(&ctx, g).root().is_none());
3819 assert!(FunctionBody::from_id(&ctx, h).root().is_none());
3820
3821 assert_eq!(ctx.split_function_at(split), g);
3822 assert_eq!(addrs(&ctx, g), vec![0x2000]);
3823 assert_eq!(addrs(&ctx, f), vec![0x1000, 0x3000, 0x3005]);
3824 assert!(FunctionBody::from_id(&ctx, h).root().is_none());
3825
3826 let g_entry = block_at_addr(&ctx, g, 0x2000);
3827 let term = BasicBlock::from_id(&ctx, g_entry)
3828 .instructions()
3829 .last()
3830 .map(|i| i.mnemonic().clone());
3831 assert!(
3832 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(h)),
3833 "split tail must stop and tail-call rootless stub H, got {term:?}",
3834 );
3835 }
3836
3837 #[test]
3838 fn split_rehomes_block_param_origin_into_destination_arena() {
3839 let mut ctx = Context::new();
3840 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3841 let entry = block_at(&mut ctx, f, 0x1000);
3842 let tail = block_at(&mut ctx, f, 0x2000);
3843 let param = BasicBlock::from_id_mut(&mut ctx, tail).push_param(8).id;
3844 crate::value::BlockParam::from_id_mut(&mut ctx, param)
3845 .set_origin(ValueId::BlockParam(param));
3846
3847 let arg = ctx.get_const(7, 8).id();
3848 let branch = ctx.builder(entry).push_branch_with_args(tail, vec![arg]).id;
3849 Instruction::from_id_mut(&mut ctx, branch).set_address(0x1000);
3850 let ret = ctx.builder(tail).push_return(ValueId::BlockParam(param)).id;
3851 Instruction::from_id_mut(&mut ctx, ret).set_address(0x2000);
3852 FunctionBody::from_id_mut(&mut ctx, f)
3853 .set_root(entry)
3854 .unwrap();
3855
3856 let g = ctx.split_function_at(tail);
3857 let new_tail = block_at_addr(&ctx, g, 0x2000);
3858 let new_param = BasicBlock::from_id(&ctx, new_tail).params().next().unwrap();
3859 assert_eq!(new_param.origin(), Some(ValueId::BlockParam(new_param.id)));
3860 }
3861
3862 #[test]
3870 fn split_rehomes_symbolic_block_literals() {
3871 use crate::value::literal::SymbolicRef;
3872
3873 let mut ctx = Context::new();
3874 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3875 let entry = block_at(&mut ctx, f, 0x1000);
3876 let tail = block_at(&mut ctx, f, 0x2000);
3877 let landing = block_at(&mut ctx, f, 0x2008);
3878
3879 let lit = ctx.get_const(0x2008, 8).id();
3882 let ValueId::Literal(lit_id) = lit else {
3883 panic!("expected a literal");
3884 };
3885 ctx.shared.values.literals[lit_id].symbolic = Some(SymbolicRef::Block(landing));
3886
3887 branch_at(&mut ctx, entry, tail, 0x1000);
3888 let ind = ctx.builder(tail).push_branchind(lit).id;
3890 Instruction::from_id_mut(&mut ctx, ind).set_address(0x2000);
3891 ctx.add_cfg_edge(tail, landing);
3892 return_at(&mut ctx, landing, 0x2008);
3893 FunctionBody::from_id_mut(&mut ctx, f)
3894 .set_root(entry)
3895 .unwrap();
3896
3897 let g = ctx.split_function_at(tail);
3898
3899 let new_landing = block_at_addr(&ctx, g, 0x2008);
3900 let new_tail = block_at_addr(&ctx, g, 0x2000);
3901 let Mnemonic::BranchInd(b) = BasicBlock::from_id(&ctx, new_tail)
3902 .instructions()
3903 .last()
3904 .unwrap()
3905 .mnemonic()
3906 .clone()
3907 else {
3908 panic!("tail must still end in an indirect branch");
3909 };
3910 let crate::value::LocalValueId::Literal(new_lit) = b.ptr else {
3911 panic!("indirect branch operand must still be a literal");
3912 };
3913 assert_eq!(
3914 ctx.shared.values.literals[new_lit].symbolic,
3915 Some(SymbolicRef::Block(new_landing)),
3916 "the relocated literal must name the clone, not the deleted original",
3917 );
3918 assert_eq!(
3919 ctx.shared.values.literals[new_lit].value, 0x2008,
3920 "re-pointing the symbol must not disturb the numeric value",
3921 );
3922 }
3923
3924 #[test]
3928 fn conditional_arm_into_split_block_uses_a_trampoline() {
3929 let mut ctx = Context::new();
3930 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3931 let entry = block_at(&mut ctx, f, 0x1000);
3932 let cont = block_at(&mut ctx, f, 0x1008);
3933 let tail = block_at(&mut ctx, f, 0x2000);
3934 cbranch_at(&mut ctx, entry, tail, cont, 0x1000);
3935 return_at(&mut ctx, cont, 0x1008);
3936 return_at(&mut ctx, tail, 0x2000);
3937 FunctionBody::from_id_mut(&mut ctx, f)
3938 .set_root(entry)
3939 .unwrap();
3940
3941 let g = ctx.split_function_at(tail);
3942
3943 let entry = block_at_addr(&ctx, f, 0x1000);
3944 let cont = block_at_addr(&ctx, f, 0x1008);
3945 assert_eq!(addrs(&ctx, g), vec![0x2000]);
3946
3947 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, entry)
3948 .instructions()
3949 .last()
3950 .unwrap()
3951 .mnemonic()
3952 .clone()
3953 else {
3954 panic!("entry must still end in a cbranch");
3955 };
3956 assert_eq!(cb.failure_block, cont.local, "fall-through arm untouched");
3957 let tramp = BlockId::new(entry.func, cb.success_block);
3958 assert_eq!(
3959 BasicBlock::from_id(&ctx, tramp).parent().map(|f| f.id),
3960 Some(f),
3961 "trampoline lives in F",
3962 );
3963 let term = BasicBlock::from_id(&ctx, tramp)
3964 .instructions()
3965 .last()
3966 .map(|i| i.mnemonic().clone());
3967 assert!(
3968 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
3969 "trampoline must tail-call G, got {term:?}",
3970 );
3971 for (_, s) in BasicBlock::from_id(&ctx, entry).successors() {
3973 assert_eq!(BasicBlock::from_id(&ctx, s).parent().map(|f| f.id), Some(f));
3974 }
3975 }
3976
3977 #[test]
3980 fn mints_a_conventional_function_when_no_stub_exists() {
3981 let mut ctx = Context::new();
3982 wazabin_qcode_macro::qcode!(
3983 ctx,
3984 "
3985 fn f:
3986 <entry>
3987 goto <0x1008>;
3988 <0x1008>
3989 return 0x0;
3990 "
3991 );
3992
3993 let mid = block_at_addr(&ctx, f, 0x1008);
3994 let g = ctx.split_function_at(mid);
3995 assert_eq!(FunctionBody::from_id(&ctx, g).name(), "fn_1008");
3996 assert_eq!(FunctionBody::from_id(&ctx, f).block_ids().len(), 1);
3998 assert_eq!(addrs(&ctx, g), vec![0x1008]);
3999 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
4000 assert_eq!(addresses.function_at(0x1008), Some(g));
4001 for b in FunctionBody::from_id(&ctx, g).block_ids() {
4002 assert_eq!(b.func, g);
4003 }
4004 }
4005
4006 fn assert_no_dangling_terminators(ctx: &Context) {
4010 for b in ctx.block_ids() {
4011 let Some(mnemonic) = BasicBlock::from_id(ctx, b)
4012 .instructions()
4013 .last()
4014 .map(|t| t.mnemonic().clone())
4015 else {
4016 continue;
4017 };
4018 let targets = match &mnemonic {
4019 Mnemonic::Branch(crate::value::insn::Branch { target, .. }) => vec![*target],
4020 Mnemonic::CBranch(crate::value::insn::CBranch {
4021 success_block,
4022 failure_block,
4023 ..
4024 }) => vec![*success_block, *failure_block],
4025 _ => vec![],
4026 };
4027 let succs: std::collections::HashSet<BlockId> = BasicBlock::from_id(ctx, b)
4028 .successors()
4029 .map(|(_, s)| s)
4030 .collect();
4031 for t in targets {
4032 let tid = BlockId::new(b.func, t);
4033 assert!(
4034 ctx.contains_block(tid),
4035 "block {b:?} terminator names dead block {tid:?}"
4036 );
4037 assert!(
4038 succs.contains(&tid),
4039 "block {b:?} terminator target {tid:?} has no CFG edge (operand/edge desync)"
4040 );
4041 }
4042 }
4043 }
4044
4045 #[test]
4050 fn retained_predecessor_into_mid_tail_promotes_the_landing() {
4051 let mut ctx = Context::new();
4052 wazabin_qcode_macro::qcode!(
4055 ctx,
4056 "
4057 fn f:
4058 <entry @c:i8>
4059 if @c goto <0x2000> else goto <0x1008>;
4060 <0x1008>
4061 goto <0x2008>;
4062 <0x2000>
4063 goto <0x2008>;
4064 <0x2008>
4065 return 0x0;
4066 "
4067 );
4068
4069 let tail = block_at_addr(&ctx, f, 0x2000);
4070 let g = ctx.split_function_at(tail);
4071
4072 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
4075 let landing_fn = addresses
4076 .function_at(0x2008)
4077 .expect("mid-tail landing must be promoted to a function");
4078 assert_ne!(landing_fn, g);
4079 assert_eq!(addrs(&ctx, g), vec![0x2000]);
4080 assert_no_dangling_terminators(&ctx);
4081
4082 for (holder, addr) in [(f, 0x1008u64), (g, 0x2000u64)] {
4083 let block = block_at_addr(&ctx, holder, addr);
4084 let term = BasicBlock::from_id(&ctx, block)
4085 .instructions()
4086 .last()
4087 .map(|i| i.mnemonic().clone());
4088 assert!(
4089 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(landing_fn)),
4090 "branch at {addr:#x} into the landing must tail-call it, got {term:?}",
4091 );
4092 }
4093 }
4094
4095 #[test]
4103 fn tail_conditional_to_own_registered_entry_uses_a_trampoline() {
4104 let mut ctx = Context::new();
4105 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4106 let entry = block_at(&mut ctx, f, 0x1000);
4107 let tail = block_at(&mut ctx, f, 0x2000);
4108 let cont = block_at(&mut ctx, f, 0x2008);
4109 branch_at(&mut ctx, entry, tail, 0x1000);
4110 cbranch_at(&mut ctx, tail, entry, cont, 0x2000);
4112 return_at(&mut ctx, cont, 0x2008);
4113 FunctionBody::from_id_mut(&mut ctx, f)
4114 .set_root(entry)
4115 .unwrap();
4116
4117 let g = ctx.split_function_at(tail);
4118
4119 assert_no_dangling_terminators(&ctx);
4120 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4121 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4122
4123 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4126 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4127 .instructions()
4128 .last()
4129 .unwrap()
4130 .mnemonic()
4131 .clone()
4132 else {
4133 panic!("moved tail must still end in a cbranch");
4134 };
4135 let tramp = BlockId::new(g, cb.success_block);
4136 assert_eq!(tramp.func, g, "trampoline must have relocated into g");
4137 let term = BasicBlock::from_id(&ctx, tramp)
4138 .instructions()
4139 .last()
4140 .map(|i| i.mnemonic().clone());
4141 assert!(
4142 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(f)),
4143 "back-edge trampoline must tail-call f, got {term:?}",
4144 );
4145 }
4146
4147 #[test]
4152 fn tail_conditional_to_foreign_entry_relocates_its_trampoline() {
4153 let mut ctx = Context::new();
4154 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4155 let entry = block_at(&mut ctx, f, 0x1000);
4156 let tail = block_at(&mut ctx, f, 0x2000);
4157 let cont = block_at(&mut ctx, f, 0x2008);
4158 let foreign = block_at(&mut ctx, f, 0x3000);
4159 branch_at(&mut ctx, entry, tail, 0x1000);
4160 cbranch_at(&mut ctx, tail, foreign, cont, 0x2000);
4162 return_at(&mut ctx, cont, 0x2008);
4163 return_at(&mut ctx, foreign, 0x3000);
4164 FunctionBody::from_id_mut(&mut ctx, f)
4165 .set_root(entry)
4166 .unwrap();
4167 let h = FunctionBody::make_at_addr(&mut ctx, 0x3000, Some(Cow::Borrowed("h"))).id;
4168
4169 let g = ctx.split_function_at(tail);
4170
4171 assert_no_dangling_terminators(&ctx);
4172 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4173 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4174
4175 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4178 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4179 .instructions()
4180 .last()
4181 .unwrap()
4182 .mnemonic()
4183 .clone()
4184 else {
4185 panic!("moved tail must still end in a cbranch");
4186 };
4187 let tramp = BlockId::new(g, cb.success_block);
4188 assert_eq!(tramp.func, g, "trampoline must have relocated into g");
4189 let term = BasicBlock::from_id(&ctx, tramp)
4190 .instructions()
4191 .last()
4192 .map(|i| i.mnemonic().clone());
4193 assert!(
4194 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(h)),
4195 "relocated trampoline must tail-call H, got {term:?}",
4196 );
4197 }
4198
4199 #[test]
4203 fn conditional_failure_arm_into_split_block_uses_a_trampoline() {
4204 let mut ctx = Context::new();
4205 wazabin_qcode_macro::qcode!(
4208 ctx,
4209 "
4210 fn f:
4211 <entry @c:i8>
4212 if @c goto <0x1008> else goto <0x2000>;
4213 <0x1008>
4214 return 0x0;
4215 <0x2000>
4216 return 0x0;
4217 "
4218 );
4219
4220 let tail = block_at_addr(&ctx, f, 0x2000);
4221 let g = ctx.split_function_at(tail);
4222
4223 assert_no_dangling_terminators(&ctx);
4224 let entry = BlockId::new(f, ctx.bodies[f].root_id().unwrap());
4225 let cont = block_at_addr(&ctx, f, 0x1008);
4226 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, entry)
4227 .instructions()
4228 .last()
4229 .unwrap()
4230 .mnemonic()
4231 .clone()
4232 else {
4233 panic!("entry must still end in a cbranch");
4234 };
4235 assert_eq!(
4236 cb.success_block, cont.local,
4237 "success (fall-through) untouched"
4238 );
4239 let tramp = BlockId::new(entry.func, cb.failure_block);
4240 let term = BasicBlock::from_id(&ctx, tramp)
4241 .instructions()
4242 .last()
4243 .map(|i| i.mnemonic().clone());
4244 assert!(
4245 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
4246 "failure arm must route through a trampoline tail-calling G, got {term:?}",
4247 );
4248 }
4249
4250 #[test]
4253 fn moved_tail_internal_conditional_remaps_both_arms() {
4254 let mut ctx = Context::new();
4255 wazabin_qcode_macro::qcode!(
4258 ctx,
4259 "
4260 fn f:
4261 <entry>
4262 goto <0x2000>;
4263 <0x2000>
4264 %c = 0x0 == 0x0;
4265 if %c goto <0x2008> else goto <0x2010>;
4266 <0x2008>
4267 return 0x0;
4268 <0x2010>
4269 return 0x0;
4270 "
4271 );
4272
4273 let tail = block_at_addr(&ctx, f, 0x2000);
4274 let g = ctx.split_function_at(tail);
4275
4276 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2008, 0x2010]);
4277 assert_no_dangling_terminators(&ctx);
4278 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4279 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4280 .instructions()
4281 .last()
4282 .unwrap()
4283 .mnemonic()
4284 .clone()
4285 else {
4286 panic!("moved tail must still end in a cbranch");
4287 };
4288 let a = block_at_addr(&ctx, g, 0x2008);
4289 let b = block_at_addr(&ctx, g, 0x2010);
4290 assert_eq!(cb.success_block, a.local, "success arm re-pointed to clone");
4291 assert_eq!(cb.failure_block, b.local, "failure arm re-pointed to clone");
4292 }
4293
4294 #[test]
4297 fn moved_tail_internal_branch_remaps_target() {
4298 let mut ctx = Context::new();
4299 wazabin_qcode_macro::qcode!(
4300 ctx,
4301 "
4302 fn f:
4303 <entry>
4304 goto <0x2000>;
4305 <0x2000>
4306 goto <0x2008>;
4307 <0x2008>
4308 return 0x0;
4309 "
4310 );
4311
4312 let tail = block_at_addr(&ctx, f, 0x2000);
4313 let g = ctx.split_function_at(tail);
4314
4315 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2008]);
4316 assert_no_dangling_terminators(&ctx);
4317 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4318 let Mnemonic::Branch(br) = BasicBlock::from_id(&ctx, moved_tail)
4319 .instructions()
4320 .last()
4321 .unwrap()
4322 .mnemonic()
4323 .clone()
4324 else {
4325 panic!("moved tail must still end in a branch");
4326 };
4327 let end = block_at_addr(&ctx, g, 0x2008);
4328 assert_eq!(br.target, end.local, "internal branch re-pointed to clone");
4329 }
4330
4331 #[test]
4334 fn split_rehomes_store_temporary_space() {
4335 let mut ctx = Context::new();
4336 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4337 let entry = block_at(&mut ctx, f, 0x1000);
4338 let tail = block_at(&mut ctx, f, 0x2000);
4339 branch_at(&mut ctx, entry, tail, 0x1000);
4340
4341 let slot = ctx.builder(tail).make_named_temp(Cow::Borrowed("slot"), 8);
4342 let space = ctx.bodies[f].temps[slot.local].space;
4343 let value = ctx.get_const(0x2a, 8).id();
4344 {
4345 let mut builder = ctx.builder(tail);
4346 builder.push_store(value, ValueId::Temp(slot), LocalMemorySpaceId::Temp(space));
4347 builder.push_return(value);
4348 }
4349 FunctionBody::from_id_mut(&mut ctx, f)
4350 .set_root(entry)
4351 .unwrap();
4352
4353 let g = ctx.split_function_at(tail);
4354 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4355 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4356
4357 let moved_store = FunctionBody::from_id(&ctx, g)
4358 .blocks()
4359 .flat_map(|block| block.instructions())
4360 .find(|insn| matches!(insn.mnemonic(), Mnemonic::Store(_)))
4361 .expect("store moved with the split");
4362 let Mnemonic::Store(moved) = moved_store.mnemonic() else {
4363 unreachable!()
4364 };
4365 let LocalMemorySpaceId::Temp(moved_space) = moved.space else {
4366 panic!("store lost temporary-space provenance")
4367 };
4368 assert!(usize::from(moved_space) < ctx.bodies[g].temp_spaces.len());
4369 }
4370 }
4371}