// Smoldot

// Copyright (C) 2019-2022  Parity Technologies (UK) Ltd.

// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify

// it under the terms of the GNU General Public License as published by

// the Free Software Foundation, either version 3 of the License, or

// (at your option) any later version.

// This program is distributed in the hope that it will be useful,

// but WITHOUT ANY WARRANTY; without even the implied warranty of

// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License

// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! This module provides the `Delay` struct, which implement `Future` and becomes ready after a

//! certain time.

//!

//! In order to optimize performances, we avoid invoking the FFI once per timer. Instead, the FFI

//! is only used in order to wake up when the earliest timer finishes, then restarted for the next

//! timer.

use crate::bindings;

use alloc::collections::BTreeSet;

use async_lock::Mutex;

use core::{

    cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd},

    future, mem,

    pin::Pin,

    task::{Context, Poll, Waker},

    time::Duration,

};

pub(crate) fn timer_finished() {

    process_timers();

}

/// `Future` that automatically wakes up after a certain amount of time has elapsed.

pub struct Delay {

    /// Index in `TIMERS::timers`. Guaranteed to have `is_obsolete` equal to `false`.

    /// If `None`, then this timer is already ready.

    timer_id: Option<usize>,

}

impl Delay {

    pub fn new(after: Duration) -> Self {

        let now = unsafe { Duration::from_micros(bindings::monotonic_clock_us()) };

        Self::new_inner(now + after, now)

    }

    pub fn new_at_monotonic_clock(when: Duration) -> Self {

        let now = unsafe { Duration::from_micros(bindings::monotonic_clock_us()) };

        Self::new_inner(when, now)

    }

    fn new_inner(when: Duration, now: Duration) -> Self {

        // Small optimization because sleeps of 0 seconds are frequent.

        if when <= now {

            return Delay { timer_id: None };

        }

        // Because we're in a single-threaded environment, `try_lock()` should always succeed.

        let mut lock = TIMERS.try_lock().unwrap();

        let timer_id = lock.timers.insert(Timer {

            is_finished: false,

            is_obsolete: false,

            waker: None,

        });

        lock.timers_queue.insert(QueuedTimer { when, timer_id });

        // If the timer that has just been inserted is the one that ends the soonest, then

        // actually start the callback that will process timers.

        // Ideally we would instead cancel or update the deadline of the previous call to

        // `start_timer`, but this isn't possible.

        if lock

            .timers_queue

            .first()

            .unwrap_or_else(|| unreachable!())

            .timer_id

            == timer_id

        {

            start_timer(when - now);

        }

        Delay {

            timer_id: Some(timer_id),

        }

    }

}

impl future::Future for Delay {

    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {

        let timer_id = match self.timer_id {

            Some(id) => id,

            None => return Poll::Ready(()),

        };

        // Because we're in a single-threaded environment, `try_lock()` should always succeed.

        let mut lock = TIMERS.try_lock().unwrap();

        debug_assert!(!lock.timers[timer_id].is_obsolete);

        if lock.timers[timer_id].is_finished {

            lock.timers.remove(timer_id);

            self.timer_id = None;

            return Poll::Ready(());

        }

        lock.timers[timer_id].waker = Some(cx.waker().clone());

        Poll::Pending

    }

}

impl Drop for Delay {

    fn drop(&mut self) {

        let timer_id = match self.timer_id {

            Some(id) => id,

            None => return,

        };

        // Because we're in a single-threaded environment, `try_lock()` should always succeed.

        let mut lock = TIMERS.try_lock().unwrap();

        debug_assert!(!lock.timers[timer_id].is_obsolete);

        if lock.timers[timer_id].is_finished {

            lock.timers.remove(timer_id);

            return;

        }

        lock.timers[timer_id].is_obsolete = true;

        lock.timers[timer_id].waker = None;

    }

}

static TIMERS: Mutex<Timers> = Mutex::new(Timers {

    timers_queue: BTreeSet::new(),

    timers: slab::Slab::new(),

});

struct Timers {

    /// Same entries as `timer`, but ordered based on when they're finished (from soonest to

    /// latest). Items are only ever removed from [`process_timers`] when they finish, even if

    /// the corresponding [`Delay`] is destroyed.

    timers_queue: BTreeSet<QueuedTimer>,

    /// List of all timers.

    timers: slab::Slab<Timer>,

}

struct Timer {

    /// If `true`, then this timer has elapsed.

    is_finished: bool,

    /// If `true`, then the corresponding `Delay` has been destroyed or no longer points to this

    /// item.

    is_obsolete: bool,

    /// How to wake up the `Delay`.

    waker: Option<Waker>,

}

struct QueuedTimer {

    when: Duration,

    // Entry in `TIMERS::timers`. Guaranteed to always have `is_finished` equal to `false`.

    timer_id: usize,

}

impl PartialEq for QueuedTimer {

    fn eq(&self, other: &Self) -> bool {

        matches!(self.cmp(other), Ordering::Equal)

    }

}

impl Eq for QueuedTimer {}

impl PartialOrd for QueuedTimer {

    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {

        Some(Ord::cmp(self, other))

    }

}

impl Ord for QueuedTimer {

    fn cmp(&self, other: &Self) -> Ordering {

        // `when` takes higher priority in the ordering.

        match self.when.cmp(&other.when) {

            Ordering::Equal => self.timer_id.cmp(&other.timer_id),

            ord => ord,

        }

    }

}

/// Marks as ready all the timers in `TIMERS` that are finished.

fn process_timers() {

    // Because we're in a single-threaded environment, `try_lock()` should always succeed.

    let mut lock = TIMERS.try_lock().unwrap();

    let lock = &mut *lock;

    let now = unsafe { Duration::from_micros(bindings::monotonic_clock_us()) };

    // Note that this function can be called spuriously.

    // For example, `process_timers` can be scheduled twice from two different timers, and the

    // first call leads to both timers being finished, after which the second call will be

    // spurious.

    // We remove all the queued timers whose `when` is inferior to `now`.

    let expired_timers = {

        let timers_remaining = lock.timers_queue.split_off(&QueuedTimer {

            when: now,

            // Note that `split_off` returns values greater or equal, meaning that if a timer had

            // a `timer_id` equal to `max_value()` it would erroneously be returned instead of being

            // left in the collection as expected. For obvious reasons, a `timer_id` of

            // `usize::MAX` is impossible, so this isn't a problem.

            timer_id: usize::MAX,

        });

        mem::replace(&mut lock.timers_queue, timers_remaining)

    };

    // Wake up the expired timers.

    for timer in expired_timers {

        debug_assert!(timer.when <= now);

        debug_assert!(!lock.timers[timer.timer_id].is_finished);

        lock.timers[timer.timer_id].is_finished = true;

        if let Some(waker) = lock.timers[timer.timer_id].waker.take() {

            waker.wake();

        }

    }

    // Figure out the next time we should call `process_timers`.

    //

    // This iterates through all the elements in `timers_queue` until a valid one is found.

    let next_wakeup: Option<Duration> = loop {

        let next_timer = match lock.timers_queue.first() {

            Some(t) => t,

            None => break None,

        };

        // The `Delay` corresponding to the iterated timer has been destroyed. Removing it and

        // `continue`.

        if lock.timers[next_timer.timer_id].is_obsolete {

            let next_timer_id = next_timer.timer_id;

            lock.timers.remove(next_timer_id);

            lock.timers_queue

                .pop_first()

                .unwrap_or_else(|| unreachable!());

            continue;

        }

        // Iterated timer is not ready.

        break Some(next_timer.when);

    };

    if let Some(next_wakeup) = next_wakeup {

        start_timer(next_wakeup - now);

    } else {

        // Clean up memory a bit. Hopefully this doesn't impact performances too much.

        if !lock.timers.is_empty() && lock.timers.capacity() > lock.timers.len() * 8 {

            lock.timers.shrink_to_fit();

        }

    }

}

/// Instructs the environment to call [`process_timers`] after the given duration.

fn start_timer(duration: Duration) {

    // Note that ideally `duration` should be rounded up in order to make sure that it is not

    // truncated, but the precision of an `f64` is so high and the precision of the operating

    // system generally so low that this is not worth dealing with.

    unsafe { bindings::start_timer(duration.as_secs_f64() * 1000.0) }

}