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, 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 get_const(&self, value: u64, size: usize) -> ValueId {
243 let type_id = self.types.get_or_make_int(size);
244 ValueId::Literal(self.values.get_or_make_typed_literal(value, type_id, size))
245 }
246
247 pub fn get_bool_const(&self, value: bool) -> ValueId {
250 let type_id = self.types.get_or_make_bool();
251 ValueId::Literal(
252 self.values
253 .get_or_make_typed_literal(u64::from(value), type_id, 1),
254 )
255 }
256
257 pub fn get_typed_const(&self, value: u64, type_id: crate::types::TypeId) -> ValueId {
260 let size = self.types.size_of(type_id);
261 ValueId::Literal(self.values.get_or_make_typed_literal(value, type_id, size))
262 }
263
264 pub fn get_bytes(&self, data: Vec<u8>) -> ValueId {
267 let i8_ty = self.types.get_or_make_int(1);
268 let type_id = self.types.get_or_make_array(i8_ty, data.len());
269 ValueId::Bytes(
270 self.values
271 .bytes
272 .push(crate::value::Bytes { data, type_id }),
273 )
274 }
275
276 pub fn bytes_display(&self, id: crate::value::BytesId) -> crate::value::BytesDisplay {
281 self.values
282 .bytes_display
283 .get(&id)
284 .copied()
285 .unwrap_or_default()
286 }
287
288 pub fn get_typed_bytes(&self, data: Vec<u8>, type_id: crate::types::TypeId) -> ValueId {
292 ValueId::Bytes(
293 self.values
294 .bytes
295 .push(crate::value::Bytes { data, type_id }),
296 )
297 }
298
299 pub fn truth(&self, prop: Proposition) -> Option<Truth> {
303 self.values.truths.get(&prop).copied()
304 }
305
306 pub fn assumed_call_convention(&self) -> Option<&crate::assumption::AssumedCallEffect> {
312 self.assumed_call_convention.as_ref()
313 }
314
315 pub fn varnodes(&self) -> impl Iterator<Item = crate::value::VarnodeRef<'str, '_>> + '_ {
318 self.values
319 .varnodes
320 .iter()
321 .map(move |v| crate::value::Varnode::from_id(self, v.id))
322 }
323
324 pub fn varnode_count(&self) -> usize {
327 self.values.varnodes.len()
328 }
329
330 pub fn stored_type_of(&self, id: ValueId) -> Option<crate::types::TypeId> {
337 match id {
338 ValueId::Literal(lid) => Some(self.values.literals[lid].type_id),
339 ValueId::Bytes(bid) => Some(self.values.bytes[bid].type_id),
340 ValueId::Varnode(vid) => self.values.varnode_types.get(&vid).copied(),
341 ValueId::Poison(pid) => Some(self.values.poisons[pid].type_id),
342 ValueId::Instruction(_)
343 | ValueId::BlockParam(_)
344 | ValueId::BasicBlock(_)
345 | ValueId::Temp(_)
346 | ValueId::Function(_) => None,
347 }
348 }
349}
350
351pub use wazabin_binary::TargetOs;
355
356impl<'str> Context<'str> {
357 pub fn new() -> Self {
363 let mut ctx = Self::default();
364 ctx.shared.spaces.push(Space::new(Some("const"), 1, 8));
366 let default_space = Space::new(Some("ram"), 1, 8);
368 ctx.shared.default_space = ctx.shared.spaces.push(default_space);
369 ctx
370 }
371
372 pub fn try_get_space(&self, name: &str) -> Option<SpaceId> {
375 self.shared.named_spaces.get(name).copied()
376 }
377
378 pub fn get_or_make_named_space(&mut self, name: &str) -> SpaceId {
383 if let Some(id) = self.try_get_space(name) {
384 return id;
385 }
386 let default_id = self.shared.default_space;
387 if self.shared.spaces[default_id].name.as_deref() == Some(name) {
388 return default_id;
389 }
390 let default = &self.shared.spaces[self.shared.default_space];
391 let space = Space::new(Some(name), default.word_size, default.addr_size);
392 self.add_space(space)
393 }
394
395 pub fn add_space(&mut self, space: Space) -> SpaceId {
397 let name_key: Option<Box<str>> = space.name.clone();
398 let id = self.shared.spaces.push(space);
399 if let Some(name) = name_key {
400 self.shared.named_spaces.insert(name, id);
401 }
402 id
403 }
404
405 pub fn space_count(&self) -> usize {
407 self.shared.spaces.len()
408 }
409
410 pub fn set_primary_entrypoint(&mut self, entrypoint: Option<u64>) {
411 self.shared.primary_entrypoint = entrypoint;
412 }
413
414 pub fn primary_entrypoint(&self) -> Option<u64> {
415 self.shared.primary_entrypoint
416 }
417
418 pub fn set_ignored_functions(&mut self, addrs: HashSet<u64>) {
422 self.shared.ignored_functions = addrs;
423 }
424
425 pub fn ignored_functions(&self) -> &HashSet<u64> {
427 &self.shared.ignored_functions
428 }
429
430 pub fn is_function_ignored(&self, addr: Option<u64>) -> bool {
433 addr.is_some_and(|a| self.shared.ignored_functions.contains(&a))
434 }
435
436 pub fn set_target_os(&mut self, os: TargetOs) {
439 self.shared.target_os = os;
440 }
441
442 pub fn target_os(&self) -> TargetOs {
444 self.shared.target_os
445 }
446
447 pub fn set_linked_libraries(&mut self, libs: Vec<String>) {
450 self.shared.linked_libraries = libs;
451 }
452
453 pub fn linked_libraries(&self) -> &[String] {
456 &self.shared.linked_libraries
457 }
458
459 pub fn load_spaces(&mut self, spaces: registry::Registry<SpaceId, Space>) {
461 self.shared.spaces = spaces;
462 }
463
464 pub fn mark_protections_known(&mut self) {
468 self.shared.protections_known = true;
469 }
470
471 pub fn protections_known(&self) -> bool {
473 self.shared.protections_known
474 }
475
476 pub fn assume_executable(
495 &mut self,
496 binary: &dyn wazabin_binary::BinaryFormat,
497 addr: u64,
498 ) -> bool {
499 let bounds = binary.segment_bounds(addr);
500 if let Some((start, end)) = bounds
501 && let Some(known) = self.known(Proposition::ExecutableMemory { start, end })
502 {
503 return known;
504 }
505 if !self.shared.protections_known || binary.is_executable(addr) {
506 return true;
507 }
508 if let Some((start, end)) = bounds {
509 self.set_known(Proposition::ExecutableMemory { start, end }, false);
510 }
511 false
512 }
513
514 pub fn discover(&mut self, discovery: crate::discovery::Discovery) -> bool {
517 self.shared.discoveries.insert(discovery)
518 }
519
520 pub fn discover_code(&mut self, func_entry: u64, source_block: u64, target: u64) {
529 self.shared.discoveries.insert(
530 crate::discovery::Discovery::block(target, func_entry)
531 .with_edge_kind(crate::discovery::EdgeKind::JumpTableTarget)
532 .from_block_addr(source_block)
533 .with_provenance(crate::discovery::DiscoveryProvenance::Optimization {
534 pass: "handle_jump_tables".to_string(),
535 assumption: None,
536 }),
537 );
538 }
539
540 pub fn drain_discoveries(&mut self) -> Vec<crate::discovery::Discovery> {
542 self.shared.discoveries.drain()
543 }
544
545 pub fn discoveries(&self) -> impl Iterator<Item = &crate::discovery::Discovery> + '_ {
547 self.shared.discoveries.iter()
548 }
549
550 pub fn has_no_discoveries(&self) -> bool {
552 self.shared.discoveries.is_empty()
553 }
554
555 pub fn lifted_code_seeds(&self) -> Vec<crate::discovery::CodeSeed> {
560 self.shared.discoveries.lifted_seeds()
561 }
562
563 pub fn seed_code(&mut self, seeds: impl IntoIterator<Item = crate::discovery::CodeSeed>) {
569 for seed in seeds {
570 self.shared.discoveries.insert(seed.into_discovery());
571 }
572 }
573
574 pub fn mark_discovery_lifted(&mut self, key: crate::discovery::DiscoveryKey) {
575 self.shared.discoveries.mark_lifted(key);
576 }
577
578 pub fn mark_discovery_failed(
579 &mut self,
580 key: crate::discovery::DiscoveryKey,
581 reason: impl Into<String>,
582 ) {
583 self.shared.discoveries.mark_failed(key, reason);
584 }
585
586 pub fn mark_discovery_skipped(
587 &mut self,
588 key: crate::discovery::DiscoveryKey,
589 reason: impl Into<String>,
590 ) {
591 self.shared.discoveries.mark_skipped(key, reason);
592 }
593
594 pub fn get_or_make_block(&mut self, addr: u64, func: FunctionId) -> BlockId {
599 let mut addresses = crate::address_index::AddressIndex::analyze(self);
600 self.get_or_make_block_indexed(&mut addresses, addr, func)
601 }
602
603 #[track_caller]
607 pub fn get_or_make_block_indexed(
608 &mut self,
609 addresses: &mut crate::address_index::AddressIndex,
610 addr: u64,
611 func: FunctionId,
612 ) -> BlockId {
613 use crate::address_index::AddressTarget;
614
615 if let Some(AddressTarget::Function(owner)) = addresses.get(addr) {
616 assert_eq!(
617 owner, func,
618 "cannot create a block at an address owned by another function"
619 );
620 }
621 let existing = match addresses.get(addr) {
622 Some(AddressTarget::Block(block)) => Some(block),
623 Some(AddressTarget::Function(function)) => FunctionBody::from_id(self, function)
624 .root()
625 .map(|root| root.id),
626 None => None,
627 };
628 match existing {
629 Some(block) => {
630 if self.block(block).address != Some(addr)
640 && block.func == func
641 && self.block(block).extra_addresses.contains(&addr)
642 {
643 return self.split_block_at_address(addresses, block, addr);
644 }
645 if block.func != func {
646 let stored = FunctionBody::from_id(self, block.func);
647 let requested = FunctionBody::from_id(self, func);
648 let parent = Some(block.func);
652 let caller = std::panic::Location::caller();
653 let detail = format!(
654 "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}",
655 block.func,
656 stored.name(),
657 stored.address(),
658 func,
659 requested.name(),
660 requested.address(),
661 );
662 log::error!(
663 target: "qcode::arena",
664 "{detail}\nbacktrace:\n{}",
665 std::backtrace::Backtrace::force_capture()
666 );
667 panic!("{detail}");
668 }
669 block
670 }
671 None => {
672 BasicBlock::make(self, func)
673 .with_address_indexed(addresses, addr)
674 .id
675 }
676 }
677 }
678
679 pub fn split_block_at_address(
699 &mut self,
700 addresses: &mut crate::address_index::AddressIndex,
701 block: BlockId,
702 addr: u64,
703 ) -> BlockId {
704 addresses.mark_boundary(addr);
707 let tail = BasicBlock::make(self, block.func).id;
708
709 self.bodies[block.func].clear_block_instructions(block);
712
713 let absorbed = std::mem::take(&mut self.block_mut(block).extra_addresses);
717 for absorbed_addr in absorbed {
718 addresses.forget(absorbed_addr);
719 }
720 BasicBlock::from_id_mut(self, tail)
721 .in_function(block.func)
722 .with_address_indexed(addresses, addr);
723 tail
724 }
725
726 pub fn builder(&mut self, block: BlockId) -> crate::builder::Builder<'str, '_> {
729 let body = &mut self.bodies[block.func];
730 crate::builder::Builder::new(body, &self.shared, &self.interfaces, block)
731 }
732
733 pub fn builder_at(&mut self, address: u64) -> crate::builder::Builder<'str, '_> {
736 use crate::address_index::AddressTarget;
737
738 let mut addresses = crate::address_index::AddressIndex::analyze(self);
739 let block = match addresses.get(address) {
740 Some(AddressTarget::Function(function)) => self.bodies[function]
741 .root_id()
742 .map(|local| BlockId::new(function, local))
743 .unwrap_or_else(|| {
744 self.get_or_make_block_indexed(&mut addresses, address, function)
745 }),
746 Some(AddressTarget::Block(block)) => block,
747 None => {
748 let function = FunctionBody::make(self, Cow::Owned(format!("blk_{address:x}")))
749 .expect("anonymous host function")
750 .id;
751 self.get_or_make_block_indexed(&mut addresses, address, function)
752 }
753 };
754 let mut builder = self.builder(block);
755 builder.set_address(address);
756 builder
757 }
758
759 pub fn bytes_display(&self, id: crate::value::BytesId) -> crate::value::BytesDisplay {
762 self.shared
763 .values
764 .bytes_display
765 .get(&id)
766 .copied()
767 .unwrap_or_default()
768 }
769
770 pub fn set_bytes_display(
774 &mut self,
775 id: crate::value::BytesId,
776 mode: crate::value::BytesDisplay,
777 ) {
778 if mode == crate::value::BytesDisplay::Auto {
779 self.shared.values.bytes_display.remove(&id);
780 } else {
781 self.shared.values.bytes_display.insert(id, mode);
782 }
783 }
784
785 pub fn block_ids(&self) -> Vec<BlockId> {
787 self.functions().flat_map(|f| f.block_ids()).collect()
788 }
789
790 pub fn instruction_ids(&self) -> Vec<InstructionId> {
794 let mut ids: Vec<_> = self.functions().flat_map(|f| f.instruction_ids()).collect();
795 ids.sort_unstable();
796 ids
797 }
798
799 pub fn function_ids(&self) -> Vec<FunctionId> {
801 self.interfaces.iter().map(|i| i.id).collect()
802 }
803
804 pub fn anon_function(&mut self) -> FunctionId {
810 let name = self.get_unique_name(std::borrow::Cow::Borrowed("anon"));
811 crate::value::FunctionBody::make(self, name)
812 .expect("unique anon function name")
813 .id
814 }
815
816 pub fn instruction_arena_stats(&self) -> (usize, usize) {
818 let mut total = 0;
819 let mut dead = 0;
820 for f in self.bodies.iter() {
821 total += f.insns.issued_len();
822 dead += f.insns.issued_len() - f.insns.len();
823 }
824 (total, dead)
825 }
826
827 pub fn body_arena_stats(&self) -> crate::value::BodyArenaStats {
833 let mut total = crate::value::BodyArenaStats::default();
834 for body in self.bodies.iter() {
835 total.add_assign(body.arena_stats());
836 }
837 total
838 }
839
840 pub fn shrink_bodies_to_fit(&mut self) {
847 for mut body in self.bodies.iter_mut() {
848 body.shrink_to_fit();
849 }
850 }
851
852 pub fn instructions(&self) -> impl Iterator<Item = InstructionRef<'str, '_>> + '_ {
854 self.instruction_ids()
855 .into_iter()
856 .map(move |id| Instruction::from_id(self, id))
857 }
858
859 pub fn blocks(&self) -> impl Iterator<Item = BlockRef<'str, '_>> + '_ {
861 self.block_ids()
862 .into_iter()
863 .map(move |id| BlockRef::from_id(self, id))
864 }
865
866 pub fn functions(&self) -> FunctionIter<'str, '_> {
868 FunctionIter {
869 ctx: self,
870 inner: self.bodies.iter(),
871 }
872 }
873
874 pub fn iter(&self) -> FunctionIter<'str, '_> {
877 self.functions()
878 }
879
880 pub fn varnodes(&self) -> impl Iterator<Item = VarnodeRef<'str, '_>> + '_ {
881 self.shared.varnodes()
882 }
883
884 pub fn varnode_count(&self) -> usize {
889 self.shared.varnode_count()
890 }
891
892 pub fn remove_cfg_edge(&mut self, func: FunctionId, edge_id: EdgeId) {
896 self.bodies[func].remove_cfg_edge(edge_id);
897 }
898
899 pub fn rehome_owned_blocks(
919 &mut self,
920 addresses: &mut crate::address_index::AddressIndex,
921 target: FunctionId,
922 olds: &[BlockId],
923 ) -> HashMap<BlockId, BlockId> {
924 let mut needed_temps: HashSet<TempId> = HashSet::default();
928 let mut needed_temp_spaces: HashSet<TempSpaceId> = HashSet::default();
929 for &old in olds {
930 for ¶m_local in &self.block(old).params {
931 let param = self.block_param(BlockParamId::new(old.func, param_local));
932 if let Some(crate::value::LocalValueId::Temp(temp)) = param.origin {
933 needed_temps.insert(TempId::new(old.func, temp));
934 }
935 if let Some(MemorySpaceId::Temp(space)) = self.shared.types.space_of(param.type_id)
936 {
937 needed_temp_spaces.insert(space);
938 }
939 }
940 for &insn_local in &self.block(old).instructions {
941 let insn = self.instruction(InstructionId::new(old.func, insn_local));
942 for arg in insn.mnemonic().args() {
943 if let crate::value::LocalValueId::Temp(temp) = arg {
944 needed_temps.insert(TempId::new(old.func, temp));
945 }
946 }
947 let explicit_space = match insn.mnemonic() {
948 Mnemonic::Load(load) => Some(load.space),
949 Mnemonic::Store(store) => Some(store.space),
950 _ => None,
951 };
952 if let Some(LocalMemorySpaceId::Temp(space)) = explicit_space {
953 needed_temp_spaces.insert(TempSpaceId::new(old.func, space));
954 }
955 if let Some(MemorySpaceId::Temp(space)) = self.shared.types.space_of(insn.type_id) {
956 needed_temp_spaces.insert(space);
957 }
958 }
959 }
960 for &temp in &needed_temps {
961 let data = &self.bodies[temp.func].temps[temp.local];
962 needed_temp_spaces.insert(TempSpaceId::new(temp.func, data.space));
963 }
964
965 let mut needed_temp_spaces: Vec<_> = needed_temp_spaces.into_iter().collect();
966 needed_temp_spaces.sort_unstable();
967 let mut temp_space_map: HashMap<TempSpaceId, TempSpaceId> = HashMap::default();
968 for old in needed_temp_spaces {
969 if old.func == target {
970 continue;
971 }
972 let space = self.bodies[old.func].temp_spaces[old.local].clone();
973 let new = self.bodies[target].push_temp_space(space);
974 temp_space_map.insert(old, new);
975 }
976
977 let mut needed_temps: Vec<_> = needed_temps.into_iter().collect();
978 needed_temps.sort_unstable();
979 let mut value_map: HashMap<ValueId, ValueId> = HashMap::default();
980 for old in needed_temps {
981 if old.func == target {
982 continue;
983 }
984 let mut temp = self.bodies[old.func].temps[old.local].clone();
985 temp.space = temp_space_map[&TempSpaceId::new(old.func, temp.space)].local;
986 if let Some(name) = temp.name.take() {
987 temp.name = Some(self.bodies[target].names.unique(name));
988 }
989 let new = self.bodies[target].push_temp(temp);
990 value_map.insert(ValueId::Temp(old), ValueId::Temp(new));
991 }
992
993 let mut block_map: HashMap<BlockId, BlockId> = HashMap::default();
996 for &old in olds {
997 let new = BasicBlock::clone_block_into(self, old, target, &mut value_map);
998 block_map.insert(old, new);
999 }
1000
1001 for (&old, &new) in &block_map {
1006 let old_params = self.block(old).params.clone();
1007 let new_params = self.block(new).params.clone();
1008 for (old_local, new_local) in old_params.into_iter().zip(new_params) {
1009 let old_param = BlockParamId::new(old.func, old_local);
1010 let new_param = BlockParamId::new(new.func, new_local);
1011 let type_id = remap_rehomed_type(
1012 self,
1013 self.block_param(new_param).type_id,
1014 target,
1015 &temp_space_map,
1016 );
1017 self.block_param_mut(new_param).type_id = type_id;
1018 let Some(origin) = self.block_param(new_param).origin else {
1019 continue;
1020 };
1021 let qualified = origin.qualify(old.func);
1022 let remapped = value_map.get(&qualified).copied().unwrap_or(qualified);
1023 debug_assert!(
1024 remapped.owning_function().is_none_or(|f| f == target),
1025 "rehome: relocated block param {old_param:?} has an origin in another \
1026 function ({qualified:?}); the relocated set is not closed",
1027 );
1028 self.block_param_mut(new_param).origin = Some(remapped.localize(new.func));
1029 }
1030
1031 let insns = self.block(new).instructions.clone();
1032 for insn_local in insns {
1033 let insn_id = InstructionId::new(new.func, insn_local);
1034 let type_id = remap_rehomed_type(
1035 self,
1036 self.instruction(insn_id).type_id,
1037 target,
1038 &temp_space_map,
1039 );
1040 self.instruction_mut(insn_id).type_id = type_id;
1041 let mut mnemonic = self.instruction(insn_id).mnemonic().clone();
1042 let mut pairs = Vec::new();
1043 for arg in mnemonic.args() {
1044 let qualified = arg.qualify(old.func);
1048 if let Some(&new_val) = value_map.get(&qualified) {
1049 pairs.push((arg, new_val.localize(new.func)));
1050 } else if let Some(new_lit) =
1051 remap_symbolic_block_literal(&self.shared.values.literals, arg, &block_map)
1052 {
1053 pairs.push((arg, new_lit));
1054 } else {
1055 debug_assert!(
1060 qualified.owning_function().is_none_or(|f| f == target),
1061 "rehome: relocated block references a value in another \
1062 function ({qualified:?}); the relocated set is not closed",
1063 );
1064 }
1065 }
1066 crate::value::block::substitute_operands(&mut mnemonic, &pairs);
1067 remap_rehomed_memory_space(&mut mnemonic, old.func, target, &temp_space_map);
1068 remap_block_targets(&mut mnemonic, old.func, new.func, &block_map);
1069 *self.instruction_mut(insn_id).mnemonic_mut() = mnemonic;
1070 }
1071 }
1072
1073 let mut incident: HashSet<(FunctionId, EdgeId)> = HashSet::default();
1080 for &old in olds {
1081 incident.extend(self.block(old).edges.iter().map(|&e| (old.func, e)));
1082 }
1083 let mut incident: Vec<_> = incident.into_iter().collect();
1084 incident.sort_unstable();
1085 for (edge_func, edge) in incident {
1086 let EdgeData { from, to } = *self.edge(edge_func, edge);
1087 let from = BlockId::new(edge_func, from);
1088 let to = BlockId::new(edge_func, to);
1089 let new_from = block_map.get(&from).copied().unwrap_or(from);
1090 let new_to = block_map.get(&to).copied().unwrap_or(to);
1091 self.add_cfg_edge(new_from, new_to);
1092 }
1093
1094 for &old in olds {
1100 let Some(addr) = self.block(old).address else {
1101 continue;
1102 };
1103 let new = block_map[&old];
1104 let extra = self.block(old).extra_addresses.clone();
1105 addresses.rehome_block(addr, old, new);
1106 for &e in &extra {
1107 addresses.rehome_block(e, old, new);
1108 }
1109 self.block_mut(new).extra_addresses = extra;
1110 self.block_mut(new).address = Some(addr);
1111 }
1112
1113 for &old in olds {
1116 BasicBlock::from_id_mut(self, old).delete();
1117 }
1118
1119 self.rebuild_users(target);
1122 block_map
1123 }
1124
1125 fn function_registered_at_block(
1129 &self,
1130 addresses: &crate::address_index::AddressIndex,
1131 block: BlockId,
1132 ) -> Option<FunctionId> {
1133 self.block(block)
1134 .address
1135 .and_then(|addr| addresses.function_at(addr))
1136 }
1137
1138 fn split_tail(
1147 &self,
1148 addresses: &crate::address_index::AddressIndex,
1149 block: BlockId,
1150 g: FunctionId,
1151 ) -> Vec<BlockId> {
1152 let mut seen: HashSet<BlockId> = HashSet::default();
1153 seen.insert(block);
1154 let mut queue = vec![block];
1155 while let Some(b) = queue.pop() {
1156 let succs: Vec<BlockId> = BasicBlock::from_id(self, b)
1157 .successors()
1158 .map(|(_, s)| s)
1159 .collect();
1160 for s in succs {
1161 if seen.contains(&s) {
1162 continue;
1163 }
1164 if let Some(entry_func) = self.function_registered_at_block(addresses, s)
1168 && entry_func != g
1169 {
1170 continue;
1171 }
1172 seen.insert(s);
1173 queue.push(s);
1174 }
1175 }
1176 let mut tail: Vec<BlockId> = seen.into_iter().collect();
1177 tail.sort_unstable_by_key(|&b| (self.block(b).address, b.local, b.func));
1178 tail
1179 }
1180
1181 pub fn split_function_at(&mut self, block: BlockId) -> FunctionId {
1205 let mut addresses = crate::address_index::AddressIndex::analyze(self);
1206 self.split_function_at_indexed(&mut addresses, block)
1207 }
1208
1209 pub fn split_function_at_indexed(
1212 &mut self,
1213 addresses: &mut crate::address_index::AddressIndex,
1214 block: BlockId,
1215 ) -> FunctionId {
1216 use crate::value::insn::{Branch, CBranch, Callee, TailCall};
1217
1218 let addr = self
1219 .block(block)
1220 .address
1221 .expect("split_function_at: block has no machine address");
1222
1223 let g = match addresses.function_at(addr) {
1228 Some(existing) => existing,
1229 None => FunctionBody::make_at_addr_indexed(self, addresses, addr, None).id,
1230 };
1231
1232 loop {
1245 let tail_set: HashSet<BlockId> =
1246 self.split_tail(addresses, block, g).into_iter().collect();
1247 let mut promote: Option<BlockId> = None;
1248 'scan: for b in self.block_ids() {
1249 if tail_set.contains(&b) {
1250 continue;
1252 }
1253 let Some(mnemonic) = BasicBlock::from_id(self, b)
1254 .instructions()
1255 .last()
1256 .map(|t| t.mnemonic().clone())
1257 else {
1258 continue;
1259 };
1260 let targets = match &mnemonic {
1261 Mnemonic::Branch(Branch { target, .. }) => vec![*target],
1262 Mnemonic::CBranch(CBranch {
1263 success_block,
1264 failure_block,
1265 ..
1266 }) => vec![*success_block, *failure_block],
1267 _ => vec![],
1268 };
1269 for t in targets {
1270 let tid = BlockId::new(b.func, t);
1271 if tid == block || !tail_set.contains(&tid) {
1275 continue;
1276 }
1277 promote = Some(tid);
1286 break 'scan;
1287 }
1288 }
1289 match promote {
1290 Some(tid) => {
1291 self.split_function_at_indexed(addresses, tid);
1292 }
1293 None => break,
1294 }
1295 }
1296
1297 let mut tail = self.split_tail(addresses, block, g);
1300 let mut tail_set: HashSet<BlockId> = tail.iter().copied().collect();
1301
1302 let mut prev_owners: HashSet<FunctionId> = HashSet::default();
1305 for &b in &tail {
1306 prev_owners.insert(b.func);
1308 }
1309
1310 let effective_owner = |_ctx: &Context, candidate: BlockId| {
1314 if tail_set.contains(&candidate) {
1315 Some(g)
1316 } else {
1317 Some(candidate.func)
1319 }
1320 };
1321
1322 let foreign_entry =
1325 |ctx: &Context, target: BlockId, owner: FunctionId| -> Option<FunctionId> {
1326 let callee = if target == block {
1327 g
1328 } else {
1329 ctx.function_registered_at_block(addresses, target)?
1330 };
1331 (callee != owner).then_some(callee)
1332 };
1333
1334 let mut tail_calls: Vec<(InstructionId, FunctionId)> = Vec::new();
1340 let mut cond_calls: Vec<(
1347 InstructionId,
1348 BlockId,
1349 FunctionId,
1350 crate::value::LocalBlockId,
1351 )> = Vec::new();
1352 let relevant: Vec<BlockId> = self.block_ids();
1353 for b in relevant {
1354 let Some(owner) = effective_owner(self, b) else {
1355 continue;
1356 };
1357 let Some((term_id, mnemonic)) = BasicBlock::from_id(self, b)
1358 .instructions()
1359 .last()
1360 .map(|t| (t.id, t.mnemonic().clone()))
1361 else {
1362 continue;
1363 };
1364 match mnemonic {
1367 Mnemonic::Branch(Branch { target, .. }) => {
1368 if let Some(callee) = foreign_entry(self, BlockId::new(b.func, target), owner) {
1369 tail_calls.push((term_id, callee));
1370 }
1371 }
1372 Mnemonic::CBranch(CBranch {
1373 success_block,
1374 failure_block,
1375 ..
1376 }) => {
1377 if let Some(callee) =
1378 foreign_entry(self, BlockId::new(b.func, success_block), owner)
1379 {
1380 cond_calls.push((term_id, b, callee, success_block));
1381 }
1382 if let Some(callee) =
1383 foreign_entry(self, BlockId::new(b.func, failure_block), owner)
1384 {
1385 cond_calls.push((term_id, b, callee, failure_block));
1386 }
1387 }
1388 _ => {}
1389 }
1390 }
1391
1392 for (insn, callee) in tail_calls {
1393 self.replace_instruction_mnemonic(
1394 insn,
1395 Mnemonic::TailCall(TailCall {
1396 target: Callee::Real(callee),
1397 args: vec![],
1398 }),
1399 );
1400 }
1401 for (insn, owner_block, callee, arm_target) in cond_calls {
1402 let tramp = BasicBlock::make(self, owner_block.func).id;
1404 (self).builder(tramp).push_tail_call(callee);
1405 self.add_cfg_edge(owner_block, tramp);
1406
1407 if tail_set.contains(&owner_block) {
1415 tail.push(tramp);
1416 tail_set.insert(tramp);
1417 }
1418
1419 let Mnemonic::CBranch(mut cb) = self.instruction(insn).mnemonic().clone() else {
1420 continue;
1421 };
1422 let tramp_local = tramp.localize(insn.func);
1427 if cb.success_block == arm_target {
1428 cb.success_block = tramp_local;
1429 }
1430 if cb.failure_block == arm_target {
1431 cb.failure_block = tramp_local;
1432 }
1433 self.replace_instruction_mnemonic(insn, Mnemonic::CBranch(cb));
1434 }
1435
1436 let moved_owner = |candidate: BlockId| {
1445 if tail_set.contains(&candidate) {
1446 g
1447 } else {
1448 candidate.func
1449 }
1450 };
1451 let mut stale: HashSet<(FunctionId, EdgeId)> = HashSet::default();
1452 for &b in &tail {
1453 for edge in self.block(b).edges.iter().copied() {
1454 let &EdgeData { from, to } = self.edge(b.func, edge);
1455 let from = BlockId::new(b.func, from);
1456 let to = BlockId::new(b.func, to);
1457 let cross = moved_owner(from) != moved_owner(to);
1458 let touches_tail = tail_set.contains(&from) || tail_set.contains(&to);
1459 if cross && touches_tail {
1460 stale.insert((b.func, edge));
1461 }
1462 }
1463 }
1464 let mut stale: Vec<_> = stale.into_iter().collect();
1465 stale.sort_unstable();
1466 for (func, edge) in stale {
1467 self.remove_cfg_edge(func, edge);
1468 }
1469
1470 let moved = self.rehome_owned_blocks(addresses, g, &tail);
1474 self.bodies[g].set_root_id(Some(moved[&block].local));
1475
1476 self.recompute_instruction_addrs(g);
1478 for owner in prev_owners {
1479 if owner != g {
1480 self.recompute_instruction_addrs(owner);
1481 }
1482 }
1483
1484 g
1485 }
1486
1487 fn recompute_instruction_addrs(&mut self, func: FunctionId) {
1490 let blocks = FunctionBody::from_id(self, func).block_ids();
1491 let mut addrs = std::collections::BTreeSet::new();
1492 for b in blocks {
1493 for insn in BasicBlock::from_id(self, b).instructions() {
1494 if let Some(a) = insn.address() {
1495 addrs.insert(a);
1496 }
1497 }
1498 }
1499 self.bodies[func].instruction_addrs = addrs;
1500 }
1501
1502 fn rebuild_users(&mut self, func: FunctionId) {
1506 let live: Vec<InstructionId> = FunctionBody::from_id(self, func).instruction_ids();
1507 let users = &mut self.bodies[func].users;
1508 users.clear();
1509 for id in live {
1510 let args = self.bodies[func].insns[id.local].mnemonic().args();
1511 let users = &mut self.bodies[func].users;
1512 for arg in args {
1513 users.entry(arg).or_default().push(id.localize(func));
1514 }
1515 }
1516 }
1517
1518 pub fn assume_true(&mut self, prop: Proposition) -> bool {
1523 self.assume(prop, true)
1524 }
1525
1526 pub fn assume_false(&mut self, prop: Proposition) -> bool {
1528 self.assume(prop, false)
1529 }
1530
1531 fn assume(&mut self, prop: Proposition, value: bool) -> bool {
1532 match self.shared.values.truths.get(&prop) {
1533 Some(t) => t.value == value,
1534 None => {
1535 self.shared.values.truths.insert(
1536 prop,
1537 Truth {
1538 value,
1539 certainty: Certainty::Assumed,
1540 pass: PassName(pass_scope::current_pass()),
1541 },
1542 );
1543 true
1544 }
1545 }
1546 }
1547
1548 pub fn set_known(&mut self, prop: Proposition, value: bool) -> bool {
1556 let pass = PassName(pass_scope::current_pass());
1557 let novel = match self.shared.values.truths.get(&prop) {
1558 Some(prior) => {
1559 if prior.certainty == Certainty::Known && prior.value != value {
1564 self.shared
1565 .values
1566 .known_contradictions
1567 .push(KnownContradiction {
1568 prop,
1569 known: prior.value,
1570 proven: value,
1571 known_pass: prior.pass,
1572 proven_pass: pass,
1573 });
1574 return false;
1575 }
1576 if prior.certainty == Certainty::Assumed && prior.value != value {
1577 self.shared.values.violations.push(Violation {
1578 prop,
1579 assumed: prior.value,
1580 assuming_pass: prior.pass,
1581 asserting_pass: pass,
1582 });
1583 true
1584 } else {
1585 false
1586 }
1587 }
1588 None => true,
1589 };
1590 self.shared.values.truths.insert(
1591 prop,
1592 Truth {
1593 value,
1594 certainty: Certainty::Known,
1595 pass,
1596 },
1597 );
1598 novel
1599 }
1600
1601 pub fn seed_known(&mut self, prop: Proposition, value: bool, pass: PassName) {
1611 let prior = self.shared.values.truths.insert(
1612 prop,
1613 Truth {
1614 value,
1615 certainty: Certainty::Known,
1616 pass,
1617 },
1618 );
1619 debug_assert!(prior.is_none(), "seeding {prop:?} over an existing truth");
1620 }
1621
1622 pub fn truth(&self, prop: Proposition) -> Option<Truth> {
1624 self.shared.values.truths.get(&prop).copied()
1625 }
1626
1627 pub fn set_assumed_call_convention(
1633 &mut self,
1634 effect: Option<crate::assumption::AssumedCallEffect>,
1635 ) {
1636 self.shared.assumed_call_convention = effect;
1637 }
1638
1639 pub fn assumed_call_convention(&self) -> Option<&crate::assumption::AssumedCallEffect> {
1643 self.shared.assumed_call_convention.as_ref()
1644 }
1645
1646 pub fn known(&self, prop: Proposition) -> Option<bool> {
1648 self.truth(prop)
1649 .filter(|t| t.certainty == Certainty::Known)
1650 .map(|t| t.value)
1651 }
1652
1653 pub fn truths(&self) -> impl Iterator<Item = (Proposition, Truth)> + '_ {
1655 self.shared.values.truths.iter().map(|(&p, &t)| (p, t))
1656 }
1657
1658 pub fn known_facts(&self) -> impl Iterator<Item = (Proposition, bool, PassName)> + '_ {
1661 self.truths()
1662 .filter(|(_, t)| t.certainty == Certainty::Known)
1663 .map(|(p, t)| (p, t.value, t.pass))
1664 }
1665
1666 pub fn violations(&self) -> &[Violation] {
1669 &self.shared.values.violations
1670 }
1671
1672 pub fn known_contradictions(&self) -> &[KnownContradiction] {
1676 &self.shared.values.known_contradictions
1677 }
1678
1679 pub fn get_literal_value(&self, id: LiteralId) -> u64 {
1681 self.shared.values.literals[id].value
1682 }
1683
1684 pub fn get_insn(&self, id: InstructionId) -> InstructionRef<'str, '_> {
1686 InstructionRef::from_id(self, id)
1687 }
1688
1689 pub fn body(&self, fid: FunctionId) -> &crate::value::FunctionBody<'str> {
1699 &self.bodies[fid]
1700 }
1701
1702 pub fn body_mut(&mut self, fid: FunctionId) -> &mut crate::value::FunctionBody<'str> {
1704 &mut self.bodies[fid]
1705 }
1706
1707 pub fn push_insn(&mut self, func: FunctionId, insn: Instruction<'str>) -> InstructionId {
1722 let args = insn.mnemonic().args();
1723 let local = self.bodies[func].insns.push(insn);
1724 let id = InstructionId::new(func, local);
1725 for arg in args {
1726 self.bodies[func]
1727 .users
1728 .entry(arg)
1729 .or_default()
1730 .push(id.localize(func));
1731 }
1732 id
1733 }
1734
1735 pub fn instruction(&self, id: InstructionId) -> &Instruction<'str> {
1737 &self.bodies[id.func].insns[id.local]
1738 }
1739
1740 pub fn instruction_mut(&mut self, id: InstructionId) -> &mut Instruction<'str> {
1742 &mut self.bodies[id.func].insns[id.local]
1743 }
1744
1745 pub fn contains_instruction(&self, id: InstructionId) -> bool {
1747 Into::<usize>::into(id.func) < self.bodies.len()
1748 && self.bodies[id.func].insns.contains(id.local)
1749 }
1750
1751 pub fn block(&self, id: BlockId) -> &BasicBlock<'str> {
1753 &self.bodies[id.func].blocks[id.local]
1754 }
1755
1756 pub fn block_mut(&mut self, id: BlockId) -> &mut BasicBlock<'str> {
1758 &mut self.bodies[id.func].blocks[id.local]
1759 }
1760
1761 pub fn contains_block(&self, id: BlockId) -> bool {
1763 Into::<usize>::into(id.func) < self.bodies.len()
1764 && self.bodies[id.func].blocks.contains(id.local)
1765 }
1766
1767 pub fn block_param(&self, id: BlockParamId) -> &BlockParam<'str> {
1769 &self.bodies[id.func].params[id.local]
1770 }
1771
1772 pub fn block_param_mut(&mut self, id: BlockParamId) -> &mut BlockParam<'str> {
1774 &mut self.bodies[id.func].params[id.local]
1775 }
1776
1777 pub fn contains_block_param(&self, id: BlockParamId) -> bool {
1779 Into::<usize>::into(id.func) < self.bodies.len()
1780 && self.bodies[id.func].params.contains(id.local)
1781 }
1782
1783 pub fn edge(&self, func: FunctionId, id: EdgeId) -> &EdgeData {
1785 &self.bodies[func].edges[id]
1786 }
1787
1788 pub fn edge_mut(&mut self, func: FunctionId, id: EdgeId) -> &mut EdgeData {
1790 &mut self.bodies[func].edges[id]
1791 }
1792
1793 pub fn users_of(&self, value: ValueId) -> Vec<InstructionId> {
1798 match value.owning_function() {
1799 Some(func) => self.bodies[func].users_of(value),
1800 None => Vec::new(),
1801 }
1802 }
1803
1804 pub fn has_users(&self, value: ValueId) -> bool {
1806 match value.owning_function() {
1807 Some(func) => self.bodies[func].has_users(value),
1808 None => false,
1809 }
1810 }
1811
1812 pub fn push_block(&mut self, func: FunctionId, block: BasicBlock<'str>) -> BlockId {
1813 let local = self.bodies[func].blocks.push(block);
1814 let id = BlockId::new(func, local);
1815 self.bodies[func].roster.push(local);
1817 id
1818 }
1819
1820 pub fn push_block_param(&mut self, func: FunctionId, param: BlockParam<'str>) -> BlockParamId {
1821 let local = self.bodies[func].params.push(param);
1822 BlockParamId::new(func, local)
1823 }
1824
1825 pub fn push_edge(&mut self, func: FunctionId, edge: EdgeData) -> EdgeId {
1826 self.bodies[func].edges.push(edge)
1827 }
1828
1829 pub fn push_function(
1832 &mut self,
1833 interface: crate::value::function::FunctionInterface<'str>,
1834 body: FunctionBody<'str>,
1835 ) -> FunctionId {
1836 let expected = FunctionId::from(self.bodies.len());
1837 assert_eq!(
1838 body.id(),
1839 expected,
1840 "function body id does not match its registry slot"
1841 );
1842 let id = self.bodies.push(body);
1843 let iid = self.interfaces.push(interface);
1844 debug_assert_eq!(
1845 Into::<usize>::into(id),
1846 Into::<usize>::into(iid),
1847 "function body/interface registries drifted"
1848 );
1849 id
1850 }
1851
1852 pub fn get_register(&self, id: RegisterId) -> VarnodeRef<'str, '_> {
1855 Varnode::from_id(self, self.shared.registers[&id])
1856 }
1857
1858 pub fn get_const(&self, value: u64, size: usize) -> LiteralRef<'str, '_> {
1860 let type_id = self.shared.types.get_or_make_int(size);
1861 let id = self
1862 .shared
1863 .values
1864 .get_or_make_typed_literal(value, type_id, size);
1865 LiteralRef::from_id(self, id)
1866 }
1867
1868 pub fn get_bool_const(&self, value: bool) -> LiteralRef<'str, '_> {
1871 let type_id = self.shared.types.get_or_make_bool();
1872 let id = self
1873 .shared
1874 .values
1875 .get_or_make_typed_literal(u64::from(value), type_id, 1);
1876 LiteralRef::from_id(self, id)
1877 }
1878
1879 pub fn get_poison(&self, type_id: crate::types::TypeId) -> ValueId {
1883 ValueId::Poison(self.shared.values.push_poison(type_id))
1884 }
1885
1886 pub fn get_typed_const(
1892 &self,
1893 value: u64,
1894 type_id: crate::types::TypeId,
1895 ) -> LiteralRef<'str, '_> {
1896 let size = self.shared.types.size_of(type_id);
1897 let id = self
1898 .shared
1899 .values
1900 .get_or_make_typed_literal(value, type_id, size);
1901 LiteralRef::from_id(self, id)
1902 }
1903
1904 pub fn get_bytes(&self, data: Vec<u8>) -> crate::value::BytesRef<'str, '_> {
1912 let i8_ty = self.shared.types.get_or_make_int(1);
1913 let type_id = self.shared.types.get_or_make_array(i8_ty, data.len());
1914 self.get_typed_bytes(data, type_id)
1915 }
1916
1917 pub fn get_typed_bytes(
1923 &self,
1924 data: Vec<u8>,
1925 type_id: crate::types::TypeId,
1926 ) -> crate::value::BytesRef<'str, '_> {
1927 let id = self
1928 .shared
1929 .values
1930 .bytes
1931 .push(crate::value::Bytes { data, type_id });
1932 crate::value::BytesRef::from_id(self, id)
1933 }
1934
1935 pub fn type_of(&self, id: ValueId) -> crate::types::TypeId {
1940 match id {
1941 ValueId::Literal(lid) => self.shared.values.literals[lid].type_id,
1942 ValueId::Bytes(bid) => self.shared.values.bytes[bid].type_id,
1943 ValueId::Instruction(iid) => self.instruction(iid).type_id,
1944 ValueId::BlockParam(pid) => self.block_param(pid).type_id,
1945 ValueId::Varnode(vid) => {
1946 if let Some(&ty) = self.shared.values.varnode_types.get(&vid) {
1947 return ty;
1948 }
1949 let size = self.shared.values.varnodes[vid].size_bytes();
1950 self.shared.types.get_or_make_int(size)
1951 }
1952 ValueId::Temp(id) => self
1953 .shared
1954 .types
1955 .get_or_make_int(self.bodies[id.func].temps[id.local].size),
1956 ValueId::Poison(pid) => self.shared.values.poisons[pid].type_id,
1957 ValueId::BasicBlock(_) | ValueId::Function(_) => self.shared.types.get_or_make_int(0),
1960 }
1961 }
1962
1963 pub fn stored_type_of(&self, id: ValueId) -> Option<crate::types::TypeId> {
1969 match id {
1970 ValueId::Literal(lid) => Some(self.shared.values.literals[lid].type_id),
1971 ValueId::Bytes(bid) => Some(self.shared.values.bytes[bid].type_id),
1972 ValueId::Instruction(iid) => Some(self.instruction(iid).type_id),
1973 ValueId::BlockParam(pid) => Some(self.block_param(pid).type_id),
1974 ValueId::Varnode(vid) => self.shared.values.varnode_types.get(&vid).copied(),
1975 ValueId::Poison(pid) => Some(self.shared.values.poisons[pid].type_id),
1976 ValueId::Temp(_) => None,
1977 ValueId::BasicBlock(_) | ValueId::Function(_) => None,
1978 }
1979 }
1980
1981 pub fn set_varnode_type(&mut self, varnode: VarnodeId, type_id: crate::types::TypeId) {
1986 self.shared.values.varnode_types.insert(varnode, type_id);
1987 }
1988
1989 pub fn users(&self, value: impl Into<ValueId>) -> Vec<InstructionId> {
1997 self.users_of(value.into())
1998 }
1999
2000 pub fn users_across_functions(&self, value: impl Into<ValueId>) -> Vec<InstructionId> {
2005 let value = value.into();
2006 if value.owning_function().is_some() {
2007 self.users_of(value)
2008 } else {
2009 self.functions().flat_map(|f| f.users_of(value)).collect()
2010 }
2011 }
2012
2013 pub fn view(&self) -> ModuleView<'_, 'str> {
2022 ModuleView::new(self)
2023 }
2024 pub fn shr(&self) -> &Shared<'str> {
2029 &self.shared
2030 }
2031 pub fn function(&self, f: FunctionId) -> &FunctionBody<'str> {
2033 &self.bodies[f]
2034 }
2035 pub fn function_mut(&mut self, f: FunctionId) -> &mut FunctionBody<'str> {
2037 &mut self.bodies[f]
2038 }
2039
2040 pub fn block_ref(&self, id: BlockId) -> BlockRef<'str, '_, ModuleView<'_, 'str>> {
2042 self.view().block_ref(id)
2043 }
2044 pub fn insn_ref(&self, id: InstructionId) -> InstructionRef<'str, '_, ModuleView<'_, 'str>> {
2046 self.view().insn_ref(id)
2047 }
2048 pub fn param_ref(&self, id: BlockParamId) -> BlockParamRef<'str, '_, ModuleView<'_, 'str>> {
2050 self.view().param_ref(id)
2051 }
2052 pub fn function_ref(&self, id: FunctionId) -> FunctionRef<'str, '_, ModuleView<'_, 'str>> {
2054 self.view().function_ref(id)
2055 }
2056
2057 pub fn push_mnemonic(
2059 &mut self,
2060 func: FunctionId,
2061 mnemonic: Mnemonic,
2062 size: usize,
2063 ) -> InstructionId {
2064 let type_id = self.shared.types.get_or_make_int(size);
2065 self.push_insn(func, Instruction::new(type_id, mnemonic))
2066 }
2067
2068 pub fn push_mnemonic_with_type(
2071 &mut self,
2072 func: FunctionId,
2073 mnemonic: Mnemonic,
2074 type_id: crate::types::TypeId,
2075 ) -> InstructionId {
2076 self.push_insn(func, Instruction::new(type_id, mnemonic))
2077 }
2078
2079 pub fn make_block(&mut self, func: FunctionId) -> BlockId {
2082 self.push_block(func, BasicBlock::detached())
2083 }
2084
2085 pub fn register_local_name(
2088 &mut self,
2089 id: ValueId,
2090 name: Cow<'str, str>,
2091 old_name: Option<&str>,
2092 ) -> Result<()> {
2093 let existing = match id.name_scope_function() {
2094 Some(func) => self
2095 .function(func)
2096 .names
2097 .get(&name)
2098 .map(|id| id.qualify(func)),
2099 None => self.get_named(&name),
2100 };
2101 if let Some(existing) = existing {
2102 return if existing == id {
2103 Ok(())
2104 } else {
2105 Err(Error::spanless(ErrorTy::DuplicateName(name.to_string())))
2106 };
2107 }
2108 match id.name_scope_function() {
2109 Some(func) => self
2110 .function_mut(func)
2111 .names
2112 .register(name, id.localize(func), old_name),
2113 None => self.update_name(name, id, old_name),
2114 }
2115 }
2116
2117 pub(crate) fn set_address_indexed(
2119 &mut self,
2120 addresses: &mut crate::address_index::AddressIndex,
2121 addr: u64,
2122 id: ValueId,
2123 ) -> crate::error::Result<()> {
2124 let target = match id {
2125 ValueId::Function(id) => crate::address_index::AddressTarget::Function(id),
2126 ValueId::BasicBlock(id) => crate::address_index::AddressTarget::Block(id),
2127 _ => unreachable!("only functions and blocks have module addresses"),
2128 };
2129 addresses.register(self, addr, target)
2130 }
2131
2132 pub fn update_name(
2135 &mut self,
2136 name: Cow<'str, str>,
2137 id: ValueId,
2138 old_name: Option<&str>,
2139 ) -> Result<()> {
2140 match id.name_scope_function() {
2141 Some(func) => self.bodies[func]
2142 .names
2143 .register(name, id.localize(func), old_name),
2144 None => self.shared.name_map.register(name, id, old_name),
2145 }
2146 }
2147
2148 pub fn get_named_in_scope(&self, id: ValueId, name: &str) -> Option<ValueId> {
2153 match id.name_scope_function() {
2154 Some(func) => self.bodies[func].names.get(name).map(|id| id.qualify(func)),
2155 None => self.shared.name_map.get(name),
2156 }
2157 }
2158
2159 pub fn get_named(&self, name: &str) -> Option<ValueId> {
2168 self.shared.name_map.get(name)
2169 }
2170
2171 pub fn get_unique_name(&mut self, name: Cow<'str, str>) -> Cow<'str, str> {
2176 self.shared.name_map.unique(name)
2177 }
2178
2179 pub fn get_unique_name_in(&mut self, func: FunctionId, name: Cow<'str, str>) -> Cow<'str, str> {
2183 self.bodies[func].names.unique(name)
2184 }
2185}
2186
2187#[derive(Clone, serde::Serialize, serde::Deserialize)]
2198pub struct NameTable<'str, Id = ValueId> {
2199 map: HashMap<Cow<'str, str>, Id>,
2201 #[serde(skip)]
2205 suffix_hint: HashMap<String, u32>,
2206}
2207
2208impl<Id> Default for NameTable<'_, Id> {
2209 fn default() -> Self {
2210 Self {
2211 map: HashMap::default(),
2212 suffix_hint: HashMap::default(),
2213 }
2214 }
2215}
2216
2217impl<'str, Id: Copy + Eq> NameTable<'str, Id> {
2218 pub(crate) fn entries(&self) -> impl Iterator<Item = (&str, Id)> + '_ {
2219 self.map.iter().map(|(name, &value)| (name.as_ref(), value))
2220 }
2221
2222 pub fn get(&self, name: &str) -> Option<Id> {
2224 self.map.get(name).copied()
2225 }
2226
2227 pub fn contains(&self, name: &str) -> bool {
2229 self.map.contains_key(name)
2230 }
2231
2232 pub fn register(&mut self, name: Cow<'str, str>, id: Id, old_name: Option<&str>) -> Result<()> {
2236 if let Some(old_name) = old_name {
2237 self.forget(old_name);
2238 }
2239 match self.map.insert(name.clone(), id) {
2240 Some(_) => Err(Error::spanless(ErrorTy::DuplicateName(name.to_string()))),
2241 None => Ok(()),
2242 }
2243 }
2244
2245 pub fn forget(&mut self, name: &str) {
2249 self.map.remove(name);
2250 if let Some((base, suffix)) = split_generated_suffix(name)
2251 && let Some(hint) = self.suffix_hint.get_mut(base)
2252 {
2253 *hint = (*hint).min(suffix);
2254 }
2255 }
2256
2257 pub fn unique(&mut self, name: Cow<'str, str>) -> Cow<'str, str> {
2262 use std::fmt::Write as _;
2263
2264 if !self.map.contains_key(&name) {
2265 return name;
2266 }
2267 let base: &str = &name;
2268 let mut suffix = self.suffix_hint.get(base).copied().unwrap_or(1).max(1);
2269 let mut unique_name = format!("{base}_{suffix}");
2270 while self.map.contains_key(unique_name.as_str()) {
2271 suffix += 1;
2272 unique_name.clear();
2273 let _ = write!(unique_name, "{base}_{suffix}");
2274 }
2275 self.suffix_hint.insert(base.to_string(), suffix);
2276 Cow::Owned(unique_name)
2277 }
2278}
2279
2280fn split_generated_suffix(name: &str) -> Option<(&str, u32)> {
2285 let (base, digits) = name.rsplit_once('_')?;
2286 if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) {
2287 return None;
2288 }
2289 Some((base, digits.parse().ok()?))
2290}
2291
2292fn remap_rehomed_type(
2295 ctx: &Context<'_>,
2296 type_id: crate::types::TypeId,
2297 target: FunctionId,
2298 temp_space_map: &HashMap<TempSpaceId, TempSpaceId>,
2299) -> crate::types::TypeId {
2300 let Some(MemorySpaceId::Temp(old_space)) = ctx.shared.types.space_of(type_id) else {
2301 return type_id;
2302 };
2303 let Some(&new_space) = temp_space_map.get(&old_space) else {
2304 debug_assert_eq!(
2305 old_space.func, target,
2306 "rehome: result type references unmapped foreign temporary space {old_space:?}"
2307 );
2308 return type_id;
2309 };
2310 ctx.shared.types.get_or_make_space_address(
2311 ctx.shared.types.size_of(type_id),
2312 MemorySpaceId::Temp(new_space),
2313 )
2314}
2315
2316fn remap_rehomed_memory_space(
2319 mnemonic: &mut Mnemonic,
2320 old_func: FunctionId,
2321 target: FunctionId,
2322 temp_space_map: &HashMap<TempSpaceId, TempSpaceId>,
2323) {
2324 let remap = |space: &mut LocalMemorySpaceId| {
2325 let LocalMemorySpaceId::Temp(old_local) = *space else {
2326 return;
2327 };
2328 let old = TempSpaceId::new(old_func, old_local);
2329 if let Some(&new) = temp_space_map.get(&old) {
2330 *space = LocalMemorySpaceId::Temp(new.local);
2331 } else {
2332 debug_assert_eq!(
2333 old_func, target,
2334 "rehome: mnemonic references unmapped foreign temporary space {old:?}"
2335 );
2336 }
2337 };
2338 match mnemonic {
2339 Mnemonic::Load(load) => remap(&mut load.space),
2340 Mnemonic::Store(store) => remap(&mut store.space),
2341 _ => {}
2342 }
2343}
2344
2345fn remap_symbolic_block_literal(
2371 literals: &crate::value::interner::LiteralInterner,
2372 arg: crate::value::LocalValueId,
2373 block_map: &HashMap<BlockId, BlockId>,
2374) -> Option<crate::value::LocalValueId> {
2375 use crate::value::literal::SymbolicRef;
2376
2377 let crate::value::LocalValueId::Literal(lid) = arg else {
2378 return None;
2379 };
2380 let literal = literals[lid].clone();
2381 let Some(SymbolicRef::Block(old_block)) = literal.symbolic else {
2382 return None;
2383 };
2384 let &new_block = block_map.get(&old_block)?;
2385 let new_lit = literals.push_literal(crate::value::literal::Literal {
2386 symbolic: Some(SymbolicRef::Block(new_block)),
2387 ..literal
2388 });
2389 Some(crate::value::LocalValueId::Literal(new_lit))
2390}
2391
2392fn remap_block_targets(
2393 mnemonic: &mut Mnemonic,
2394 old_func: FunctionId,
2395 new_func: FunctionId,
2396 block_map: &HashMap<BlockId, BlockId>,
2397) {
2398 let remap = |b: &mut crate::value::LocalBlockId| {
2399 if let Some(&new) = block_map.get(&BlockId::new(old_func, *b)) {
2400 *b = new.localize(new_func);
2401 }
2402 };
2403 match mnemonic {
2404 Mnemonic::Branch(branch) => remap(&mut branch.target),
2405 Mnemonic::CBranch(cbranch) => {
2406 remap(&mut cbranch.success_block);
2407 remap(&mut cbranch.failure_block);
2408 }
2409 _ => {}
2410 }
2411}
2412
2413impl Display for Context<'_> {
2414 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2415 self.functions().try_for_each(|fun| fun.fmt(f))?;
2416
2417 self.blocks()
2418 .filter(|block| block.parent().is_none())
2419 .try_for_each(|block| block.fmt(f))
2420 }
2421}
2422
2423pub struct FunctionIter<'str, 'ctx> {
2424 ctx: &'ctx Context<'str>,
2425 inner: registry::Iter<'ctx, FunctionId, FunctionBody<'str>>,
2426}
2427
2428impl<'str, 'ctx> Iterator for FunctionIter<'str, 'ctx> {
2429 type Item = FunctionRef<'str, 'ctx>;
2430
2431 fn next(&mut self) -> Option<Self::Item> {
2432 let ctx = self.ctx;
2433 self.inner.next().map(|f| FunctionRef::from_id(ctx, f.id))
2434 }
2435}
2436
2437impl<'str, 'ctx> IntoIterator for &'ctx Context<'str> {
2438 type Item = FunctionRef<'str, 'ctx>;
2439 type IntoIter = FunctionIter<'str, 'ctx>;
2440
2441 fn into_iter(self) -> Self::IntoIter {
2442 self.iter()
2443 }
2444}
2445
2446#[cfg(test)]
2447mod tests {
2448 use super::*;
2449 use crate::value::{
2450 BasicBlock, FunctionBody, ValueId,
2451 insn::{Binary, Binop, Call, Callee, IntBinop, Load, Mnemonic},
2452 };
2453 use wazabin_qcode_macro::qcode;
2454
2455 fn make_fn_with_blocks(ctx: &mut Context<'static>, name: &'static str, n: usize) -> FunctionId {
2456 let f = FunctionBody::make(ctx, name.into()).unwrap().id;
2458 for _ in 0..n {
2459 BasicBlock::make(ctx, f);
2460 }
2461 f
2462 }
2463
2464 #[test]
2465 #[should_panic(expected = "cannot reuse a block stored in another function arena")]
2466 fn get_or_make_block_rejects_foreign_storage_at_address() {
2467 let mut ctx = Context::new();
2468 let a = FunctionBody::make(&mut ctx, "address_owner".into())
2469 .unwrap()
2470 .id;
2471 let b = FunctionBody::make(&mut ctx, "address_requester".into())
2472 .unwrap()
2473 .id;
2474 BasicBlock::make(&mut ctx, a).with_address(0x1000);
2475
2476 ctx.get_or_make_block(0x1000, b);
2477 }
2478
2479 #[test]
2480 #[should_panic(expected = "cannot create a block at an address owned by another function")]
2481 fn get_or_make_block_rejects_foreign_function_address_without_root() {
2482 let mut ctx = Context::new();
2483 FunctionBody::make_at_addr(&mut ctx, 0x1000, None);
2484 let requester = FunctionBody::make(&mut ctx, "address_requester".into())
2485 .unwrap()
2486 .id;
2487
2488 ctx.get_or_make_block(0x1000, requester);
2489 }
2490
2491 #[test]
2492 fn functions_iter_yields_all_functions() {
2493 let mut ctx = Context::new();
2494 let alpha = make_fn_with_blocks(&mut ctx, "alpha", 1);
2495 let beta = make_fn_with_blocks(&mut ctx, "beta", 1);
2496
2497 let names: Vec<_> = ctx.functions().map(|f| f.name().to_string()).collect();
2498 assert!(names.contains(&"alpha".to_string()));
2499 assert!(names.contains(&"beta".to_string()));
2500 assert_eq!(names.len(), 2);
2501 assert_eq!(ctx.function_ids(), vec![alpha, beta]);
2502 assert_eq!(ctx.function_ids().len(), ctx.interfaces.len());
2503 }
2504
2505 #[test]
2506 fn body_view_reads_match_module_reads() {
2507 use crate::value::{BodyView, FunctionId, FunctionRef, ModuleView, QCodeView};
2508
2509 let mut ctx = Context::new();
2510 qcode!(
2511 ctx,
2512 "
2513 fn foo:
2514 <bb1>
2515 if i8 1 goto <bb2> else goto <bb3>;
2516 <bb2>
2517 goto <bb3>;
2518 <bb3>
2519 return at 0;
2520 "
2521 );
2522 let fid = FunctionBody::from_name(&ctx, "foo").unwrap().id();
2523 let fid = ValueId::as_function(fid).unwrap();
2524
2525 type Snap = (String, Vec<(String, Vec<String>, Vec<String>, usize)>);
2530 fn snapshot<'a, 'str: 'a>(view: impl QCodeView<'a, 'str>, fid: FunctionId) -> Snap {
2531 let f = FunctionRef::new(view, fid);
2532 let blocks = f
2533 .blocks()
2534 .map(|b| {
2535 let name = b.name().unwrap_or("?").to_string();
2536 let mut succ: Vec<String> = b
2537 .successors()
2538 .map(|(_, s)| BlockRef::new(view, s).name().unwrap_or("?").to_string())
2539 .collect();
2540 succ.sort();
2541 let ops: Vec<String> =
2542 b.instructions().map(|i| i.opcode().to_string()).collect();
2543 (name, succ, ops, b.num_params())
2544 })
2545 .collect();
2546 (f.name().to_string(), blocks)
2547 }
2548
2549 let module_snap = snapshot(ModuleView::new(&ctx), fid);
2550 assert!(!module_snap.1.is_empty(), "sanity: foo has blocks");
2551
2552 let checked = BodyView::new(&ctx.bodies[fid], &ctx.shared, &ctx.interfaces);
2555 let checked_snap = snapshot(checked, fid);
2556 assert_eq!(
2557 module_snap, checked_snap,
2558 "reads through BodyView must match the module reads"
2559 );
2560 }
2561
2562 #[test]
2563 fn body_mut_mut_matches_module_mut() {
2564 use crate::value::{
2565 BlockParam, FunctionId, FunctionRef, InstructionId, Renameable,
2566 block::BlockId,
2567 block_param::BlockParamId,
2568 util::{base_ref::BaseRef, body_mut::BodyMut},
2569 };
2570
2571 fn build(mut ctx: &mut Context<'static>) -> (FunctionId, BlockId, BlockId, InstructionId) {
2572 qcode!(
2573 ctx,
2574 "
2575 varnode i64 x;
2576 fn foo:
2577 <entry>
2578 %a = load(x:8, &x);
2579 %b = load(x:8, &x);
2580 goto <bb1>;
2581 <bb1>
2582 return at %a;
2583 "
2584 );
2585 let fid = foo;
2586 let entry = FunctionRef::from_id(ctx, fid).root().unwrap().id;
2587 let bb1 = FunctionRef::from_id(ctx, fid)
2588 .blocks()
2589 .map(|b| b.id)
2590 .find(|&b| b != entry)
2591 .unwrap();
2592 let insns = BasicBlock::from_id(ctx, entry).instruction_ids();
2593 (fid, entry, bb1, insns[0])
2594 }
2595
2596 fn add_param(ctx: &mut Context<'static>, bb1: BlockId) -> BlockParamId {
2598 BasicBlock::from_id_mut(ctx, bb1).push_param(8).id
2599 }
2600
2601 type MSnap = Vec<(String, Option<String>, Vec<usize>, Vec<String>, Vec<String>)>;
2604 fn snap(ctx: &Context, fid: FunctionId) -> MSnap {
2605 FunctionRef::from_id(ctx, fid)
2606 .blocks()
2607 .map(|b| {
2608 let name = b.name().unwrap_or("?").to_string();
2609 let comment = b.comment().map(str::to_string);
2610 let params: Vec<usize> = b.params().map(|p| p.size()).collect();
2611 let ops: Vec<String> =
2612 b.instructions().map(|i| i.opcode().to_string()).collect();
2613 let mut succ: Vec<String> = b
2614 .successors()
2615 .map(|(_, s)| {
2616 BasicBlock::from_id(ctx, s)
2617 .name()
2618 .unwrap_or("?")
2619 .to_string()
2620 })
2621 .collect();
2622 succ.sort();
2623 (name, comment, params, ops, succ)
2624 })
2625 .collect()
2626 }
2627
2628 let mut ctx_a = Context::new();
2630 let (fid, entry, bb1, a) = build(&mut ctx_a);
2631 let param = add_param(&mut ctx_a, bb1);
2632 let b = BasicBlock::from_id(&ctx_a, entry).instruction_ids()[1];
2633 BasicBlock::from_id_mut(&mut ctx_a, entry).set_comment(Some("c".into()));
2634 BasicBlock::from_id_mut(&mut ctx_a, entry)
2635 .rename("start".into())
2636 .unwrap();
2637 let e = ctx_a.add_cfg_edge(entry, bb1);
2638 ctx_a.remove_cfg_edge(entry.func, e);
2639 ctx_a.replace_instruction(a, ValueId::Instruction(b));
2640 BlockParam::from_id_mut(&mut ctx_a, param).set_size(4);
2641 let snap_a = snap(&ctx_a, fid);
2642
2643 let mut ctx_b = Context::new();
2645 let (fid_b, entry_b, bb1_b, a_b) = build(&mut ctx_b);
2646 let param_b = add_param(&mut ctx_b, bb1_b);
2647 let b_b = BasicBlock::from_id(&ctx_b, entry_b).instruction_ids()[1];
2648
2649 {
2650 let mut host = BodyMut::new(&mut ctx_b.bodies[fid_b], &ctx_b.shared, &ctx_b.interfaces);
2651 let mut r = BaseRef::new(host.reborrow(), entry_b);
2652 r.set_comment(Some("c".into()));
2653 let mut r = BaseRef::new(host.reborrow(), entry_b);
2654 r.rename("start".into()).unwrap();
2655 let e = host.add_cfg_edge(entry_b, bb1_b);
2656 host.remove_cfg_edge(e);
2657 host.replace_instruction(a_b, ValueId::Instruction(b_b));
2658 let mut r = BaseRef::new(host.reborrow(), param_b);
2659 r.set_size(4);
2660 }
2661 let snap_b = snap(&ctx_b, fid_b);
2662
2663 assert_eq!(
2664 snap_a, snap_b,
2665 "mutations through a pass-scoped host must match the module-path mutations"
2666 );
2667 }
2668
2669 #[test]
2670 fn into_iterator_for_context_matches_functions() {
2671 let mut ctx = Context::new();
2672 make_fn_with_blocks(&mut ctx, "f1", 1);
2673 make_fn_with_blocks(&mut ctx, "f2", 1);
2674
2675 let via_method: Vec<_> = ctx.functions().map(|f| f.id()).collect();
2676 let via_into: Vec<_> = (&ctx).into_iter().map(|f| f.id()).collect();
2677 assert_eq!(via_method, via_into);
2678 }
2679
2680 #[test]
2681 fn blocks_iter_yields_all_blocks() {
2682 let mut ctx = Context::new();
2683 make_fn_with_blocks(&mut ctx, "g", 3);
2684
2685 let count = ctx.blocks().count();
2686 assert_eq!(count, 3);
2687 }
2688
2689 #[test]
2690 fn instructions_iter_yields_all_instructions() {
2691 let mut ctx = Context::new();
2692
2693 qcode!(
2694 ctx,
2695 "
2696 varnode i64 ptr;
2697
2698 <block>
2699 store(ptr:8, &ptr <- i64 0x1234);
2700 return at ptr;
2701 "
2702 );
2703
2704 let count = ctx.instructions().count();
2705 assert!(count >= 1, "expected at least one instruction, got {count}");
2706 }
2707
2708 #[test]
2709 fn move_insn_before_preserves_id_and_supports_arbitrary_anchors() {
2710 let mut ctx = Context::new();
2711 qcode!(
2712 ctx,
2713 "
2714 fn f:
2715 <source>
2716 %a = i64 0x1 + i64 0x2;
2717 %free = i64 0x5 + i64 0x6;
2718 goto <target>;
2719 <target>
2720 %b = i64 0x3 + i64 0x4;
2721 %consumer = %a + %b;
2722 return %consumer;
2723 "
2724 );
2725
2726 assert!(ctx.users(a).contains(&consumer));
2727 ctx.move_insn_before(a, b);
2728
2729 assert!(ctx.contains_instruction(a), "moving keeps the ID live");
2730 assert_eq!(ctx.get_insn(a).parent().map(|block| block.id), Some(target));
2731 assert!(
2732 !BasicBlock::from_id(&ctx, source)
2733 .instruction_ids()
2734 .contains(&a)
2735 );
2736 assert_eq!(
2737 BasicBlock::from_id(&ctx, target).instruction_ids()[..3],
2738 [a, b, consumer]
2739 );
2740 assert!(
2741 ctx.users(a).contains(&consumer),
2742 "moving preserves use-map entries"
2743 );
2744
2745 ctx.move_insn_before(b, a);
2747 assert_eq!(
2748 BasicBlock::from_id(&ctx, target).instruction_ids()[..3],
2749 [b, a, consumer]
2750 );
2751
2752 let return_id = *BasicBlock::from_id(&ctx, target)
2754 .instruction_ids()
2755 .last()
2756 .unwrap();
2757 ctx.move_insn_before(free, return_id);
2758 assert_eq!(
2759 BasicBlock::from_id(&ctx, target).instruction_ids()[..4],
2760 [b, a, consumer, free]
2761 );
2762 }
2763
2764 #[test]
2765 fn remove_instruction_removes_from_block() {
2766 let mut ctx = Context::new();
2767 qcode!(
2768 ctx,
2769 "
2770 varnode i64 x;
2771 <block>
2772 %a = load(x:8, &x);
2773 %b = load(x:8, &x);
2774 return at %a;
2775 "
2776 );
2777 let block_ref = BasicBlock::from_id(&ctx, block);
2778 let ids = block_ref.instruction_ids();
2779 let load_a = ids[0];
2780 let original_len = ids.len();
2781
2782 ctx.remove_instruction(load_a);
2783
2784 let remaining = BasicBlock::from_id(&ctx, block).instruction_ids();
2785 assert_eq!(remaining.len(), original_len - 1);
2786 assert!(!remaining.contains(&load_a));
2787 }
2788
2789 #[test]
2790 fn remove_instruction_drops_payload() {
2791 let mut ctx = Context::new();
2792 qcode!(
2793 ctx,
2794 "
2795 varnode i64 x;
2796 <block>
2797 %a = load(x:8, &x);
2798 return at %a;
2799 "
2800 );
2801 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2802
2803 ctx.remove_instruction(load_id);
2804
2805 assert!(!ctx.contains_instruction(load_id));
2806 }
2807
2808 #[test]
2809 fn remove_instruction_frees_name() {
2810 let mut ctx = Context::new();
2811 qcode!(
2812 ctx,
2813 "
2814 varnode i64 x;
2815 <block>
2816 %a = load(x:8, &x);
2817 return at %a;
2818 "
2819 );
2820 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2821 assert!(
2823 ctx.get_named_in_scope(load_id.into(), "a").is_some(),
2824 "name should be in map before removal"
2825 );
2826
2827 ctx.remove_instruction(load_id);
2828
2829 assert!(
2830 ctx.get_named_in_scope(load_id.into(), "a").is_none(),
2831 "name should be gone after removal"
2832 );
2833 assert!(!ctx.contains_instruction(load_id));
2834 }
2835
2836 #[test]
2837 fn remove_instruction_frees_name_for_reuse() {
2838 let mut ctx = Context::new();
2839 qcode!(
2840 ctx,
2841 "
2842 varnode i64 x;
2843 <block>
2844 %a = load(x:8, &x);
2845 return at %a;
2846 "
2847 );
2848 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2849
2850 ctx.remove_instruction(load_id);
2851
2852 qcode!(
2854 ctx,
2855 "
2856 varnode i64 y;
2857 <block2>
2858 %a = load(y:8, &y);
2859 return at %a;
2860 "
2861 );
2862 let a2 = BasicBlock::from_id(&ctx, block2).instruction_ids()[0];
2863 assert!(
2864 ctx.get_named_in_scope(a2.into(), "a").is_some(),
2865 "name should be reusable after removal"
2866 );
2867 }
2868
2869 #[test]
2870 fn remove_instruction_updates_users_map() {
2871 let mut ctx = Context::new();
2872 qcode!(
2873 ctx,
2874 "
2875 varnode i64 x;
2876 <block>
2877 %a = load(x:8, &x);
2878 %b = %a + i64 1;
2879 return at %b;
2880 "
2881 );
2882 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
2883 let load_id = ids[0];
2884 let add_id = ids[1];
2885
2886 assert!(
2887 ctx.users(load_id).contains(&add_id),
2888 "add should be a user of load before removal"
2889 );
2890
2891 ctx.remove_instruction(add_id);
2892
2893 assert!(
2894 ctx.users(load_id).is_empty(),
2895 "load should have no users after add is removed"
2896 );
2897 }
2898
2899 #[test]
2900 fn removed_instruction_is_absent_and_not_iterated() {
2901 let mut ctx = Context::new();
2906 qcode!(
2907 ctx,
2908 "
2909 varnode i64 x;
2910 <block>
2911 %a = load(x:8, &x);
2912 %dead = %a + i64 1;
2913 return at i64 0;
2914 "
2915 );
2916 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
2917 let dead_id = ids[1]; assert!(
2920 ctx.instructions().any(|i| i.id == dead_id),
2921 "the instruction is iterated while live"
2922 );
2923
2924 ctx.remove_instruction(dead_id);
2925
2926 assert!(!ctx.contains_instruction(dead_id));
2927 assert!(
2928 !ctx.instructions().any(|i| i.id == dead_id),
2929 "a deleted instruction must not be yielded by ctx.instructions()"
2930 );
2931 }
2932
2933 #[test]
2934 fn replace_instruction_mnemonic_rewrites_callind_users() {
2935 let mut ctx = Context::new();
2936 qcode!(
2937 ctx,
2938 "
2939 varnode i64 ptr;
2940 <block>
2941 call [ptr];
2942 "
2943 );
2944 let call_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
2945 let ptr = match ctx.get_insn(call_id).mnemonic() {
2946 Mnemonic::CallInd(call) => call.ptr.qualify(call_id.func),
2947 other => panic!("expected CallInd, got {other:?}"),
2948 };
2949 assert_eq!(ctx.users_across_functions(ptr), vec![call_id]);
2951
2952 let target = FunctionBody::make(&mut ctx, "target".into()).unwrap().id;
2953 ctx.replace_instruction_mnemonic(
2954 call_id,
2955 Mnemonic::Call(Call {
2956 target: Callee::Real(target),
2957 args: vec![],
2958 clobbers: vec![],
2959 tag: Default::default(),
2960 }),
2961 );
2962
2963 assert!(
2964 ctx.users_across_functions(ptr).is_empty(),
2965 "old indirect pointer should no longer list the rewritten call"
2966 );
2967 assert!(matches!(
2968 ctx.get_insn(call_id).mnemonic(),
2969 Mnemonic::Call(Call {
2970 target: actual,
2971 args,
2972 ..
2973 }) if *actual == Callee::Real(target) && args.is_empty()
2974 ));
2975 }
2976
2977 #[test]
2978 fn users_across_functions_keeps_ssa_users_in_the_owning_function() {
2979 let mut ctx = Context::new();
2980 qcode!(
2981 ctx,
2982 "
2983 fn f:
2984 <f_entry>
2985 %fx = i64 1 + i64 2;
2986 %fuse = %fx + i64 3;
2987 return at %fuse;
2988 fn g:
2989 <g_entry>
2990 %gx = i64 4 + i64 5;
2991 %guse = %gx + i64 6;
2992 return at %guse;
2993 "
2994 );
2995 let f_ids = BasicBlock::from_id(&ctx, f_entry).instruction_ids();
2996 let g_ids = BasicBlock::from_id(&ctx, g_entry).instruction_ids();
2997 assert_eq!(
2998 f_ids[0].local, g_ids[0].local,
2999 "precondition: arena-local ids collide"
3000 );
3001 assert_eq!(
3002 ctx.users_across_functions(ValueId::Instruction(f_ids[0])),
3003 vec![f_ids[1]],
3004 "an SSA query must not pick up the same local key from another function"
3005 );
3006 }
3007
3008 #[test]
3009 fn replace_instruction_mnemonic_moves_operand_users() {
3010 let mut ctx = Context::new();
3011 qcode!(
3012 ctx,
3013 "
3014 varnode i64 x;
3015 varnode i64 y;
3016 <block>
3017 %a = load(x:8, x);
3018 return at %a;
3019 "
3020 );
3021 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3022 let old_ptr = ValueId::Varnode(x);
3023 let new_ptr = ValueId::Varnode(y);
3024 assert_eq!(ctx.users_across_functions(old_ptr), vec![load_id]);
3026 assert!(ctx.users_across_functions(new_ptr).is_empty());
3027
3028 ctx.replace_instruction_mnemonic(
3029 load_id,
3030 Mnemonic::Load(Load {
3031 space: ctx.shared.default_space.into(),
3032 ptr: new_ptr.localize(load_id.func),
3033 size: 8,
3034 }),
3035 );
3036
3037 assert!(ctx.users_across_functions(old_ptr).is_empty());
3038 assert_eq!(ctx.users_across_functions(new_ptr), vec![load_id]);
3039 }
3040
3041 #[test]
3042 fn replace_instruction_mnemonic_tracks_repeated_operands() {
3043 let mut ctx = Context::new();
3044 qcode!(
3045 ctx,
3046 "
3047 varnode i64 x;
3048 varnode i64 y;
3049 <block>
3050 %a = load(x:8, x);
3051 return at %a;
3052 "
3053 );
3054 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3055 let old_ptr = ValueId::Varnode(x);
3056 let new_arg = ValueId::Varnode(y);
3057
3058 ctx.replace_instruction_mnemonic(
3059 load_id,
3060 Mnemonic::Binop(Binary {
3061 op: Binop::Int(IntBinop::Add),
3062 lhs: new_arg.localize(load_id.func),
3063 rhs: new_arg.localize(load_id.func),
3064 }),
3065 );
3066
3067 assert!(ctx.users_across_functions(old_ptr).is_empty());
3068 assert_eq!(
3069 ctx.users_across_functions(new_arg),
3070 vec![load_id, load_id],
3071 "a mnemonic using the same operand twice should record both uses"
3072 );
3073 }
3074
3075 #[test]
3076 fn remove_instruction_unparented_noop() {
3077 let mut ctx = Context::new();
3078 qcode!(
3079 ctx,
3080 "
3081 varnode i64 x;
3082 <block>
3083 %a = load(x:8, &x);
3084 return at %a;
3085 "
3086 );
3087 let load_id = BasicBlock::from_id(&ctx, block).instruction_ids()[0];
3088
3089 ctx.instruction_mut(load_id).parent = None;
3092
3093 ctx.remove_instruction(load_id);
3095
3096 assert!(ctx.get_named("a").is_none());
3097 }
3098
3099 #[test]
3100 fn add_cfg_edge_returns_id_and_remove_unlinks_both_blocks() {
3101 let mut ctx = Context::new();
3102 let f = ctx.anon_function();
3104 let a = BasicBlock::make(&mut ctx, f).id;
3105 let b = BasicBlock::make(&mut ctx, f).id;
3106 let c = BasicBlock::make(&mut ctx, f).id;
3107
3108 let edge = ctx.add_cfg_edge(a, b);
3109 let surviving_edge = ctx.add_cfg_edge(b, c);
3110 assert_eq!(
3111 BasicBlock::from_id(&ctx, a)
3112 .successors()
3113 .collect::<Vec<_>>(),
3114 vec![(edge, b)]
3115 );
3116 assert_eq!(
3117 BasicBlock::from_id(&ctx, b)
3118 .predecessors()
3119 .collect::<Vec<_>>(),
3120 vec![(edge, a)]
3121 );
3122
3123 ctx.remove_cfg_edge(a.func, edge);
3124 assert!(BasicBlock::from_id(&ctx, a).successors().next().is_none());
3125 assert!(BasicBlock::from_id(&ctx, b).predecessors().next().is_none());
3126 assert!(!ctx.bodies[a.func].edges.contains(edge));
3127 let surviving = ctx.edge(a.func, surviving_edge);
3128 assert_eq!(
3129 surviving.from, b.local,
3130 "swap removal must preserve the source"
3131 );
3132 assert_eq!(
3133 surviving.to, c.local,
3134 "swap removal must preserve the target"
3135 );
3136 assert_eq!(ctx.bodies[a.func].edges.len(), 1);
3137
3138 let self_edge = ctx.add_cfg_edge(a, a);
3139 ctx.remove_cfg_edge(a.func, self_edge);
3140 assert!(!ctx.bodies[a.func].edges.contains(self_edge));
3141 assert!(ctx.block(a).edges.is_empty());
3142
3143 let parallel_a = ctx.add_cfg_edge(a, b);
3144 let parallel_b = ctx.add_cfg_edge(a, b);
3145 ctx.remove_cfg_edge(a.func, parallel_a);
3146 assert!(!ctx.bodies[a.func].edges.contains(parallel_a));
3147 assert!(ctx.bodies[a.func].edges.contains(parallel_b));
3148 assert_eq!(
3149 BasicBlock::from_id(&ctx, a)
3150 .successors()
3151 .collect::<Vec<_>>(),
3152 vec![(parallel_b, b)],
3153 );
3154 }
3155
3156 #[test]
3157 fn truth_map_tracks_four_states_and_conflicts() {
3158 let mut ctx = Context::new();
3159 let callee = FunctionBody::make(&mut ctx, "callee".into()).unwrap().id;
3160 let prop = Proposition::FunctionReturns(callee);
3161
3162 assert!(ctx.assume_true(prop));
3164 assert!(ctx.assume_true(prop));
3165 assert!(!ctx.assume_false(prop));
3166 assert_eq!(ctx.known(prop), None, "assumed is not known");
3167
3168 let snapshot = ctx.clone();
3170
3171 let scope = pass_scope::enter("verifier");
3174 assert!(ctx.set_known(prop, false), "overturning is novel");
3175 drop(scope);
3176 assert_eq!(ctx.known(prop), Some(false));
3177 let [v] = ctx.violations() else {
3178 panic!("expected one violation")
3179 };
3180 assert_eq!(v.prop, prop);
3181 assert!(v.assumed);
3182 assert_eq!(v.asserting_pass, "verifier");
3183
3184 assert!(!ctx.set_known(prop, false));
3186
3187 assert!(snapshot.violations().is_empty());
3189 assert_eq!(snapshot.known(prop), None);
3190
3191 assert!(!ctx.assume_true(prop));
3193 assert!(ctx.assume_false(prop));
3194 }
3195
3196 #[test]
3197 fn seeded_facts_are_not_novel() {
3198 let mut ctx = Context::new();
3199 let callee = FunctionBody::make(&mut ctx, "exit".into()).unwrap().id;
3200 let prop = Proposition::FunctionReturns(callee);
3201
3202 ctx.seed_known(prop, false, PassName("seed"));
3203 assert_eq!(ctx.known(prop), Some(false));
3204 assert!(!ctx.assume_true(prop), "seeded fact blocks opposite assume");
3205 assert!(
3206 !ctx.set_known(prop, false),
3207 "re-proving a seed is not novel"
3208 );
3209 assert!(ctx.violations().is_empty());
3210 }
3211
3212 #[test]
3213 fn discovered_code_records_and_survives_round_trip() {
3214 let mut ctx = Context::new();
3215 ctx.discover_code(0x1000, 0x10f0, 0x1100);
3216 ctx.discover_code(0x1000, 0x10f0, 0x1200);
3217 ctx.discover_code(0x1000, 0x10f0, 0x1100); let targets: Vec<u64> = ctx.discoveries().map(|d| d.target).collect();
3220 assert_eq!(targets, vec![0x1100, 0x1200]);
3221
3222 let config = bincode::config::standard();
3223 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3224 let (restored, _): (Context<'static>, usize) =
3225 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3226 assert_eq!(
3227 restored.discoveries().map(|d| d.target).collect::<Vec<_>>(),
3228 targets
3229 );
3230 }
3231
3232 #[test]
3233 fn assume_executable_narrows_once_protections_known() {
3234 let mut ctx = Context::new();
3235 let mut image = crate::memory_image::MemoryImage::default();
3236 image.add_segment(0x1000, vec![0u8; 4], true, false); image.add_segment(0x2000, vec![0u8; 4], false, true); let binary: &dyn wazabin_binary::BinaryFormat = ℑ
3239
3240 assert!(ctx.assume_executable(binary, 0x1000));
3243 assert!(ctx.assume_executable(binary, 0x2000));
3244 assert!(ctx.assume_executable(binary, 0x9999));
3245
3246 ctx.mark_protections_known();
3247 assert!(
3248 ctx.assume_executable(binary, 0x1000),
3249 "code region stays liftable"
3250 );
3251 assert!(
3252 !ctx.assume_executable(binary, 0x2000),
3253 "data region is skipped once protections are known"
3254 );
3255 assert!(
3256 !ctx.assume_executable(binary, 0x9999),
3257 "unmapped is skipped once known"
3258 );
3259 assert_eq!(
3261 ctx.known(Proposition::ExecutableMemory {
3262 start: 0x2000,
3263 end: 0x2004,
3264 }),
3265 Some(false),
3266 );
3267 }
3268
3269 #[test]
3270 fn assume_executable_honors_region_override() {
3271 let mut ctx = Context::new();
3272 let mut image = crate::memory_image::MemoryImage::default();
3273 image.add_segment(0x1000, vec![0u8; 4], true, false); image.add_segment(0x2000, vec![0u8; 4], false, true); let binary: &dyn wazabin_binary::BinaryFormat = ℑ
3276 ctx.mark_protections_known();
3277
3278 ctx.seed_known(
3280 Proposition::ExecutableMemory {
3281 start: 0x2000,
3282 end: 0x2004,
3283 },
3284 true,
3285 PassName("override"),
3286 );
3287 ctx.seed_known(
3288 Proposition::ExecutableMemory {
3289 start: 0x1000,
3290 end: 0x1004,
3291 },
3292 false,
3293 PassName("override"),
3294 );
3295
3296 assert!(
3297 ctx.assume_executable(binary, 0x2000),
3298 "override wins over the non-executable segment flag"
3299 );
3300 assert!(
3301 !ctx.assume_executable(binary, 0x1000),
3302 "override wins over the executable segment flag"
3303 );
3304 }
3305
3306 #[test]
3307 fn context_survives_bincode_round_trip() {
3308 let mut ctx = Context::new();
3309 qcode!(
3310 ctx,
3311 "
3312 varnode i64 ptr;
3313 <block>
3314 %a = load(ptr:8, &ptr);
3315 %b = %a + i64 0x10;
3316 store(ptr:8, &ptr <- i64 0x1234);
3317 return at %b;
3318 "
3319 );
3320
3321 let some_space = ctx.get_or_make_named_space("scratch");
3323 let sa = ctx.shared.types.get_or_make_space_address(8, some_space);
3324 let sa_size = ctx.shared.types.size_of(sa);
3325
3326 let blocks_before = ctx.block_ids().len();
3327 let insns_before = ctx.instruction_ids().len();
3328 let funcs_before = ctx.function_ids().len();
3329
3330 let config = bincode::config::standard();
3331 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3332 let (restored, _): (Context<'static>, usize) =
3333 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3334
3335 assert_eq!(restored.block_ids().len(), blocks_before);
3336 assert_eq!(restored.instruction_ids().len(), insns_before);
3337 assert_eq!(restored.function_ids().len(), funcs_before);
3338 for function_id in restored.function_ids() {
3339 assert_eq!(restored.bodies[function_id].id(), function_id);
3340 }
3341 assert_eq!(restored.shared.types.size_of(sa), sa_size);
3343 assert_eq!(
3344 restored.shared.types.space_of(sa),
3345 Some(crate::space::MemorySpaceId::Shared(some_space))
3346 );
3347 }
3348
3349 #[test]
3350 fn compact_edge_arena_preserves_ids_across_round_trip() {
3351 let mut ctx = Context::new();
3352 let function = ctx.anon_function();
3353 let a = BasicBlock::make(&mut ctx, function).id;
3354 let b = BasicBlock::make(&mut ctx, function).id;
3355 let c = BasicBlock::make(&mut ctx, function).id;
3356 let d = BasicBlock::make(&mut ctx, function).id;
3357 let first = ctx.add_cfg_edge(a, b);
3358 let removed = ctx.add_cfg_edge(b, c);
3359 let last = ctx.add_cfg_edge(c, d);
3360 ctx.remove_cfg_edge(function, removed);
3361
3362 let physical_order: Vec<_> = ctx.bodies[function]
3363 .edges
3364 .iter()
3365 .map(|edge| edge.id)
3366 .collect();
3367 assert_eq!(physical_order, vec![first, last]);
3368
3369 let config = bincode::config::standard();
3370 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3371 let (mut restored, _): (Context<'static>, usize) =
3372 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3373
3374 assert!(!restored.bodies[function].edges.contains(removed));
3375 assert_eq!(
3376 restored.bodies[function]
3377 .edges
3378 .iter()
3379 .map(|edge| edge.id)
3380 .collect::<Vec<_>>(),
3381 physical_order,
3382 );
3383 assert_eq!(restored.edge(function, first).to, b.local);
3384 assert_eq!(restored.edge(function, last).from, c.local);
3385
3386 let fresh = restored.add_cfg_edge(a, d);
3387 assert!(fresh > last);
3388 assert_ne!(fresh, removed, "removed edge IDs must never be reused");
3389 }
3390
3391 #[test]
3392 fn compact_instruction_arena_preserves_ids_across_round_trip() {
3393 let mut ctx = Context::new();
3394 qcode!(
3395 ctx,
3396 "
3397 <block>
3398 %first = i64 1 + i64 2;
3399 %removed = i64 3 + i64 4;
3400 return at %first;
3401 "
3402 );
3403 let ids = BasicBlock::from_id(&ctx, block).instruction_ids();
3404 let first = ids[0];
3405 let removed = ids[1];
3406 let last = ids[2];
3407 ctx.remove_instruction(removed);
3408
3409 let physical_order: Vec<_> = ctx.bodies[first.func]
3410 .insns
3411 .iter()
3412 .map(|insn| insn.id)
3413 .collect();
3414 assert_eq!(physical_order, vec![first.local, last.local]);
3415
3416 let config = bincode::config::standard();
3417 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3418 let (mut restored, _): (Context<'static>, usize) =
3419 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3420
3421 assert!(!restored.contains_instruction(removed));
3422 assert_eq!(
3423 restored.bodies[first.func]
3424 .insns
3425 .iter()
3426 .map(|insn| insn.id)
3427 .collect::<Vec<_>>(),
3428 physical_order,
3429 );
3430 assert!(restored.contains_instruction(first));
3431 assert!(restored.contains_instruction(last));
3432
3433 let template = restored.instruction(last).clone();
3434 let fresh = restored.push_insn(first.func, template);
3435 assert!(fresh.local > last.local);
3436 assert_ne!(
3437 fresh, removed,
3438 "removed instruction IDs must never be reused"
3439 );
3440 }
3441
3442 #[test]
3443 fn compact_param_arena_preserves_ids_across_round_trip() {
3444 let mut ctx = Context::new();
3445 let function = ctx.anon_function();
3446 let block = BasicBlock::make(&mut ctx, function).id;
3447 let first = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3448 let removed = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3449 let last = BasicBlock::from_id_mut(&mut ctx, block).push_param(8).id;
3450
3451 ctx.block_mut(block).params.remove(1);
3452 ctx.block_param_mut(last).index = 1;
3453 ctx.remove_block_param(removed);
3454
3455 let physical_order: Vec<_> = ctx.bodies[function]
3456 .params
3457 .iter()
3458 .map(|param| param.id)
3459 .collect();
3460 assert_eq!(physical_order, vec![first.local, last.local]);
3461 assert_eq!(ctx.block_param(first).index, 0);
3462 assert_eq!(ctx.block_param(last).index, 1);
3463
3464 let config = bincode::config::standard();
3465 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3466 let (mut restored, _): (Context<'static>, usize) =
3467 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3468
3469 assert!(!restored.contains_block_param(removed));
3470 assert_eq!(
3471 restored.bodies[function]
3472 .params
3473 .iter()
3474 .map(|param| param.id)
3475 .collect::<Vec<_>>(),
3476 physical_order,
3477 );
3478 assert!(restored.contains_block_param(first));
3479 assert!(restored.contains_block_param(last));
3480
3481 let fresh = BasicBlock::from_id_mut(&mut restored, block)
3482 .push_param(8)
3483 .id;
3484 assert!(fresh.local > last.local);
3485 assert_ne!(fresh, removed, "removed parameter IDs must never be reused");
3486 }
3487
3488 #[test]
3489 fn compact_block_arena_preserves_ids_across_round_trip() {
3490 let mut ctx = Context::new();
3491 let function = ctx.anon_function();
3492 let first = BasicBlock::make(&mut ctx, function).id;
3493 let removed = BasicBlock::make(&mut ctx, function).id;
3494 let last = BasicBlock::make(&mut ctx, function).id;
3495 FunctionBody::from_id_mut(&mut ctx, function)
3496 .set_root(first)
3497 .expect("set root");
3498
3499 ctx.delete_block(removed);
3500
3501 let physical_order: Vec<_> = ctx.bodies[function]
3502 .blocks
3503 .iter()
3504 .map(|block| block.id)
3505 .collect();
3506 assert_eq!(physical_order, vec![first.local, last.local]);
3507 assert_eq!(ctx.block_ids(), vec![first, last]);
3508
3509 let config = bincode::config::standard();
3510 let bytes = bincode::serde::encode_to_vec(&ctx, config).expect("encode");
3511 let (mut restored, _): (Context<'static>, usize) =
3512 bincode::serde::decode_from_slice(&bytes, config).expect("decode");
3513
3514 assert!(!restored.contains_block(removed));
3515 assert_eq!(
3516 restored.bodies[function]
3517 .blocks
3518 .iter()
3519 .map(|block| block.id)
3520 .collect::<Vec<_>>(),
3521 physical_order,
3522 );
3523 assert!(restored.contains_block(first));
3524 assert!(restored.contains_block(last));
3525 assert_eq!(
3526 FunctionBody::from_id(&restored, function)
3527 .root()
3528 .map(|block| block.id),
3529 Some(first),
3530 );
3531
3532 let fresh = BasicBlock::make(&mut restored, function).id;
3533 assert!(fresh.local > last.local);
3534 assert_ne!(fresh, removed, "removed block IDs must never be reused");
3535 }
3536
3537 #[test]
3538 fn deleting_root_clears_function_root() {
3539 let mut ctx = Context::new();
3540 let function = ctx.anon_function();
3541 let root = BasicBlock::make(&mut ctx, function).id;
3542 FunctionBody::from_id_mut(&mut ctx, function)
3543 .set_root(root)
3544 .expect("set root");
3545
3546 ctx.delete_block(root);
3547
3548 assert!(!ctx.contains_block(root));
3549 assert!(FunctionBody::from_id(&ctx, function).root().is_none());
3550 assert!(ctx.block_ids().is_empty());
3551 }
3552
3553 #[test]
3554 fn get_unique_name_resumes_probe_and_reuses_freed_suffixes() {
3555 use crate::value::VarnodeId;
3556
3557 let mut ctx = Context::new();
3558 let id = ValueId::Varnode(VarnodeId::from(0usize));
3559
3560 fn take(ctx: &mut Context<'static>, id: ValueId, base: &str) -> String {
3562 let name = ctx
3563 .get_unique_name(Cow::Owned(base.to_string()))
3564 .to_string();
3565 ctx.update_name(Cow::Owned(name.clone()), id, None).unwrap();
3566 name
3567 }
3568
3569 assert_eq!(take(&mut ctx, id, "tmp"), "tmp");
3571 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_1");
3572 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_2");
3573 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_3");
3574
3575 assert_eq!(take(&mut ctx, id, "x"), "x");
3577 assert_eq!(take(&mut ctx, id, "x"), "x_1");
3578
3579 ctx.update_name(Cow::Borrowed("relocated"), id, Some("tmp_1"))
3582 .unwrap();
3583 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_1");
3584 assert_eq!(take(&mut ctx, id, "tmp"), "tmp_4");
3586 }
3587
3588 mod split_function_at {
3591 use super::*;
3592
3593 use crate::value::insn::{Callee, Mnemonic, TailCall};
3594 use crate::value::{BasicBlock, FunctionBody, Instruction, Value};
3595 use std::borrow::Cow;
3596
3597 fn block_at(ctx: &mut Context<'static>, func: FunctionId, addr: u64) -> BlockId {
3598 BasicBlock::make(ctx, func).with_address(addr).id
3599 }
3600
3601 fn branch_at(ctx: &mut Context<'static>, block: BlockId, target: BlockId, addr: u64) {
3602 let id = (ctx).builder(block).push_branch(target).id;
3603 Instruction::from_id_mut(ctx, id).set_address(addr);
3604 }
3605
3606 fn cbranch_at(
3607 ctx: &mut Context<'static>,
3608 block: BlockId,
3609 success: BlockId,
3610 failure: BlockId,
3611 addr: u64,
3612 ) {
3613 let cond = ctx.get_const(1, 1).id();
3614 let id = (ctx).builder(block).push_cbranch(cond, success, failure).id;
3615 Instruction::from_id_mut(ctx, id).set_address(addr);
3616 }
3617
3618 fn return_at(ctx: &mut Context<'static>, block: BlockId, addr: u64) {
3619 let zero = ctx.get_const(0, 8).id();
3620 let id = (ctx).builder(block).push_return(zero).id;
3621 Instruction::from_id_mut(ctx, id).set_address(addr);
3622 }
3623
3624 fn block_at_addr(ctx: &Context, func: FunctionId, addr: u64) -> BlockId {
3625 FunctionBody::from_id(ctx, func)
3626 .block_ids()
3627 .into_iter()
3628 .find(|b| ctx.block(*b).address == Some(addr))
3629 .unwrap_or_else(|| panic!("{func:?} has no block at {addr:#x}"))
3630 }
3631
3632 fn addrs(ctx: &Context, func: FunctionId) -> Vec<u64> {
3633 let mut got: Vec<u64> = FunctionBody::from_id(ctx, func)
3634 .block_ids()
3635 .into_iter()
3636 .filter_map(|b| ctx.block(b).address)
3637 .collect();
3638 got.sort_unstable();
3639 got
3640 }
3641
3642 #[test]
3647 fn splits_absorbed_body_reusing_the_stub() {
3648 let mut ctx = Context::new();
3649 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("thunk"))).id;
3650 let b0 = block_at(&mut ctx, f, 0x1000);
3651 let b1 = block_at(&mut ctx, f, 0x2000);
3652 let b2 = block_at(&mut ctx, f, 0x2005);
3653 branch_at(&mut ctx, b0, b1, 0x1000);
3654 branch_at(&mut ctx, b1, b2, 0x2000);
3655 return_at(&mut ctx, b2, 0x2005);
3656 {
3657 let mut func = FunctionBody::from_id_mut(&mut ctx, f);
3658 func.set_root(b0).unwrap();
3659 }
3660 let g = FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("real"))).id;
3662
3663 let split_g = ctx.split_function_at(b1);
3664 assert_eq!(
3665 split_g, g,
3666 "the split must reuse the existing stub at 0x2000"
3667 );
3668
3669 assert_eq!(addrs(&ctx, f), vec![0x1000]);
3670 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2005]);
3671 let g_entry = block_at_addr(&ctx, g, 0x2000);
3672 assert_eq!(ctx.bodies[g].root_id(), Some(g_entry.local));
3673
3674 for b in FunctionBody::from_id(&ctx, g).block_ids() {
3676 assert_eq!(b.func, g);
3677 }
3678
3679 let f_entry = block_at_addr(&ctx, f, 0x1000);
3681 assert_eq!(BasicBlock::from_id(&ctx, f_entry).successors().count(), 0);
3682 let term = BasicBlock::from_id(&ctx, f_entry)
3683 .instructions()
3684 .last()
3685 .map(|i| i.mnemonic().clone());
3686 assert!(
3687 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
3688 "thunk branch must become TailCall(G), got {term:?}",
3689 );
3690 }
3691
3692 #[test]
3693 fn split_rehomes_temporary_values_spaces_and_pointer_types() {
3694 let mut ctx = Context::new();
3695 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3696 let entry = block_at(&mut ctx, f, 0x1000);
3697 let tail = block_at(&mut ctx, f, 0x2000);
3698 branch_at(&mut ctx, entry, tail, 0x1000);
3699 FunctionBody::from_id_mut(&mut ctx, f)
3700 .set_root(entry)
3701 .unwrap();
3702
3703 let temp = ctx
3704 .builder(tail)
3705 .make_named_temp(Cow::Borrowed("scratch"), 8);
3706 ctx.builder(entry)
3707 .make_named_temp(Cow::Borrowed("unused"), 4);
3708 let temp_space = ctx.bodies[f].temps[temp.local].space;
3709 let load = {
3710 let mut builder = ctx.builder(tail);
3711 let ValueId::Instruction(load) = builder
3712 .push_load::<false>(
3713 ValueId::Temp(temp),
3714 8,
3715 LocalMemorySpaceId::Temp(temp_space),
3716 )
3717 .id()
3718 else {
3719 unreachable!()
3720 };
3721 builder.push_return(ValueId::Instruction(load));
3722 load
3723 };
3724 let pointer_type = ctx
3725 .shared
3726 .types
3727 .get_or_make_space_address(8, MemorySpaceId::Temp(TempSpaceId::new(f, temp_space)));
3728 ctx.instruction_mut(load).type_id = pointer_type;
3729
3730 let g =
3731 FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("discovered"))).id;
3732 assert_eq!(ctx.split_function_at(tail), g);
3733
3734 let diagnostics = crate::verify_body_arena_integrity(&ctx);
3735 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
3736 assert_eq!(ctx.bodies[g].temp_spaces.len(), 1);
3737 assert_eq!(ctx.bodies[g].temps.len(), 1);
3738 assert_eq!(ctx.bodies[f].temps.len(), 2, "source arenas remain intact");
3739
3740 let moved_load = FunctionBody::from_id(&ctx, g)
3741 .blocks()
3742 .flat_map(|block| block.instructions())
3743 .find(|insn| matches!(insn.mnemonic(), Mnemonic::Load(_)))
3744 .expect("load moved with the split");
3745 let Mnemonic::Load(moved) = moved_load.mnemonic() else {
3746 unreachable!()
3747 };
3748 let LocalMemorySpaceId::Temp(moved_space) = moved.space else {
3749 panic!("load lost temporary-space provenance")
3750 };
3751 assert!(matches!(moved.ptr, crate::value::LocalValueId::Temp(_)));
3752 assert!(usize::from(moved_space) < ctx.bodies[g].temp_spaces.len());
3753 assert_eq!(
3754 ctx.shared.types.space_of(moved_load.type_id()),
3755 Some(MemorySpaceId::Temp(TempSpaceId::new(g, moved_space)))
3756 );
3757
3758 let rendered = FunctionBody::from_id(&ctx, g).to_string();
3760 assert!(rendered.contains("scratch"));
3761 }
3762
3763 #[test]
3764 fn split_stops_at_a_foreign_rootless_stub_address() {
3765 let mut ctx = Context::new();
3766 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3767 let entry = block_at(&mut ctx, f, 0x1000);
3768 let split = block_at(&mut ctx, f, 0x2000);
3769 let foreign_entry = block_at(&mut ctx, f, 0x3000);
3770 let foreign_body = block_at(&mut ctx, f, 0x3005);
3771 branch_at(&mut ctx, entry, split, 0x1000);
3772 branch_at(&mut ctx, split, foreign_entry, 0x2000);
3773 branch_at(&mut ctx, foreign_entry, foreign_body, 0x3000);
3774 return_at(&mut ctx, foreign_body, 0x3005);
3775 FunctionBody::from_id_mut(&mut ctx, f)
3776 .set_root(entry)
3777 .unwrap();
3778
3779 let g = FunctionBody::make_at_addr(&mut ctx, 0x2000, Some(Cow::Borrowed("g"))).id;
3780 let h = FunctionBody::make_at_addr(&mut ctx, 0x3000, Some(Cow::Borrowed("h"))).id;
3781 assert!(FunctionBody::from_id(&ctx, g).root().is_none());
3782 assert!(FunctionBody::from_id(&ctx, h).root().is_none());
3783
3784 assert_eq!(ctx.split_function_at(split), g);
3785 assert_eq!(addrs(&ctx, g), vec![0x2000]);
3786 assert_eq!(addrs(&ctx, f), vec![0x1000, 0x3000, 0x3005]);
3787 assert!(FunctionBody::from_id(&ctx, h).root().is_none());
3788
3789 let g_entry = block_at_addr(&ctx, g, 0x2000);
3790 let term = BasicBlock::from_id(&ctx, g_entry)
3791 .instructions()
3792 .last()
3793 .map(|i| i.mnemonic().clone());
3794 assert!(
3795 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(h)),
3796 "split tail must stop and tail-call rootless stub H, got {term:?}",
3797 );
3798 }
3799
3800 #[test]
3801 fn split_rehomes_block_param_origin_into_destination_arena() {
3802 let mut ctx = Context::new();
3803 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3804 let entry = block_at(&mut ctx, f, 0x1000);
3805 let tail = block_at(&mut ctx, f, 0x2000);
3806 let param = BasicBlock::from_id_mut(&mut ctx, tail).push_param(8).id;
3807 crate::value::BlockParam::from_id_mut(&mut ctx, param)
3808 .set_origin(ValueId::BlockParam(param));
3809
3810 let arg = ctx.get_const(7, 8).id();
3811 let branch = ctx.builder(entry).push_branch_with_args(tail, vec![arg]).id;
3812 Instruction::from_id_mut(&mut ctx, branch).set_address(0x1000);
3813 let ret = ctx.builder(tail).push_return(ValueId::BlockParam(param)).id;
3814 Instruction::from_id_mut(&mut ctx, ret).set_address(0x2000);
3815 FunctionBody::from_id_mut(&mut ctx, f)
3816 .set_root(entry)
3817 .unwrap();
3818
3819 let g = ctx.split_function_at(tail);
3820 let new_tail = block_at_addr(&ctx, g, 0x2000);
3821 let new_param = BasicBlock::from_id(&ctx, new_tail).params().next().unwrap();
3822 assert_eq!(new_param.origin(), Some(ValueId::BlockParam(new_param.id)));
3823 }
3824
3825 #[test]
3833 fn split_rehomes_symbolic_block_literals() {
3834 use crate::value::literal::SymbolicRef;
3835
3836 let mut ctx = Context::new();
3837 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3838 let entry = block_at(&mut ctx, f, 0x1000);
3839 let tail = block_at(&mut ctx, f, 0x2000);
3840 let landing = block_at(&mut ctx, f, 0x2008);
3841
3842 let lit = ctx.get_const(0x2008, 8).id();
3845 let ValueId::Literal(lit_id) = lit else {
3846 panic!("expected a literal");
3847 };
3848 ctx.shared.values.literals[lit_id].symbolic = Some(SymbolicRef::Block(landing));
3849
3850 branch_at(&mut ctx, entry, tail, 0x1000);
3851 let ind = ctx.builder(tail).push_branchind(lit).id;
3853 Instruction::from_id_mut(&mut ctx, ind).set_address(0x2000);
3854 ctx.add_cfg_edge(tail, landing);
3855 return_at(&mut ctx, landing, 0x2008);
3856 FunctionBody::from_id_mut(&mut ctx, f)
3857 .set_root(entry)
3858 .unwrap();
3859
3860 let g = ctx.split_function_at(tail);
3861
3862 let new_landing = block_at_addr(&ctx, g, 0x2008);
3863 let new_tail = block_at_addr(&ctx, g, 0x2000);
3864 let Mnemonic::BranchInd(b) = BasicBlock::from_id(&ctx, new_tail)
3865 .instructions()
3866 .last()
3867 .unwrap()
3868 .mnemonic()
3869 .clone()
3870 else {
3871 panic!("tail must still end in an indirect branch");
3872 };
3873 let crate::value::LocalValueId::Literal(new_lit) = b.ptr else {
3874 panic!("indirect branch operand must still be a literal");
3875 };
3876 assert_eq!(
3877 ctx.shared.values.literals[new_lit].symbolic,
3878 Some(SymbolicRef::Block(new_landing)),
3879 "the relocated literal must name the clone, not the deleted original",
3880 );
3881 assert_eq!(
3882 ctx.shared.values.literals[new_lit].value, 0x2008,
3883 "re-pointing the symbol must not disturb the numeric value",
3884 );
3885 }
3886
3887 #[test]
3891 fn conditional_arm_into_split_block_uses_a_trampoline() {
3892 let mut ctx = Context::new();
3893 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
3894 let entry = block_at(&mut ctx, f, 0x1000);
3895 let cont = block_at(&mut ctx, f, 0x1008);
3896 let tail = block_at(&mut ctx, f, 0x2000);
3897 cbranch_at(&mut ctx, entry, tail, cont, 0x1000);
3898 return_at(&mut ctx, cont, 0x1008);
3899 return_at(&mut ctx, tail, 0x2000);
3900 FunctionBody::from_id_mut(&mut ctx, f)
3901 .set_root(entry)
3902 .unwrap();
3903
3904 let g = ctx.split_function_at(tail);
3905
3906 let entry = block_at_addr(&ctx, f, 0x1000);
3907 let cont = block_at_addr(&ctx, f, 0x1008);
3908 assert_eq!(addrs(&ctx, g), vec![0x2000]);
3909
3910 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, entry)
3911 .instructions()
3912 .last()
3913 .unwrap()
3914 .mnemonic()
3915 .clone()
3916 else {
3917 panic!("entry must still end in a cbranch");
3918 };
3919 assert_eq!(cb.failure_block, cont.local, "fall-through arm untouched");
3920 let tramp = BlockId::new(entry.func, cb.success_block);
3921 assert_eq!(
3922 BasicBlock::from_id(&ctx, tramp).parent().map(|f| f.id),
3923 Some(f),
3924 "trampoline lives in F",
3925 );
3926 let term = BasicBlock::from_id(&ctx, tramp)
3927 .instructions()
3928 .last()
3929 .map(|i| i.mnemonic().clone());
3930 assert!(
3931 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
3932 "trampoline must tail-call G, got {term:?}",
3933 );
3934 for (_, s) in BasicBlock::from_id(&ctx, entry).successors() {
3936 assert_eq!(BasicBlock::from_id(&ctx, s).parent().map(|f| f.id), Some(f));
3937 }
3938 }
3939
3940 #[test]
3943 fn mints_a_conventional_function_when_no_stub_exists() {
3944 let mut ctx = Context::new();
3945 wazabin_qcode_macro::qcode!(
3946 ctx,
3947 "
3948 fn f:
3949 <entry>
3950 goto <0x1008>;
3951 <0x1008>
3952 return 0x0;
3953 "
3954 );
3955
3956 let mid = block_at_addr(&ctx, f, 0x1008);
3957 let g = ctx.split_function_at(mid);
3958 assert_eq!(FunctionBody::from_id(&ctx, g).name(), "fn_1008");
3959 assert_eq!(FunctionBody::from_id(&ctx, f).block_ids().len(), 1);
3961 assert_eq!(addrs(&ctx, g), vec![0x1008]);
3962 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
3963 assert_eq!(addresses.function_at(0x1008), Some(g));
3964 for b in FunctionBody::from_id(&ctx, g).block_ids() {
3965 assert_eq!(b.func, g);
3966 }
3967 }
3968
3969 fn assert_no_dangling_terminators(ctx: &Context) {
3973 for b in ctx.block_ids() {
3974 let Some(mnemonic) = BasicBlock::from_id(ctx, b)
3975 .instructions()
3976 .last()
3977 .map(|t| t.mnemonic().clone())
3978 else {
3979 continue;
3980 };
3981 let targets = match &mnemonic {
3982 Mnemonic::Branch(crate::value::insn::Branch { target, .. }) => vec![*target],
3983 Mnemonic::CBranch(crate::value::insn::CBranch {
3984 success_block,
3985 failure_block,
3986 ..
3987 }) => vec![*success_block, *failure_block],
3988 _ => vec![],
3989 };
3990 let succs: std::collections::HashSet<BlockId> = BasicBlock::from_id(ctx, b)
3991 .successors()
3992 .map(|(_, s)| s)
3993 .collect();
3994 for t in targets {
3995 let tid = BlockId::new(b.func, t);
3996 assert!(
3997 ctx.contains_block(tid),
3998 "block {b:?} terminator names dead block {tid:?}"
3999 );
4000 assert!(
4001 succs.contains(&tid),
4002 "block {b:?} terminator target {tid:?} has no CFG edge (operand/edge desync)"
4003 );
4004 }
4005 }
4006 }
4007
4008 #[test]
4013 fn retained_predecessor_into_mid_tail_promotes_the_landing() {
4014 let mut ctx = Context::new();
4015 wazabin_qcode_macro::qcode!(
4018 ctx,
4019 "
4020 fn f:
4021 <entry @c:i8>
4022 if @c goto <0x2000> else goto <0x1008>;
4023 <0x1008>
4024 goto <0x2008>;
4025 <0x2000>
4026 goto <0x2008>;
4027 <0x2008>
4028 return 0x0;
4029 "
4030 );
4031
4032 let tail = block_at_addr(&ctx, f, 0x2000);
4033 let g = ctx.split_function_at(tail);
4034
4035 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
4038 let landing_fn = addresses
4039 .function_at(0x2008)
4040 .expect("mid-tail landing must be promoted to a function");
4041 assert_ne!(landing_fn, g);
4042 assert_eq!(addrs(&ctx, g), vec![0x2000]);
4043 assert_no_dangling_terminators(&ctx);
4044
4045 for (holder, addr) in [(f, 0x1008u64), (g, 0x2000u64)] {
4046 let block = block_at_addr(&ctx, holder, addr);
4047 let term = BasicBlock::from_id(&ctx, block)
4048 .instructions()
4049 .last()
4050 .map(|i| i.mnemonic().clone());
4051 assert!(
4052 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(landing_fn)),
4053 "branch at {addr:#x} into the landing must tail-call it, got {term:?}",
4054 );
4055 }
4056 }
4057
4058 #[test]
4066 fn tail_conditional_to_own_registered_entry_uses_a_trampoline() {
4067 let mut ctx = Context::new();
4068 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4069 let entry = block_at(&mut ctx, f, 0x1000);
4070 let tail = block_at(&mut ctx, f, 0x2000);
4071 let cont = block_at(&mut ctx, f, 0x2008);
4072 branch_at(&mut ctx, entry, tail, 0x1000);
4073 cbranch_at(&mut ctx, tail, entry, cont, 0x2000);
4075 return_at(&mut ctx, cont, 0x2008);
4076 FunctionBody::from_id_mut(&mut ctx, f)
4077 .set_root(entry)
4078 .unwrap();
4079
4080 let g = ctx.split_function_at(tail);
4081
4082 assert_no_dangling_terminators(&ctx);
4083 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4084 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4085
4086 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4089 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4090 .instructions()
4091 .last()
4092 .unwrap()
4093 .mnemonic()
4094 .clone()
4095 else {
4096 panic!("moved tail must still end in a cbranch");
4097 };
4098 let tramp = BlockId::new(g, cb.success_block);
4099 assert_eq!(tramp.func, g, "trampoline must have relocated into g");
4100 let term = BasicBlock::from_id(&ctx, tramp)
4101 .instructions()
4102 .last()
4103 .map(|i| i.mnemonic().clone());
4104 assert!(
4105 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(f)),
4106 "back-edge trampoline must tail-call f, got {term:?}",
4107 );
4108 }
4109
4110 #[test]
4115 fn tail_conditional_to_foreign_entry_relocates_its_trampoline() {
4116 let mut ctx = Context::new();
4117 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4118 let entry = block_at(&mut ctx, f, 0x1000);
4119 let tail = block_at(&mut ctx, f, 0x2000);
4120 let cont = block_at(&mut ctx, f, 0x2008);
4121 let foreign = block_at(&mut ctx, f, 0x3000);
4122 branch_at(&mut ctx, entry, tail, 0x1000);
4123 cbranch_at(&mut ctx, tail, foreign, cont, 0x2000);
4125 return_at(&mut ctx, cont, 0x2008);
4126 return_at(&mut ctx, foreign, 0x3000);
4127 FunctionBody::from_id_mut(&mut ctx, f)
4128 .set_root(entry)
4129 .unwrap();
4130 let h = FunctionBody::make_at_addr(&mut ctx, 0x3000, Some(Cow::Borrowed("h"))).id;
4131
4132 let g = ctx.split_function_at(tail);
4133
4134 assert_no_dangling_terminators(&ctx);
4135 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4136 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4137
4138 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4141 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4142 .instructions()
4143 .last()
4144 .unwrap()
4145 .mnemonic()
4146 .clone()
4147 else {
4148 panic!("moved tail must still end in a cbranch");
4149 };
4150 let tramp = BlockId::new(g, cb.success_block);
4151 assert_eq!(tramp.func, g, "trampoline must have relocated into g");
4152 let term = BasicBlock::from_id(&ctx, tramp)
4153 .instructions()
4154 .last()
4155 .map(|i| i.mnemonic().clone());
4156 assert!(
4157 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(h)),
4158 "relocated trampoline must tail-call H, got {term:?}",
4159 );
4160 }
4161
4162 #[test]
4166 fn conditional_failure_arm_into_split_block_uses_a_trampoline() {
4167 let mut ctx = Context::new();
4168 wazabin_qcode_macro::qcode!(
4171 ctx,
4172 "
4173 fn f:
4174 <entry @c:i8>
4175 if @c goto <0x1008> else goto <0x2000>;
4176 <0x1008>
4177 return 0x0;
4178 <0x2000>
4179 return 0x0;
4180 "
4181 );
4182
4183 let tail = block_at_addr(&ctx, f, 0x2000);
4184 let g = ctx.split_function_at(tail);
4185
4186 assert_no_dangling_terminators(&ctx);
4187 let entry = BlockId::new(f, ctx.bodies[f].root_id().unwrap());
4188 let cont = block_at_addr(&ctx, f, 0x1008);
4189 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, entry)
4190 .instructions()
4191 .last()
4192 .unwrap()
4193 .mnemonic()
4194 .clone()
4195 else {
4196 panic!("entry must still end in a cbranch");
4197 };
4198 assert_eq!(
4199 cb.success_block, cont.local,
4200 "success (fall-through) untouched"
4201 );
4202 let tramp = BlockId::new(entry.func, cb.failure_block);
4203 let term = BasicBlock::from_id(&ctx, tramp)
4204 .instructions()
4205 .last()
4206 .map(|i| i.mnemonic().clone());
4207 assert!(
4208 matches!(term, Some(Mnemonic::TailCall(TailCall { target, .. })) if target == Callee::Real(g)),
4209 "failure arm must route through a trampoline tail-calling G, got {term:?}",
4210 );
4211 }
4212
4213 #[test]
4216 fn moved_tail_internal_conditional_remaps_both_arms() {
4217 let mut ctx = Context::new();
4218 wazabin_qcode_macro::qcode!(
4221 ctx,
4222 "
4223 fn f:
4224 <entry>
4225 goto <0x2000>;
4226 <0x2000>
4227 %c = 0x0 == 0x0;
4228 if %c goto <0x2008> else goto <0x2010>;
4229 <0x2008>
4230 return 0x0;
4231 <0x2010>
4232 return 0x0;
4233 "
4234 );
4235
4236 let tail = block_at_addr(&ctx, f, 0x2000);
4237 let g = ctx.split_function_at(tail);
4238
4239 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2008, 0x2010]);
4240 assert_no_dangling_terminators(&ctx);
4241 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4242 let Mnemonic::CBranch(cb) = BasicBlock::from_id(&ctx, moved_tail)
4243 .instructions()
4244 .last()
4245 .unwrap()
4246 .mnemonic()
4247 .clone()
4248 else {
4249 panic!("moved tail must still end in a cbranch");
4250 };
4251 let a = block_at_addr(&ctx, g, 0x2008);
4252 let b = block_at_addr(&ctx, g, 0x2010);
4253 assert_eq!(cb.success_block, a.local, "success arm re-pointed to clone");
4254 assert_eq!(cb.failure_block, b.local, "failure arm re-pointed to clone");
4255 }
4256
4257 #[test]
4260 fn moved_tail_internal_branch_remaps_target() {
4261 let mut ctx = Context::new();
4262 wazabin_qcode_macro::qcode!(
4263 ctx,
4264 "
4265 fn f:
4266 <entry>
4267 goto <0x2000>;
4268 <0x2000>
4269 goto <0x2008>;
4270 <0x2008>
4271 return 0x0;
4272 "
4273 );
4274
4275 let tail = block_at_addr(&ctx, f, 0x2000);
4276 let g = ctx.split_function_at(tail);
4277
4278 assert_eq!(addrs(&ctx, g), vec![0x2000, 0x2008]);
4279 assert_no_dangling_terminators(&ctx);
4280 let moved_tail = block_at_addr(&ctx, g, 0x2000);
4281 let Mnemonic::Branch(br) = BasicBlock::from_id(&ctx, moved_tail)
4282 .instructions()
4283 .last()
4284 .unwrap()
4285 .mnemonic()
4286 .clone()
4287 else {
4288 panic!("moved tail must still end in a branch");
4289 };
4290 let end = block_at_addr(&ctx, g, 0x2008);
4291 assert_eq!(br.target, end.local, "internal branch re-pointed to clone");
4292 }
4293
4294 #[test]
4297 fn split_rehomes_store_temporary_space() {
4298 let mut ctx = Context::new();
4299 let f = FunctionBody::make_at_addr(&mut ctx, 0x1000, Some(Cow::Borrowed("f"))).id;
4300 let entry = block_at(&mut ctx, f, 0x1000);
4301 let tail = block_at(&mut ctx, f, 0x2000);
4302 branch_at(&mut ctx, entry, tail, 0x1000);
4303
4304 let slot = ctx.builder(tail).make_named_temp(Cow::Borrowed("slot"), 8);
4305 let space = ctx.bodies[f].temps[slot.local].space;
4306 let value = ctx.get_const(0x2a, 8).id();
4307 {
4308 let mut builder = ctx.builder(tail);
4309 builder.push_store(value, ValueId::Temp(slot), LocalMemorySpaceId::Temp(space));
4310 builder.push_return(value);
4311 }
4312 FunctionBody::from_id_mut(&mut ctx, f)
4313 .set_root(entry)
4314 .unwrap();
4315
4316 let g = ctx.split_function_at(tail);
4317 let diagnostics = crate::verify_body_arena_integrity(&ctx);
4318 assert!(diagnostics.is_empty(), "{diagnostics:#?}");
4319
4320 let moved_store = FunctionBody::from_id(&ctx, g)
4321 .blocks()
4322 .flat_map(|block| block.instructions())
4323 .find(|insn| matches!(insn.mnemonic(), Mnemonic::Store(_)))
4324 .expect("store moved with the split");
4325 let Mnemonic::Store(moved) = moved_store.mnemonic() else {
4326 unreachable!()
4327 };
4328 let LocalMemorySpaceId::Temp(moved_space) = moved.space else {
4329 panic!("store lost temporary-space provenance")
4330 };
4331 assert!(usize::from(moved_space) < ctx.bodies[g].temp_spaces.len());
4332 }
4333 }
4334}