BurritOS/src/kernel/thread_manager.rs

use std::{rc::Rc, cell::{RefCell, Ref}};

use crate::{utility::{list::List, objaddr::ObjAddr}, simulator::{machine::{NUM_INT_REGS, NUM_FP_REGS, Machine}, interrupt::InterruptStatus, error::{MachineOk, MachineError}}};

use super::{thread::Thread, process::Process};

pub const SIMULATORSTACKSIZE: usize = 32 * 1024;

/// # Thread manager
/// 
/// An instance of this struct is responsible for managing threads on behalf of the system
#[derive(PartialEq)]
pub struct ThreadManager {
    /// Current running thread
    pub g_current_thread: Option<Rc<RefCell<Thread>>>,
    /// The thread to be destroyed next
    pub g_thread_to_be_destroyed: Option<Rc<RefCell<Thread>>>,
    /// The list of alive threads
    pub g_alive: List<Rc<RefCell<Thread>>>,
    /// Thread in ready state waiting to become active
    ready_list: List<Rc<RefCell<Thread>>>,
    obj_addrs: ObjAddr,
    debug: bool
}

impl ThreadManager {

    /// Thread manager constructor
    pub fn new(debug: bool) -> Self {
        Self { 
            g_current_thread: Option::None, 
            g_thread_to_be_destroyed: Option::None, 
            g_alive: List::default(),
            ready_list: List::default(),
            obj_addrs: ObjAddr::init(),
            debug
        }
    }

    /// Mark a thread as aready, but not necessarily running yet.
    /// 
    /// Put it in the ready list, for later scheduling onto the CPU.
    /// 
    /// ## Pamameter
    /// 
    /// **thread** is the thread to be put on the read list
    pub fn ready_to_run(&mut self, thread: Rc<RefCell<Thread>>) {
        self.ready_list.push(thread);
    }

    /// Return the next thread to be scheduled onto the CPU.
    /// If there are no ready threads, return Option::None
    /// 
    /// Thread is removed from the ready list.
    /// 
    /// **return** Thread thread to be scheduled
    pub fn find_next_to_run(&mut self) -> Option<Rc<RefCell<Thread>>> {
        self.ready_list.pop()
    }

    /// Dispatch the CPU to next_thread. Save the state of the old thread 
    /// and load the state of the new thread.
    /// 
    /// We assume the state of the previously running thread has already been changed from running to blocked or ready.
    /// 
    /// Global variable g_current_thread become next_thread
    /// 
    /// ## Parameter
    /// 
    /// **next_thread** thread to dispatch to the CPU
    pub fn switch_to(&mut self, machine: &mut Machine, next_thread: Rc<RefCell<Thread>>) {
        match self.get_g_current_thread() {
            Some(old) => {
                let old1 = Rc::clone(old);
                let old2 = Rc::clone(old);
                self.thread_save_processor_state(machine, old1);
                // old_thread.save_simulator_state();
                if old2 != next_thread {
                    self.debug(format!("switching from \"{}\" to \"{}\"", old2.borrow().get_name(), next_thread.borrow().get_name()));
                    self.thread_restore_processor_state(machine, Rc::clone(&next_thread));
                    // next_thread.restore_simulator_state();
                    debug_assert!(!self.ready_list.contains(&next_thread));
                    self.set_g_current_thread(Some(next_thread));
                }
            },
            None => {
                self.thread_restore_processor_state(machine, Rc::clone(&next_thread));
                // next_thread.restore_simulator_state();
                self.set_g_current_thread(Some(next_thread));
            }
        }

        if let Some(th) = self.get_g_thread_to_be_destroyed() {
            drop(th);
            self.g_thread_to_be_destroyed = None
        }
    }

    /// Start a thread, attaching it to a process
    pub fn start_thread(&mut self, thread: Rc<RefCell<Thread>>, owner: Rc<RefCell<Process>>, func_pc: u64, sp_loc: u64, argument: i64) {
        let mut thread_m = thread.borrow_mut();
        assert_eq!(thread_m.process, Option::None);
        thread_m.process = Option::Some(Rc::clone(&owner));
        let ptr = sp_loc; // todo addrspace
        thread_m.init_thread_context(func_pc, ptr, argument);
        owner.borrow_mut().num_thread += 1;
        self.get_g_alive().push(Rc::clone(&thread));
        self.ready_to_run(Rc::clone(&thread));
    }

    /// Wait for another thread to finish its execution
    pub fn thread_join(&mut self, machine: &mut Machine, waiter: Rc<RefCell<Thread>>, waiting_for: Rc<RefCell<Thread>>) {
        let waiting_for = Rc::clone(&waiting_for);
        while self.get_g_alive().contains(&waiting_for) {
            self.debug(format!("Joining \"{}\" to \"{}\"", waiter.borrow().get_name(), waiting_for.borrow().get_name()));
            self.thread_yield(machine, Rc::clone(&waiter));
        }
    }

    /// Relinquish the CPU if any other thread is runnable.
    /// 
    /// Cannot use yield as a function name -> reserved name in rust
    pub fn thread_yield(&mut self, machine: &mut Machine, thread: Rc<RefCell<Thread>>) {
        let old_status = machine.interrupt.set_status(crate::simulator::interrupt::InterruptStatus::InterruptOff);
            
        
        self.debug(format!("Yeilding thread: {}", thread.borrow().get_name()));
        debug_assert_eq!(&Option::Some(Rc::clone(&thread)), self.get_g_current_thread());
        let next_thread = self.find_next_to_run();
        if let Some(next_thread) = next_thread {
            self.ready_to_run(thread);
            self.switch_to(machine, next_thread);
        }
        machine.interrupt.set_status(old_status);    
    }

    /// Put the thread to sleep and relinquish the processor
    pub fn thread_sleep(&mut self, machine: &mut Machine, thread: Rc<RefCell<Thread>>) {
        debug_assert_eq!(Option::Some(Rc::clone(&thread)), self.g_current_thread);
        debug_assert_eq!(machine.interrupt.get_status(), InterruptStatus::InterruptOff);

        let mut next_thread = self.find_next_to_run();
        while next_thread.is_none() {
            eprintln!("Nobody to run => idle");
            machine.interrupt.idle();
            next_thread = self.find_next_to_run();
        }
        self.switch_to(machine, Rc::clone(&next_thread.unwrap()));
        
    }

    /// Finish the execution of the thread and prepare its deallocation
    pub fn thread_finish(&mut self, machine: &mut Machine, thread: Rc<RefCell<Thread>>) {
        let old_status = machine.interrupt.set_status(InterruptStatus::InterruptOff);
        self.set_g_thread_to_be_destroyed(Option::Some(Rc::clone(&thread)));
        self.g_alive.remove(Rc::clone(&thread));
        self.debug(format!("Sleeping thread {}", thread.borrow().get_name()));
        // g_objets_addrs->removeObject(self.thread) // a ajouté plus tard
        self.thread_sleep(machine, Rc::clone(&thread));
        machine.interrupt.set_status(old_status);
    }

	/// Save the CPU state of a user program on a context switch.
    pub fn thread_save_processor_state(&mut self, machine: &mut Machine, thread: Rc<RefCell<Thread>>) {
        let mut t = thread.borrow_mut();
        for i in 0..NUM_INT_REGS {
            t.thread_context.int_registers[i] = machine.read_int_register(i);
        }
        for i in 0..NUM_FP_REGS {
            t.thread_context.float_registers[i] = machine.read_fp_register(i);
        }
        t.thread_context.pc = machine.pc;
    }

    /// Restore the CPU state of a user program on a context switch.
    pub fn thread_restore_processor_state(&self, machine: &mut Machine, thread: Rc<RefCell<Thread>>) {
        let t: Ref<_> = thread.borrow();
        for i in 0..NUM_INT_REGS {
            machine.write_int_register(i, t.thread_context.int_registers[i]);
        }
        machine.pc = t.thread_context.pc;
    }

    /// Decrement the value, and wait if it becomes < 0. Checking the
    ///	value and decrementing must be done atomically, so we
    ///	need to disable interrupts before checking the value.
    ///
    ///	Note that thread_manager::thread_sleep assumes that interrupts are disabled
    ///	when it is called.
    ///
    /// ### Parameters
    /// - *id_sema* id of the semaphore, stored in [`ObjAddr`], id given by user program thought exceptions
    /// - *machine* Current state of the machine
    pub fn sem_p(&mut self, id_sema: i32, machine: &mut Machine) -> Result<MachineOk, MachineError> {
        match self.get_g_current_thread() {
            Some(thread) => {
                let rc1_thread = Rc::clone(thread);
                let rc2_thread = Rc::clone(thread);
                let sema = self.get_obj_addrs().search_semaphore(id_sema as i32);
                if let Some(sema) = sema {
                    let old_status = machine.interrupt.set_status(InterruptStatus::InterruptOff);
                    sema.counter -= 1;
                    if sema.counter < 0 {
                        sema.waiting_queue.push(rc1_thread);
                        self.thread_sleep(machine, rc2_thread);
                    }
                    machine.interrupt.set_status(old_status);
                    Ok(MachineOk::Ok)
                } else {
                    Err("Cannot find semaphore")?
                }
            },
            None => unreachable!("Current thread should not be None")
        }   
    }

    /// Increment semaphore value, waking up a waiting thread if any.
    /// As with P(), this operation must be atomic, so we need to disable
    /// interrupts.
    ///
    /// scheduler::ready_to_run() assumes that interrupts
    /// are disabled when it is called.
    ///
    /// ### Parameters
    /// - *id_sema* id of the semaphore, stored in [`ObjAddr`], id given by user program thought exceptions
    /// - **machine** the machine where the threads are executed
    pub fn sem_v(&mut self, id_sema: i32, machine: &mut Machine) -> Result<MachineOk, MachineError> {
        let sema = self.get_obj_addrs().search_semaphore(id_sema as i32);
        match sema {
            Some(sema) => {
                let old_status = machine.interrupt.set_status(InterruptStatus::InterruptOff);
                sema.counter += 1;
                match sema.waiting_queue.pop() {
                    Some(thread) => self.ready_to_run(thread),
                    None => ()
                }
                machine.interrupt.set_status(old_status);
                Ok(MachineOk::Ok)
            },
            None => {
                Err("Cannot find semaphore")?
            }
        }
        
    }


    ///    Wait until the lock become free.  Checking the
    ///	state of the lock (free or busy) and modify it must be done
    ///	atomically, so we need to disable interrupts before checking
    ///	the value of free.
    ///
    ///	Note that thread_manager::thread_seep assumes that interrupts are disabled
    ///	when it is called.
    ///
    /// ### Parameters
    /// - **id** id of the lock, stored in [`ObjAddr`], id given by user program thought exceptions
    /// - **machine** the machine where the threads are executed
    pub fn lock_acquire(&mut self, id: i32, machine: &mut Machine) -> Result<MachineOk, MachineError> {
        match self.get_g_current_thread() {
            Some(thread) => {
                let rc_thread = Rc::clone(thread);
                let lock = self.get_obj_addrs().search_lock(id);
                match lock {
                    Some(lock) => {
                        let old_status = machine.interrupt.set_status(InterruptStatus::InterruptOff);
                        if lock.free {
                            lock.free = false;
                            lock.owner = Some(Rc::clone(&rc_thread))
                        } else {
                            let thread = Rc::clone(&rc_thread);
                            lock.waiting_queue.push(Rc::clone(&thread));
                            self.thread_sleep(machine, Rc::clone(&thread));
                        }
                        machine.interrupt.set_status(old_status);
                        Ok(MachineOk::Ok)
                    },
                    None => Err("Cannot find lock")?
                }
            },
            None => unreachable!("Current thread should not be None")
        }
    }

    ///     Wake up a waiter if necessary, or release it if no thread is waiting.
    pub fn lock_release(&mut self, id: i32, machine: &mut Machine) -> Result<MachineOk, MachineError> {
        let current_thread = self.get_g_current_thread();
        match current_thread {
            Some(thread) => {
                let current_thread = Rc::clone(&thread);

                let mut lock_option = self.get_obj_addrs().search_lock(id).cloned();
                match lock_option {
                    Some(mut lock) => {
                        let old_status = machine.interrupt.set_status(InterruptStatus::InterruptOff);

                        match &lock.owner {
                            Some(lock_owner) => {
                                if Rc::ptr_eq(&current_thread, lock_owner) {
                                    match lock.waiting_queue.pop() {
                                        Some(thread) => {
                                            lock.owner = Some(thread);
                                            match &lock.owner {
                                                Some(x) => {
                                                    self.ready_to_run(Rc::clone(x));
                                                    lock.free = true;
                                                },
                                                None => ()
                                            }
                                        },
                                        None => {
                                            lock.free = true;
                                            lock.owner = None;
                                        }
                                    }
                                }
                            },
                            None => ()
                        }
                        machine.interrupt.set_status(old_status);
                        Ok(MachineOk::Ok)
                    },
                    None => Err("Cannot find lock")?
                }
            },
            None => unreachable!("Current thread should not be None")
        }

    }

    /// Currently running thread
    pub fn get_g_current_thread(&mut self) -> &Option<Rc<RefCell<Thread>>> {
        &self.g_current_thread
    }

    /// Thread to be destroyed by [...]
    /// 
    /// TODO: Finish the comment with the relevant value
    pub fn get_g_thread_to_be_destroyed(&mut self) -> &Option<Rc<RefCell<Thread>>> {
        &self.g_thread_to_be_destroyed
    }

    /// List of alive threads
    pub fn get_g_alive(&mut self) -> &mut List<Rc<RefCell<Thread>>> {
        &mut self.g_alive
    }

    /// Set currently running thread
    pub fn set_g_current_thread(&mut self, thread: Option<Rc<RefCell<Thread>>>) {
        self.g_current_thread = thread
    }

    /// Set thread to be destroyed next
    pub fn set_g_thread_to_be_destroyed(&mut self, thread: Option<Rc<RefCell<Thread>>>) {
        self.g_thread_to_be_destroyed = thread
    }

    pub fn get_obj_addrs(&mut self) -> &mut ObjAddr {
        &mut self.obj_addrs 
    }

    fn debug(&self, message: String) {
        if self.debug {
            println!("{}", message);
        }
    }

}

#[cfg(test)]
mod test {
    use std::{rc::Rc, cell::RefCell};

    use crate::{simulator::{machine::Machine, loader}, kernel::{system::System, thread::Thread, process::Process, thread_manager::ThreadManager, synch::Semaphore}};
    use crate::kernel::synch::Lock;

    #[test]
    fn test_thread_context() {
        let mut machine = Machine::new(true);

        let (loader, ptr) = loader::Loader::new("./target/guac/halt.guac", &mut machine, 0).expect("IO Error");
        let start_pc = loader.elf_header.entrypoint;
        let system = &mut System::new(true);

        let thread1 = Thread::new("th1");
        let thread1 = Rc::new(RefCell::new(thread1));
        system.get_thread_manager().get_g_alive().push(Rc::clone(&thread1));

        let owner1 = Process { num_thread: 0 };
        let owner1 = Rc::new(RefCell::new(owner1));
        system.get_thread_manager().start_thread(Rc::clone(&thread1), owner1, loader.elf_header.entrypoint, ptr, -1);
        debug_assert_eq!(thread1.borrow_mut().thread_context.pc, start_pc);
        debug_assert!(system.get_thread_manager().get_g_alive().contains(&Rc::clone(&thread1)));

        let to_run = system.get_thread_manager().find_next_to_run().unwrap();
        debug_assert_eq!(to_run, Rc::clone(&thread1));

        system.get_thread_manager().switch_to(&mut machine, Rc::clone(&to_run));
        debug_assert_eq!(system.get_thread_manager().g_current_thread, Option::Some(Rc::clone(&thread1)));
        debug_assert_eq!(machine.pc, loader.elf_header.entrypoint);

        machine.run(system);
    }

    #[test]
    fn test_lock(){
        let mut machine = Machine::new(true);
        let mut thread_manager = ThreadManager::new(true);
        let lock = Lock::new();
        let lock_id = thread_manager.get_obj_addrs().add_lock(lock);
        let thread = Rc::new(RefCell::new(Thread::new("test_lock")));
        thread_manager.ready_to_run(Rc::clone(&thread));
        thread_manager.set_g_current_thread(Some(thread));

        thread_manager.lock_acquire(lock_id, &mut machine).expect("lock acquire return an error: ");
        {
            let lock = thread_manager.get_obj_addrs().search_lock(lock_id).unwrap();
            assert!(!lock.free);
            assert!(lock.waiting_queue.is_empty());
        }

        thread_manager.lock_release(lock_id, &mut machine).expect("lock release return an error: ");
        {
            let lock = thread_manager.get_obj_addrs().search_lock(lock_id).unwrap();
            assert!(!lock.free);
            assert!(lock.waiting_queue.is_empty());
        }
    }

    #[test]
    fn test_semaphore_single() {
        // Init
        let mut machine = Machine::new(true);
        let mut thread_manager = ThreadManager::new(true);
        let semaphore = Semaphore::new(1);
        let sema_id = thread_manager.get_obj_addrs().add_semaphore(semaphore);
        let thread = Rc::new(RefCell::new(Thread::new("test_semaphore")));
        thread_manager.ready_to_run(Rc::clone(&thread));
        thread_manager.set_g_current_thread(Some(thread));
        // P
        thread_manager.sem_p(sema_id, &mut machine).expect("semaphore P return an error: ");
        {
            let semaphore = thread_manager.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 0);
            assert!(semaphore.waiting_queue.is_empty());
        }
        // V
        thread_manager.sem_v(sema_id, &mut machine).expect("semaphore V return an error: ");
        {
            let semaphore = thread_manager.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 1);
            assert!(semaphore.waiting_queue.is_empty());
        }
    }

    #[test]
    fn test_semaphore_multiple() {
        // Init
        let mut tm = ThreadManager::new(true);
        let mut machine = Machine::new(true);
        let semaphore = Semaphore::new(2);
        let sema_id = tm.get_obj_addrs().add_semaphore(semaphore);
        let thread1 = Rc::new(RefCell::new(Thread::new("test_semaphore_1")));
        let thread2 = Rc::new(RefCell::new(Thread::new("test_semaphore_2")));
        let thread3 = Rc::new(RefCell::new(Thread::new("test_semaphore_3")));

        // let mut borrow_tm = tm.borrow_mut();
        // let scheduler = &mut tm.g_scheduler;
        tm.ready_to_run(Rc::clone(&thread1));
        tm.ready_to_run(Rc::clone(&thread2));
        tm.ready_to_run(Rc::clone(&thread3));
        // P
        tm.set_g_current_thread(Some(Rc::clone(&thread1)));
        tm.sem_p(sema_id, &mut machine).expect("semaphore P return an error: ");
        {
            let semaphore = tm.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 1);
            assert!(semaphore.waiting_queue.is_empty());
        }
        
        tm.set_g_current_thread(Some(Rc::clone(&thread2)));
        tm.sem_p(sema_id, &mut machine).expect("semaphore P return an error: ");
        {
            let semaphore = tm.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 0);
            assert!(semaphore.waiting_queue.is_empty());
        }

        tm.set_g_current_thread(Some(Rc::clone(&thread3)));
        tm.sem_p( sema_id, &mut machine).expect("semaphore P return an error: ");
        {
            let semaphore = tm.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, -1);
            assert!(semaphore.waiting_queue.iter().count() == 1);
        }

        // V
        tm.sem_v(sema_id, &mut machine).expect("semaphore V return an error: ");
        {
            let semaphore = tm.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 0);
            assert!(semaphore.waiting_queue.is_empty());
        }

        tm.sem_v(sema_id, &mut machine).expect("semaphore V return an error: ");
        {
            let semaphore = tm.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 1);
            assert!(semaphore.waiting_queue.is_empty());
        }

        tm.sem_v(sema_id, &mut machine).expect("semaphore V return an error: ");
        {
            let semaphore = tm.get_obj_addrs().search_semaphore(sema_id).unwrap();
            assert_eq!(semaphore.counter, 2);
            assert!(semaphore.waiting_queue.is_empty());
        }
        
    }

}