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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use cairo_lang_casm::instructions::{
    AssertEqInstruction, CallInstruction, Instruction, InstructionBody, JnzInstruction,
    JumpInstruction,
};
use cairo_lang_casm::operand::{BinOpOperand, DerefOrImmediate, ResOperand};
use cairo_lang_sierra::ids::ConcreteTypeId;
use cairo_lang_sierra::program::StatementIdx;
use cairo_lang_sierra_gas::objects::ConstCost;

use crate::compiler::ConstSegmentInfo;

pub type CodeOffset = usize;

#[derive(Debug, Eq, PartialEq)]
pub enum Relocation {
    /// Adds program_offset(StatementIdx) and subtracts the program offset of the casm instruction
    /// that is being relocated.
    RelativeStatementId(StatementIdx),
    /// Adds the offset between the current statement index and the end of the program code
    /// segment.
    EndOfProgram,
    /// Adds the offset of the const type in the const segment, assuming the const segment is at
    /// the end of the program.
    Const(ConcreteTypeId),
}

impl Relocation {
    pub fn apply(
        &self,
        instruction_offset: CodeOffset,
        statement_offsets: &[CodeOffset],
        const_segment_info: &ConstSegmentInfo,
        instruction: &mut Instruction,
    ) {
        let target_pc = match self {
            Relocation::RelativeStatementId(statement_idx) => statement_offsets[statement_idx.0],
            Relocation::EndOfProgram => {
                *statement_offsets.last().unwrap() + const_segment_info.const_segment_size
            }
            Relocation::Const(ty) => {
                *statement_offsets.last().unwrap()
                    + const_segment_info
                        .const_allocations
                        .get(ty)
                        .expect("Const type not found in the const segment.")
                        .offset
            }
        };
        match instruction {
            Instruction {
                body:
                    InstructionBody::Call(CallInstruction {
                        target: DerefOrImmediate::Immediate(value),
                        relative: true,
                    }),
                inc_ap: false,
                ..
            }
            | Instruction {
                body:
                    InstructionBody::Jnz(JnzInstruction {
                        jump_offset: DerefOrImmediate::Immediate(value),
                        condition: _,
                    }),
                ..
            }
            | Instruction {
                body:
                    InstructionBody::Jump(JumpInstruction {
                        target: DerefOrImmediate::Immediate(value),
                        relative: true,
                    }),
                ..
            }
            | Instruction {
                body:
                    InstructionBody::AssertEq(AssertEqInstruction {
                        b:
                            ResOperand::BinOp(BinOpOperand {
                                b: DerefOrImmediate::Immediate(value),
                                ..
                            }),
                        ..
                    }),
                ..
            } => {
                value.value += target_pc as i128 - instruction_offset as i128;
            }
            _ => panic!("Bad relocation."),
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
pub struct RelocationEntry {
    /// The index of the casm instruction that needs to be relocated.
    pub instruction_idx: CodeOffset,
    /// The relocation the needs to be applied.
    pub relocation: Relocation,
}

/// Applies 'relocations' to 'instructions'.
///
/// This is currently O(instruction.len()) rather then O(relocations.len()),
/// But another pass is required anyhow to generate the bytecode and the relocations
/// can be applied during that pass.
pub fn relocate_instructions(
    relocations: &[RelocationEntry],
    statement_offsets: &[usize],
    const_segment_info: &ConstSegmentInfo,
    instructions: &mut [Instruction],
) {
    let mut program_offset = 0;
    let mut relocations_iter = relocations.iter();
    let mut relocation_entry = relocations_iter.next();
    for (instruction_idx, instruction) in instructions.iter_mut().enumerate() {
        if let Some(RelocationEntry { instruction_idx: relocation_idx, relocation }) =
            relocation_entry
        {
            if *relocation_idx == instruction_idx {
                relocation.apply(
                    program_offset,
                    statement_offsets,
                    const_segment_info,
                    instruction,
                );
                relocation_entry = relocations_iter.next();
            } else {
                assert!(
                    *relocation_idx > instruction_idx,
                    "Relocation for already handled instruction #{relocation_idx} - currently \
                     handling #{instruction_idx}."
                );
            }
        }

        program_offset += instruction.body.op_size();
    }
    assert!(
        relocation_entry.is_none(),
        "No relocations should be left when done with all instructions."
    );
}

/// Represents a list of instructions with their relocations.
#[derive(Default)]
pub struct InstructionsWithRelocations {
    pub instructions: Vec<Instruction>,
    pub relocations: Vec<RelocationEntry>,
    /// The gas cost of executing the instructions.
    pub cost: ConstCost,
}