alloc_cat 1.1.1

use cfg_if::cfg_if;

cfg_if! {
    if #[cfg(feature = "audit_sizes")] {
        use arr_macro::arr;
    }
}
use core::alloc::Layout;
use core::sync::atomic::{AtomicUsize, Ordering};

use crate::AllocatorStats;


cfg_if! {
    if #[cfg(feature = "audit_sizes")] {
        #[derive(Debug)]
        pub struct AtomicSizeInfo {
            pub count: AtomicUsize,
            pub high_water: AtomicUsize,
            pub trip_high_water: AtomicUsize,
        }

        impl AtomicSizeInfo {
            const fn new() -> AtomicSizeInfo {
                AtomicSizeInfo {
                    count: AtomicUsize::new(0),
                    high_water: AtomicUsize::new(0),
                    trip_high_water: AtomicUsize::new(0),
                }
            }

            fn unpack(&self, size: usize) -> SizeInfo {
                SizeInfo {
                    size,
                    count: self.count.load(Ordering::Relaxed),
                    high_water: self.high_water.load(Ordering::Relaxed),
                    trip_high_water: self.trip_high_water.load(Ordering::Relaxed),
                }
            }
        }

        impl Default for AtomicSizeInfo {
            fn default() -> AtomicSizeInfo {
                AtomicSizeInfo::new()
            }
        }
    }
}


#[derive(Clone, Copy, Debug)]
pub struct SizeInfo {
    pub size: usize,
    pub count: usize,
    pub high_water: usize,
    pub trip_high_water: usize,
}

impl SizeInfo {
    pub const fn new() -> SizeInfo {
        SizeInfo { size: 0, count: 0, high_water: 0, trip_high_water: 0 }
    }
}

impl Default for SizeInfo {
    fn default() -> SizeInfo {
        SizeInfo::new()
    }
}


cfg_if! {
    if #[cfg(feature = "audit_sizes")] {
        #[derive(Debug)]
        pub struct Metrics {
            // bytes currently in use, and highest ever in use
            in_use: AtomicUsize,
            high_water: AtomicUsize,
            trip_high_water: AtomicUsize,

            // count of requests for each alignment
            align_table: [AtomicUsize; 16],

            // buckets per allocation granularity
            size_table: [AtomicSizeInfo; 257],
        }
    } else {
        #[derive(Debug)]
        pub struct Metrics {
            // bytes currently in use, and highest ever in use
            in_use: AtomicUsize,
            high_water: AtomicUsize,
            trip_high_water: AtomicUsize,
        }
    }
}

impl Metrics {
    pub const fn new() -> Metrics {
        cfg_if! {
            if #[cfg(feature = "audit_sizes")] {
                Metrics {
                    in_use: AtomicUsize::new(0),
                    high_water: AtomicUsize::new(0),
                    trip_high_water: AtomicUsize::new(0),
                    align_table: arr![AtomicUsize::new(0); 16],
                    size_table: arr![AtomicSizeInfo::new(); 257],
                }
            } else {
                Metrics {
                    in_use: AtomicUsize::new(0),
                    high_water: AtomicUsize::new(0),
                    trip_high_water: AtomicUsize::new(0),
                }
            }
        }
    }

    pub fn reset_trip(&self) {
        self.trip_high_water.store(0, Ordering::Relaxed);
    }

    pub fn get_stats(&self, stats: &mut AllocatorStats) {
        stats.in_use = self.in_use.load(Ordering::Relaxed);
        stats.high_water = self.high_water.load(Ordering::Relaxed);
        stats.trip_high_water = self.trip_high_water.load(Ordering::Relaxed);

        cfg_if! {
            if #[cfg(feature = "audit_sizes")] {
                for (i, atom) in self.align_table.iter().enumerate() {
                    stats.align_table[i] = atom.load(Ordering::Relaxed);
                }

                for (i, bucket) in self.size_table.iter().enumerate() {
                    stats.size_table[i] = bucket.unpack(slot_to_bucket(i));
                }
            }
        }
    }

    pub fn add(&self, layout: Layout) {
        let in_use = self.in_use.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
        self.high_water.fetch_max(in_use, Ordering::Relaxed);
        self.trip_high_water.fetch_max(in_use, Ordering::Relaxed);

        cfg_if! {
            if #[cfg(feature = "audit_sizes")] {
                let i = u32::BITS - (layout.align() as u32).leading_zeros();
                assert!(i < 16);
                self.align_table[i as usize].fetch_add(1, Ordering::Relaxed);

                let bucket = &self.size_table[bucket_to_slot(layout.size())];
                let prev = bucket.count.fetch_add(1, Ordering::Relaxed);
                bucket.high_water.fetch_max(prev + 1, Ordering::Relaxed);
            }
        }
    }

    pub fn subtract(&self, layout: Layout) {
        self.in_use.fetch_sub(layout.size(), Ordering::Relaxed);
        cfg_if! {
            if #[cfg(feature = "audit_sizes")] {
                let bucket = &self.size_table[bucket_to_slot(layout.size())];
                bucket.count.fetch_sub(1, Ordering::Relaxed);
            }
        }
    }
}

// 256 buckets are split into 8 ranges, each with different granularity:
//
//            bucket   bucket
//     count  range    size  span  range
//     -----  --       ----  ----  -----
//     48     0-47     1     48    1-48
//     16     48-63    2     32    49-80
//     32     64-95    4     128   81-208
//     32     96-127   8     256   209-464
//     32     128-159  16    512   465-976
//     32     160-191  32    1024  977-2000
//     32     192-223  128   4096  2001-6096
//     32     224-255  512   16384 6097-22480
//
// bucket 257 is reserved for "bigger than that".

cfg_if! {
    if #[cfg(feature = "audit_sizes")] {
        const fn bucket_to_slot(n: usize) -> usize {
            match n {
                _ if n <= 48 => n - 1,
                _ if n <= 80 => ((n - 49) >> 1) + 48,
                _ if n <= 208 => ((n - 81) >> 2) + 64,
                _ if n <= 464 => ((n - 209) >> 3) + 96,
                _ if n <= 976 => ((n - 465) >> 4) + 128,
                _ if n <= 2000 => ((n - 977) >> 5) + 160,
                _ if n <= 6096 => ((n - 2001) >> 7) + 192,
                _ if n <= 22480 => ((n - 6097) >> 9) + 224,
                _ => 256
            }
        }

        const fn slot_to_bucket(n: usize) -> usize {
            match n {
                _ if n < 48 => n + 1,
                _ if n < 64 => ((n - 48) << 1) + 50,
                _ if n < 96 => ((n - 64) << 2) + 84,
                _ if n < 128 => ((n - 96) << 3) + 216,
                _ if n < 160 => ((n - 128) << 4) + 480,
                _ if n < 192 => ((n - 160) << 5) + 1008,
                _ if n < 224 => ((n - 192) << 7) + 2128,
                _ if n < 256 => ((n - 224) << 9) + 6608,
                _ => usize::MAX,
            }
        }
    }
}