lowlet 0.1.2

use std::marker::PhantomData;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::sync::asm::fence::pause;

const PTR_MASK: u64 = 0x0000_FFFF_FFFF_FFFF;
const TAG_SHIFT: u32 = 48;

#[repr(C)]
pub struct Node<T> {
    pub value: T,
    next: AtomicU64,
}

impl<T> Node<T> {
    pub const fn new(value: T) -> Self {
        Self {
            value,
            next: AtomicU64::new(0),
        }
    }
}

#[repr(C, align(64))]
pub struct Stack<T> {
    head: AtomicU64,
    _marker: PhantomData<T>,
}

unsafe impl<T: Send> Send for Stack<T> {}
unsafe impl<T: Send> Sync for Stack<T> {}

impl<T> Stack<T> {
    pub const fn new() -> Self {
        Self {
            head: AtomicU64::new(0),
            _marker: PhantomData,
        }
    }

    fn pack(ptr: *mut Node<T>, tag: u16) -> u64 {
        (u64::from(tag) << TAG_SHIFT) | (ptr as u64 & PTR_MASK)
    }

    fn unpack(val: u64) -> (*mut Node<T>, u16) {
        let ptr = (val & PTR_MASK) as *mut Node<T>;
        let tag = (val >> TAG_SHIFT) as u16;
        (ptr, tag)
    }

    pub unsafe fn push(&self, node: *mut Node<T>) {
        loop {
            let head = self.head.load(Ordering::Acquire);
            let (_, tag) = Self::unpack(head);

            (*node).next.store(head, Ordering::Relaxed);

            let new = Self::pack(node, tag.wrapping_add(1));

            if self
                .head
                .compare_exchange_weak(head, new, Ordering::Release, Ordering::Relaxed)
                .is_ok()
            {
                return;
            }

            pause();
        }
    }

    pub unsafe fn pop(&self) -> Option<*mut Node<T>> {
        loop {
            let head = self.head.load(Ordering::Acquire);
            let (ptr, tag) = Self::unpack(head);

            if ptr.is_null() {
                return None;
            }

            let next = (*ptr).next.load(Ordering::Acquire);
            let new = Self::pack(Self::unpack(next).0, tag.wrapping_add(1));

            if self
                .head
                .compare_exchange_weak(head, new, Ordering::Release, Ordering::Relaxed)
                .is_ok()
            {
                return Some(ptr);
            }

            pause();
        }
    }

    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::Acquire);
        Self::unpack(head).0.is_null()
    }
}

impl<T> Default for Stack<T> {
    fn default() -> Self {
        Self::new()
    }
}