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
153
154
155
156
157
158
159
160
161
162

//! This crate contains structures and routines for context switching 
//! when SSE/SIMD extensions are not active. 

use zerocopy::FromBytes;

/// The registers saved before a context switch and restored after a context switch.
///
/// Note: the order of the registers here MUST MATCH the order of 
/// registers popped in the [`restore_registers_regular!`] macro. 
#[derive(FromBytes)]
#[repr(C, packed)]
pub struct ContextRegular {
    rflags: usize,
    r15: usize, 
    r14: usize,
    r13: usize,
    r12: usize,
    rbp: usize,
    rbx: usize,
    /// The instruction pointer.
    ///
    /// `rip` is implicitly pushed onto the stack when a function is called, and
    /// popped off when returning. When a task's stack is set to an instance of
    /// [`ContextRegular`], [`context_switch_regular`] will execute `ret` when
    /// the stack pointer is pointing to the value of `rip`. Hence, the program
    /// will "return" to that address and continue executing.
    rip: usize,
}

impl ContextRegular {
    /// Creates a new [`ContextRegular`] struct that will cause the
    /// Task containing it to begin its execution at the given `rip`.
    pub fn new(rip: usize) -> ContextRegular {
        ContextRegular {
            // The ninth bit is the interrupt enable flag. When a task is first
            // run, interrupts should already be enabled.
            rflags: 1 << 9,
            r15: 0,
            r14: 0,
            r13: 0,
            r12: 0,
            rbp: 0,
            rbx: 0,
            rip,
        }
    }

    /// Sets the value of the first register to the given `value`.
    /// 
    /// This is useful for storing a value (e.g., task ID) in that register
    /// and then recovering it later with [`read_first_register()`].
    /// 
    /// On x86_64, this sets the `r15` register.
    pub fn set_first_register(&mut self, value: usize) {
        self.r15 = value;
    }
}

/// Reads the value of the first register from the actual CPU register hardware.
/// 
/// This can be called at any time, but is intended for use as the second half
/// of "saving and restoring" a register value.
/// The first half was a previous call to [`ContextRegular::set_first_register()`],
/// and the second half is a call to this function immediately after the original
/// `ContextRegular` has been used for switching to a new task for the first time.
/// 
/// Returns the current value held in the specified CPU register.
/// On x86_64, this reads the `r15` register.
#[naked]
pub extern "C" fn read_first_register() -> usize {
    // SAFE: simply reads and returns the value of `r15`.
    unsafe {
        core::arch::asm!(
            "mov rax, r15", // rax is used for return values on x86_64
            "ret",
            options(noreturn)
        )
    }
}


/// An assembly block for saving regular x86_64 registers
/// by pushing them onto the stack.
#[macro_export]
macro_rules! save_registers_regular {
    () => (
        // Save all general purpose registers into the previous task.
        r#"
            push rbx
            push rbp
            push r12
            push r13
            push r14
            push r15
            pushfq
        "#
    );
}


/// An assembly block for switching stacks,
/// which is the integral part of the actual context switching routine.
/// 
/// * The `rdi` register must contain a pointer to the previous task's stack pointer.
/// * The `rsi` register must contain the value of the next task's stack pointer.
#[macro_export]
macro_rules! switch_stacks {
    () => (
        // switch the stack pointers
        r#"
            mov [rdi], rsp
            mov rsp, rsi
        "#
    );
}


/// An assembly block for restoring regular x86_64 registers
/// by popping them off of the stack.
/// 
/// This assembly statement ends with an explicit `ret` instruction at the end,
/// which is the final component of a context switch operation. 
/// Note that this is intentional and required in order to accommodate 
/// the `noreturn` option is required by Rust's naked functions.
#[macro_export]
macro_rules! restore_registers_regular {
    () => (
        // Restore the next task's general purpose registers.
        r#" 
            popfq
            pop r15
            pop r14
            pop r13
            pop r12
            pop rbp
            pop rbx
            ret
        "#
    );
}


/// Switches context from a regular Task to another regular Task.
/// 
/// # Arguments
/// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
/// * Second argument (in `RSI`): the value of the next task's stack pointer
/// 
/// # Safety
/// This function is unsafe because it changes the content on both task's stacks. 
#[naked]
pub unsafe extern "C" fn context_switch_regular(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
    // Since this is a naked function that expects its arguments in two registers,
    // you CANNOT place any log statements or other instructions here
    // before, in between, or after anything below.
    core::arch::asm!(
        save_registers_regular!(),
        switch_stacks!(),
        restore_registers_regular!(),
        options(noreturn)
    );
}