1use crate::value::QCodeMut;
2use crate::{
3 context::Context,
4 error::Result,
5 value::{
6 FunctionBody, Instruction, LocalValueId, ModuleView, QCodeView, Value, ValueId,
7 block_param::{BlockParam, BlockParamId, BlockParamMutRef, BlockParamRef, LocalParamId},
8 function::{FunctionId, FunctionMutRef, FunctionRef},
9 insn::{InstructionId, InstructionRef, LocalInsnId, Mnemonic},
10 util::{
11 base_ref::{BaseRef, WithCtx, WithCtxMut},
12 body_mut::BodyMut,
13 named::{Named, Renameable},
14 },
15 },
16};
17use core::slice;
18use jstd::graph::FxBuildHasher;
19use std::{
20 borrow::Cow,
21 collections::HashSet,
22 fmt::{Display, Formatter},
23 marker::PhantomData,
24};
25
26use rustc_hash::FxHashMap as HashMap;
27
28pub(crate) use self::cfg::EdgeData;
29pub use self::cfg::{BlockId, EdgeId};
30pub mod cfg;
31
32pub(crate) fn substitute_operands(mnemonic: &mut Mnemonic, pairs: &[(LocalValueId, LocalValueId)]) {
35 let mut occupied: Vec<LocalValueId> = mnemonic
36 .args()
37 .into_iter()
38 .chain(pairs.iter().map(|&(_, new)| new))
39 .collect();
40 let mut sentinels = Vec::with_capacity(pairs.len());
41 let mut next = 0usize;
42 for _ in pairs {
43 let sentinel = loop {
44 let candidate = LocalValueId::Varnode(crate::value::VarnodeId::from(next));
45 next += 1;
46 if !occupied.contains(&candidate) {
47 occupied.push(candidate);
48 break candidate;
49 }
50 };
51 sentinels.push(sentinel);
52 }
53 for (&(old, _), &sentinel) in pairs.iter().zip(&sentinels) {
54 mnemonic.replace_value(old, sentinel);
55 }
56 for (&(_, new), &sentinel) in pairs.iter().zip(&sentinels) {
57 mnemonic.replace_value(sentinel, new);
58 }
59}
60
61#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
64pub struct BasicBlock<'str> {
65 name: Option<Cow<'str, str>>,
67
68 comment: Option<String>,
70
71 pub params: Vec<LocalParamId>,
74
75 pub instructions: Vec<LocalInsnId>,
77
78 pub edges: HashSet<EdgeId, FxBuildHasher>,
90
91 pub address: Option<u64>,
93
94 pub extra_addresses: Vec<u64>,
96}
97
98impl<'str> BasicBlock<'str> {
99 pub fn instruction_ids(&self) -> &[LocalInsnId] {
104 &self.instructions
105 }
106
107 pub fn param_ids(&self) -> &[LocalParamId] {
111 &self.params
112 }
113
114 pub fn from_id<'ctx>(ctx: &'ctx Context<'str>, id: BlockId) -> BlockRef<'str, 'ctx> {
116 BlockRef::new(ModuleView::new(ctx), id)
117 }
118
119 pub fn from_id_mut<'ctx>(ctx: &'ctx mut Context<'str>, id: BlockId) -> BlockMutRef<'str, 'ctx> {
121 BlockMutRef::new(ctx, id)
122 }
123
124 pub fn from_name<'ctx>(ctx: &'ctx Context<'str>, name: &str) -> Option<BlockRef<'str, 'ctx>> {
130 ctx.functions()
131 .find_map(|f| f.local_named(name))
132 .and_then(ValueId::as_block)
133 .map(|id| BasicBlock::from_id(ctx, id))
134 }
135
136 pub fn make<'ctx>(ctx: &'ctx mut Context<'str>, func: FunctionId) -> BlockMutRef<'str, 'ctx> {
139 let id = ctx.push_block(func, BasicBlock::default());
140 BlockMutRef::new(ctx, id)
141 }
142
143 pub(crate) fn set_name(&mut self, name: Option<Cow<'str, str>>) {
147 self.name = name;
148 }
149
150 pub(crate) fn local_name(&self) -> Option<&str> {
152 self.name.as_deref()
153 }
154
155 pub(crate) fn detached() -> Self {
160 BasicBlock::default()
161 }
162
163 pub fn clone_block_into(
177 ctx: &mut Context<'str>,
178 orig: BlockId,
179 target: FunctionId,
180 value_map: &mut HashMap<ValueId, ValueId>,
181 ) -> BlockId {
182 let new_block_id = BasicBlock::make(ctx, target).id;
183
184 let name = ctx.block(orig).name.clone().unwrap_or_else(|| {
187 Cow::Owned(format!("clone_{:x}", ctx.block(orig).address.unwrap_or(0)))
188 });
189 let unique_name = ctx.get_unique_name_in(target, name);
190 BasicBlock::from_id_mut(ctx, new_block_id)
191 .rename(unique_name)
192 .expect("name was deduplicated");
193
194 for old_param_local in ctx.block(orig).params.clone() {
196 let old_param_id = BlockParamId::new(orig.func, old_param_local);
197 let old_param = ctx.block_param(old_param_id).clone();
198 let new_param_id = ctx.push_block_param(
199 target,
200 BlockParam {
201 parent: Some(new_block_id.local),
202 ..old_param
203 },
204 );
205 BasicBlock::from_id_mut(ctx, new_block_id).push_existing_param(new_param_id);
206 value_map.insert(
207 ValueId::BlockParam(old_param_id),
208 ValueId::BlockParam(new_param_id),
209 );
210 }
211
212 let orig_insns = ctx.block(orig).instructions.clone();
215 for old_insn_local in orig_insns {
216 let old_insn_id = InstructionId::new(orig.func, old_insn_local);
217 let (mnemonic, type_id, address) = {
218 let insn = Instruction::from_id(&*ctx, old_insn_id);
219 (insn.mnemonic().clone(), insn.type_id(), insn.address())
220 };
221 let new_insn_id =
222 InstructionRef::from_mnemonic_with_type(ctx, target, mnemonic, type_id).id;
223 if let Some(addr) = address {
224 ctx.instruction_mut(new_insn_id).set_address(addr);
225 }
226 BasicBlock::from_id_mut(ctx, new_block_id).push_insn(new_insn_id);
227 value_map.insert(
228 ValueId::Instruction(old_insn_id),
229 ValueId::Instruction(new_insn_id),
230 );
231 }
232
233 new_block_id
234 }
235
236 pub fn clone_into_ctx(
239 ctx: &mut Context<'str>,
240 orig: BlockId,
241 value_map: &mut HashMap<ValueId, ValueId>,
242 ) -> BlockId {
243 let new_block_id = BasicBlock::make(ctx, orig.func).id;
245
246 let name = Cow::Owned(format!("clone_{:x}", ctx.block(orig).address.unwrap_or(0)));
247 let unique_name = ctx.get_unique_name_in(new_block_id.func, name);
248 BasicBlock::from_id_mut(ctx, new_block_id)
249 .rename(unique_name)
250 .expect("name was deduplicated");
251
252 for old_param_local in ctx.block(orig).params.clone() {
254 let old_param_id = BlockParamId::new(orig.func, old_param_local);
255 let old_param = ctx.block_param(old_param_id).clone();
256 let new_param_id = ctx.push_block_param(
257 new_block_id.func,
258 BlockParam {
259 parent: Some(new_block_id.local),
260 ..old_param
261 },
262 );
263
264 BasicBlock::from_id_mut(ctx, new_block_id).push_existing_param(new_param_id);
265 value_map.insert(
266 ValueId::BlockParam(old_param_id),
267 ValueId::BlockParam(new_param_id),
268 );
269 }
270
271 let orig_insns = ctx.block(orig).instructions.clone();
273 for old_insn_local in orig_insns {
274 let old_insn_id = InstructionId::new(orig.func, old_insn_local);
275 let insn_ref = Instruction::from_id(&*ctx, old_insn_id);
277 let size = insn_ref.size();
278 let space = insn_ref.space().map(|s| s.id);
279
280 let mut new_mnemonic = insn_ref.mnemonic().clone();
282
283 let pairs: Vec<_> = new_mnemonic
288 .args()
289 .into_iter()
290 .filter_map(|old| {
291 let qualified = old.qualify(orig.func);
295 value_map
296 .get(&qualified)
297 .map(|&new| (old, new.localize(new_block_id.func)))
298 })
299 .collect();
300 substitute_operands(&mut new_mnemonic, &pairs);
301
302 let new_insn_id = InstructionRef::from_mnemonic_with_space(
303 ctx,
304 new_block_id.func,
305 new_mnemonic,
306 size,
307 space,
308 )
309 .id;
310
311 BasicBlock::from_id_mut(ctx, new_block_id).push_insn(new_insn_id);
312
313 value_map.insert(
314 ValueId::Instruction(old_insn_id),
315 ValueId::Instruction(new_insn_id),
316 );
317 }
318
319 new_block_id
320 }
321}
322
323impl<'s, 'ctx: 's, 'str: 'ctx, R> BlockRef<'str, 'ctx, R>
325where
326 R: QCodeView<'ctx, 'str>,
327{
328 fn inner(&'s self) -> &'ctx BasicBlock<'str> {
329 self.view.block(self.id)
330 }
331
332 pub fn successors(&'s self) -> impl Iterator<Item = (EdgeId, BlockId)> + 's {
340 let view = self.view;
341 let id = self.id;
342 self.inner().edges.iter().copied().filter_map(move |edge| {
343 let e = view.edge(id.func, edge);
344 (e.from == id.local).then_some((edge, BlockId::new(id.func, e.to)))
345 })
346 }
347
348 pub fn predecessors(&'s self) -> impl Iterator<Item = (EdgeId, BlockId)> + 's {
351 let view = self.view;
352 let id = self.id;
353 self.inner().edges.iter().copied().filter_map(move |edge| {
354 let e = view.edge(id.func, edge);
355 (e.to == id.local).then_some((edge, BlockId::new(id.func, e.from)))
356 })
357 }
358
359 pub fn name(&'s self) -> Option<&'ctx str> {
360 self.inner().name.as_deref()
361 }
362
363 pub fn address(&'s self) -> Option<u64> {
365 self.inner().address
366 }
367
368 pub fn comment(&'s self) -> Option<&'ctx str> {
369 self.inner().comment.as_deref()
370 }
371
372 pub fn params(&'s self) -> impl Iterator<Item = BlockParamRef<'str, 'ctx, R>> + 's {
374 let func = self.id.func;
375 self.inner()
376 .params
377 .iter()
378 .map(move |&local| BlockParamRef::new(self.view, BlockParamId::new(func, local)))
379 }
380
381 pub fn num_params(&'s self) -> usize {
383 self.inner().params.len()
384 }
385
386 pub fn instructions(&'s self) -> InstructionIter<'str, 'ctx, R> {
388 let inner = self.inner();
389 InstructionIter {
390 view: self.view,
391 func: self.id.func,
392 inner: inner.instructions.iter(),
393 marker: PhantomData,
394 }
395 }
396
397 pub fn iter(&'s self) -> InstructionIter<'str, 'ctx, R> {
400 self.instructions()
401 }
402
403 pub fn instruction_ids(&'s self) -> Vec<InstructionId> {
404 let func = self.id.func;
405 self.inner()
406 .instructions
407 .iter()
408 .map(|&local| InstructionId::new(func, local))
409 .collect()
410 }
411
412 pub fn instruction_count(&'s self) -> usize {
419 self.inner().instructions.len()
420 }
421
422 pub fn is_empty(&'s self) -> bool {
424 self.inner().instructions.is_empty()
425 }
426
427 pub fn is_terminated(&'s self) -> bool {
429 self.iter().last().is_some_and(|insn| insn.is_terminator())
430 }
431
432 pub fn parent(&'s self) -> Option<FunctionRef<'str, 'ctx, R>> {
433 Some(FunctionRef::new(self.view, self.id.func))
436 }
437
438 pub fn function(&'s self) -> Option<FunctionRef<'str, 'ctx, R>> {
439 self.parent()
440 }
441
442 fn fmt(&'s self, f: &mut Formatter<'_>) -> std::fmt::Result {
443 let name = self.name().unwrap_or("unnamed");
444 write!(f, "<{name}")?;
445 for param in self.params() {
446 write!(f, " ")?;
447 param.fmt_decl(f)?;
448 }
449 writeln!(f, ">")?;
450
451 if let Some(comment) = self.comment() {
452 for line in comment.lines() {
453 writeln!(f, "\t// {line}")?;
454 }
455 }
456
457 self.iter().try_for_each(|instr| {
458 write!(f, "\t")?;
459 instr.as_statement().fmt(f)?;
460 writeln!(f)
461 })?;
462
463 use crate::value::insn::Mnemonic;
469 if let Some(term) = self.iter().last()
470 && matches!(
471 term.mnemonic(),
472 Mnemonic::Call(_) | Mnemonic::CallInd(_) | Mnemonic::BranchInd(_)
473 )
474 {
475 let mut succ: Vec<&str> = self
476 .successors()
477 .map(|(_, b)| BlockRef::new(self.view, b).name().unwrap_or("unnamed"))
478 .collect();
479 if !succ.is_empty() {
480 succ.sort_unstable();
481 write!(f, "\t// ->")?;
482 for (i, name) in succ.iter().enumerate() {
483 write!(f, "{} <{name}>", if i == 0 { "" } else { "," })?;
484 }
485 writeln!(f)?;
486 }
487 }
488
489 Ok(())
490 }
491}
492
493#[derive(Clone, Copy)]
494pub struct BlockRef<'str, 'ctx, R = ModuleView<'ctx, 'str>> {
495 pub id: BlockId,
496 pub(in crate::value) view: R,
497 marker: PhantomData<&'ctx &'str ()>,
498}
499
500impl<'str, 'ctx, R> BlockRef<'str, 'ctx, R> {
501 pub fn new(view: R, id: BlockId) -> Self {
502 Self {
503 id,
504 view,
505 marker: PhantomData,
506 }
507 }
508
509 pub fn id(&self) -> ValueId {
510 self.id.into()
511 }
512}
513
514impl<'str, 'ctx> BlockRef<'str, 'ctx> {
515 pub fn from_id(ctx: &'ctx Context<'str>, id: BlockId) -> Self {
516 Self::new(ModuleView::new(ctx), id)
517 }
518}
519
520impl<'s, 'ctx: 's, 'str: 'ctx> WithCtx<'s, 'ctx, 'str> for BlockRef<'str, 'ctx> {
521 fn ctx(&'s self) -> &'ctx Context<'str> {
522 self.view.context()
526 }
527}
528
529impl<'str: 'ctx, 'ctx, R> Named for BlockRef<'str, 'ctx, R>
530where
531 R: QCodeView<'ctx, 'str>,
532{
533 fn name(&self) -> Option<&str> {
534 self.view.block(self.id).name.as_deref()
535 }
536}
537
538impl<'str: 'ctx, 'ctx, R> Display for BlockRef<'str, 'ctx, R>
539where
540 R: QCodeView<'ctx, 'str>,
541{
542 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
543 BlockRef::fmt(self, f)
544 }
545}
546
547impl<'str: 'ctx, 'ctx, R> Value<'str, 'ctx> for BlockRef<'str, 'ctx, R>
548where
549 R: QCodeView<'ctx, 'str>,
550{
551 fn id(&self) -> ValueId {
552 self.id()
553 }
554
555 fn size(&self) -> usize {
556 0
557 }
558}
559
560pub struct InstructionIter<'str, 'ctx, R = ModuleView<'ctx, 'str>> {
561 view: R,
562 func: FunctionId,
563 inner: slice::Iter<'ctx, LocalInsnId>,
564 marker: PhantomData<&'str ()>,
565}
566
567impl<'str: 'ctx, 'ctx, R> Iterator for InstructionIter<'str, 'ctx, R>
568where
569 R: QCodeView<'ctx, 'str>,
570{
571 type Item = InstructionRef<'str, 'ctx, R>;
572
573 fn next(&mut self) -> Option<Self::Item> {
574 self.inner
575 .next()
576 .map(|&local| InstructionRef::new(self.view, InstructionId::new(self.func, local)))
577 }
578}
579
580impl<'str: 'ctx, 'ctx, R> IntoIterator for &BlockRef<'str, 'ctx, R>
581where
582 R: QCodeView<'ctx, 'str>,
583{
584 type Item = InstructionRef<'str, 'ctx, R>;
585 type IntoIter = InstructionIter<'str, 'ctx, R>;
586
587 fn into_iter(self) -> Self::IntoIter {
588 self.iter()
589 }
590}
591
592pub type BlockMutRef<'str, 'ctx> = BaseRef<&'ctx mut Context<'str>, BlockId>;
593
594impl<'s, 'ctx: 's, 'str: 'ctx> WithCtxMut<'s, 'str> for BlockMutRef<'str, 'ctx> {
595 fn ctx_mut(&'s mut self) -> &'s mut Context<'str> {
596 self.ctx
597 }
598}
599
600impl<'s, 'str> WithCtx<'s, 's, 'str> for BaseRef<&mut Context<'str>, BlockId>
604where
605 'str: 's,
606{
607 fn ctx(&'s self) -> &'s Context<'str> {
608 self.ctx
609 }
610}
611
612impl Named for BlockMutRef<'_, '_> {
616 fn name(&self) -> Option<&str> {
617 self.ctx.block(self.id).name.as_deref()
618 }
619}
620
621impl<'a, 'str> Named for BaseRef<BodyMut<'a, 'str>, BlockId> {
622 fn name(&self) -> Option<&str> {
623 self.ctx.fun.blocks[self.id.local].name.as_deref()
624 }
625}
626
627impl<'str, 'ctx, H: QCodeMut<'str>> Renameable<'str, 'ctx> for BaseRef<H, BlockId>
629where
630 Self: Named,
631{
632 fn rename(&mut self, name: Cow<'str, str>) -> Result<()> {
633 self.rename_local(name)
634 }
635}
636
637impl<'str, H: QCodeMut<'str>> BaseRef<H, BlockId> {
640 pub fn is_terminated(&self) -> bool {
642 let body = self.ctx.body(self.id.func);
643 body.block(self.id)
644 .instructions
645 .last()
646 .is_some_and(|&local| {
647 body.insn(InstructionId::new(self.id.func, local))
648 .mnemonic()
649 .is_terminator()
650 })
651 }
652
653 pub fn set_comment(&mut self, comment: Option<String>) {
655 self.ctx.block_mut(self.id).comment = comment;
656 }
657
658 pub fn rename_local(&mut self, name: Cow<'str, str>) -> crate::error::Result<()> {
663 let old_name = self
664 .ctx
665 .body(self.id.func)
666 .block(self.id)
667 .name
668 .as_deref()
669 .map(str::to_owned);
670 self.ctx
671 .register_body_name(self.id.into(), name.clone(), old_name.as_deref())?;
672 self.ctx.block_mut(self.id).name = Some(name);
673 Ok(())
674 }
675
676 fn insert_insn(&mut self, index: usize, insn_id: InstructionId) {
677 self.ctx.instruction_mut(insn_id).parent = Some(self.id.local);
678 self.ctx
679 .block_mut(self.id)
680 .instructions
681 .insert(index, insn_id.localize(self.id.func));
682 }
683
684 pub fn insert_insn_at_index(&mut self, index: usize, insn_id: InstructionId) {
687 self.insert_insn(index, insn_id);
688 }
689
690 pub fn push_insn(&mut self, id: InstructionId) {
692 let len = self
693 .ctx
694 .body(self.id.func)
695 .block(self.id)
696 .instructions
697 .len();
698 self.insert_insn(len, id);
699 }
700
701 pub fn insert_insn_before(&mut self, before_id: InstructionId, insn_id: InstructionId) {
704 let id = self.id;
705 self.ctx.insert_insn_before(id, before_id, insn_id);
706 }
707
708 pub fn delete(&mut self) {
711 let id = self.id;
712 self.ctx.delete_block(id);
713 }
714
715 pub fn absorb_block(&mut self, other: BlockId, edge_ab: EdgeId) {
718 let id = self.id;
719 self.ctx.absorb_block(id, other, edge_ab);
720 }
721}
722
723impl Display for BlockMutRef<'_, '_> {
724 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
725 self.as_ref().fmt(f)
726 }
727}
728
729impl<'str, 'ctx> Value<'str, 'ctx> for BlockMutRef<'str, 'ctx> {
730 fn id(&self) -> ValueId {
731 self.id()
732 }
733
734 fn size(&self) -> usize {
735 0
736 }
737}
738
739impl<'str, 'ctx> BlockMutRef<'str, 'ctx> {
740 fn inner(&self) -> &BasicBlock<'str> {
741 self.ctx.block(self.id)
742 }
743
744 pub fn num_params(&self) -> usize {
745 self.as_ref().num_params()
746 }
747
748 pub fn instruction_ids(&self) -> Vec<InstructionId> {
749 self.as_ref().instruction_ids()
750 }
751
752 pub fn instructions(&self) -> impl Iterator<Item = InstructionRef<'str, '_>> {
753 self.as_ref().instructions().collect::<Vec<_>>().into_iter()
754 }
755
756 pub fn address(&self) -> Option<u64> {
757 self.as_ref().address()
758 }
759
760 pub fn successors(&self) -> impl Iterator<Item = (EdgeId, BlockId)> {
761 self.as_ref().successors().collect::<Vec<_>>().into_iter()
762 }
763
764 pub fn with_address(mut self, addr: u64) -> Self {
767 self.set_address(addr)
768 .expect("address is already mapped to a value");
769 self
770 }
771
772 pub fn with_address_indexed(
774 mut self,
775 addresses: &mut crate::address_index::AddressIndex,
776 addr: u64,
777 ) -> Self {
778 self.set_address_indexed(addresses, addr)
779 .expect("address is already mapped to a value");
780 self
781 }
782
783 #[allow(unused_mut)]
784 pub fn in_function(mut self, fun_id: FunctionId) -> Self {
785 FunctionBody::from_id_mut(self.ctx, fun_id).add_block(self.id);
786 self
787 }
788
789 pub fn with_id(&mut self, id: BlockId) -> &mut Self {
790 self.id = id;
791 self
792 }
793
794 pub fn reborrow(&mut self) -> BlockMutRef<'str, '_> {
796 BlockMutRef::from_id(self.ctx, self.id)
797 }
798
799 pub(in crate::value) fn inner_mut(&mut self) -> &mut BasicBlock<'str> {
800 self.ctx.block_mut(self.id)
801 }
802
803 pub fn parent_mut(&mut self) -> Option<FunctionMutRef<'str, '_>> {
804 Some(FunctionBody::from_id_mut(self.ctx, self.id.func))
806 }
807
808 pub fn as_ref(&self) -> BlockRef<'str, '_> {
809 BlockRef::new(ModuleView::new(self.ctx), self.id)
810 }
811
812 pub fn push_param(&mut self, size: usize) -> BlockParamMutRef<'str, '_> {
818 let block_id = self.id;
819 let index = self.inner().params.len();
820 let type_id = self.ctx.shared.types.get_or_make_int(size);
821 let id = self.ctx.push_block_param(
822 block_id.func,
823 BlockParam {
824 index,
825 type_id,
826 parent: Some(block_id.local),
827 name: None,
828 origin: None,
829 },
830 );
831 self.inner_mut().params.push(id.localize(block_id.func));
832 BlockParamMutRef::from_id(self.ctx, id)
833 }
834
835 pub fn push_existing_param(&mut self, id: BlockParamId) {
837 let func = self.id.func;
838 self.inner_mut().params.push(id.localize(func));
839 }
840
841 pub fn insert_insn_after(&mut self, after_id: InstructionId, insn_id: InstructionId) {
844 let index = self
845 .inner()
846 .instructions
847 .iter()
848 .position(|&local| InstructionId::new(self.id.func, local) == after_id)
849 .expect("after_id not found in block");
850 self.insert_insn(index + 1, insn_id);
851 }
852
853 pub fn retain_insns(&mut self, mut f: impl FnMut(&InstructionId) -> bool) {
856 let func = self.id.func;
857 let mut removed = Vec::new();
858 self.inner_mut().instructions.retain(|&local| {
859 let id = InstructionId::new(func, local);
860 if f(&id) {
861 true
862 } else {
863 removed.push(id);
864 false
865 }
866 });
867 for id in removed {
868 self.ctx.remove_instruction(id);
869 }
870 }
871
872 pub fn pop_insn(&mut self) {
874 if let Some(&local) = self.inner().instructions.last() {
875 self.ctx
876 .remove_instruction(InstructionId::new(self.id.func, local));
877 }
878 }
879
880 pub fn extend_insns(&mut self, insns: &[InstructionId]) {
882 let func = self.id.func;
883 self.inner_mut()
884 .instructions
885 .extend(insns.iter().map(|&id| id.localize(func)));
886 }
887
888 pub fn set_address(&mut self, addr: u64) -> Result<()> {
892 let mut addresses = crate::address_index::AddressIndex::analyze(&*self.ctx);
893 self.set_address_indexed(&mut addresses, addr)
894 }
895
896 pub fn set_address_indexed(
898 &mut self,
899 addresses: &mut crate::address_index::AddressIndex,
900 addr: u64,
901 ) -> Result<()> {
902 let old_address = self.inner().address;
903 self.inner_mut().address = Some(addr);
904 if let Err(error) = self
905 .ctx
906 .set_address_indexed(addresses, addr, self.id.into())
907 {
908 self.inner_mut().address = old_address;
909 return Err(error);
910 }
911
912 if self.name().is_none() {
913 let label = self
914 .ctx
915 .get_unique_name_in(self.id.func, Cow::Owned(format!("{addr:x}")));
916 self.rename(label)?;
917 }
918 Ok(())
919 }
920}
921
922#[cfg(test)]
923mod tests {
924
925 use super::*;
926 use crate::value::insn::{Binary, Binop, IntBinop, LocalInsnId};
927 use wazabin_qcode_macro::qcode;
928
929 #[test]
930 fn simultaneous_operand_substitution_does_not_chain_local_ids() {
931 let a = LocalValueId::Instruction(LocalInsnId::from(1));
932 let b = LocalValueId::Instruction(LocalInsnId::from(2));
933 let c = LocalValueId::Instruction(LocalInsnId::from(3));
934 let mut mnemonic = Mnemonic::Binop(Binary {
935 op: Binop::Int(IntBinop::Add),
936 lhs: a,
937 rhs: b,
938 });
939
940 substitute_operands(&mut mnemonic, &[(a, b), (b, c)]);
941
942 let Mnemonic::Binop(binary) = mnemonic else {
943 unreachable!();
944 };
945 assert_eq!((binary.lhs, binary.rhs), (b, c));
946 }
947
948 #[test]
949 fn test_create_block_at_address() {
950 let mut ctx = Context::new();
951 let id = {
952 let __f = ctx.anon_function();
953 BasicBlock::make(&mut ctx, __f)
954 }
955 .with_address(0x2000)
956 .id;
957 let addresses = crate::address_index::AddressIndex::analyze(&ctx);
958 let block_by_addr = BasicBlock::from_id(
959 &ctx,
960 addresses
961 .block_at(0x2000)
962 .expect("block not found by address"),
963 );
964 assert_eq!(id, block_by_addr.id);
965 assert_eq!(block_by_addr.address(), Some(0x2000));
966 }
967
968 #[test]
969 #[should_panic(expected = "address is already mapped to a value")]
970 fn test_create_block_at_duplicate_address() {
971 let mut ctx = Context::new();
972 {
973 let __f = ctx.anon_function();
974 BasicBlock::make(&mut ctx, __f)
975 }
976 .with_address(0x2000);
977 {
978 let __f = ctx.anon_function();
979 BasicBlock::make(&mut ctx, __f)
980 }
981 .with_address(0x2000);
982 }
983
984 #[test]
985 fn test_block_child() {
986 let mut ctx = Context::new();
987 qcode!(
988 ctx,
989 "
990 <entry>
991 goto <body>;
992 <body>
993 "
994 );
995 let entry = BasicBlock::from_name(&ctx, "entry").unwrap();
996 let children: Vec<_> = entry
997 .successors()
998 .map(|(_, b)| {
999 BasicBlock::from_id(&ctx, b)
1000 .name()
1001 .unwrap_or("")
1002 .to_string()
1003 })
1004 .collect();
1005 assert_eq!(children, ["body"]);
1006 }
1007
1008 #[test]
1009 fn iter_yields_all_instructions() {
1010 let mut ctx = Context::new();
1011 qcode!(
1012 ctx,
1013 "
1014 varnode i64 X;
1015 varnode i64 Y;
1016
1017 <block>
1018 %x = load(X:8, &X);
1019 %y = load(Y:8, &Y);
1020 %sum = i64 %x + i64 %y;
1021 return at i64 0;
1022 "
1023 );
1024
1025 let block = BasicBlock::from_id(&ctx, block);
1026 let count = block.iter().count();
1027 assert_eq!(count, 4);
1028
1029 let mut iter = block.iter();
1030
1031 assert_eq!(
1032 iter.next().unwrap().as_statement().to_string(),
1033 "i64 %x = load(X:8, i64 X);"
1034 );
1035 assert_eq!(
1036 iter.next().unwrap().as_statement().to_string(),
1037 "i64 %y = load(Y:8, i64 Y);"
1038 );
1039 assert_eq!(
1040 iter.next().unwrap().as_statement().to_string(),
1041 "i64 %sum = i64 %x + i64 %y;"
1042 );
1043 assert_eq!(
1044 iter.next().unwrap().as_statement().to_string(),
1045 "return at i64 0x0;"
1046 );
1047 }
1048
1049 #[test]
1050 fn into_iterator_for_block_ref_matches_iter() {
1051 let mut ctx = Context::new();
1052
1053 qcode!(
1054 ctx,
1055 "
1056 varnode i64 X;
1057 varnode i64 Y;
1058
1059 <block>
1060 %x = load(X:8, &X);
1061 %y = load(Y:8, &Y);
1062 %sum = i64 %x + i64 %y;
1063 return at i64 0;
1064 "
1065 );
1066
1067 let block = BasicBlock::from_id(&ctx, block);
1068 let via_iter: Vec<_> = block.iter().map(|i| i.id).collect();
1069 let via_into: Vec<_> = (&block).into_iter().map(|i| i.id).collect();
1070 assert_eq!(via_iter, via_into);
1071 }
1072
1073 #[test]
1074 fn push_param_adds_to_params_not_instructions() {
1075 let mut ctx = Context::new();
1076 let mut block = {
1077 let __f = ctx.anon_function();
1078 BasicBlock::make(&mut ctx, __f)
1079 };
1080
1081 assert_eq!(block.num_params(), 0);
1082 assert_eq!(block.as_ref().instruction_ids().len(), 0);
1083
1084 block.push_param(8);
1085 assert_eq!(block.num_params(), 1);
1086 assert_eq!(block.as_ref().instruction_ids().len(), 0);
1087
1088 block.push_param(4);
1089 assert_eq!(block.num_params(), 2);
1090 assert_eq!(block.as_ref().instruction_ids().len(), 0);
1091 }
1092
1093 #[test]
1094 fn params_iter_yields_in_order() {
1095 let mut ctx = Context::new();
1096 let mut block = {
1097 let __f = ctx.anon_function();
1098 BasicBlock::make(&mut ctx, __f)
1099 };
1100
1101 let p0_id = block.push_param(8).id;
1102 let p1_id = block.push_param(4).id;
1103 let block_ref = block.as_ref();
1104
1105 let param_ids: Vec<_> = block_ref.params().map(|p| p.id).collect();
1106 assert_eq!(param_ids, [p0_id, p1_id]);
1107 assert_eq!(block_ref.params().next().unwrap().index(), 0);
1108 assert_eq!(block_ref.params().nth(1).unwrap().index(), 1);
1109 }
1110
1111 #[test]
1112 fn wazabin_qcode_macro_block_with_params() {
1113 use crate::context::Context;
1114 use wazabin_qcode_macro::qcode;
1115
1116 let mut ctx = Context::new();
1117 qcode!(
1118 ctx,
1119 "
1120 <entry @v1:i64 @v2:i32>
1121 goto <done @x=@v1 @y=@v2>;
1122
1123 <done @x:i64 @y:i32>
1124 goto <0x1001>;
1125 "
1126 );
1127
1128 let entry = BasicBlock::from_id(&ctx, entry);
1129 assert_eq!(entry.num_params(), 2, "entry should have 2 params");
1130
1131 let params = entry.params().collect::<Vec<_>>();
1132 assert_eq!(params[0].name(), Some("v1"));
1133 assert_eq!(params[0].size(), 8);
1134 assert_eq!(params[1].name(), Some("v2"));
1135 assert_eq!(params[1].size(), 4);
1136
1137 let done_block = BasicBlock::from_id(&ctx, done);
1138 assert_eq!(done_block.num_params(), 2, "done should have 2 params");
1139 let done_params = done_block.params().collect::<Vec<_>>();
1140 assert_eq!(done_params[0].size(), 8);
1141 assert_eq!(done_params[1].size(), 4);
1142
1143 let branch_insn = entry.iter().last().expect("entry has instructions");
1145 let crate::value::insn::Mnemonic::Branch(branch) = branch_insn.mnemonic() else {
1146 panic!("expected branch");
1147 };
1148 assert_eq!(branch.args.len(), 2);
1149 }
1150
1151 #[test]
1152 fn block_display_uses_qcode_param_syntax() {
1153 use crate::context::Context;
1154 use wazabin_qcode_macro::qcode;
1155
1156 let mut ctx = Context::new();
1157 qcode!(
1158 ctx,
1159 "
1160 <entry @a @b>
1161 goto <0x1001>;
1162 "
1163 );
1164
1165 let entry = BasicBlock::from_id(&ctx, entry);
1166 assert!(entry.to_string().starts_with("<entry @a @b>\n"));
1167 }
1168
1169 #[test]
1170 fn param_value_id_usable_in_instruction() {
1171 use crate::value::ValueId;
1172 let mut ctx = Context::new();
1173 let block_id = {
1174 let __f = ctx.anon_function();
1175 BasicBlock::make(&mut ctx, __f)
1176 }
1177 .id;
1178
1179 let param_id: ValueId = {
1180 let mut block = BasicBlock::from_id_mut(&mut ctx, block_id);
1181 block.push_param(8).id()
1182 };
1183
1184 let mut builder = ctx.builder(block_id);
1185 let sum = builder.push_add(param_id, param_id);
1186 assert_eq!(sum.size(), 8);
1187 }
1188
1189 #[test]
1190 fn remove_terminator_branch() {
1191 let mut ctx = Context::new();
1192 qcode!(
1193 ctx,
1194 "
1195 <entry>
1196 goto <done>;
1197 <done>
1198 "
1199 );
1200
1201 let mut entry_block = BasicBlock::from_id_mut(&mut ctx, entry);
1202 assert_eq!(entry_block.instruction_ids().len(), 1);
1203 assert_eq!(entry_block.successors().count(), 1);
1204 assert!(entry_block.is_terminated());
1205
1206 entry_block.pop_insn();
1207 assert_eq!(entry_block.instruction_ids().len(), 0);
1208 assert_eq!(entry_block.successors().count(), 0);
1209 assert!(!entry_block.is_terminated());
1210 }
1211
1212 #[test]
1213 fn remove_terminator_cbranch() {
1214 let mut ctx = Context::new();
1215 qcode!(
1216 ctx,
1217 "
1218 varnode i8 cond;
1219
1220 <entry>
1221 %c = load(cond:1, cond);
1222 if %c goto <then_lbl> else goto <else_lbl>;
1223
1224 <then_lbl>
1225
1226 <else_lbl>
1227 "
1228 );
1229
1230 let mut entry_block = BasicBlock::from_id_mut(&mut ctx, entry);
1231 assert_eq!(entry_block.successors().count(), 2);
1232 assert_eq!(entry_block.instruction_ids().len(), 2);
1233 assert!(entry_block.is_terminated());
1234
1235 entry_block.pop_insn();
1236
1237 assert_eq!(entry_block.successors().count(), 0);
1238 assert_eq!(entry_block.instruction_ids().len(), 1);
1239 assert!(!entry_block.is_terminated());
1240
1241 assert_eq!(
1242 BasicBlock::from_id(&ctx, then_lbl).predecessors().count(),
1243 0
1244 );
1245 assert_eq!(
1246 BasicBlock::from_id(&ctx, else_lbl).predecessors().count(),
1247 0
1248 );
1249 }
1250
1251 #[test]
1252 fn remove_terminator_return() {
1253 let mut ctx = Context::new();
1254 qcode!(
1255 ctx,
1256 "
1257 <entry>
1258 return at i64 0;
1259 "
1260 );
1261
1262 let mut entry_block = BasicBlock::from_id_mut(&mut ctx, entry);
1263 assert_eq!(entry_block.instruction_ids().len(), 1);
1264 assert_eq!(entry_block.successors().count(), 0);
1265
1266 entry_block.pop_insn();
1267 assert_eq!(entry_block.instruction_ids().len(), 0);
1268 assert_eq!(entry_block.successors().count(), 0);
1269 }
1270
1271 #[test]
1272 fn clone_into_ctx_produces_distinct_ids() {
1273 let mut ctx = Context::new();
1274 qcode!(
1275 ctx,
1276 "
1277 varnode i64 X;
1278 varnode i64 Y;
1279
1280 <block>
1281 %x = load(X:8, &X);
1282 %y = load(Y:8, &Y);
1283 %sum = i64 %x + i64 %y;
1284 return at i64 0;
1285 "
1286 );
1287
1288 let mut value_map = HashMap::default();
1289 let cloned_id = BasicBlock::clone_into_ctx(&mut ctx, block, &mut value_map);
1290
1291 let orig = BasicBlock::from_id(&ctx, block);
1292 let cloned = BasicBlock::from_id(&ctx, cloned_id);
1293
1294 assert_ne!(block, cloned_id, "cloned block must have a different id");
1295 assert_ne!(
1296 orig.name(),
1297 cloned.name(),
1298 "cloned block must have a different name"
1299 );
1300
1301 assert_eq!(orig.instruction_ids().len(), cloned.instruction_ids().len());
1302 for (orig_id, clone_id) in orig
1303 .instruction_ids()
1304 .into_iter()
1305 .zip(cloned.instruction_ids())
1306 {
1307 assert_ne!(
1308 orig_id, clone_id,
1309 "cloned instruction must have a different id"
1310 );
1311 }
1312 }
1313
1314 #[test]
1315 fn clone_into_ctx_remaps_operands() {
1316 let mut ctx = Context::new();
1317 qcode!(
1318 ctx,
1319 "
1320 varnode i64 X;
1321 varnode i64 Y;
1322
1323 <block>
1324 %x = load(X:8, &X);
1325 %y = load(Y:8, &Y);
1326 %sum = i64 %x + i64 %y;
1327 return at i64 0;
1328 "
1329 );
1330
1331 let mut value_map = HashMap::default();
1332 let cloned_id = BasicBlock::clone_into_ctx(&mut ctx, block, &mut value_map);
1333 let cloned = BasicBlock::from_id(&ctx, cloned_id);
1334
1335 let orig_value_ids: HashSet<ValueId> = value_map.keys().copied().collect();
1336
1337 for arg in cloned.iter().flat_map(|i| i.operands()) {
1340 assert!(
1341 !orig_value_ids.contains(&arg),
1342 "cloned instruction still references original value {arg:?}"
1343 );
1344 }
1345 }
1346
1347 #[test]
1352 fn delete_unlinks_incident_edges() {
1353 let mut ctx = Context::new();
1354 qcode!(
1355 ctx,
1356 "
1357 fn f:
1358 <a>
1359 goto <exit>;
1360 <b>
1361 goto <exit>;
1362 <exit>
1363 return at i64 0;
1364 "
1365 );
1366
1367 assert_eq!(BasicBlock::from_id(&ctx, exit).predecessors().count(), 2);
1368
1369 BasicBlock::from_id_mut(&mut ctx, b).delete();
1370
1371 assert_eq!(
1372 BasicBlock::from_id(&ctx, exit).predecessors().count(),
1373 1,
1374 "deleted block's edge must not linger as a phantom predecessor"
1375 );
1376 assert!(
1377 !ctx.contains_block(b),
1378 "deleted block payload must be absent"
1379 );
1380 }
1381
1382 #[test]
1385 fn delete_unlinks_self_loop() {
1386 let mut ctx = Context::new();
1387 qcode!(
1388 ctx,
1389 "
1390 fn f:
1391 <a>
1392 goto <loop_hdr>;
1393 <loop_hdr>
1394 goto <loop_hdr>;
1395 "
1396 );
1397
1398 assert!(
1399 BasicBlock::from_id(&ctx, loop_hdr)
1400 .successors()
1401 .any(|(_, s)| s == loop_hdr)
1402 );
1403
1404 BasicBlock::from_id_mut(&mut ctx, loop_hdr).delete();
1405
1406 assert!(
1407 !ctx.contains_block(loop_hdr),
1408 "deleted self-loop block payload must be absent"
1409 );
1410 }
1411
1412 #[test]
1419 fn delete_removes_instructions_from_use_lists() {
1420 let mut ctx = Context::new();
1421 qcode!(
1422 ctx,
1423 "
1424 fn f:
1425 <a>
1426 %x = i64 1 + i64 2;
1427 goto <b>;
1428 <b>
1429 %y = %x + i64 3;
1430 goto <exit>;
1431 <exit>
1432 return at i64 0;
1433 "
1434 );
1435
1436 assert!(
1438 ctx.users(crate::value::ValueId::Instruction(x))
1439 .to_vec()
1440 .contains(&y),
1441 "precondition: %x is used by %y"
1442 );
1443
1444 BasicBlock::from_id_mut(&mut ctx, b).delete();
1445
1446 assert!(
1447 !ctx.users(crate::value::ValueId::Instruction(x))
1448 .to_vec()
1449 .contains(&y),
1450 "deleting b must unregister %y from %x's use-list, not orphan it"
1451 );
1452 assert!(
1453 !ctx.contains_instruction(y),
1454 "deleted instruction payload must be physically absent"
1455 );
1456 }
1457}