ahrightos-process 0.1.0

use crate::{ProcessContext, ProcessId};
use alloc::vec::Vec;
use core::fmt;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProcessState {
    Ready,
    Running,
    Blocked,
    Zombie,
}

impl ProcessState {
    pub fn can_schedule(self) -> bool {
        self == Self::Ready
    }

    pub fn can_transition_to(self, new_state: Self) -> bool {
        match (self, new_state) {
            (Self::Ready, Self::Running)
            | (Self::Running, Self::Ready)
            | (Self::Running, Self::Blocked)
            | (Self::Running, Self::Zombie)
            | (Self::Blocked, Self::Ready) => true,
            (old, new) if old == new => true,
            _ => false,
        }
    }
}

impl fmt::Display for ProcessState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let state = match self {
            Self::Ready => "Ready",
            Self::Running => "Running",
            Self::Blocked => "Blocked",
            Self::Zombie => "Zombie",
        };
        write!(f, "{}", state)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProcessStateError {
    pub from: ProcessState,
    pub to: ProcessState,
}

impl fmt::Display for ProcessStateError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "invalid process transition {} -> {}", self.from, self.to)
    }
}

#[derive(Clone)]
pub struct ProcessControlBlock {
    pid: ProcessId,
    name: &'static str,
    state: ProcessState,
    exit_code: Option<i32>,
    parent_pid: Option<ProcessId>,
    children: Vec<ProcessId>,
    context: ProcessContext,
    time_slice: usize,
    priority: usize,
}

impl ProcessControlBlock {
    pub fn new(pid: ProcessId, name: &'static str, parent_pid: Option<ProcessId>) -> Self {
        Self {
            pid,
            name,
            state: ProcessState::Ready,
            exit_code: None,
            parent_pid,
            children: Vec::new(),
            context: ProcessContext::new(),
            time_slice: 5,
            priority: 1,
        }
    }

    pub fn add_child(&mut self, pid: ProcessId) {
        if !self.children.contains(&pid) {
            self.children.push(pid);
        }
    }

    pub fn remove_child(&mut self, pid: ProcessId) {
        self.children.retain(|&p| p != pid);
    }

    pub fn children(&self) -> &[ProcessId] {
        &self.children
    }

    pub fn parent_pid(&self) -> Option<ProcessId> {
        self.parent_pid
    }

    pub fn state(&self) -> ProcessState {
        self.state
    }

    pub fn set_state(&mut self, state: ProcessState) -> Result<(), ProcessStateError> {
        if !self.state.can_transition_to(state) {
            return Err(ProcessStateError {
                from: self.state,
                to: state,
            });
        }
        self.state = state;
        Ok(())
    }

    #[cfg(test)]
    pub(crate) fn set_state_unchecked(&mut self, state: ProcessState) {
        self.state = state;
    }

    pub fn is_runnable(&self) -> bool {
        self.state == ProcessState::Ready
    }

    pub fn set_exit_code(&mut self, code: i32) {
        self.exit_code = Some(code);
        self.state = ProcessState::Zombie;
    }

    pub fn exit_code(&self) -> Option<i32> {
        self.exit_code
    }

    pub fn context_mut(&mut self) -> &mut ProcessContext {
        &mut self.context
    }

    pub fn context(&self) -> &ProcessContext {
        &self.context
    }

    pub fn pid(&self) -> ProcessId {
        self.pid
    }

    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn tick(&mut self) -> usize {
        if self.time_slice > 0 {
            self.time_slice -= 1;
        }
        self.time_slice
    }

    pub fn reset_time_slice(&mut self) {
        self.time_slice = 5;
    }

    pub fn time_slice(&self) -> usize {
        self.time_slice
    }

    pub fn set_priority(&mut self, priority: usize) {
        self.priority = priority;
    }

    pub fn priority(&self) -> usize {
        self.priority
    }

    pub fn is_init(&self) -> bool {
        self.pid.is_init()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pcb_tracks_children_and_zombie_exit_code() {
        let parent = ProcessId::from_usize(1);
        let child = ProcessId::from_usize(2);
        let mut pcb = ProcessControlBlock::new(parent, "init", None);
        pcb.add_child(child);
        pcb.add_child(child);
        assert_eq!(pcb.children(), &[child]);
        pcb.remove_child(child);
        assert!(pcb.children().is_empty());
        pcb.set_exit_code(42);
        assert_eq!(pcb.state(), ProcessState::Zombie);
        assert_eq!(pcb.exit_code(), Some(42));
    }

    #[test]
    fn invalid_state_transition_is_rejected() {
        let mut pcb = ProcessControlBlock::new(ProcessId::from_usize(3), "worker", None);
        let err = pcb.set_state(ProcessState::Blocked).unwrap_err();
        assert_eq!(err.from, ProcessState::Ready);
        assert_eq!(err.to, ProcessState::Blocked);
        assert_eq!(pcb.state(), ProcessState::Ready);
    }
}