    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

    use super::UPSafeCell;
    use crate::task::TaskControlBlock;
    use crate::task::{add_task, current_task};
    use crate::task::{block_current_and_run_next, suspend_current_and_run_next};
    use alloc::{collections::VecDeque, sync::Arc};

    pub trait Mutex: Sync + Send {
        fn lock(&self);
        fn unlock(&self);
    }

    pub struct MutexSpin {
        locked: UPSafeCell<bool>,
    }

    impl MutexSpin {
        pub fn new() -> Self {
            Self {
                locked: unsafe { UPSafeCell::new(false) },
            }
        }
    }

    impl Mutex for MutexSpin {
        fn lock(&self) {
            loop {
                let mut locked = self.locked.exclusive_access();
                if *locked {
                    drop(locked);
                    suspend_current_and_run_next();
                    continue;
                } else {
                    *locked = true;
                    return;
                }
            }
        }

        fn unlock(&self) {
            let mut locked = self.locked.exclusive_access();
            *locked = false;
        }
    }

    pub struct MutexBlocking {
        inner: UPSafeCell<MutexBlockingInner>,
    }

    pub struct MutexBlockingInner {
        locked: bool,
        wait_queue: VecDeque<Arc<TaskControlBlock>>,
    }

    impl MutexBlocking {
        pub fn new() -> Self {
            Self {
                inner: unsafe {
                    UPSafeCell::new(MutexBlockingInner {
                        locked: false,
                        wait_queue: VecDeque::new(),
                    })
                },
            }
        }
    }

    impl Mutex for MutexBlocking {
        fn lock(&self) {
            let mut mutex_inner = self.inner.exclusive_access();
            if mutex_inner.locked {
                mutex_inner.wait_queue.push_back(current_task().unwrap());
                drop(mutex_inner);
                block_current_and_run_next();
            } else {
                mutex_inner.locked = true;
            }
        }

        fn unlock(&self) {
            let mut mutex_inner = self.inner.exclusive_access();
            assert!(mutex_inner.locked);
            if let Some(waking_task) = mutex_inner.wait_queue.pop_front() {
                add_task(waking_task);
            } else {
                mutex_inner.locked = false;
            }
        }
    }