1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
//! Program counter handlers.

use crate::{
    jump::{relocate, Jump, JumpTable},
    Error, Result, BUFFER_LIMIT,
};

impl JumpTable {
    /// Shift program counter for all items.
    pub fn shift_pc(&mut self, start: u16, offset: u16) -> Result<()> {
        tracing::trace!("shift pc from 0x{start:x} with offset={offset}");
        self.shift_label_pc(start, offset)?;
        self.shift_label_target(start, offset)?;
        self.shift_func_target(start, offset)?;

        Ok(())
    }

    /// Shift program counter for labels.
    pub fn shift_label_pc(&mut self, start: u16, offset: u16) -> Result<()> {
        self.jump = self
            .jump
            .iter()
            .map(|(k, v)| {
                let mut k = *k;
                if k > start {
                    tracing::trace!(
                        "shift {v:x?} pc with offset={offset}: 0x{k:x}(0x{start:x}) -> 0x{:x}",
                        k + offset
                    );
                    k += offset;
                    if k > BUFFER_LIMIT as u16 {
                        return Err(Error::InvalidPC(k as usize));
                    }
                }

                Ok((k, v.clone()))
            })
            .collect::<Result<_>>()?;

        Ok(())
    }

    /// Shift the target program counters.
    ///
    /// Calculating target pc from the offset of original pc.
    pub fn shift_targets(&mut self) -> Result<()> {
        let mut total_offset = 0;
        self.jump
            .clone()
            .keys()
            .try_for_each(|original_pc| -> Result<()> {
                let pc = original_pc + total_offset;
                let offset = relocate::offset(pc)?;
                total_offset += offset;
                self.shift_target(pc, offset)
            })
    }

    /// Shift the program counter of targets with given ptr and offset.
    ///
    /// 1. shift code section.
    /// 2. shift label targets.
    /// 3. shift function targets.
    pub fn shift_target(&mut self, ptr: u16, offset: u16) -> Result<()> {
        self.code.shift(offset);
        self.shift_label_target(ptr, offset)?;
        self.shift_func_target(ptr, offset)
    }

    /// Shift program counter for functions.
    pub fn shift_func_target(&mut self, start: u16, offset: u16) -> Result<()> {
        self.func.iter_mut().try_for_each(|(k, v)| {
            if *v > start {
                tracing::trace!(
                    "shift Func({k}) target with offset={offset}: 0x{v:x}(0x{start:x}) -> 0x{:x}",
                    *v + offset
                );
                *v += offset;
                if *v > BUFFER_LIMIT as u16 {
                    return Err(Error::InvalidPC(*v as usize));
                }
            }

            Ok(())
        })?;

        Ok(())
    }

    /// Shift target program counter for labels.
    pub fn shift_label_target(&mut self, ptr: u16, offset: u16) -> Result<()> {
        self.jump.iter_mut().try_for_each(|(pc, target)| {
            if let Jump::Label(label) = target {
                if *label > ptr {
                    tracing::trace!(
                        "shift Label(pc=0x{pc:x}) target with offset={offset} 0x{label:x}(0x{ptr:x}) -> 0x{:x}",
                        *label + offset
                    );
                    *label += offset;
                }
            }

            Ok(())
        })
    }
}