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 get_bool_const(&self, value: bool) -> ValueId {
256 let type_id = self.types.get_or_make_bool();
257 ValueId::Literal(
258 self.values
259 .get_or_make_typed_literal(u64::from(value), type_id, 1),
260 )
261 }
262
263 pub fn get_typed_const(&self, value: u64, type_id: crate::types::TypeId) -> ValueId {
266 let size = self.types.size_of(type_id);
267 ValueId::Literal(self.values.get_or_make_typed_literal(value, type_id, size))
268 }
269
270 pub fn get_bytes(&self, data: Vec<u8>) -> ValueId {
273 let i8_ty = self.types.get_or_make_int(1);
274 let type_id = self.types.get_or_make_array(i8_ty, data.len());
275 ValueId::Bytes(
276 self.values
277 .bytes
278 .push(crate::value::Bytes { data, type_id }),
279 )
280 }
281
282 pub fn bytes_display(&self, id: crate::value::BytesId) -> crate::value::BytesDisplay {
287 self.values
288 .bytes_display
289 .get(&id)
290 .copied()
291 .unwrap_or_default()
292 }
293
294 pub fn get_typed_bytes(&self, data: Vec<u8>, type_id: crate::types::TypeId) -> ValueId {
298 ValueId::Bytes(
299 self.values
300 .bytes
301 .push(crate::value::Bytes { data, type_id }),
302 )
303 }
304
305 pub fn truth(&self, prop: Proposition) -> Option<Truth> {
309 self.values.truths.get(&prop).copied()
310 }
311
312 pub fn assumed_call_convention(&self) -> Option<&crate::assumption::AssumedCallEffect> {
318 self.assumed_call_convention.as_ref()
319 }
320
321 pub fn varnodes(&self) -> impl Iterator<Item = crate::value::VarnodeRef<'str, '_>> + '_ {
324 self.values
325 .varnodes
326 .iter()
327 .map(move |v| crate::value::Varnode::from_id(self, v.id))
328 }
329
330 pub fn varnode_count(&self) -> usize {
333 self.values.varnodes.len()
334 }
335
336 pub fn stored_type_of(&self, id: ValueId) -> Option<crate::types::TypeId> {
343 match id {
344 ValueId::Literal(lid) => Some(self.values.literals[lid].type_id),
345 ValueId::Bytes(bid) => Some(self.values.bytes[bid].type_id),
346 ValueId::Varnode(vid) => self.values.varnode_types.get(&vid).copied(),
347 ValueId::Poison(pid) => Some(self.values.poisons[pid].type_id),
348 ValueId::Instruction(_)
349 | ValueId::BlockParam(_)
350 | ValueId::BasicBlock(_)
351 | ValueId::Temp(_)
352 | ValueId::Function(_) => None,
353 }
354 }
355}
356
357pub use wazabin_binary::TargetOs;
361
362impl<'str> Context<'str> {
363 pub fn new() -> Self {
369 let mut ctx = Self::default();
370 ctx.shared.spaces.push(Space::new(Some("const"), 1, 8));
372 let default_space = Space::new(Some("ram"), 1, 8);
374 ctx.shared.default_space = ctx.shared.spaces.push(default_space);
375 ctx
376 }
377
378 pub fn try_get_space(&self, name: &str) -> Option<SpaceId> {
381 self.shared.named_spaces.get(name).copied()
382 }
383
384 pub fn get_or_make_named_space(&mut self, name: &str) -> SpaceId {
389 if let Some(id) = self.try_get_space(name) {
390 return id;
391 }
392 let default_id = self.shared.default_space;
393 if self.shared.spaces[default_id].name.as_deref() == Some(name) {
394 return default_id;
395 }
396 let default = &self.shared.spaces[self.shared.default_space];
397 let space = Space::new(Some(name), default.word_size, default.addr_size);
398 self.add_space(space)
399 }
400
401 pub fn add_space(&mut self, space: Space) -> SpaceId {
403 let name_key: Option<Box<str>> = space.name.clone();
404 let id = self.shared.spaces.push(space);
405 if let Some(name) = name_key {
406 self.shared.named_spaces.insert(name, id);
407 }
408 id
409 }
410
411 pub fn space_count(&self) -> usize {
413 self.shared.spaces.len()
414 }
415
416 pub fn spaces(&self) -> impl Iterator<Item = Identified<SpaceId, &Space>> + '_ {
419 self.shared.spaces()
420 }
421
422 pub fn set_primary_entrypoint(&mut self, entrypoint: Option<u64>) {
423 self.shared.primary_entrypoint = entrypoint;
424 }
425
426 pub fn primary_entrypoint(&self) -> Option<u64> {
427 self.shared.primary_entrypoint
428 }
429
430 pub fn set_ignored_functions(&mut self, addrs: HashSet<u64>) {
434 self.shared.ignored_functions = addrs;
435 }
436
437 pub fn ignored_functions(&self) -> &HashSet<u64> {
439 &self.shared.ignored_functions
440 }
441
442 pub fn is_function_ignored(&self, addr: Option<u64>) -> bool {
445 addr.is_some_and(|a| self.shared.ignored_functions.contains(&a))
446 }
447
448 pub fn set_target_os(&mut self, os: TargetOs) {
451 self.shared.target_os = os;
452 }
453
454 pub fn target_os(&self) -> TargetOs {
456 self.shared.target_os
457 }
458
459 pub fn set_linked_libraries(&mut self, libs: Vec<String>) {
462 self.shared.linked_libraries = libs;
463 }
464
465 pub fn linked_libraries(&self) -> &[String] {
468 &self.shared.linked_libraries
469 }
470
471 pub fn load_spaces(&mut self, spaces: registry::Registry<SpaceId, Space>) {
473 self.shared.spaces = spaces;
474 }
475
476 pub fn mark_protections_known(&mut self) {
480 self.shared.protections_known = true;
481 }
482
483 pub fn protections_known(&self) -> bool {
485 self.shared.protections_known
486 }
487
488 pub fn assume_executable(
507 &mut self,
508 binary: &dyn wazabin_binary::BinaryFormat,
509 addr: u64,
510 ) -> bool {
511 let bounds = binary.segment_bounds(addr);
512 if let Some((start, end)) = bounds
513 && let Some(known) = self.known(Proposition::ExecutableMemory { start, end })
514 {
515 return known;
516 }
517 if !self.shared.protections_known || binary.is_executable(addr) {
518 return true;
519 }
520 if let Some((start, end)) = bounds {
521 self.set_known(Proposition::ExecutableMemory { start, end }, false);
522 }
523 false
524 }
525
526 pub fn discover(&mut self, discovery: crate::discovery::Discovery) -> bool {
529 self.shared.discoveries.insert(discovery)
530 }
531
532 pub fn discover_code(&mut self, func_entry: u64, source_block: u64, target: u64) {
541 self.shared.discoveries.insert(
542 crate::discovery::Discovery::block(target, func_entry)
543 .with_edge_kind(crate::discovery::EdgeKind::JumpTableTarget)
544 .from_block_addr(source_block)
545 .with_provenance(crate::discovery::DiscoveryProvenance::Optimization {
546 pass: "handle_jump_tables".to_string(),
547 assumption: None,
548 }),
549 );
550 }
551
552 pub fn drain_discoveries(&mut self) -> Vec<crate::discovery::Discovery> {
554 self.shared.discoveries.drain()
555 }
556
557 pub fn discoveries(&self) -> impl Iterator<Item = &crate::discovery::Discovery> + '_ {
559 self.shared.discoveries.iter()
560 }
561
562 pub fn has_no_discoveries(&self) -> bool {
564 self.shared.discoveries.is_empty()
565 }
566
567 pub fn lifted_code_seeds(&self) -> Vec<crate::discovery::CodeSeed> {
572 self.shared.discoveries.lifted_seeds()
573 }
574
575 pub fn seed_code(&mut self, seeds: impl IntoIterator<Item = crate::discovery::CodeSeed>) {
581 for seed in seeds {
582 self.shared.discoveries.insert(seed.into_discovery());
583 }
584 }
585
586 pub fn mark_discovery_lifted(&mut self, key: crate::discovery::DiscoveryKey) {
587 self.shared.discoveries.mark_lifted(key);
588 }
589
590 pub fn mark_discovery_failed(
591 &mut self,
592 key: crate::discovery::DiscoveryKey,
593 reason: impl Into<String>,
594 ) {
595 self.shared.discoveries.mark_failed(key, reason);
596 }
597
598 pub fn mark_discovery_skipped(
599 &mut self,
600 key: crate::discovery::DiscoveryKey,
601 reason: impl Into<String>,
602 ) {
603 self.shared.discoveries.mark_skipped(key, reason);
604 }
605
606 pub fn get_or_make_block(&mut self, addr: u64, func: FunctionId) -> BlockId {
611 let mut addresses = crate::address_index::AddressIndex::analyze(self);
612 self.get_or_make_block_indexed(&mut addresses, addr, func)
613 }
614
615 #[track_caller]
619 pub fn get_or_make_block_indexed(
620 &mut self,
621 addresses: &mut crate::address_index::AddressIndex,
622 addr: u64,
623 func: FunctionId,
624 ) -> BlockId {
625 use crate::address_index::AddressTarget;
626
627 if let Some(AddressTarget::Function(owner)) = addresses.get(addr) {
628 assert_eq!(
629 owner, func,
630 "cannot create a block at an address owned by another function"
631 );
632 }
633 let existing = match addresses.get(addr) {
634 Some(AddressTarget::Block(block)) => Some(block),
635 Some(AddressTarget::Function(function)) => FunctionBody::from_id(self, function)
636 .root()
637 .map(|root| root.id),
638 None => None,
639 };
640 match existing {
641 Some(block) => {
642 if self.block(block).address != Some(addr)
652 && block.func == func
653 && self.block(block).extra_addresses.contains(&addr)
654 {
655 return self.split_block_at_address(addresses, block, addr);
656 }
657 if block.func != func {
658 let stored = FunctionBody::from_id(self, block.func);
659 let requested = FunctionBody::from_id(self, func);
660 let parent = Some(block.func);
664 let caller = std::panic::Location::caller();
665 let detail = format!(
666 "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}",
667 block.func,
668 stored.name(),
669 stored.address(),
670 func,
671 requested.name(),
672 requested.address(),
673 );
674 log::error!(
675 target: "qcode::arena",
676 "{detail}\nbacktrace:\n{}",
677 std::backtrace::Backtrace::force_capture()
678 );
679 panic!("{detail}");
680 }
681 block
682 }
683 None => {
684 BasicBlock::make(self, func)
685 .with_address_indexed(addresses, addr)
686 .id
687 }
688 }
689 }
690
691 pub fn split_block_at_address(
711 &mut self,
712 addresses: &mut crate::address_index::AddressIndex,
713 block: BlockId,
714 addr: u64,
715 ) -> BlockId {
716 addresses.mark_boundary(addr);
719 let tail = BasicBlock::make(self, block.func).id;
720
721 self.bodies[block.func].clear_block_instructions(block);
724
725 let absorbed = std::mem::take(&mut self.block_mut(block).extra_addresses);
729 for absorbed_addr in absorbed {
730 addresses.forget(absorbed_addr);
731 }
732 BasicBlock::from_id_mut(self, tail)
733 .in_function(block.func)
734 .with_address_indexed(addresses, addr);
735 tail
736 }
737
738 pub fn builder(&mut self, block: BlockId) -> crate::builder::Builder<'str, '_> {
741 let body = &mut self.bodies[block.func];
742 crate::builder::Builder::new(body, &self.shared, &self.interfaces, block)
743 }
744
745 pub fn builder_at(&mut self, address: u64) -> crate::builder::Builder<'str, '_> {
748 use crate::address_index::AddressTarget;
749
750 let mut addresses = crate::address_index::AddressIndex::analyze(self);
751 let block = match addresses.get(address) {
752 Some(AddressTarget::Function(function)) => self.bodies[function]
753 .root_id()
754 .map(|local| BlockId::new(function, local))
755 .unwrap_or_else(|| {
756 self.get_or_make_block_indexed(&mut addresses, address, function)
757 }),
758 Some(AddressTarget::Block(block)) => block,
759 None => {
760 let function = FunctionBody::make(self, Cow::Owned(format!("blk_{address:x}")))
761 .expect("anonymous host function")
762 .id;
763 self.get_or_make_block_indexed(&mut addresses, address, function)
764 }
765 };
766 let mut builder = self.builder(block);
767 builder.set_address(address);
768 builder
769 }
770
771 pub fn bytes_display(&self, id: crate::value::BytesId) -> crate::value::BytesDisplay {
774 self.shared
775 .values
776 .bytes_display
777 .get(&id)
778 .copied()
779 .unwrap_or_default()
780 }
781
782 pub fn set_bytes_display(
786 &mut self,
787 id: crate::value::BytesId,
788 mode: crate::value::BytesDisplay,
789 ) {
790 if mode == crate::value::BytesDisplay::Auto {
791 self.shared.values.bytes_display.remove(&id);
792 } else {
793 self.shared.values.bytes_display.insert(id, mode);
794 }
795 }
796
797 pub fn block_ids(&self) -> Vec<BlockId> {
799 self.functions().flat_map(|f| f.block_ids()).collect()
800 }
801
802 pub fn instruction_ids(&self) -> Vec<InstructionId> {
806 let mut ids: Vec<_> = self.functions().flat_map(|f| f.instruction_ids()).collect();
807 ids.sort_unstable();
808 ids
809 }
810
811 pub fn function_ids(&self) -> Vec<FunctionId> {
813 self.interfaces.iter().map(|i| i.id).collect()
814 }
815
816 pub fn anon_function(&mut self) -> FunctionId {
822 let name = self.get_unique_name(std::borrow::Cow::Borrowed("anon"));
823 crate::value::FunctionBody::make(self, name)
824 .expect("unique anon function name")
825 .id
826 }
827
828 pub fn instruction_arena_stats(&self) -> (usize, usize) {
830 let mut total = 0;
831 let mut dead = 0;
832 for f in self.bodies.iter() {
833 total += f.insns.issued_len();
834 dead += f.insns.issued_len() - f.insns.len();
835 }
836 (total, dead)
837 }
838
839 pub fn body_arena_stats(&self) -> crate::value::BodyArenaStats {
845 let mut total = crate::value::BodyArenaStats::default();
846 for body in self.bodies.iter() {
847 total.add_assign(body.arena_stats());
848 }
849 total
850 }
851
852 pub fn shrink_bodies_to_fit(&mut self) {
859 for mut body in self.bodies.iter_mut() {
860 body.shrink_to_fit();
861 }
862 }
863
864 pub fn instructions(&self) -> impl Iterator<Item = InstructionRef<'str, '_>> + '_ {
866 self.instruction_ids()
867 .into_iter()
868 .map(move |id| Instruction::from_id(self, id))
869 }
870
871 pub fn blocks(&self) -> impl Iterator<Item = BlockRef<'str, '_>> + '_ {
873 self.block_ids()
874 .into_iter()
875 .map(move |id| BlockRef::from_id(self, id))
876 }
877
878 pub fn functions(&self) -> FunctionIter<'str, '_> {
880 FunctionIter {
881 ctx: self,
882 inner: self.bodies.iter(),
883 }
884 }
885
886 pub fn iter(&self) -> FunctionIter<'str, '_> {
889 self.functions()
890 }
891
892 pub fn varnodes(&self) -> impl Iterator<Item = VarnodeRef<'str, '_>> + '_ {
893 self.shared.varnodes()
894 }
895
896 pub fn varnode_count(&self) -> usize {
901 self.shared.varnode_count()
902 }
903
904 pub fn remove_cfg_edge(&mut self, func: FunctionId, edge_id: EdgeId) {
908 self.bodies[func].remove_cfg_edge(edge_id);
909 }
910
911 pub fn rehome_owned_blocks(
931 &mut self,
932 addresses: &mut crate::address_index::AddressIndex,
933 target: FunctionId,
934 olds: &[BlockId],
935 ) -> HashMap<BlockId, BlockId> {
936 let mut needed_temps: HashSet<TempId> = HashSet::default();
940 let mut needed_temp_spaces: HashSet<TempSpaceId> = HashSet::default();
941 for &old in olds {
942 for ¶m_local in &self.block(old).params {
943 let param = self.block_param(BlockParamId::new(old.func, param_local));
944 if let Some(crate::value::LocalValueId::Temp(temp)) = param.origin {
945 needed_temps.insert(TempId::new(old.func, temp));
946 }
947 if let Some(MemorySpaceId::Temp(space)) = self.shared.types.space_of(param.type_id)
948 {
949 needed_temp_spaces.insert(space);
950 }
951 }
952 for &insn_local in &self.block(old).instructions {
953 let insn = self.instruction(InstructionId::new(old.func, insn_local));
954 for arg in insn.mnemonic().args() {
955 if let crate::value::LocalValueId::Temp(temp) = arg {
956 needed_temps.insert(TempId::new(old.func, temp));
957 }
958 }
959 let explicit_space = match insn.mnemonic() {
960 Mnemonic::Load(load) => Some(load.space),
961 Mnemonic::Store(store) => Some(store.space),
962 _ => None,
963 };
964 if let Some(LocalMemorySpaceId::Temp(space)) = explicit_space {
965 needed_temp_spaces.insert(TempSpaceId::new(old.func, space));
966 }
967 if let Some(MemorySpaceId::Temp(space)) = self.shared.types.space_of(insn.type_id) {
968 needed_temp_spaces.insert(space);
969 }
970 }
971 }
972 for &temp in &needed_temps {
973 let data = &self.bodies[temp.func].temps[temp.local];
974 needed_temp_spaces.insert(TempSpaceId::new(temp.func, data.space));
975 }
976
977 let mut needed_temp_spaces: Vec<_> = needed_temp_spaces.into_iter().collect();
978 needed_temp_spaces.sort_unstable();
979 let mut temp_space_map: HashMap<TempSpaceId, TempSpaceId> = HashMap::default();
980 for old in needed_temp_spaces {
981 if old.func == target {
982 continue;
983 }
984 let space = self.bodies[old.func].temp_spaces[old.local].clone();
985 let new = self.bodies[target].push_temp_space(space);
986 temp_space_map.insert(old, new);
987 }
988
989 let mut needed_temps: Vec<_> = needed_temps.into_iter().collect();
990 needed_temps.sort_unstable();
991 let mut value_map: HashMap<ValueId, ValueId> = HashMap::default();
992 for old in needed_temps {
993 if old.func == target {
994 continue;
995 }
996 let mut temp = self.bodies[old.func].temps[old.local].clone();
997 temp.space = temp_space_map[&TempSpaceId::new(old.func, temp.space)].local;
998 if let Some(name) = temp.name.take() {
999 temp.name = Some(self.bodies[target].names.unique(name));
1000 }
1001 let new = self.bodies[target].push_temp(temp);
1002 value_map.insert(ValueId::Temp(old), ValueId::Temp(new));
1003 }
1004
1005 let mut block_map: HashMap<BlockId, BlockId> = HashMap::default();
1008 for &old in olds {
1009 let new = BasicBlock::clone_block_into(self, old, target, &mut value_map);
1010 block_map.insert(old, new);
1011 }
1012
1013 for (&old, &new) in &block_map {
1018 let old_params = self.block(old).params.clone();
1019 let new_params = self.block(new).params.clone();
1020 for (old_local, new_local) in old_params.into_iter().zip(new_params) {
1021 let old_param = BlockParamId::new(old.func, old_local);
1022 let new_param = BlockParamId::new(new.func, new_local);
1023 let type_id = remap_rehomed_type(
1024 self,
1025 self.block_param(new_param).type_id,
1026 target,
1027 &temp_space_map,
1028 );
1029 self.block_param_mut(new_param).type_id = type_id;
1030 let Some(origin) = self.block_param(new_param).origin else {
1031 continue;
1032 };
1033 let qualified = origin.qualify(old.func);
1034 let remapped = value_map.get(&qualified).copied().unwrap_or(qualified);
1035 debug_assert!(
1036 remapped.owning_function().is_none_or(|f| f == target),
1037 "rehome: relocated block param {old_param:?} has an origin in another \
1038 function ({qualified:?}); the relocated set is not closed",
1039 );
1040 self.block_param_mut(new_param).origin = Some(remapped.localize(new.func));
1041 }
1042
1043 let insns = self.block(new).instructions.clone();
1044 for insn_local in insns {
1045 let insn_id = InstructionId::new(new.func, insn_local);
1046 let type_id = remap_rehomed_type(
1047 self,
1048 self.instruction(insn_id).type_id,
1049 target,
1050 &temp_space_map,
1051 );
1052 self.instruction_mut(insn_id).type_id = type_id;
1053 let mut mnemonic = self.instruction(insn_id).mnemonic().clone();
1054 let mut pairs = Vec::new();
1055 for arg in mnemonic.args() {
1056 let qualified = arg.qualify(old.func);
1060 if let Some(&new_val) = value_map.get(&qualified) {
1061 pairs.push((arg, new_val.localize(new.func)));
1062 } else if let Some(new_lit) =
1063 remap_symbolic_block_literal(&self.shared.values.literals, arg, &block_map)
1064 {
1065 pairs.push((arg, new_lit));
1066 } else {
1067 debug_assert!(
1072 qualified.owning_function().is_none_or(|f| f == target),
1073 "rehome: relocated block references a value in another \
1074 function ({qualified:?}); the relocated set is not closed",
1075 );
1076 }
1077 }
1078 crate::value::block::substitute_operands(&mut mnemonic, &pairs);
1079 remap_rehomed_memory_space(&mut mnemonic, old.func, target, &temp_space_map);
1080 remap_block_targets(&mut mnemonic, old.func, new.func, &block_map);
1081 *self.instruction_mut(insn_id).mnemonic_mut() = mnemonic;
1082 }
1083 }
1084
1085 let mut incident: HashSet<(FunctionId, EdgeId)> = HashSet::default();
1092 for &old in olds {
1093 incident.extend(self.block(old).edges.iter().map(|&e| (old.func, e)));
1094 }
1095 let mut incident: Vec<_> = incident.into_iter().collect();
1096 incident.sort_unstable();
1097 for (edge_func, edge) in incident {
1098 let EdgeData { from, to } = *self.edge(edge_func, edge);
1099 let from = BlockId::new(edge_func, from);
1100 let to = BlockId::new(edge_func, to);
1101 let new_from = block_map.get(&from).copied().unwrap_or(from);
1102 let new_to = block_map.get(&to).copied().unwrap_or(to);
1103 self.add_cfg_edge(new_from, new_to);
1104 }
1105
1106 for &old in olds {
1112 let Some(addr) = self.block(old).address else {
1113 continue;
1114 };
1115 let new = block_map[&old];
1116 let extra = self.block(old).extra_addresses.clone();
1117 addresses.rehome_block(addr, old, new);
1118 for &e in &extra {
1119 addresses.rehome_block(e, old, new);
1120 }
1121 self.block_mut(new).extra_addresses = extra;
1122 self.block_mut(new).address = Some(addr);
1123 }
1124
1125 for &old in olds {
1128 BasicBlock::from_id_mut(self, old).delete();
1129 }
1130
1131 self.rebuild_users(target);
1134 block_map
1135 }
1136
1137 fn function_registered_at_block(
1141 &self,
1142 addresses: &crate::address_index::AddressIndex,
1143 block: BlockId,
1144 ) -> Option<FunctionId> {
1145 self.block(block)
1146 .address
1147 .and_then(|addr| addresses.function_at(addr))
1148 }
1149
1150 fn split_tail(
1159 &self,
1160 addresses: &crate::address_index::AddressIndex,
1161 block: BlockId,
1162 g: FunctionId,
1163 ) -> Vec<BlockId> {
1164 let mut seen: HashSet<BlockId> = HashSet::default();
1165 seen.insert(block);
1166 let mut queue = vec![block];
1167 while let Some(b) = queue.pop() {
1168 let succs: Vec<BlockId> = BasicBlock::from_id(self, b)
1169 .successors()
1170 .map(|(_, s)| s)
1171 .collect();
1172 for s in succs {
1173 if seen.contains(&s) {
1174 continue;
1175 }
1176 if let Some(entry_func) = self.function_registered_at_block(addresses, s)
1180 && entry_func != g
1181 {
1182 continue;
1183 }
1184 seen.insert(s);
1185 queue.push(s);
1186 }
1187 }
1188 let mut tail: Vec<BlockId> = seen.into_iter().collect();
1189 tail.sort_unstable_by_key(|&b| (self.block(b).address, b.local, b.func));
1190 tail
1191 }
1192
1193 pub fn split_function_at(&mut self, block: BlockId) -> FunctionId {
1217 let mut addresses = crate::address_index::AddressIndex::analyze(self);
1218 self.split_function_at_indexed(&mut addresses, block)
1219 }
1220
1221 pub fn split_function_at_indexed(
1224 &mut self,
1225 addresses: &mut crate::address_index::AddressIndex,
1226 block: BlockId,
1227 ) -> FunctionId {
1228 use crate::value::insn::{Branch, CBranch, Callee, TailCall};
1229
1230 let addr = self
1231 .block(block)
1232 .address
1233 .expect("split_function_at: block has no machine address");
1234
1235 let g = match addresses.function_at(addr) {
1240 Some(existing) => existing,
1241 None => FunctionBody::make_at_addr_indexed(self, addresses, addr, None).id,
1242 };
1243
1244 loop {
1257 let tail_set: HashSet<BlockId> =
1258 self.split_tail(addresses, block, g).into_iter().collect();
1259 let mut promote: Option<BlockId> = None;
1260 'scan: for b in self.block_ids() {
1261 if tail_set.contains(&b) {
1262 continue;
1264 }
1265 let Some(mnemonic) = BasicBlock::from_id(self, b)
1266 .instructions()
1267 .last()
1268 .map(|t| t.mnemonic().clone())
1269 else {
1270 continue;
1271 };
1272 let targets = match &mnemonic {
1273 Mnemonic::Branch(Branch { target, .. }) => vec![*target],
1274 Mnemonic::CBranch(CBranch {
1275 success_block,
1276 failure_block,
1277 ..
1278 }) => vec![*success_block, *failure_block],
1279 _ => vec![],
1280 };
1281 for t in targets {
1282 let tid = BlockId::new(b.func, t);
1283 if tid == block || !tail_set.contains(&tid) {
1287 continue;
1288 }
1289 promote = Some(tid);
1298 break 'scan;
1299 }
1300 }
1301 match promote {
1302 Some(tid) => {
1303 self.split_function_at_indexed(addresses, tid);
1304 }
1305 None => break,
1306 }
1307 }
1308
1309 let mut tail = self.split_tail(addresses, block, g);
1312 let mut tail_set: HashSet<BlockId> = tail.iter().copied().collect();
1313
1314 let mut prev_owners: HashSet<FunctionId> = HashSet::default();
1317 for &b in &tail {
1318 prev_owners.insert(b.func);
1320 }
1321
1322 let effective_owner = |_ctx: &Context, candidate: BlockId| {
1326 if tail_set.contains(&candidate) {
1327 Some(g)
1328 } else {
1329 Some(candidate.func)
1331 }
1332 };
1333
1334 let foreign_entry =
1337 |ctx: &Context, target: BlockId, owner: FunctionId| -> Option<FunctionId> {
1338 let callee = if target == block {
1339 g
1340 } else {
1341 ctx.function_registered_at_block(addresses, target)?
1342 };
1343 (callee != owner).then_some(callee)
1344 };
1345
1346 let mut tail_calls: Vec<(InstructionId, FunctionId)> = Vec::new();
1352 let mut cond_calls: Vec<(
1359 InstructionId,
1360 BlockId,
1361 FunctionId,
1362 crate::value::LocalBlockId,
1363 )> = Vec::new();
1364 let relevant: Vec<BlockId> = self.block_ids();
1365 for b in relevant {
1366 let Some(owner) = effective_owner(self, b) else {
1367 continue;
1368 };
1369 let Some((term_id, mnemonic)) = BasicBlock::from_id(self, b)
1370 .instructions()
1371 .last()
1372 .map(|t| (t.id, t.mnemonic().clone()))
1373 else {
1374 continue;
1375 };
1376 match mnemonic {
1379 Mnemonic::Branch(Branch { target, .. }) => {
1380 if let Some(callee) = foreign_entry(self, BlockId::new(b.func, target), owner) {
1381 tail_calls.push((term_id, callee));
1382 }
1383 }
1384 Mnemonic::CBranch(CBranch {
1385 success_block,
1386 failure_block,
1387 ..
1388 }) => {
1389 if let Some(callee) =
1390 foreign_entry(self, BlockId::new(b.func, success_block), owner)
1391 {
1392 cond_calls.push((term_id, b, callee, success_block));
1393 }
1394 if let Some(callee) =
1395 foreign_entry(self, BlockId::new(b.func, failure_block), owner)
1396 {
1397 cond_calls.push((term_id, b, callee, failure_block));
1398 }
1399 }
1400 _ => {}
1401 }
1402 }
1403
1404 for (insn, callee) in tail_calls {
1405 self.replace_instruction_mnemonic(
1406 insn,
1407 Mnemonic::TailCall(TailCall {
1408 target: Callee::Real(callee),
1409 args: vec![],
1410 }),
1411 );
1412 }
1413 for (insn, owner_block, callee, arm_target) in cond_calls {
1414 let tramp = BasicBlock::make(self, owner_block.func).id;
1416 (self).builder(tramp).push_tail_call(callee);
1417 self.add_cfg_edge(owner_block, tramp);
1418
1419 if tail_set.contains(&owner_block) {
1427 tail.push(tramp);
1428 tail_set.insert(tramp);
1429 }
1430
1431 let Mnemonic::CBranch(mut cb) = self.instruction(insn).mnemonic().clone() else {
1432 continue;
1433 };
1434 let tramp_local = tramp.localize(insn.func);
1439 if cb.success_block == arm_target {
1440 cb.success_block = tramp_local;
1441 }
1442 if cb.failure_block == arm_target {
1443 cb.failure_block = tramp_local;
1444 }
1445 self.replace_instruction_mnemonic(insn, Mnemonic::CBranch(cb));
1446 }
1447
1448 let moved_owner = |candidate: BlockId| {
1457 if tail_set.contains(&candidate) {
1458 g
1459 } else {
1460 candidate.func
1461 }
1462 };
1463 let mut stale: HashSet<(FunctionId, EdgeId)> = HashSet::default();
1464 for &b in &tail {
1465 for edge in self.block(b).edges.iter().copied() {
1466 let &EdgeData { from, to } = self.edge(b.func, edge);
1467 let from = BlockId::new(b.func, from);
1468 let to = BlockId::new(b.func, to);
1469 let cross = moved_owner(from) != moved_owner(to);
1470 let touches_tail = tail_set.contains(&from) || tail_set.contains(&to);
1471 if cross && touches_tail {
1472 stale.insert((b.func, edge));
1473 }
1474 }
1475 }
1476 let mut stale: Vec<_> = stale.into_iter().collect();
1477 stale.sort_unstable();
1478 for (func, edge) in stale {
1479 self.remove_cfg_edge(func, edge);
1480 }
1481
1482 let moved = self.rehome_owned_blocks(addresses, g, &tail);
1486 self.bodies[g].set_root_id(Some(moved[&block].local));
1487
1488 self.recompute_instruction_addrs(g);
1490 for owner in prev_owners {
1491 if owner != g {
1492 self.recompute_instruction_addrs(owner);
1493 }
1494 }
1495
1496 g
1497 }
1498
1499 fn recompute_instruction_addrs(&mut self, func: FunctionId) {
1502 let blocks = FunctionBody::from_id(self, func).block_ids();
1503 let mut addrs = std::collections::BTreeSet::new();
1504 for b in blocks {
1505 for insn in BasicBlock::from_id(self, b).instructions() {
1506 if let Some(a) = insn.address() {
1507 addrs.insert(a);
1508 }
1509 }
1510 }
1511 self.bodies[func].instruction_addrs = addrs;
1512 }
1513
1514 fn rebuild_users(&mut self, func: FunctionId) {
1518 let live: Vec<InstructionId> = FunctionBody::from_id(self, func).instruction_ids();
1519 let users = &mut self.bodies[func].users;
1520 users.clear();
1521 for id in live {
1522 let args = self.bodies[func].insns[id.local].mnemonic().args();
1523 let users = &mut self.bodies[func].users;
1524 for arg in args {
1525 users.entry(arg).or_default().push(id.localize(func));
1526 }
1527 }
1528 }
1529
1530 pub fn assume_true(&mut self, prop: Proposition) -> bool {
1535 self.assume(prop, true)
1536 }
1537
1538 pub fn assume_false(&mut self, prop: Proposition) -> bool {
1540 self.assume(prop, false)
1541 }
1542
1543 fn assume(&mut self, prop: Proposition, value: bool) -> bool {
1544 match self.shared.values.truths.get(&prop) {
1545 Some(t) => t.value == value,
1546 None => {
1547 self.shared.values.truths.insert(
1548 prop,
1549 Truth {
1550 value,
1551 certainty: Certainty::Assumed,
1552 pass: PassName(pass_scope::current_pass()),
1553 },
1554 );
1555 true
1556 }
1557 }
1558 }
1559
1560 pub fn set_known(&mut self, prop: Proposition, value: bool) -> bool {
1568 let pass = PassName(pass_scope::current_pass());
1569 let novel = match self.shared.values.truths.get(&prop) {
1570 Some(prior) => {
1571 if prior.certainty == Certainty::Known && prior.value != value {
1576 self.shared
1577 .values
1578 .known_contradictions
1579 .push(KnownContradiction {
1580 prop,
1581 known: prior.value,
1582 proven: value,
1583 known_pass: prior.pass,
1584 proven_pass: pass,
1585 });
1586 return false;
1587 }
1588 if prior.certainty == Certainty::Assumed && prior.value != value {
1589 self.shared.values.violations.push(Violation {
1590 prop,
1591 assumed: prior.value,
1592 assuming_pass: prior.pass,
1593 asserting_pass: pass,
1594 });
1595 true
1596 } else {
1597 false
1598 }
1599 }
1600 None => true,
1601 };
1602 self.shared.values.truths.insert(
1603 prop,
1604 Truth {
1605 value,
1606 certainty: Certainty::Known,
1607 pass,
1608 },
1609 );
1610 novel
1611 }
1612
1613 pub fn seed_known(&mut self, prop: Proposition, value: bool, pass: PassName) {
1623 let prior = self.shared.values.truths.insert(
1624 prop,
1625 Truth {
1626 value,
1627 certainty: Certainty::Known,
1628 pass,
1629 },
1630 );
1631 debug_assert!(prior.is_none(), "seeding {prop:?} over an existing truth");
1632 }
1633
1634 pub fn truth(&self, prop: Proposition) -> Option<Truth> {
1636 self.shared.values.truths.get(&prop).copied()
1637 }
1638
1639 pub fn set_assumed_call_convention(
1645 &mut self,
1646 effect: Option<crate::assumption::AssumedCallEffect>,
1647 ) {
1648 self.shared.assumed_call_convention = effect;
1649 }
1650
1651 pub fn assumed_call_convention(&self) -> Option<&crate::assumption::AssumedCallEffect> {
1655 self.shared.assumed_call_convention.as_ref()
1656 }
1657
1658 pub fn known(&self, prop: Proposition) -> Option<bool> {
1660 self.truth(prop)
1661 .filter(|t| t.certainty == Certainty::Known)
1662 .map(|t| t.value)
1663 }
1664
1665 pub fn truths(&self) -> impl Iterator<Item = (Proposition, Truth)> + '_ {
1667 self.shared.values.truths.iter().map(|(&p, &t)| (p, t))
1668 }
1669
1670 pub fn known_facts(&self) -> impl Iterator<Item = (Proposition, bool, PassName)> + '_ {
1673 self.truths()
1674 .filter(|(_, t)| t.certainty == Certainty::Known)
1675 .map(|(p, t)| (p, t.value, t.pass))
1676 }
1677
1678 pub fn violations(&self) -> &[Violation] {
1681 &self.shared.values.violations
1682 }
1683
1684 pub fn known_contradictions(&self) -> &[KnownContradiction] {
1688 &self.shared.values.known_contradictions
1689 }
1690
1691 pub fn get_literal_value(&self, id: LiteralId) -> u64 {
1693 self.shared.values.literals[id].value
1694 }
1695
1696 pub fn get_insn(&self, id: InstructionId) -> InstructionRef<'str, '_> {
1698 InstructionRef::from_id(self, id)
1699 }
1700
1701 pub fn body(&self, fid: FunctionId) -> &crate::value::FunctionBody<'str> {
1711 &self.bodies[fid]
1712 }
1713
1714 pub fn body_mut(&mut self, fid: FunctionId) -> &mut crate::value::FunctionBody<'str> {
1716 &mut self.bodies[fid]
1717 }
1718
1719 pub fn push_insn(&mut self, func: FunctionId, insn: Instruction<'str>) -> InstructionId {
1734 let args = insn.mnemonic().args();
1735 let local = self.bodies[func].insns.push(insn);
1736 let id = InstructionId::new(func, local);
1737 for arg in args {
1738 self.bodies[func]
1739 .users
1740 .entry(arg)
1741 .or_default()
1742 .push(id.localize(func));
1743 }
1744 id
1745 }
1746
1747 pub fn instruction(&self, id: InstructionId) -> &Instruction<'str> {
1749 &self.bodies[id.func].insns[id.local]
1750 }
1751
1752 pub fn instruction_mut(&mut self, id: InstructionId) -> &mut Instruction<'str> {
1754 &mut self.bodies[id.func].insns[id.local]
1755 }
1756
1757 pub fn contains_instruction(&self, id: InstructionId) -> bool {
1759 Into::<usize>::into(id.func) < self.bodies.len()
1760 && self.bodies[id.func].insns.contains(id.local)
1761 }
1762
1763 pub fn block(&self, id: BlockId) -> &BasicBlock<'str> {
1765 &self.bodies[id.func].blocks[id.local]
1766 }
1767
1768 pub fn block_mut(&mut self, id: BlockId) -> &mut BasicBlock<'str> {
1770 &mut self.bodies[id.func].blocks[id.local]
1771 }
1772
1773 pub fn contains_block(&self, id: BlockId) -> bool {
1775 Into::<usize>::into(id.func) < self.bodies.len()
1776 && self.bodies[id.func].blocks.contains(id.local)
1777 }
1778
1779 pub fn block_param(&self, id: BlockParamId) -> &BlockParam<'str> {
1781 &self.bodies[id.func].params[id.local]
1782 }
1783
1784 pub fn block_param_mut(&mut self, id: BlockParamId) -> &mut BlockParam<'str> {
1786 &mut self.bodies[id.func].params[id.local]
1787 }
1788
1789 pub fn contains_block_param(&self, id: BlockParamId) -> bool {
1791 Into::<usize>::into(id.func) < self.bodies.len()
1792 && self.bodies[id.func].params.contains(id.local)
1793 }
1794
1795 pub fn edge(&self, func: FunctionId, id: EdgeId) -> &EdgeData {
1797 &self.bodies[func].edges[id]
1798 }
1799
1800 pub fn edge_mut(&mut self, func: FunctionId, id: EdgeId) -> &mut EdgeData {
1802 &mut self.bodies[func].edges[id]
1803 }
1804
1805 pub fn users_of(&self, value: ValueId) -> Vec<InstructionId> {
1810 match value.owning_function() {
1811 Some(func) => self.bodies[func].users_of(value),
1812 None => Vec::new(),
1813 }
1814 }
1815
1816 pub fn has_users(&self, value: ValueId) -> bool {
1818 match value.owning_function() {
1819 Some(func) => self.bodies[func].has_users(value),
1820 None => false,
1821 }
1822 }
1823
1824 pub fn push_block(&mut self, func: FunctionId, block: BasicBlock<'str>) -> BlockId {
1825 let local = self.bodies[func].blocks.push(block);
1826 let id = BlockId::new(func, local);
1827 self.bodies[func].roster.push(local);
1829 id
1830 }
1831
1832 pub fn push_block_param(&mut self, func: FunctionId, param: BlockParam<'str>) -> BlockParamId {
1833 let local = self.bodies[func].params.push(param);
1834 BlockParamId::new(func, local)
1835 }
1836
1837 pub fn push_edge(&mut self, func: FunctionId, edge: EdgeData) -> EdgeId {
1838 self.bodies[func].edges.push(edge)
1839 }
1840
1841 pub fn push_function(
1844 &mut self,
1845 interface: crate::value::function::FunctionInterface<'str>,
1846 body: FunctionBody<'str>,
1847 ) -> FunctionId {
1848 let expected = FunctionId::from(self.bodies.len());
1849 assert_eq!(
1850 body.id(),
1851 expected,
1852 "function body id does not match its registry slot"
1853 );
1854 let id = self.bodies.push(body);
1855 let iid = self.interfaces.push(interface);
1856 debug_assert_eq!(
1857 Into::<usize>::into(id),
1858 Into::<usize>::into(iid),
1859 "function body/interface registries drifted"
1860 );
1861 id
1862 }
1863
1864 pub fn get_register(&self, id: RegisterId) -> VarnodeRef<'str, '_> {
1867 Varnode::from_id(self, self.shared.registers[&id])
1868 }
1869
1870 pub fn get_const(&self, value: u64, size: usize) -> LiteralRef<'str, '_> {
1872 let type_id = self.shared.types.get_or_make_int(size);
1873 let id = self
1874 .shared
1875 .values
1876 .get_or_make_typed_literal(value, type_id, size);
1877 LiteralRef::from_id(self, id)
1878 }
1879
1880 pub fn get_bool_const(&self, value: bool) -> LiteralRef<'str, '_> {
1883 let type_id = self.shared.types.get_or_make_bool();
1884 let id = self
1885 .shared
1886 .values
1887 .get_or_make_typed_literal(u64::from(value), type_id, 1);
1888 LiteralRef::from_id(self, id)
1889 }
1890
1891 pub fn get_poison(&self, type_id: crate::types::TypeId) -> ValueId {
1895 ValueId::Poison(self.shared.values.push_poison(type_id))
1896 }
1897
1898 pub fn get_typed_const(
1904 &self,
1905 value: u64,
1906 type_id: crate::types::TypeId,
1907 ) -> LiteralRef<'str, '_> {
1908 let size = self.shared.types.size_of(type_id);
1909 let id = self
1910 .shared
1911 .values
1912 .get_or_make_typed_literal(value, type_id, size);
1913 LiteralRef::from_id(self, id)
1914 }
1915
1916 pub fn get_bytes(&self, data: Vec<u8>) -> crate::value::BytesRef<'str, '_> {
1924 let i8_ty = self.shared.types.get_or_make_int(1);
1925 let type_id = self.shared.types.get_or_make_array(i8_ty, data.len());
1926 self.get_typed_bytes(data, type_id)
1927 }
1928
1929 pub fn get_typed_bytes(
1935 &self,
1936 data: Vec<u8>,
1937 type_id: crate::types::TypeId,
1938 ) -> crate::value::BytesRef<'str, '_> {
1939 let id = self
1940 .shared
1941 .values
1942 .bytes
1943 .push(crate::value::Bytes { data, type_id });
1944 crate::value::BytesRef::from_id(self, id)
1945 }
1946
1947 pub fn type_of(&self, id: ValueId) -> crate::types::TypeId {
1952 match id {
1953 ValueId::Literal(lid) => self.shared.values.literals[lid].type_id,
1954 ValueId::Bytes(bid) => self.shared.values.bytes[bid].type_id,
1955 ValueId::Instruction(iid) => self.instruction(iid).type_id,
1956 ValueId::BlockParam(pid) => self.block_param(pid).type_id,
1957 ValueId::Varnode(vid) => {
1958 if let Some(&ty) = self.shared.values.varnode_types.get(&vid) {
1959 return ty;
1960 }
1961 let size = self.shared.values.varnodes[vid].size_bytes();
1962 self.shared.types.get_or_make_int(size)
1963 }
1964 ValueId::Temp(id) => self
1965 .shared
1966 .types
1967 .get_or_make_int(self.bodies[id.func].temps[id.local].size),
1968 ValueId::Poison(pid) => self.shared.values.poisons[pid].type_id,
1969 ValueId::BasicBlock(_) | ValueId::Function(_) => self.shared.types.get_or_make_int(0),
1972 }
1973 }
1974
1975 pub fn stored_type_of(&self, id: ValueId) -> Option<crate::types::TypeId> {
1981 match id {
1982 ValueId::Literal(lid) => Some(self.shared.values.literals[lid].type_id),
1983 ValueId::Bytes(bid) => Some(self.shared.values.bytes[bid].type_id),
1984 ValueId::Instruction(iid) => Some(self.instruction(iid).type_id),
1985 ValueId::BlockParam(pid) => Some(self.block_param(pid).type_id),
1986 ValueId::Varnode(vid) => self.shared.values.varnode_types.get(&vid).copied(),
1987 ValueId::Poison(pid) => Some(self.shared.values.poisons[pid].type_id),
1988 ValueId::Temp(_) => None,
1989 ValueId::BasicBlock(_) | ValueId::Function(_) => None,
1990 }
1991 }
1992
1993 pub fn set_varnode_type(&mut self, varnode: VarnodeId, type_id: crate::types::TypeId) {
1998 self.shared.values.varnode_types.insert(varnode, type_id);
1999 }
2000
2001 pub fn users(&self, value: impl Into<ValueId>) -> Vec<InstructionId> {
2009 self.users_of(value.into())
2010 }
2011
2012 pub fn users_across_functions(&self, value: impl Into<ValueId>) -> Vec<InstructionId> {
2017 let value = value.into();
2018 if value.owning_function().is_some() {
2019 self.users_of(value)
2020 } else {
2021 self.functions().flat_map(|f| f.users_of(value)).collect()
2022 }
2023 }
2024
2025 pub fn view(&self) -> ModuleView<'_, 'str> {
2034 ModuleView::new(self)
2035 }
2036 pub fn shr(&self) -> &Shared<'str> {
2041 &self.shared
2042 }
2043 pub fn function(&self, f: FunctionId) -> &FunctionBody<'str> {
2045 &self.bodies[f]
2046 }
2047 pub fn function_mut(&mut self, f: FunctionId) -> &mut FunctionBody<'str> {
2049 &mut self.bodies[f]
2050 }
2051
2052 pub fn block_ref(&self, id: BlockId) -> BlockRef<'str, '_, ModuleView<'_, 'str>> {
2054 self.view().block_ref(id)
2055 }
2056 pub fn insn_ref(&self, id: InstructionId) -> InstructionRef<'str, '_, ModuleView<'_, 'str>> {
2058 self.view().insn_ref(id)
2059 }
2060 pub fn param_ref(&self, id: BlockParamId) -> BlockParamRef<'str, '_, ModuleView<'_, 'str>> {
2062 self.view().param_ref(id)
2063 }
2064 pub fn function_ref(&self, id: FunctionId) -> FunctionRef<'str, '_, ModuleView<'_, 'str>> {
2066 self.view().function_ref(id)
2067 }
2068
2069 pub fn push_mnemonic(
2071 &mut self,
2072 func: FunctionId,
2073 mnemonic: Mnemonic,
2074 size: usize,
2075 ) -> InstructionId {
2076 let type_id = self.shared.types.get_or_make_int(size);
2077 self.push_insn(func, Instruction::new(type_id, mnemonic))
2078 }
2079
2080 pub fn push_mnemonic_with_type(
2083 &mut self,
2084 func: FunctionId,
2085 mnemonic: Mnemonic,
2086 type_id: crate::types::TypeId,
2087 ) -> InstructionId {
2088 self.push_insn(func, Instruction::new(type_id, mnemonic))
2089 }
2090
2091 pub fn make_block(&mut self, func: FunctionId) -> BlockId {
2094 self.push_block(func, BasicBlock::detached())
2095 }
2096
2097 pub fn register_local_name(
2100 &mut self,
2101 id: ValueId,
2102 name: Cow<'str, str>,
2103 old_name: Option<&str>,
2104 ) -> Result<()> {
2105 let existing = match id.name_scope_function() {
2106 Some(func) => self
2107 .function(func)
2108 .names
2109 .get(&name)
2110 .map(|id| id.qualify(func)),
2111 None => self.get_named(&name),
2112 };
2113 if let Some(existing) = existing {
2114 return if existing == id {
2115 Ok(())
2116 } else {
2117 Err(Error::spanless(ErrorTy::DuplicateName(name.to_string())))
2118 };
2119 }
2120 match id.name_scope_function() {
2121 Some(func) => self
2122 .function_mut(func)
2123 .names
2124 .register(name, id.localize(func), old_name),
2125 None => self.update_name(name, id, old_name),
2126 }
2127 }
2128
2129 pub(crate) fn set_address_indexed(
2131 &mut self,
2132 addresses: &mut crate::address_index::AddressIndex,
2133 addr: u64,
2134 id: ValueId,
2135 ) -> crate::error::Result<()> {
2136 let target = match id {
2137 ValueId::Function(id) => crate::address_index::AddressTarget::Function(id),
2138 ValueId::BasicBlock(id) => crate::address_index::AddressTarget::Block(id),
2139 _ => unreachable!("only functions and blocks have module addresses"),
2140 };
2141 addresses.register(self, addr, target)
2142 }
2143
2144 pub fn update_name(
2147 &mut self,
2148 name: Cow<'str, str>,
2149 id: ValueId,
2150 old_name: Option<&str>,
2151 ) -> Result<()> {
2152 match id.name_scope_function() {
2153 Some(func) => self.bodies[func]
2154 .names
2155 .register(name, id.localize(func), old_name),
2156 None => self.shared.name_map.register(name, id, old_name),
2157 }
2158 }
2159
2160 pub fn get_named_in_scope(&self, id: ValueId, name: &str) -> Option<ValueId> {
2165 match id.name_scope_function() {
2166 Some(func) => self.bodies[func].names.get(name).map(|id| id.qualify(func)),
2167 None => self.shared.name_map.get(name),
2168 }
2169 }
2170
2171 pub fn get_named(&self, name: &str) -> Option<ValueId> {
2180 self.shared.name_map.get(name)
2181 }
2182
2183 pub fn get_unique_name(&mut self, name: Cow<'str, str>) -> Cow<'str, str> {
2188 self.shared.name_map.unique(name)
2189 }
2190
2191 pub fn get_unique_name_in(&mut self, func: FunctionId, name: Cow<'str, str>) -> Cow<'str, str> {
2195 self.bodies[func].names.unique(name)
2196 }
2197}
2198
2199#[derive(Clone, serde::Serialize, serde::Deserialize)]
2210pub struct NameTable<'str, Id = ValueId> {
2211 map: HashMap<Cow<'str, str>, Id>,
2213 #[serde(skip)]
2217 suffix_hint: HashMap<String, u32>,
2218}
2219
2220impl<Id> Default for NameTable<'_, Id> {
2221 fn default() -> Self {
2222 Self {
2223 map: HashMap::default(),
2224 suffix_hint: HashMap::default(),
2225 }
2226 }
2227}
2228
2229impl<'str, Id: Copy + Eq> NameTable<'str, Id> {
2230 pub(crate) fn entries(&self) -> impl Iterator<Item = (&str, Id)> + '_ {
2231 self.map.iter().map(|(name, &value)| (name.as_ref(), value))
2232 }
2233
2234 pub fn get(&self, name: &str) -> Option<Id> {
2236 self.map.get(name).copied()
2237 }
2238
2239 pub fn contains(&self, name: &str) -> bool {
2241 self.map.contains_key(name)
2242 }
2243
2244 pub fn register(&mut self, name: Cow<'str, str>, id: Id, old_name: Option<&str>) -> Result<()> {
2248 if let Some(old_name) = old_name {
2249 self.forget(old_name);
2250 }
2251 match self.map.insert(name.clone(), id) {
2252 Some(_) => Err(Error::spanless(ErrorTy::DuplicateName(name.to_string()))),
2253 None => Ok(()),
2254 }
2255 }
2256
2257 pub fn forget(&mut self, name: &str) {
2261 self.map.remove(name);
2262 if let Some((base, suffix)) = split_generated_suffix(name)
2263 && let Some(hint) = self.suffix_hint.get_mut(base)
2264 {
2265 *hint = (*hint).min(suffix);
2266 }
2267 }
2268
2269 pub fn unique(&mut self, name: Cow<'str, str>) -> Cow<'str, str> {
2274 use std::fmt::Write as _;
2275
2276 if !self.map.contains_key(&name) {
2277 return name;
2278 }
2279 let base: &str = &name;
2280 let mut suffix = self.suffix_hint.get(base).copied().unwrap_or(1).max(1);
2281 let mut unique_name = format!("{base}_{suffix}");
2282 while self.map.contains_key(unique_name.as_str()) {
2283 suffix += 1;
2284 unique_name.clear();
2285 let _ = write!(unique_name, "{base}_{suffix}");
2286 }
2287 self.suffix_hint.insert(base.to_string(), suffix);
2288 Cow::Owned(unique_name)
2289 }
2290}
2291
2292fn split_generated_suffix(name: &str) -> Option<(&str, u32)> {
2297 let (base, digits) = name.rsplit_once('_')?;
2298 if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) {
2299 return None;
2300 }
2301 Some((base, digits.parse().ok()?))
2302}
2303
2304fn remap_rehomed_type(
2307 ctx: &Context<'_>,
2308 type_id: crate::types::TypeId,
2309 target: FunctionId,
2310 temp_space_map: &HashMap<TempSpaceId, TempSpaceId>,
2311) -> crate::types::TypeId {
2312 let Some(MemorySpaceId::Temp(old_space)) = ctx.shared.types.space_of(type_id) else {
2313 return type_id;
2314 };
2315 let Some(&new_space) = temp_space_map.get(&old_space) else {
2316 debug_assert_eq!(
2317 old_space.func, target,
2318 "rehome: result type references unmapped foreign temporary space {old_space:?}"
2319 );
2320 return type_id;
2321 };
2322 ctx.shared.types.get_or_make_space_address(
2323 ctx.shared.types.size_of(type_id),
2324 MemorySpaceId::Temp(new_space),
2325 )
2326}
2327
2328fn remap_rehomed_memory_space(
2331 mnemonic: &mut Mnemonic,
2332 old_func: FunctionId,
2333 target: FunctionId,
2334 temp_space_map: &HashMap<TempSpaceId, TempSpaceId>,
2335) {
2336 let remap = |space: &mut LocalMemorySpaceId| {
2337 let LocalMemorySpaceId::Temp(old_local) = *space else {
2338 return;
2339 };
2340 let old = TempSpaceId::new(old_func, old_local);
2341 if let Some(&new) = temp_space_map.get(&old) {
2342 *space = LocalMemorySpaceId::Temp(new.local);
2343 } else {
2344 debug_assert_eq!(
2345 old_func, target,
2346 "rehome: mnemonic references unmapped foreign temporary space {old:?}"
2347 );
2348 }
2349 };
2350 match mnemonic {
2351 Mnemonic::Load(load) => remap(&mut load.space),
2352 Mnemonic::Store(store) => remap(&mut store.space),
2353 _ => {}
2354 }
2355}
2356
2357fn remap_symbolic_block_literal(
2383 literals: &crate::value::interner::LiteralInterner,
2384 arg: crate::value::LocalValueId,
2385 block_map: &HashMap<BlockId, BlockId>,
2386) -> Option<crate::value::LocalValueId> {
2387 use crate::value::literal::SymbolicRef;
2388
2389 let crate::value::LocalValueId::Literal(lid) = arg else {
2390 return None;
2391 };
2392 let literal = literals[lid].clone();
2393 let Some(SymbolicRef::Block(old_block)) = literal.symbolic else {
2394 return None;
2395 };
2396 let &new_block = block_map.get(&old_block)?;
2397 let new_lit = literals.push_literal(crate::value::literal::Literal {
2398 symbolic: Some(SymbolicRef::Block(new_block)),
2399 ..literal
2400 });
2401 Some(crate::value::LocalValueId::Literal(new_lit))
2402}
2403
2404fn remap_block_targets(
2405 mnemonic: &mut Mnemonic,
2406 old_func: FunctionId,
2407 new_func: FunctionId,
2408 block_map: &HashMap<BlockId, BlockId>,
2409) {
2410 let remap = |b: &mut crate::value::LocalBlockId| {
2411 if let Some(&new) = block_map.get(&BlockId::new(old_func, *b)) {
2412 *b = new.localize(new_func);
2413 }
2414 };
2415 match mnemonic {
2416 Mnemonic::Branch(branch) => remap(&mut branch.target),
2417 Mnemonic::CBranch(cbranch) => {
2418 remap(&mut cbranch.success_block);
2419 remap(&mut cbranch.failure_block);
2420 }
2421 _ => {}
2422 }
2423}
2424
2425impl Display for Context<'_> {
2426 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2427 self.functions().try_for_each(|fun| fun.fmt(f))?;
2428
2429 self.blocks()
2430 .filter(|block| block.parent().is_none())
2431 .try_for_each(|block| block.fmt(f))
2432 }
2433}
2434
2435pub struct FunctionIter<'str, 'ctx> {
2436 ctx: &'ctx Context<'str>,
2437 inner: registry::Iter<'ctx, FunctionId, FunctionBody<'str>>,
2438}
2439
2440impl<'str, 'ctx> Iterator for FunctionIter<'str, 'ctx> {
2441 type Item = FunctionRef<'str, 'ctx>;
2442
2443 fn next(&mut self) -> Option<Self::Item> {
2444 let ctx = self.ctx;
2445 self.inner.next().map(|f| FunctionRef::from_id(ctx, f.id))
2446 }
2447}
2448
2449impl<'str, 'ctx> IntoIterator for &'ctx Context<'str> {
2450 type Item = FunctionRef<'str, 'ctx>;
2451 type IntoIter = FunctionIter<'str, 'ctx>;
2452
2453 fn into_iter(self) -> Self::IntoIter {
2454 self.iter()
2455 }
2456}
2457
2458#[cfg(test)]
2459mod tests {
2460 use super::*;
2461 use crate::value::{
2462 BasicBlock, FunctionBody, ValueId,
2463 insn::{Binary, Binop, Call, Callee, IntBinop, Load, Mnemonic},
2464 };
2465 use wazabin_qcode_macro::qcode;
2466
2467 fn make_fn_with_blocks(ctx: &mut Context<'static>, name: &'static str, n: usize) -> FunctionId {
2468 let f = FunctionBody::make(ctx, name.into()).unwrap().id;
2470 for _ in 0..n {
2471 BasicBlock::make(ctx, f);
2472 }
2473 f
2474 }
2475
2476 #[test]
2477 #[should_panic(expected = "cannot reuse a block stored in another function arena")]
2478 fn get_or_make_block_rejects_foreign_storage_at_address() {
2479 let mut ctx = Context::new();
2480 let a = FunctionBody::make(&mut ctx, "address_owner".into())
2481 .unwrap()
2482 .id;
2483 let b = FunctionBody::make(&mut ctx, "address_requester".into())
2484 .unwrap()
2485 .id;
2486 BasicBlock::make(&mut ctx, a).with_address(0x1000);
2487
2488 ctx.get_or_make_block(0x1000, b);
2489 }
2490
2491 #[test]
2492 #[should_panic(expected = "cannot create a block at an address owned by another function")]
2493 fn get_or_make_block_rejects_foreign_function_address_without_root() {
2494 let mut ctx = Context::new();
2495 FunctionBody::make_at_addr(&mut ctx, 0x1000, None);
2496 let requester = FunctionBody::make(&mut ctx, "address_requester".into())
2497 .unwrap()
2498 .id;
2499
2500 ctx.get_or_make_block(0x1000, requester);
2501 }
2502
2503 #[test]
2504 fn functions_iter_yields_all_functions() {
2505 let mut ctx = Context::new();
2506 let alpha = make_fn_with_blocks(&mut ctx, "alpha", 1);
2507 let beta = make_fn_with_blocks(&mut ctx, "beta", 1);
2508
2509 let names: Vec<_> = ctx.functions().map(|f| f.name().to_string()).collect();
2510 assert!(names.contains(&"alpha".to_string()));
2511 assert!(names.contains(&"beta".to_string()));
2512 assert_eq!(names.len(), 2);
2513 assert_eq!(ctx.function_ids(), vec![alpha, beta]);
2514 assert_eq!(ctx.function_ids().len(), ctx.interfaces.len());
2515 }
2516
2517 #[test]
2518 fn body_view_reads_match_module_reads() {
2519 use crate::value::{BodyView, FunctionId, FunctionRef, ModuleView, QCodeView};
2520
2521 let mut ctx = Context::new();
2522 qcode!(
2523 ctx,
2524 "
2525 fn foo:
2526 <bb1>
2527 if i8 1 goto <bb2> else goto <bb3>;
2528 <bb2>
2529 goto <bb3>;
2530 <bb3>
2531 return at 0;
2532 "
2533 );
2534 let fid = FunctionBody::from_name(&ctx, "foo").unwrap().id();
2535 let fid = ValueId::as_function(fid).unwrap();
2536
2537 type Snap = (String, Vec<(String, Vec<String>, Vec<String>, usize)>);
2542 fn snapshot<'a, 'str: 'a>(view: impl QCodeView<'a, 'str>, fid: FunctionId) -> Snap {
2543 let f = FunctionRef::new(view, fid);
2544 let blocks = f
2545 .blocks()
2546 .map(|b| {
2547 let name = b.name().unwrap_or("?").to_string();
2548 let mut succ: Vec<String> = b
2549 .successors()
2550 .map(|(_, s)| BlockRef::new(view, s).name().unwrap_or("?").to_string())
2551 .collect();
2552 succ.sort();
2553 let ops: Vec<String> =
2554 b.instructions().map(|i| i.opcode().to_string()).collect();
2555 (name, succ, ops, b.num_params())
2556 })
2557 .collect();
2558 (f.name().to_string(), blocks)
2559 }
2560
2561 let module_snap = snapshot(ModuleView::new(&ctx), fid);
2562 assert!(!module_snap.1.is_empty(), "sanity: foo has blocks");
2563
2564 let checked = BodyView::new(&ctx.bodies[fid], &ctx.shared, &ctx.interfaces);
2567 let checked_snap = snapshot(checked, fid);
2568 assert_eq!(
2569 module_snap, checked_snap,
2570 "reads through BodyView must match the module reads"
2571 );
2572 }
2573
2574 #[test]
2575 fn body_mut_mut_matches_module_mut() {
2576 use crate::value::{
2577 BlockParam, FunctionId, FunctionRef, InstructionId, Renameable,
2578 block::BlockId,
2579 block_param::BlockParamId,
2580 util::{base_ref::BaseRef, body_mut::BodyMut},
2581 };
2582
2583 fn build(mut ctx: &mut Context<'static>) -> (FunctionId, BlockId, BlockId, InstructionId) {
2584 qcode!(
2585 ctx,
2586 "
2587 varnode i64 x;
2588 fn foo:
2589 <entry>
2590 %a = load(x:8, &x);
2591 %b = load(x:8, &x);
2592 goto <bb1>;
2593 <bb1>
2594 return at %a;
2595 "
2596 );
2597 let fid = foo;
2598 let entry = FunctionRef::from_id(ctx, fid).root().unwrap().id;
2599 let bb1 = FunctionRef::from_id(ctx, fid)
2600 .blocks()
2601 .map(|b| b.id)
2602 .find(|&b| b != entry)
2603 .unwrap();
2604 let insns = BasicBlock::from_id(ctx, entry).instruction_ids();
2605 (fid, entry, bb1, insns[0])
2606 }
2607
2608 fn add_param(ctx: &mut Context<'static>, bb1: BlockId) -> BlockParamId {
2610 BasicBlock::from_id_mut(ctx, bb1).push_param(8).id
2611 }
2612
2613 type MSnap = Vec<(String, Option<String>, Vec<usize>, Vec<String>, Vec<String>)>;
2616 fn snap(ctx: &Context, fid: FunctionId) -> MSnap {
2617 FunctionRef::from_id(ctx, fid)
2618 .blocks()
2619 .map(|b| {
2620 let name = b.name().unwrap_or("?").to_string();
2621 let comment = b.comment().map(str::to_string);
2622 let params: Vec<usize> = b.params().map(|p| p.size()).collect();
2623 let ops: Vec<String> =
2624 b.instructions().map(|i| i.opcode().to_string()).collect();
2625 let mut succ: Vec<String> = b
2626 .successors()
2627 .map(|(_, s)| {
2628 BasicBlock::from_id(ctx, s)
2629 .name()
2630 .unwrap_or("?")
2631 .to_string()
2632 })
2633 .collect();
2634 succ.sort();
2635 (name, comment, params, ops, succ)
2636 })
2637 .collect()
2638 }
2639
2640 let mut ctx_a = Context::new();
2642 let (fid, entry, bb1, a) = build(&mut ctx_a);
2643 let param = add_param(&mut ctx_a, bb1);
2644 let b = BasicBlock::from_id(&ctx_a, entry).instruction_ids()[1];
2645 BasicBlock::from_id_mut(&mut ctx_a, entry).set_comment(Some("c".into()));
2646 BasicBlock::from_id_mut(&mut ctx_a, entry)
2647 .rename("start".into())
2648 .unwrap();
2649 let e = ctx_a.add_cfg_edge(entry, bb1);
2650 ctx_a.remove_cfg_edge(entry.func, e);
2651 ctx_a.replace_instruction(a, ValueId::Instruction(b));
2652 BlockParam::from_id_mut(&mut ctx_a, param).set_size(4);
2653 let snap_a = snap(&ctx_a, fid);
2654
2655 let mut ctx_b = Context::new();
2657 let (fid_b, entry_b, bb1_b, a_b) = build(&mut ctx_b);
2658 let param_b = add_param(&mut ctx_b, bb1_b);
2659 let b_b = BasicBlock::from_id(&ctx_b, entry_b).instruction_ids()[1];
2660
2661 {
2662 let mut host = BodyMut::new(&mut ctx_b.bodies[fid_b], &ctx_b.shared, &ctx_b.interfaces);
2663 let mut r = BaseRef::new(host.reborrow(), entry_b);
2664 r.set_comment(Some("c".into()));
2665 let mut r = BaseRef::new(host.reborrow(), entry_b);
2666 r.rename("start".into()).unwrap();
2667 let e = host.add_cfg_edge(entry_b, bb1_b);
2668 host.remove_cfg_edge(e);
2669 host.replace_instruction(a_b, ValueId::Instruction(b_b));
2670 let mut r = BaseRef::new(host.reborrow(), param_b);
2671 r.set_size(4);
2672 }
2673 let snap_b = snap(&ctx_b, fid_b);
2674
2675 assert_eq!(
2676 snap_a, snap_b,
2677 "mutations through a pass-scoped host must match the module-path mutations"
2678 );
2679 }
2680
2681 #[test]
2682 fn into_iterator_for_context_matches_functions() {
2683 let mut ctx = Context::new();
2684 make_fn_with_blocks(&mut ctx, "f1", 1);
2685 make_fn_with_blocks(&mut ctx, "f2", 1);
2686
2687 let via_method: Vec<_> = ctx.functions().map(|f| f.id()).collect();
2688 let via_into: Vec<_> = (&ctx).into_iter().map(|f| f.id()).collect();
2689 assert_eq!(via_method, via_into);
2690 }
2691
2692 #[test]
2693 fn blocks_iter_yields_all_blocks() {
2694 let mut ctx = Context::new();
2695 make_fn_with_blocks(&mut ctx, "g", 3);
2696
2697 let count = ctx.blocks().count();
2698 assert_eq!(count, 3);
2699 }
2700
2701 #[test]
2702 fn instructions_iter_yields_all_instructions() {
2703 let mut ctx = Context::new();
2704
2705 qcode!(
2706 ctx,
2707 "
2708 varnode i64 ptr;
2709
2710 <block>
2711 store(ptr:8, &ptr <- i64 0x1234);
2712 return at ptr;
2713 "
2714 );
2715
2716 let count = ctx.instructions().count();
2717 assert!(count >= 1, "expected at least one instruction, got {count}");
2718 }
2719
2720 #[test]
2721 fn move_insn_before_preserves_id_and_supports_arbitrary_anchors() {
2722 let mut ctx = Context::new();
2723 qcode!(
2724 ctx,
2725 "
2726 fn f:
2727 <source>
2728 %a = i64 0x1 + i64 0x2;
2729 %free = i64 0x5 + i64 0x6;
2730 goto <target>;
2731 <target>
2732 %b = i64 0x3 + i64 0x4;
2733 %consumer = %a + %b;
2734 return %consumer;
2735 "
2736 );
2737
2738 assert!(ctx.users(a).contains(&consumer));
2739 ctx.move_insn_before(a, b);
2740
2741 assert!(ctx.contains_instruction(a), "moving keeps the ID live");
2742 assert_eq!(ctx.get_insn(a).parent().map(|block| block.id), Some(target));
2743 assert!(
2744 !BasicBlock::from_id(&ctx, source)
2745 .instruction_ids()
2746 .contains(&a)
2747 );
2748 assert_eq!(
2749 BasicBlock::from_id(&ctx, target).instruction_ids()[..3],
2750 [a, b, consumer]
2751 );
2752 assert!(
2753 ctx.users(a).contains(&consumer),
2754 "moving preserves use-map entries"
2755 );
2756
2757 ctx.move_insn_before(b, a);
2759 assert_eq!(
2760 BasicBlock::from_id(&ctx, target).instruction_ids()[..3],
2761 [b, a, consumer]
2762 );
2763
2764 let return_id = *BasicBlock::from_id(&ctx, target)
2766 .instruction_ids()
2767 .last()
2768 .unwrap();
2769 ctx.move_insn_before(free, return_id);
2770 assert_eq!(
2771 BasicBlock::from_id(&ctx, target).instruction_ids()[..4],
2772 [b, a, consumer, free]
2773 );
2774 }
2775
2776 #[test]
2777 fn remove_instruction_removes_from_block() {
2778 let mut ctx = Context::new();
2779 qcode!(
2780 ctx,
2781 "
2782 varnode i64 x;
2783 <block>
2784 %a = load(x:8, &x);
2785 %b = load(x:8, &x);
2786 return at %a;
2787 "
2788 );
2789 let block_ref = BasicBlock::from_id(&ctx, block);
2790 let ids = block_ref.instruction_ids();
2791 let load_a = ids[0];
2792 let original_len = ids.len();
2793
2794 ctx.remove_instruction(load_a);
2795
2796 let remaining = BasicBlock::from_id(&ctx, block).instruction_ids();
2797 assert_eq!(remaining.len(), original_len - 1);
2798 assert!(!remaining.contains(&load_a));
2799 }
2800
2801 #[test]
2802 fn remove_instruction_drops_payload() {
2803 let mut ctx = Context::new();
2804 qcode!(
2805 ctx,
2806 "
2807 varnode i64 x;
2808 <block>
2809 %a = load(x:8, &x);
2810 return at %a;
2811 "
2812 );
2813 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2814
2815 ctx.remove_instruction(load_id);
2816
2817 assert!(!ctx.contains_instruction(load_id));
2818 }
2819
2820 #[test]
2821 fn remove_instruction_frees_name() {
2822 let mut ctx = Context::new();
2823 qcode!(
2824 ctx,
2825 "
2826 varnode i64 x;
2827 <block>
2828 %a = load(x:8, &x);
2829 return at %a;
2830 "
2831 );
2832 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2833 assert!(
2835 ctx.get_named_in_scope(load_id.into(), "a").is_some(),
2836 "name should be in map before removal"
2837 );
2838
2839 ctx.remove_instruction(load_id);
2840
2841 assert!(
2842 ctx.get_named_in_scope(load_id.into(), "a").is_none(),
2843 "name should be gone after removal"
2844 );
2845 assert!(!ctx.contains_instruction(load_id));
2846 }
2847
2848 #[test]
2849 fn remove_instruction_frees_name_for_reuse() {
2850 let mut ctx = Context::new();
2851 qcode!(
2852 ctx,
2853 "
2854 varnode i64 x;
2855 <block>
2856 %a = load(x:8, &x);
2857 return at %a;
2858 "
2859 );
2860 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2861
2862 ctx.remove_instruction(load_id);
2863
2864 qcode!(
2866 ctx,
2867 "
2868 varnode i64 y;
2869 <block2>
2870 %a = load(y:8, &y);
2871 return at %a;
2872 "
2873 );
2874 let a2 = BasicBlock::from_id(&ctx, block2).instruction_ids()[0];
2875 assert!(
2876 ctx.get_named_in_scope(a2.into(), "a").is_some(),
2877 "name should be reusable after removal"
2878 );
2879 }
2880
2881 #[test]
2882 fn remove_instruction_updates_users_map() {
2883 let mut ctx = Context::new();
2884 qcode!(
2885 ctx,
2886 "
2887 varnode i64 x;
2888 <block>
2889 %a = load(x:8, &x);
2890 %b = %a + i64 1;
2891 return at %b;
2892 "
2893 );
2894 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
2895 let load_id = ids[0];
2896 let add_id = ids[1];
2897
2898 assert!(
2899 ctx.users(load_id).contains(&add_id),
2900 "add should be a user of load before removal"
2901 );
2902
2903 ctx.remove_instruction(add_id);
2904
2905 assert!(
2906 ctx.users(load_id).is_empty(),
2907 "load should have no users after add is removed"
2908 );
2909 }
2910
2911 #[test]
2912 fn removed_instruction_is_absent_and_not_iterated() {
2913 let mut ctx = Context::new();
2918 qcode!(
2919 ctx,
2920 "
2921 varnode i64 x;
2922 <block>
2923 %a = load(x:8, &x);
2924 %dead = %a + i64 1;
2925 return at i64 0;
2926 "
2927 );
2928 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
2929 let dead_id = ids[1]; assert!(
2932 ctx.instructions().any(|i| i.id == dead_id),
2933 "the instruction is iterated while live"
2934 );
2935
2936 ctx.remove_instruction(dead_id);
2937
2938 assert!(!ctx.contains_instruction(dead_id));
2939 assert!(
2940 !ctx.instructions().any(|i| i.id == dead_id),
2941 "a deleted instruction must not be yielded by ctx.instructions()"
2942 );
2943 }
2944
2945 #[test]
2946 fn replace_instruction_mnemonic_rewrites_callind_users() {
2947 let mut ctx = Context::new();
2948 qcode!(
2949 ctx,
2950 "
2951 varnode i64 ptr;
2952 <block>
2953 call [ptr];
2954 "
2955 );
2956 let call_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2957 let ptr = match ctx.get_insn(call_id).mnemonic() {
2958 Mnemonic::CallInd(call) => call.ptr.qualify(call_id.func),
2959 other => panic!("expected CallInd, got {other:?}"),
2960 };
2961 assert_eq!(ctx.users_across_functions(ptr), vec![call_id]);
2963
2964 let target = FunctionBody::make(&mut ctx, "target".into()).unwrap().id;
2965 ctx.replace_instruction_mnemonic(
2966 call_id,
2967 Mnemonic::Call(Call {
2968 target: Callee::Real(target),
2969 args: vec![],
2970 clobbers: vec![],
2971 tag: Default::default(),
2972 }),
2973 );
2974
2975 assert!(
2976 ctx.users_across_functions(ptr).is_empty(),
2977 "old indirect pointer should no longer list the rewritten call"
2978 );
2979 assert!(matches!(
2980 ctx.get_insn(call_id).mnemonic(),
2981 Mnemonic::Call(Call {
2982 target: actual,
2983 args,
2984 ..
2985 }) if *actual == Callee::Real(target) && args.is_empty()
2986 ));
2987 }
2988
2989 #[test]
2990 fn users_across_functions_keeps_ssa_users_in_the_owning_function() {
2991 let mut ctx = Context::new();
2992 qcode!(
2993 ctx,
2994 "
2995 fn f:
2996 <f_entry>
2997 %fx = i64 1 + i64 2;
2998 %fuse = %fx + i64 3;
2999 return at %fuse;
3000 fn g:
3001 <g_entry>
3002 %gx = i64 4 + i64 5;
3003 %guse = %gx + i64 6;
3004 return at %guse;
3005 "
3006 );
3007 let f_ids = BasicBlock::from_id(&ctx, f_entry).instruction_ids();
3008 let g_ids = BasicBlock::from_id(&ctx, g_entry).instruction_ids();
3009 assert_eq!(
3010 f_ids[0].local, g_ids[0].local,
3011 "precondition: arena-local ids collide"
3012 );
3013 assert_eq!(
3014 ctx.users_across_functions(ValueId::Instruction(f_ids[0])),
3015 vec![f_ids[1]],
3016 "an SSA query must not pick up the same local key from another function"
3017 );
3018 }
3019
3020 #[test]
3021 fn replace_instruction_mnemonic_moves_operand_users() {
3022 let mut ctx = Context::new();
3023 qcode!(
3024 ctx,
3025 "
3026 varnode i64 x;
3027 varnode i64 y;
3028 <block>
3029 %a = load(x:8, x);
3030 return at %a;
3031 "
3032 );
3033 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3034 let old_ptr = ValueId::Varnode(x);
3035 let new_ptr = ValueId::Varnode(y);
3036 assert_eq!(ctx.users_across_functions(old_ptr), vec![load_id]);
3038 assert!(ctx.users_across_functions(new_ptr).is_empty());
3039
3040 ctx.replace_instruction_mnemonic(
3041 load_id,
3042 Mnemonic::Load(Load {
3043 space: ctx.shared.default_space.into(),
3044 ptr: new_ptr.localize(load_id.func),
3045 size: 8,
3046 }),
3047 );
3048
3049 assert!(ctx.users_across_functions(old_ptr).is_empty());
3050 assert_eq!(ctx.users_across_functions(new_ptr), vec![load_id]);
3051 }
3052
3053 #[test]
3054 fn replace_instruction_mnemonic_tracks_repeated_operands() {
3055 let mut ctx = Context::new();
3056 qcode!(
3057 ctx,
3058 "
3059 varnode i64 x;
3060 varnode i64 y;
3061 <block>
3062 %a = load(x:8, x);
3063 return at %a;
3064 "
3065 );
3066 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3067 let old_ptr = ValueId::Varnode(x);
3068 let new_arg = ValueId::Varnode(y);
3069
3070 ctx.replace_instruction_mnemonic(
3071 load_id,
3072 Mnemonic::Binop(Binary {
3073 op: Binop::Int(IntBinop::Add),
3074 lhs: new_arg.localize(load_id.func),
3075 rhs: new_arg.localize(load_id.func),
3076 }),
3077 );
3078
3079 assert!(ctx.users_across_functions(old_ptr).is_empty());
3080 assert_eq!(
3081 ctx.users_across_functions(new_arg),
3082 vec![load_id, load_id],
3083 "a mnemonic using the same operand twice should record both uses"
3084 );
3085 }
3086
3087 #[test]
3088 fn remove_instruction_unparented_noop() {
3089 let mut ctx = Context::new();
3090 qcode!(
3091 ctx,
3092 "
3093 varnode i64 x;
3094 <block>
3095 %a = load(x:8, &x);
3096 return at %a;
3097 "
3098 );
3099 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3100
3101 ctx.instruction_mut(load_id).parent = None;
3104
3105 ctx.remove_instruction(load_id);
3107
3108 assert!(ctx.get_named("a").is_none());
3109 }
3110
3111 #[test]
3112 fn add_cfg_edge_returns_id_and_remove_unlinks_both_blocks() {
3113 let mut ctx = Context::new();
3114 let f = ctx.anon_function();
3116 let a = BasicBlock::make(&mut ctx, f).id;
3117 let b = BasicBlock::make(&mut ctx, f).id;
3118 let c = BasicBlock::make(&mut ctx, f).id;
3119
3120 let edge = ctx.add_cfg_edge(a, b);
3121 let surviving_edge = ctx.add_cfg_edge(b, c);
3122 assert_eq!(
3123 BasicBlock::from_id(&ctx, a)
3124 .successors()
3125 .collect::<Vec<_>>(),
3126 vec![(edge, b)]
3127 );
3128 assert_eq!(
3129 BasicBlock::from_id(&ctx, b)
3130 .predecessors()
3131 .collect::<Vec<_>>(),
3132 vec![(edge, a)]
3133 );
3134
3135 ctx.remove_cfg_edge(a.func, edge);
3136 assert!(BasicBlock::from_id(&ctx, a).successors().next().is_none());
3137 assert!(BasicBlock::from_id(&ctx, b).predecessors().next().is_none());
3138 assert!(!ctx.bodies[a.func].edges.contains(edge));
3139 let surviving = ctx.edge(a.func, surviving_edge);
3140 assert_eq!(
3141 surviving.from, b.local,
3142 "swap removal must preserve the source"
3143 );
3144 assert_eq!(
3145 surviving.to, c.local,
3146 "swap removal must preserve the target"
3147 );
3148 assert_eq!(ctx.bodies[a.func].edges.len(), 1);
3149
3150 let self_edge = ctx.add_cfg_edge(a, a);
3151 ctx.remove_cfg_edge(a.func, self_edge);
3152 assert!(!ctx.bodies[a.func].edges.contains(self_edge));
3153 assert!(ctx.block(a).edges.is_empty());
3154
3155 let parallel_a = ctx.add_cfg_edge(a, b);
3156 let parallel_b = ctx.add_cfg_edge(a, b);
3157 ctx.remove_cfg_edge(a.func, parallel_a);
3158 assert!(!ctx.bodies[a.func].edges.contains(parallel_a));
3159 assert!(ctx.bodies[a.func].edges.contains(parallel_b));
3160 assert_eq!(
3161 BasicBlock::from_id(&ctx, a)
3162 .successors()
3163 .collect::<Vec<_>>(),
3164 vec![(parallel_b, b)],
3165 );
3166 }
3167
3168 #[test]
3169 fn truth_map_tracks_four_states_and_conflicts() {
3170 let mut ctx = Context::new();
3171 let callee = FunctionBody::make(&mut ctx, "callee".into()).unwrap().id;
3172 let prop = Proposition::FunctionReturns(callee);
3173
3174 assert!(ctx.assume_true(prop));
3176 assert!(ctx.assume_true(prop));
3177 assert!(!ctx.assume_false(prop));
3178 assert_eq!(ctx.known(prop), None, "assumed is not known");
3179
3180 let snapshot = ctx.clone();
3182
3183 let scope = pass_scope::enter("verifier");
3186 assert!(ctx.set_known(prop, false), "overturning is novel");
3187 drop(scope);
3188 assert_eq!(ctx.known(prop), Some(false));
3189 let [v] = ctx.violations() else {
3190 panic!("expected one violation")
3191 };
3192 assert_eq!(v.prop, prop);
3193 assert!(v.assumed);
3194 assert_eq!(v.asserting_pass, "verifier");
3195
3196 assert!(!ctx.set_known(prop, false));
3198
3199 assert!(snapshot.violations().is_empty());
3201 assert_eq!(snapshot.known(prop), None);
3202
3203 assert!(!ctx.assume_true(prop));
3205 assert!(ctx.assume_false(prop));
3206 }
3207
3208 #[test]
3209 fn seeded_facts_are_not_novel() {
3210 let mut ctx = Context::new();
3211 let callee = FunctionBody::make(&mut ctx, "exit".into()).unwrap().id;
3212 let prop = Proposition::FunctionReturns(callee);
3213
3214 ctx.seed_known(prop, false, PassName("seed"));
3215 assert_eq!(ctx.known(prop), Some(false));
3216 assert!(!ctx.assume_true(prop), "seeded fact blocks opposite assume");
3217 assert!(
3218 !ctx.set_known(prop, false),
3219 "re-proving a seed is not novel"
3220 );
3221 assert!(ctx.violations().is_empty());
3222 }
3223
3224 #[test]
3225 fn discovered_code_records_and_survives_round_trip() {
3226 let mut ctx = Context::new();
3227 ctx.discover_code(0x1000, 0x10f0, 0x1100);
3228 ctx.discover_code(0x1000, 0x10f0, 0x1200);
3229 ctx.discover_code(0x1000, 0x10f0, 0x1100); let targets: Vec<u64> = ctx.discoveries().map(|d| d.target).collect();
3232 assert_eq!(targets, vec![0x1100, 0x1200]);
3233
3234 let config = bincode::config::standard();
3235 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3236 let (restored, _): (Context<'static>, usize) =
3237 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3238 assert_eq!(
3239 restored.discoveries().map(|d| d.target).collect::<Vec<_>>(),
3240 targets
3241 );
3242 }
3243
3244 #[test]
3245 fn assume_executable_narrows_once_protections_known() {
3246 let mut ctx = Context::new();
3247 let mut image = crate::memory_image::MemoryImage::default();
3248 image.add_segment(0x1000, vec![0u8; 4], true, false); image.add_segment(0x2000, vec![0u8; 4], false, true); let binary: &dyn wazabin_binary::BinaryFormat = ℑ
3251
3252 assert!(ctx.assume_executable(binary, 0x1000));
3255 assert!(ctx.assume_executable(binary, 0x2000));
3256 assert!(ctx.assume_executable(binary, 0x9999));
3257
3258 ctx.mark_protections_known();
3259 assert!(
3260 ctx.assume_executable(binary, 0x1000),
3261 "code region stays liftable"
3262 );
3263 assert!(
3264 !ctx.assume_executable(binary, 0x2000),
3265 "data region is skipped once protections are known"
3266 );
3267 assert!(
3268 !ctx.assume_executable(binary, 0x9999),
3269 "unmapped is skipped once known"
3270 );
3271 assert_eq!(
3273 ctx.known(Proposition::ExecutableMemory {
3274 start: 0x2000,
3275 end: 0x2004,
3276 }),
3277 Some(false),
3278 );
3279 }
3280
3281 #[test]
3282 fn assume_executable_honors_region_override() {
3283 let mut ctx = Context::new();
3284 let mut image = crate::memory_image::MemoryImage::default();
3285 image.add_segment(0x1000, vec![0u8; 4], true, false); image.add_segment(0x2000, vec![0u8; 4], false, true); let binary: &dyn wazabin_binary::BinaryFormat = ℑ
3288 ctx.mark_protections_known();
3289
3290 ctx.seed_known(
3292 Proposition::ExecutableMemory {
3293 start: 0x2000,
3294 end: 0x2004,
3295 },
3296 true,
3297 PassName("override"),
3298 );
3299 ctx.seed_known(
3300 Proposition::ExecutableMemory {
3301 start: 0x1000,
3302 end: 0x1004,
3303 },
3304 false,
3305 PassName("override"),
3306 );
3307
3308 assert!(
3309 ctx.assume_executable(binary, 0x2000),
3310 "override wins over the non-executable segment flag"
3311 );
3312 assert!(
3313 !ctx.assume_executable(binary, 0x1000),
3314 "override wins over the executable segment flag"
3315 );
3316 }
3317
3318 #[test]
3319 fn context_survives_bincode_round_trip() {
3320 let mut ctx = Context::new();
3321 qcode!(
3322 ctx,
3323 "
3324 varnode i64 ptr;
3325 <block>
3326 %a = load(ptr:8, &ptr);
3327 %b = %a + i64 0x10;
3328 store(ptr:8, &ptr <- i64 0x1234);
3329 return at %b;
3330 "
3331 );
3332
3333 let some_space = ctx.get_or_make_named_space("scratch");
3335 let sa = ctx.shared.types.get_or_make_space_address(8, some_space);
3336 let sa_size = ctx.shared.types.size_of(sa);
3337
3338 let blocks_before = ctx.block_ids().len();
3339 let insns_before = ctx.instruction_ids().len();
3340 let funcs_before = ctx.function_ids().len();
3341
3342 let config = bincode::config::standard();
3343 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3344 let (restored, _): (Context<'static>, usize) =
3345 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3346
3347 assert_eq!(restored.block_ids().len(), blocks_before);
3348 assert_eq!(restored.instruction_ids().len(), insns_before);
3349 assert_eq!(restored.function_ids().len(), funcs_before);
3350 for function_id in restored.function_ids() {
3351 assert_eq!(restored.bodies[function_id].id(), function_id);
3352 }
3353 assert_eq!(restored.shared.types.size_of(sa), sa_size);
3355 assert_eq!(
3356 restored.shared.types.space_of(sa),
3357 Some(crate::space::MemorySpaceId::Shared(some_space))
3358 );
3359 }
3360
3361 #[test]
3362 fn compact_edge_arena_preserves_ids_across_round_trip() {
3363 let mut ctx = Context::new();
3364 let function = ctx.anon_function();
3365 let a = BasicBlock::make(&mut ctx, function).id;
3366 let b = BasicBlock::make(&mut ctx, function).id;
3367 let c = BasicBlock::make(&mut ctx, function).id;
3368 let d = BasicBlock::make(&mut ctx, function).id;
3369 let first = ctx.add_cfg_edge(a, b);
3370 let removed = ctx.add_cfg_edge(b, c);
3371 let last = ctx.add_cfg_edge(c, d);
3372 ctx.remove_cfg_edge(function, removed);
3373
3374 let physical_order: Vec<_> = ctx.bodies[function]
3375 .edges
3376 .iter()
3377 .map(|edge| edge.id)
3378 .collect();
3379 assert_eq!(physical_order, vec![first, last]);
3380
3381 let config = bincode::config::standard();
3382 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3383 let (mut restored, _): (Context<'static>, usize) =
3384 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3385
3386 assert!(!restored.bodies[function].edges.contains(removed));
3387 assert_eq!(
3388 restored.bodies[function]
3389 .edges
3390 .iter()
3391 .map(|edge| edge.id)
3392 .collect::<Vec<_>>(),
3393 physical_order,
3394 );
3395 assert_eq!(restored.edge(function, first).to, b.local);
3396 assert_eq!(restored.edge(function, last).from, c.local);
3397
3398 let fresh = restored.add_cfg_edge(a, d);
3399 assert!(fresh > last);
3400 assert_ne!(fresh, removed, "removed edge IDs must never be reused");
3401 }
3402
3403 #[test]
3404 fn compact_instruction_arena_preserves_ids_across_round_trip() {
3405 let mut ctx = Context::new();
3406 qcode!(
3407 ctx,
3408 "
3409 <block>
3410 %first = i64 1 + i64 2;
3411 %removed = i64 3 + i64 4;
3412 return at %first;
3413 "
3414 );
3415 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
3416 let first = ids[0];
3417 let removed = ids[1];
3418 let last = ids[2];
3419 ctx.remove_instruction(removed);
3420
3421 let physical_order: Vec<_> = ctx.bodies[first.func]
3422 .insns
3423 .iter()
3424 .map(|insn| insn.id)
3425 .collect();
3426 assert_eq!(physical_order, vec![first.local, last.local]);
3427
3428 let config = bincode::config::standard();
3429 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3430 let (mut restored, _): (Context<'static>, usize) =
3431 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3432
3433 assert!(!restored.contains_instruction(removed));
3434 assert_eq!(
3435 restored.bodies[first.func]
3436 .insns
3437 .iter()
3438 .map(|insn| insn.id)
3439 .collect::<Vec<_>>(),
3440 physical_order,
3441 );
3442 assert!(restored.contains_instruction(first));
3443 assert!(restored.contains_instruction(last));
3444
3445 let template = restored.instruction(last).clone();
3446 let fresh = restored.push_insn(first.func, template);
3447 assert!(fresh.local > last.local);
3448 assert_ne!(
3449 fresh, removed,
3450 "removed instruction IDs must never be reused"
3451 );
3452 }
3453
3454 #[test]
3455 fn compact_param_arena_preserves_ids_across_round_trip() {
3456 let mut ctx = Context::new();
3457 let function = ctx.anon_function();
3458 let block = BasicBlock::make(&mut ctx, function).id;
3459 let first = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3460 let removed = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3461 let last = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3462
3463 ctx.block_mut(block).params.remove(1);
3464 ctx.block_param_mut(last).index = 1;
3465 ctx.remove_block_param(removed);
3466
3467 let physical_order: Vec<_> = ctx.bodies[function]
3468 .params
3469 .iter()
3470 .map(|param| param.id)
3471 .collect();
3472 assert_eq!(physical_order, vec![first.local, last.local]);
3473 assert_eq!(ctx.block_param(first).index, 0);
3474 assert_eq!(ctx.block_param(last).index, 1);
3475
3476 let config = bincode::config::standard();
3477 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3478 let (mut restored, _): (Context<'static>, usize) =
3479 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3480
3481 assert!(!restored.contains_block_param(removed));
3482 assert_eq!(
3483 restored.bodies[function]
3484 .params
3485 .iter()
3486 .map(|param| param.id)
3487 .collect::<Vec<_>>(),
3488 physical_order,
3489 );
3490 assert!(restored.contains_block_param(first));
3491 assert!(restored.contains_block_param(last));
3492
3493 let fresh = BasicBlock::from_id_mut(&mut restored, block)
3494 .push_param(8)
3495 .id;
3496 assert!(fresh.local > last.local);
3497 assert_ne!(fresh, removed, "removed parameter IDs must never be reused");
3498 }
3499
3500 #[test]
3501 fn compact_block_arena_preserves_ids_across_round_trip() {
3502 let mut ctx = Context::new();
3503 let function = ctx.anon_function();
3504 let first = BasicBlock::make(&mut ctx, function).id;
3505 let removed = BasicBlock::make(&mut ctx, function).id;
3506 let last = BasicBlock::make(&mut ctx, function).id;
3507 FunctionBody::from_id_mut(&mut ctx, function)
3508 .set_root(first)
3509 .expect("set root");
3510
3511 ctx.delete_block(removed);
3512
3513 let physical_order: Vec<_> = ctx.bodies[function]
3514 .blocks
3515 .iter()
3516 .map(|block| block.id)
3517 .collect();
3518 assert_eq!(physical_order, vec![first.local, last.local]);
3519 assert_eq!(ctx.block_ids(), vec![first, last]);
3520
3521 let config = bincode::config::standard();
3522 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3523 let (mut restored, _): (Context<'static>, usize) =
3524 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3525
3526 assert!(!restored.contains_block(removed));
3527 assert_eq!(
3528 restored.bodies[function]
3529 .blocks
3530 .iter()
3531 .map(|block| block.id)
3532 .collect::<Vec<_>>(),
3533 physical_order,
3534 );
3535 assert!(restored.contains_block(first));
3536 assert!(restored.contains_block(last));
3537 assert_eq!(
3538 FunctionBody::from_id(&restored, function)
3539 .root()
3540 .map(|block| block.id),
3541 Some(first),
3542 );
3543
3544 let fresh = BasicBlock::make(&mut restored, function).id;
3545 assert!(fresh.local > last.local);
3546 assert_ne!(fresh, removed, "removed block IDs must never be reused");
3547 }
3548
3549 #[test]
3550 fn deleting_root_clears_function_root() {
3551 let mut ctx = Context::new();
3552 let function = ctx.anon_function();
3553 let root = BasicBlock::make(&mut ctx, function).id;
3554 FunctionBody::from_id_mut(&mut ctx, function)
3555 .set_root(root)
3556 .expect("set root");
3557
3558 ctx.delete_block(root);
3559
3560 assert!(!ctx.contains_block(root));
3561 assert!(FunctionBody::from_id(&ctx, function).root().is_none());
3562 assert!(ctx.block_ids().is_empty());
3563 }
3564
3565 #[test]
3566 fn get_unique_name_resumes_probe_and_reuses_freed_suffixes() {
3567 use crate::value::VarnodeId;
3568
3569 let mut ctx = Context::new();
3570 let id = ValueId::Varnode(VarnodeId::from(0usize));
3571
3572 fn take(ctx: &mut Context<'static>, id: ValueId, base: &str) -> String {
3574 let name = ctx
3575 .get_unique_name(Cow::Owned(base.to_string()))
3576 .to_string();
3577 ctx.update_name(Cow::Owned(name.clone()), id, None).unwrap();
3578 name
3579 }
3580
3581 assert_eq!(take(&mut ctx, id, "tmp"), "tmp");
3583 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_1");
3584 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_2");
3585 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_3");
3586
3587 assert_eq!(take(&mut ctx, id, "x"), "x");
3589 assert_eq!(take(&mut ctx, id, "x"), "x_1");
3590
3591 ctx.update_name(Cow::Borrowed("relocated"), id, Some("tmp_1"))
3594 .unwrap();
3595 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_1");
3596 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_4");
3598 }
3599
3600 mod split_function_at {
3603 use super::*;
3604
3605 use crate::value::insn::{Callee, Mnemonic, TailCall};
3606 use crate::value::{BasicBlock, FunctionBody, Instruction, Value};
3607 use std::borrow::Cow;
3608
3609 fn block_at(ctx: &mut Context<'static>, func: FunctionId, addr: u64) -> BlockId {
3610 BasicBlock::make(ctx, func).with_address(addr).id
3611 }
3612
3613 fn branch_at(ctx: &mut Context<'static>, block: BlockId, target: BlockId, addr: u64) {
3614 let id = (ctx).builder(block).push_branch(target).id;
3615 Instruction::from_id_mut(ctx, id).set_address(addr);
3616 }
3617
3618 fn cbranch_at(
3619 ctx: &mut Context<'static>,
3620 block: BlockId,
3621 success: BlockId,
3622 failure: BlockId,
3623 addr: u64,
3624 ) {
3625 let cond = ctx.get_const(1, 1).id();
3626 let id = (ctx).builder(block).push_cbranch(cond, success, failure).id;
3627 Instruction::from_id_mut(ctx, id).set_address(addr);
3628 }
3629
3630 fn return_at(ctx: &mut Context<'static>, block: BlockId, addr: u64) {
3631 let zero = ctx.get_const(0, 8).id();
3632 let id = (ctx).builder(block).push_return(zero).id;
3633 Instruction::from_id_mut(ctx, id).set_address(addr);
3634 }
3635
3636 fn block_at_addr(ctx: &Context, func: FunctionId, addr: u64) -> BlockId {
3637 FunctionBody::from_id(ctx, func)
3638 .block_ids()
3639 .into_iter()
3640 .find(|b| ctx.block(*b).address == Some(addr))
3641 .unwrap_or_else(|| panic!("{func:?} has no block at {addr:#x}"))
3642 }
3643
3644 fn addrs(ctx: &Context, func: FunctionId) -> Vec<u64> {
3645 let mut got: Vec<u64> = FunctionBody::from_id(ctx, func)
3646 .block_ids()
3647 .into_iter()
3648 .filter_map(|b| ctx.block(b).address)
3649 .collect();
3650 got.sort_unstable();
3651 got
3652 }
3653
3654 #[test]
3659 fn splits_absorbed_body_reusing_the_stub() {
3660 let mut ctx = Context::new();
3661 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("thunk"))).id;
3662 let b0 = block_at(&mut ctx, f, 0x1000);
3663 let b1 = block_at(&mut ctx, f, 0x2000);
3664 let b2 = block_at(&mut ctx, f, 0x2005);
3665 branch_at(&mut ctx, b0, b1, 0x1000);
3666 branch_at(&mut ctx, b1, b2, 0x2000);
3667 return_at(&mut ctx, b2, 0x2005);
3668 {
3669 let mut func = FunctionBody::from_id_mut(&mut ctx, f);
3670 func.set_root(b0).unwrap();
3671 }
3672 let g = FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("real"))).id;
3674
3675 let split_g = ctx.split_function_at(b1);
3676 assert_eq!(
3677 split_g, g,
3678 "the split must reuse the existing stub at 0x2000"
3679 );
3680
3681 assert_eq!(addrs(&ctx, f), vec![0x1000]);
3682 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2005]);
3683 let g_entry = block_at_addr(&ctx, g, 0x2000);
3684 assert_eq!(ctx.bodies[g].root_id(), Some(g_entry.local));
3685
3686 for b in FunctionBody::from_id(&ctx, g).block_ids() {
3688 assert_eq!(b.func, g);
3689 }
3690
3691 let f_entry = block_at_addr(&ctx, f, 0x1000);
3693 assert_eq!(BasicBlock::from_id(&ctx, f_entry).successors().count(), 0);
3694 let term = BasicBlock::from_id(&ctx, f_entry)
3695 .instructions()
3696 .last()
3697 .map(|i| i.mnemonic().clone());
3698 assert!(
3699 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
3700 "thunk branch must become TailCall(G), got {term:?}",
3701 );
3702 }
3703
3704 #[test]
3705 fn split_rehomes_temporary_values_spaces_and_pointer_types() {
3706 let mut ctx = Context::new();
3707 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3708 let entry = block_at(&mut ctx, f, 0x1000);
3709 let tail = block_at(&mut ctx, f, 0x2000);
3710 branch_at(&mut ctx, entry, tail, 0x1000);
3711 FunctionBody::from_id_mut(&mut ctx, f)
3712 .set_root(entry)
3713 .unwrap();
3714
3715 let temp = ctx
3716 .builder(tail)
3717 .make_named_temp(Cow::Borrowed("scratch"), 8);
3718 ctx.builder(entry)
3719 .make_named_temp(Cow::Borrowed("unused"), 4);
3720 let temp_space = ctx.bodies[f].temps[temp.local].space;
3721 let load = {
3722 let mut builder = ctx.builder(tail);
3723 let ValueId::Instruction(load) = builder
3724 .push_load::<false>(
3725 ValueId::Temp(temp),
3726 8,
3727 LocalMemorySpaceId::Temp(temp_space),
3728 )
3729 .id()
3730 else {
3731 unreachable!()
3732 };
3733 builder.push_return(ValueId::Instruction(load));
3734 load
3735 };
3736 let pointer_type = ctx
3737 .shared
3738 .types
3739 .get_or_make_space_address(8, MemorySpaceId::Temp(TempSpaceId::new(f, temp_space)));
3740 ctx.instruction_mut(load).type_id = pointer_type;
3741
3742 let g =
3743 FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("discovered"))).id;
3744 assert_eq!(ctx.split_function_at(tail), g);
3745
3746 let diagnostics = crate::verify_body_arena_integrity(&ctx);
3747 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
3748 assert_eq!(ctx.bodies[g].temp_spaces.len(), 1);
3749 assert_eq!(ctx.bodies[g].temps.len(), 1);
3750 assert_eq!(ctx.bodies[f].temps.len(), 2, "source arenas remain intact");
3751
3752 let moved_load = FunctionBody::from_id(&ctx, g)
3753 .blocks()
3754 .flat_map(|block| block.instructions())
3755 .find(|insn| matches!(insn.mnemonic(), Mnemonic::Load(_)))
3756 .expect("load moved with the split");
3757 let Mnemonic::Load(moved) = moved_load.mnemonic() else {
3758 unreachable!()
3759 };
3760 let LocalMemorySpaceId::Temp(moved_space) = moved.space else {
3761 panic!("load lost temporary-space provenance")
3762 };
3763 assert!(matches!(moved.ptr, crate::value::LocalValueId::Temp(_)));
3764 assert!(usize::from(moved_space) < ctx.bodies[g].temp_spaces.len());
3765 assert_eq!(
3766 ctx.shared.types.space_of(moved_load.type_id()),
3767 Some(MemorySpaceId::Temp(TempSpaceId::new(g, moved_space)))
3768 );
3769
3770 let rendered = FunctionBody::from_id(&ctx, g).to_string();
3772 assert!(rendered.contains("scratch"));
3773 }
3774
3775 #[test]
3776 fn split_stops_at_a_foreign_rootless_stub_address() {
3777 let mut ctx = Context::new();
3778 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3779 let entry = block_at(&mut ctx, f, 0x1000);
3780 let split = block_at(&mut ctx, f, 0x2000);
3781 let foreign_entry = block_at(&mut ctx, f, 0x3000);
3782 let foreign_body = block_at(&mut ctx, f, 0x3005);
3783 branch_at(&mut ctx, entry, split, 0x1000);
3784 branch_at(&mut ctx, split, foreign_entry, 0x2000);
3785 branch_at(&mut ctx, foreign_entry, foreign_body, 0x3000);
3786 return_at(&mut ctx, foreign_body, 0x3005);
3787 FunctionBody::from_id_mut(&mut ctx, f)
3788 .set_root(entry)
3789 .unwrap();
3790
3791 let g = FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("g"))).id;
3792 let h = FunctionBody::make_at_addr(&mut ctx, 0x3000, Some(Cow::Borrowed("h"))).id;
3793 assert!(FunctionBody::from_id(&ctx, g).root().is_none());
3794 assert!(FunctionBody::from_id(&ctx, h).root().is_none());
3795
3796 assert_eq!(ctx.split_function_at(split), g);
3797 assert_eq!(addrs(&ctx, g), vec![0x2000]);
3798 assert_eq!(addrs(&ctx, f), vec![0x1000, 0x3000, 0x3005]);
3799 assert!(FunctionBody::from_id(&ctx, h).root().is_none());
3800
3801 let g_entry = block_at_addr(&ctx, g, 0x2000);
3802 let term = BasicBlock::from_id(&ctx, g_entry)
3803 .instructions()
3804 .last()
3805 .map(|i| i.mnemonic().clone());
3806 assert!(
3807 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(h)),
3808 "split tail must stop and tail-call rootless stub H, got {term:?}",
3809 );
3810 }
3811
3812 #[test]
3813 fn split_rehomes_block_param_origin_into_destination_arena() {
3814 let mut ctx = Context::new();
3815 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3816 let entry = block_at(&mut ctx, f, 0x1000);
3817 let tail = block_at(&mut ctx, f, 0x2000);
3818 let param = BasicBlock::from_id_mut(&mut ctx, tail).push_param(8).id;
3819 crate::value::BlockParam::from_id_mut(&mut ctx, param)
3820 .set_origin(ValueId::BlockParam(param));
3821
3822 let arg = ctx.get_const(7, 8).id();
3823 let branch = ctx.builder(entry).push_branch_with_args(tail, vec![arg]).id;
3824 Instruction::from_id_mut(&mut ctx, branch).set_address(0x1000);
3825 let ret = ctx.builder(tail).push_return(ValueId::BlockParam(param)).id;
3826 Instruction::from_id_mut(&mut ctx, ret).set_address(0x2000);
3827 FunctionBody::from_id_mut(&mut ctx, f)
3828 .set_root(entry)
3829 .unwrap();
3830
3831 let g = ctx.split_function_at(tail);
3832 let new_tail = block_at_addr(&ctx, g, 0x2000);
3833 let new_param = BasicBlock::from_id(&ctx, new_tail).params().next().unwrap();
3834 assert_eq!(new_param.origin(), Some(ValueId::BlockParam(new_param.id)));
3835 }
3836
3837 #[test]
3845 fn split_rehomes_symbolic_block_literals() {
3846 use crate::value::literal::SymbolicRef;
3847
3848 let mut ctx = Context::new();
3849 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3850 let entry = block_at(&mut ctx, f, 0x1000);
3851 let tail = block_at(&mut ctx, f, 0x2000);
3852 let landing = block_at(&mut ctx, f, 0x2008);
3853
3854 let lit = ctx.get_const(0x2008, 8).id();
3857 let ValueId::Literal(lit_id) = lit else {
3858 panic!("expected a literal");
3859 };
3860 ctx.shared.values.literals[lit_id].symbolic = Some(SymbolicRef::Block(landing));
3861
3862 branch_at(&mut ctx, entry, tail, 0x1000);
3863 let ind = ctx.builder(tail).push_branchind(lit).id;
3865 Instruction::from_id_mut(&mut ctx, ind).set_address(0x2000);
3866 ctx.add_cfg_edge(tail, landing);
3867 return_at(&mut ctx, landing, 0x2008);
3868 FunctionBody::from_id_mut(&mut ctx, f)
3869 .set_root(entry)
3870 .unwrap();
3871
3872 let g = ctx.split_function_at(tail);
3873
3874 let new_landing = block_at_addr(&ctx, g, 0x2008);
3875 let new_tail = block_at_addr(&ctx, g, 0x2000);
3876 let Mnemonic::BranchInd(b) = BasicBlock::from_id(&ctx, new_tail)
3877 .instructions()
3878 .last()
3879 .unwrap()
3880 .mnemonic()
3881 .clone()
3882 else {
3883 panic!("tail must still end in an indirect branch");
3884 };
3885 let crate::value::LocalValueId::Literal(new_lit) = b.ptr else {
3886 panic!("indirect branch operand must still be a literal");
3887 };
3888 assert_eq!(
3889 ctx.shared.values.literals[new_lit].symbolic,
3890 Some(SymbolicRef::Block(new_landing)),
3891 "the relocated literal must name the clone, not the deleted original",
3892 );
3893 assert_eq!(
3894 ctx.shared.values.literals[new_lit].value, 0x2008,
3895 "re-pointing the symbol must not disturb the numeric value",
3896 );
3897 }
3898
3899 #[test]
3903 fn conditional_arm_into_split_block_uses_a_trampoline() {
3904 let mut ctx = Context::new();
3905 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3906 let entry = block_at(&mut ctx, f, 0x1000);
3907 let cont = block_at(&mut ctx, f, 0x1008);
3908 let tail = block_at(&mut ctx, f, 0x2000);
3909 cbranch_at(&mut ctx, entry, tail, cont, 0x1000);
3910 return_at(&mut ctx, cont, 0x1008);
3911 return_at(&mut ctx, tail, 0x2000);
3912 FunctionBody::from_id_mut(&mut ctx, f)
3913 .set_root(entry)
3914 .unwrap();
3915
3916 let g = ctx.split_function_at(tail);
3917
3918 let entry = block_at_addr(&ctx, f, 0x1000);
3919 let cont = block_at_addr(&ctx, f, 0x1008);
3920 assert_eq!(addrs(&ctx, g), vec![0x2000]);
3921
3922 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, entry)
3923 .instructions()
3924 .last()
3925 .unwrap()
3926 .mnemonic()
3927 .clone()
3928 else {
3929 panic!("entry must still end in a cbranch");
3930 };
3931 assert_eq!(cb.failure_block, cont.local, "fall-through arm untouched");
3932 let tramp = BlockId::new(entry.func, cb.success_block);
3933 assert_eq!(
3934 BasicBlock::from_id(&ctx, tramp).parent().map(|f| f.id),
3935 Some(f),
3936 "trampoline lives in F",
3937 );
3938 let term = BasicBlock::from_id(&ctx, tramp)
3939 .instructions()
3940 .last()
3941 .map(|i| i.mnemonic().clone());
3942 assert!(
3943 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
3944 "trampoline must tail-call G, got {term:?}",
3945 );
3946 for (_, s) in BasicBlock::from_id(&ctx, entry).successors() {
3948 assert_eq!(BasicBlock::from_id(&ctx, s).parent().map(|f| f.id), Some(f));
3949 }
3950 }
3951
3952 #[test]
3955 fn mints_a_conventional_function_when_no_stub_exists() {
3956 let mut ctx = Context::new();
3957 wazabin_qcode_macro::qcode!(
3958 ctx,
3959 "
3960 fn f:
3961 <entry>
3962 goto <0x1008>;
3963 <0x1008>
3964 return 0x0;
3965 "
3966 );
3967
3968 let mid = block_at_addr(&ctx, f, 0x1008);
3969 let g = ctx.split_function_at(mid);
3970 assert_eq!(FunctionBody::from_id(&ctx, g).name(), "fn_1008");
3971 assert_eq!(FunctionBody::from_id(&ctx, f).block_ids().len(), 1);
3973 assert_eq!(addrs(&ctx, g), vec![0x1008]);
3974 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
3975 assert_eq!(addresses.function_at(0x1008), Some(g));
3976 for b in FunctionBody::from_id(&ctx, g).block_ids() {
3977 assert_eq!(b.func, g);
3978 }
3979 }
3980
3981 fn assert_no_dangling_terminators(ctx: &Context) {
3985 for b in ctx.block_ids() {
3986 let Some(mnemonic) = BasicBlock::from_id(ctx, b)
3987 .instructions()
3988 .last()
3989 .map(|t| t.mnemonic().clone())
3990 else {
3991 continue;
3992 };
3993 let targets = match &mnemonic {
3994 Mnemonic::Branch(crate::value::insn::Branch { target, .. }) => vec![*target],
3995 Mnemonic::CBranch(crate::value::insn::CBranch {
3996 success_block,
3997 failure_block,
3998 ..
3999 }) => vec![*success_block, *failure_block],
4000 _ => vec![],
4001 };
4002 let succs: std::collections::HashSet<BlockId> = BasicBlock::from_id(ctx, b)
4003 .successors()
4004 .map(|(_, s)| s)
4005 .collect();
4006 for t in targets {
4007 let tid = BlockId::new(b.func, t);
4008 assert!(
4009 ctx.contains_block(tid),
4010 "block {b:?} terminator names dead block {tid:?}"
4011 );
4012 assert!(
4013 succs.contains(&tid),
4014 "block {b:?} terminator target {tid:?} has no CFG edge (operand/edge desync)"
4015 );
4016 }
4017 }
4018 }
4019
4020 #[test]
4025 fn retained_predecessor_into_mid_tail_promotes_the_landing() {
4026 let mut ctx = Context::new();
4027 wazabin_qcode_macro::qcode!(
4030 ctx,
4031 "
4032 fn f:
4033 <entry @c:i8>
4034 if @c goto <0x2000> else goto <0x1008>;
4035 <0x1008>
4036 goto <0x2008>;
4037 <0x2000>
4038 goto <0x2008>;
4039 <0x2008>
4040 return 0x0;
4041 "
4042 );
4043
4044 let tail = block_at_addr(&ctx, f, 0x2000);
4045 let g = ctx.split_function_at(tail);
4046
4047 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
4050 let landing_fn = addresses
4051 .function_at(0x2008)
4052 .expect("mid-tail landing must be promoted to a function");
4053 assert_ne!(landing_fn, g);
4054 assert_eq!(addrs(&ctx, g), vec![0x2000]);
4055 assert_no_dangling_terminators(&ctx);
4056
4057 for (holder, addr) in [(f, 0x1008u64), (g, 0x2000u64)] {
4058 let block = block_at_addr(&ctx, holder, addr);
4059 let term = BasicBlock::from_id(&ctx, block)
4060 .instructions()
4061 .last()
4062 .map(|i| i.mnemonic().clone());
4063 assert!(
4064 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(landing_fn)),
4065 "branch at {addr:#x} into the landing must tail-call it, got {term:?}",
4066 );
4067 }
4068 }
4069
4070 #[test]
4078 fn tail_conditional_to_own_registered_entry_uses_a_trampoline() {
4079 let mut ctx = Context::new();
4080 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4081 let entry = block_at(&mut ctx, f, 0x1000);
4082 let tail = block_at(&mut ctx, f, 0x2000);
4083 let cont = block_at(&mut ctx, f, 0x2008);
4084 branch_at(&mut ctx, entry, tail, 0x1000);
4085 cbranch_at(&mut ctx, tail, entry, cont, 0x2000);
4087 return_at(&mut ctx, cont, 0x2008);
4088 FunctionBody::from_id_mut(&mut ctx, f)
4089 .set_root(entry)
4090 .unwrap();
4091
4092 let g = ctx.split_function_at(tail);
4093
4094 assert_no_dangling_terminators(&ctx);
4095 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4096 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4097
4098 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4101 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4102 .instructions()
4103 .last()
4104 .unwrap()
4105 .mnemonic()
4106 .clone()
4107 else {
4108 panic!("moved tail must still end in a cbranch");
4109 };
4110 let tramp = BlockId::new(g, cb.success_block);
4111 assert_eq!(tramp.func, g, "trampoline must have relocated into g");
4112 let term = BasicBlock::from_id(&ctx, tramp)
4113 .instructions()
4114 .last()
4115 .map(|i| i.mnemonic().clone());
4116 assert!(
4117 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(f)),
4118 "back-edge trampoline must tail-call f, got {term:?}",
4119 );
4120 }
4121
4122 #[test]
4127 fn tail_conditional_to_foreign_entry_relocates_its_trampoline() {
4128 let mut ctx = Context::new();
4129 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4130 let entry = block_at(&mut ctx, f, 0x1000);
4131 let tail = block_at(&mut ctx, f, 0x2000);
4132 let cont = block_at(&mut ctx, f, 0x2008);
4133 let foreign = block_at(&mut ctx, f, 0x3000);
4134 branch_at(&mut ctx, entry, tail, 0x1000);
4135 cbranch_at(&mut ctx, tail, foreign, cont, 0x2000);
4137 return_at(&mut ctx, cont, 0x2008);
4138 return_at(&mut ctx, foreign, 0x3000);
4139 FunctionBody::from_id_mut(&mut ctx, f)
4140 .set_root(entry)
4141 .unwrap();
4142 let h = FunctionBody::make_at_addr(&mut ctx, 0x3000, Some(Cow::Borrowed("h"))).id;
4143
4144 let g = ctx.split_function_at(tail);
4145
4146 assert_no_dangling_terminators(&ctx);
4147 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4148 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4149
4150 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4153 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4154 .instructions()
4155 .last()
4156 .unwrap()
4157 .mnemonic()
4158 .clone()
4159 else {
4160 panic!("moved tail must still end in a cbranch");
4161 };
4162 let tramp = BlockId::new(g, cb.success_block);
4163 assert_eq!(tramp.func, g, "trampoline must have relocated into g");
4164 let term = BasicBlock::from_id(&ctx, tramp)
4165 .instructions()
4166 .last()
4167 .map(|i| i.mnemonic().clone());
4168 assert!(
4169 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(h)),
4170 "relocated trampoline must tail-call H, got {term:?}",
4171 );
4172 }
4173
4174 #[test]
4178 fn conditional_failure_arm_into_split_block_uses_a_trampoline() {
4179 let mut ctx = Context::new();
4180 wazabin_qcode_macro::qcode!(
4183 ctx,
4184 "
4185 fn f:
4186 <entry @c:i8>
4187 if @c goto <0x1008> else goto <0x2000>;
4188 <0x1008>
4189 return 0x0;
4190 <0x2000>
4191 return 0x0;
4192 "
4193 );
4194
4195 let tail = block_at_addr(&ctx, f, 0x2000);
4196 let g = ctx.split_function_at(tail);
4197
4198 assert_no_dangling_terminators(&ctx);
4199 let entry = BlockId::new(f, ctx.bodies[f].root_id().unwrap());
4200 let cont = block_at_addr(&ctx, f, 0x1008);
4201 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, entry)
4202 .instructions()
4203 .last()
4204 .unwrap()
4205 .mnemonic()
4206 .clone()
4207 else {
4208 panic!("entry must still end in a cbranch");
4209 };
4210 assert_eq!(
4211 cb.success_block, cont.local,
4212 "success (fall-through) untouched"
4213 );
4214 let tramp = BlockId::new(entry.func, cb.failure_block);
4215 let term = BasicBlock::from_id(&ctx, tramp)
4216 .instructions()
4217 .last()
4218 .map(|i| i.mnemonic().clone());
4219 assert!(
4220 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
4221 "failure arm must route through a trampoline tail-calling G, got {term:?}",
4222 );
4223 }
4224
4225 #[test]
4228 fn moved_tail_internal_conditional_remaps_both_arms() {
4229 let mut ctx = Context::new();
4230 wazabin_qcode_macro::qcode!(
4233 ctx,
4234 "
4235 fn f:
4236 <entry>
4237 goto <0x2000>;
4238 <0x2000>
4239 %c = 0x0 == 0x0;
4240 if %c goto <0x2008> else goto <0x2010>;
4241 <0x2008>
4242 return 0x0;
4243 <0x2010>
4244 return 0x0;
4245 "
4246 );
4247
4248 let tail = block_at_addr(&ctx, f, 0x2000);
4249 let g = ctx.split_function_at(tail);
4250
4251 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2008, 0x2010]);
4252 assert_no_dangling_terminators(&ctx);
4253 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4254 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4255 .instructions()
4256 .last()
4257 .unwrap()
4258 .mnemonic()
4259 .clone()
4260 else {
4261 panic!("moved tail must still end in a cbranch");
4262 };
4263 let a = block_at_addr(&ctx, g, 0x2008);
4264 let b = block_at_addr(&ctx, g, 0x2010);
4265 assert_eq!(cb.success_block, a.local, "success arm re-pointed to clone");
4266 assert_eq!(cb.failure_block, b.local, "failure arm re-pointed to clone");
4267 }
4268
4269 #[test]
4272 fn moved_tail_internal_branch_remaps_target() {
4273 let mut ctx = Context::new();
4274 wazabin_qcode_macro::qcode!(
4275 ctx,
4276 "
4277 fn f:
4278 <entry>
4279 goto <0x2000>;
4280 <0x2000>
4281 goto <0x2008>;
4282 <0x2008>
4283 return 0x0;
4284 "
4285 );
4286
4287 let tail = block_at_addr(&ctx, f, 0x2000);
4288 let g = ctx.split_function_at(tail);
4289
4290 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2008]);
4291 assert_no_dangling_terminators(&ctx);
4292 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4293 let Mnemonic::Branch(br) = BasicBlock::from_id(&ctx, moved_tail)
4294 .instructions()
4295 .last()
4296 .unwrap()
4297 .mnemonic()
4298 .clone()
4299 else {
4300 panic!("moved tail must still end in a branch");
4301 };
4302 let end = block_at_addr(&ctx, g, 0x2008);
4303 assert_eq!(br.target, end.local, "internal branch re-pointed to clone");
4304 }
4305
4306 #[test]
4309 fn split_rehomes_store_temporary_space() {
4310 let mut ctx = Context::new();
4311 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4312 let entry = block_at(&mut ctx, f, 0x1000);
4313 let tail = block_at(&mut ctx, f, 0x2000);
4314 branch_at(&mut ctx, entry, tail, 0x1000);
4315
4316 let slot = ctx.builder(tail).make_named_temp(Cow::Borrowed("slot"), 8);
4317 let space = ctx.bodies[f].temps[slot.local].space;
4318 let value = ctx.get_const(0x2a, 8).id();
4319 {
4320 let mut builder = ctx.builder(tail);
4321 builder.push_store(value, ValueId::Temp(slot), LocalMemorySpaceId::Temp(space));
4322 builder.push_return(value);
4323 }
4324 FunctionBody::from_id_mut(&mut ctx, f)
4325 .set_root(entry)
4326 .unwrap();
4327
4328 let g = ctx.split_function_at(tail);
4329 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4330 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4331
4332 let moved_store = FunctionBody::from_id(&ctx, g)
4333 .blocks()
4334 .flat_map(|block| block.instructions())
4335 .find(|insn| matches!(insn.mnemonic(), Mnemonic::Store(_)))
4336 .expect("store moved with the split");
4337 let Mnemonic::Store(moved) = moved_store.mnemonic() else {
4338 unreachable!()
4339 };
4340 let LocalMemorySpaceId::Temp(moved_space) = moved.space else {
4341 panic!("store lost temporary-space provenance")
4342 };
4343 assert!(usize::from(moved_space) < ctx.bodies[g].temp_spaces.len());
4344 }
4345 }
4346}