// Copyright (c) 2015 The Robigalia Project Developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.
//! See the `Bitmap` type.
#![doc(html_root_url = "https://doc.robigalia.org/")]
#![cfg_attr(all(not(test), feature = "no_std"), no_std)]
#![cfg_attr(feature = "collections", feature(collections))]

#[cfg(feature = "collections")]
extern crate collections;
#[cfg(feature = "collections")]
use collections::vec::Vec;

#[cfg(any(test, not(feature = "no_std")))]
extern crate core;

/// The "element type" of a bitmap
///
/// Types implementing this trait can be stored in a bitmap. However, you shouldn't need to
/// implement this trait yourself if you only care about semantics-free bits: use one of the
/// fixed-width types defined here, such as `OneBit`, `TwoBits`, etc, or the `DynamicWidth` type.
///
/// However, this can be used to make the `Bitmap` strongly typed storage for types such as those
/// generated by the `bitflags!` macro.
pub trait Element: Sized {
    /// Type passed to the `width` method.
    type W: Copy;
    /// Type decoded from `from_bits` and thus `Bitmap::get`.
    type R: Copy;
    /// Return the width of values of this type, in bits.
    ///
    /// This method is called somewhat frequently, on almost every access to the `Bitmap`. However,
    /// the value it is passed is only ever specified in one place: `Bitmap::from_storage`.
    ///
    /// This should never return a value greater than 64, and should always return the same value
    /// for the same `Bitmap`.
    fn width(Self::W) -> usize;
    /// Decode a value from raw bits.
    fn from_bits(bits: u64) -> Option<Self::R>;
    /// Encode a value into raw bits.
    ///
    /// No bits outside of the least significant `Element::width(w)` bits should be set.
    fn to_bits(Self::R) -> u64;
}

/// Storage for a slice of `usize`.
pub trait Storage {
    fn as_ref(&self) -> &[usize];
    fn as_mut(&mut self) -> &mut [usize];
}

//impl<T: core::ops::DerefMut<Target=[usize]>> Storage for T {
//    fn as_ref(&self) -> &[usize] { &*self }
//    fn as_mut(&mut self) -> &mut [usize] { &mut *self }
//}

impl Storage for [usize] {
    fn as_ref(&self) -> &[usize] { self }
    fn as_mut(&mut self) -> &mut [usize] { self }
}

#[cfg(any(test, not(feature = "no_std"), feature = "collections"))]
impl Storage for Vec<usize> {
    fn as_ref(&self) -> &[usize] { &*self }
    fn as_mut(&mut self) -> &mut [usize] { &mut *self }
}

macro_rules! fixed_storage_impl {
    ($n:expr) => { 
        impl Storage for [usize; $n] {
            fn as_ref(&self) -> &[usize] { self }
            fn as_mut(&mut self) -> &mut [usize] { self }
        }
    }
}

fixed_storage_impl!(0);
fixed_storage_impl!(1);
fixed_storage_impl!(2);
fixed_storage_impl!(3);
fixed_storage_impl!(4);
fixed_storage_impl!(5);
fixed_storage_impl!(6);
fixed_storage_impl!(7);
fixed_storage_impl!(8);
fixed_storage_impl!(9);
fixed_storage_impl!(10);
fixed_storage_impl!(11);
fixed_storage_impl!(12);
fixed_storage_impl!(13);
fixed_storage_impl!(14);
fixed_storage_impl!(15);
fixed_storage_impl!(16);
fixed_storage_impl!(17);
fixed_storage_impl!(18);
fixed_storage_impl!(19);
fixed_storage_impl!(20);
fixed_storage_impl!(21);
fixed_storage_impl!(22);
fixed_storage_impl!(23);
fixed_storage_impl!(24);
fixed_storage_impl!(25);
fixed_storage_impl!(26);
fixed_storage_impl!(27);
fixed_storage_impl!(28);
fixed_storage_impl!(29);
fixed_storage_impl!(30);
fixed_storage_impl!(31);
fixed_storage_impl!(32);

/// A dense bitmap, intended to store small bitslices (<= width of u64).
///
/// That is, this type represents an array of values, where each value is `E::width(instance)` bits long,
/// and the bits are stored contiguously. The first value is packed into the least significant bits
/// of the first word.
pub struct Bitmap<S: Storage, E: Element> {
    entries: usize,
    width: E::W,
    storage: S,
    phantom: core::marker::PhantomData<E>,
}

impl<S: Storage, E: Element> core::fmt::Debug for Bitmap<S, E> {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        self.storage.as_ref().fmt(f)
    }
}

#[inline(always)]
fn bits() -> usize {
    core::mem::size_of::<usize>() * 8
}

#[inline(always)]
fn get_n_bits_at(word: usize, n: u8, start: u8) -> usize {
    debug_assert!(n <= bits() as u8);
    debug_assert!(start < bits() as u8);
    debug_assert!(bits() <= 64); // y'know, just in case.
    (word >> start as usize) & (!0 >> ((bits() as u8 - n) as usize % bits()))
}

#[inline(always)]
fn set_n_bits_at(word: usize, value: usize, n: u8, start: u8) -> usize {
    debug_assert!(n <= bits() as u8);
    debug_assert!(start < bits() as u8);

    let premask = (1 << (n as usize)) - 1;
    let mask = premask << (start as usize);
    (word & !mask) | ((value & premask) << start as usize)
}

impl<S: Storage, E: Element> Bitmap<S, E> {
    #[inline(always)]
    fn width(&self) -> usize {
        E::width(self.width)
    }

    pub fn from_storage(entries: usize, w: E::W, storage: S) -> Option<Bitmap<S, E>> {
        if  E::width(w) > bits() || E::width(w) == 0 {
            None
        } else {
            entries.checked_mul(E::width(w))
            .and_then(|bits| bits.checked_add(bits % ::bits()))
            .and_then(|rbits| rbits.checked_div(::bits()))
            .and_then(|_| {
                let bits_needed = entries.wrapping_mul(E::width(w));
                if (bits_needed + (bits_needed % bits())) / bits() > storage.as_ref().len() {
                    None
                } else {
                    Some(Bitmap {
                        entries: entries,
                        width: w,
                        storage: storage,
                        phantom: core::marker::PhantomData,
                    })
                }
            })
        }
    }

    /// Get the `i`th bitslice, returning None on out-of-bounds or if `E::from_bits` returns None.
    pub fn get(&self, i: usize) -> Option<E::R> {
        if i >= self.entries {
            None
        } else {
            let mut bit_offset = i.wrapping_mul(self.width());

            let mut word_offset = bit_offset / bits();
            let mut in_word_offset = bit_offset % bits();

            let mut bits_left = self.width();

            let mut value: u64 = 0;

            while bits_left > 0 {
                // how many bits can we get from this word?
                let can_get = core::cmp::min(bits() - in_word_offset, bits_left);

                // alright, pull them out.
                let word = unsafe { *self.storage.as_ref().get_unchecked(word_offset) };
                let got = get_n_bits_at(word, can_get as u8, in_word_offset as u8);

                // make room for the bits we just read
                value <<= can_get;
                value |= got as u64;

                // update all the state
                bit_offset = bit_offset.wrapping_add(can_get);
                in_word_offset = 0;
                word_offset = word_offset.wrapping_add(1);
                bits_left = bits_left.wrapping_sub(can_get);
            }
            E::from_bits(value)
        }
    }

    /// Set the `i`th bitslice to `value`, returning false on out-of-bounds or if `value` contains
    /// bits outside of the least significant `E::width(w)` bits.
    pub fn set(&mut self, i: usize, value: E::R) -> bool {
        let mut value = E::to_bits(value);
        if i >= self.entries {
            false
        } else if (value >> self.width()) != 0 {
            debug_assert!(value >> self.width() != 0, "value contained bits outside the least\
                          significant `E::width(w)` bits");
            false
        } else {
            let mut bit_offset = i.wrapping_mul(self.width());

            let mut word_offset = bit_offset / bits();
            let mut in_word_offset = bit_offset % bits();

            let mut bits_left = self.width();

            while bits_left > 0 {
                // how many bits can we get from this word?
                let can_set = core::cmp::min(bits() - in_word_offset, bits_left);

                // alright, pull them out.
                let word = unsafe { *self.storage.as_mut().get_unchecked(word_offset) };
                let update = set_n_bits_at(word, value as usize, can_set as u8, in_word_offset as u8);
                // and put them back in
                unsafe { *self.storage.as_mut().get_unchecked_mut(word_offset) = update };

                // empty the bits we just wrote out
                value >>= can_set;

                // update all the state
                bit_offset = bit_offset.wrapping_add(can_set);
                in_word_offset = 0;
                word_offset = word_offset.wrapping_add(1);
                bits_left = bits_left.wrapping_sub(can_set);
            }
            true
        }
    }

    /// Length in number of bitslices cointained.
    pub fn len(&self) -> usize {
        self.entries
    }

    /// Size of the internal buffer, in number of `usize`s.
    pub fn usize_len(&self) -> usize {
        // can't overflow, since creation asserts that it doesn't.
        let w = self.entries.wrapping_mul(self.width());
        let r = w % bits();
        (w.wrapping_add(r)) / bits()
    }

    pub fn iter(&self) -> Slices<S, E> {
        Slices { idx: 0, bm: self }
    }

    pub fn unwrap(self) -> S {
        let Bitmap {
            storage,
            ..
        } = self;
        storage
    }
}

/// Iterator over the bitslices in the bitmap
pub struct Slices<'a, S: Storage + 'a, E: Element + 'a> {
    idx: usize,
    bm: &'a Bitmap<S, E>,
}

impl<'a, S: Storage, E: Element> Iterator for Slices<'a, S, E> {
    type Item = E::R;
    /// *NOTE*: This iterator is not "well-behaved", in that if you keep calling
    /// `next` after it returns None, eventually it will overflow and start
    /// yielding elements again. Use the `fuse` method to make this
    /// "well-behaved".
    fn next(&mut self) -> Option<E::R> {
        let rv = self.bm.get(self.idx);
        self.idx += 1;
        rv
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.bm.len(), Some(self.bm.len()))
    }
}

impl<'a, S: Storage, E: Element> core::iter::IntoIterator for &'a Bitmap<S, E> {
    type Item = E::R;
    type IntoIter = Slices<'a, S, E>;

    fn into_iter(self) -> Slices<'a, S, E> {
        self.iter()
    }
}

impl<S: Storage> Bitmap<S, OneBit> {
    /// Return the index of the first bit set
    pub fn first_set(&self) -> Option<usize> {
        for (idx, &word) in self.storage.as_ref().iter().enumerate() {
            if word.trailing_zeros() != bits() as u32 {
                return Some(idx * bits() + word.trailing_zeros() as usize);
            }
        }
        None
    }
}

pub struct DynamicSize;

impl Element for DynamicSize {
    type W = usize;
    type R = u64;
    #[inline(always)]
    fn width(w: usize) -> usize { w }
    #[inline(always)]
    fn from_bits(bits: u64) -> Option<u64> { Some(bits) }
    #[inline(always)]
    fn to_bits(value: u64) -> u64 { value }
}

macro_rules! static_size {
     ($name:ident, $width:expr) => {
         /// Helper type for semantic-free values of fixed width
         pub struct $name;
         impl Element for $name {
            type W = ();
            type R = u64;
            #[inline(always)]
            fn width(_w: ()) -> usize { $width }
            #[inline(always)]
            fn from_bits(bits: u64) -> Option<u64> { Some(bits) }
            #[inline(always)]
            fn to_bits(value: u64) -> u64 { value }
         }
     }
}

static_size!(OneBit, 1);
static_size!(TwoBits, 2);
static_size!(ThreeBits, 3);
static_size!(FourBits, 4);
static_size!(FiveBits, 5);
static_size!(SixBits, 6);
static_size!(SevenBits, 7);
static_size!(EightBits, 8);

#[cfg(test)]
mod test {
    extern crate quickcheck;

    use self::quickcheck::quickcheck;
    use super::{get_n_bits_at, Bitmap, bits, DynamicSize};

    #[test]
    fn empty() {
        let bm: Bitmap<_, DynamicSize> = Bitmap::from_storage(10, 10, vec![0; 100]).unwrap();

        for i in 0..10 {
            assert_eq!(bm.get(i), Some(0));
        }

        assert_eq!(bm.get(11), None);
    }

    #[test]
    fn get() {
        let data = [0usize];
        let mut bm: Bitmap<_, DynamicSize> = Bitmap::from_storage(8, 3, data).unwrap();

        println!("{:?}", (0..8).map(|i| bm.get(i)).collect::<Vec<_>>());
        for i in 0..8 {
            assert_eq!(bm.set(i, i as u64), true);
            assert_eq!(bm.get(i), Some(i as u64));
        }

        assert_eq!(bm.get(8), None);
        assert_eq!(bm.get(9), None);
    }

    #[test]
    fn set() {
        let mut bm: Bitmap<_, DynamicSize> = Bitmap::from_storage(10, 3, vec![0; 16]).unwrap();

        for i in 0..8 {
            assert!(bm.set(i, i as u64));
            assert_eq!(bm.get(i), Some(i as u64));
        }
        assert_eq!(bm.get(8), Some(0));
        assert_eq!(bm.get(9), Some(0));

        assert_eq!(bm.get(10), None);
    }

    #[test]
    fn get_n_bits() {
        macro_rules! t {
            ( $( $e:expr, $n:expr, $s:expr, $g:expr; )*  ) => (
                {
                    $(
                        assert_eq!(get_n_bits_at($e, $n, $s), $g);
                     )*
                }
            )
        }

        t! {
            0b00111001, 1, 0, 0b1;
            0b00111001, 8, 0, 0b00111001;
            0b11010101, 2, 0, 0b01;
            0b11010101, 2, 1, 0b10;
            0b11010101, 2, 2, 0b01;
            0b11010101, 2, 3, 0b10;
            0b11010101, 2, 4, 0b01;
            0b11010101, 3, 0, 0b101;
            0b11010101, 3, 1, 0b010;
            0b11010101, 3, 2, 0b101;
        }
    }

    #[test]
    fn iter() {
        let mut bm: Bitmap<_, DynamicSize> = Bitmap::from_storage(10, 3, vec![0; 16]).unwrap();

        bm.set(2, 0b101);
        bm.set(7, 0b110);

        let bs: Vec<u64> = bm.iter().collect();
        assert_eq!(bs, [0, 0, 0b101, 0, 0, 0, 0, 0b110, 0, 0]);
    }

    fn set_then_clear_prop(entries: usize, width: usize) -> bool {
        if width >= bits() || width == 0 { return true }
        let mut bm: Bitmap<_, DynamicSize> = Bitmap::from_storage(entries, width, vec![0; entries *
                                                                  width]).unwrap();
        let all_set = (1 << width) - 1;
        for i in 0..entries {
            assert!(bm.set(i, all_set));
        }

        for val in &bm {
            if val != all_set { return false; }
        }

        for i in 0..entries {
            assert!(bm.set(i, 0));
        }

        for val in &bm {
            if val != 0 { return false; }
        }
        true
    }

    #[test]
    fn set_then_clear_is_identity() {
        quickcheck(set_then_clear_prop as fn(usize, usize) -> bool);
    }

    #[test]
    fn get_n_bits_at_oring_in_is_same() {
        fn f(b: usize, n: u8, s: u8) -> quickcheck::TestResult {
            if s >= bits() as u8 || n > bits() as u8 {
                return quickcheck::TestResult::discard()
            }
            quickcheck::TestResult::from_bool(((get_n_bits_at(b, n, s) << s) | b) == b)
        }
        quickcheck(f as fn(usize, u8, u8) -> quickcheck::TestResult);
    }

    #[test]
    fn get_n_bits_at_one_bit() {
        fn f(b: usize, s: u8) -> quickcheck::TestResult {
            if s >= bits() as u8 {
                return quickcheck::TestResult::discard()
            }
            quickcheck::TestResult::from_bool(get_n_bits_at(b, 1, s) == (b >> s as usize) & 1)
        }
        quickcheck(f as fn(usize, u8) -> quickcheck::TestResult);
    }

    #[test]
    fn first_set_works() {
        fn f(b: usize, l: usize) -> quickcheck::TestResult {
            if b >= l*bits() || l >= 20 {
                return quickcheck::TestResult::discard()
            }
            let mut bm = Bitmap::from_storage(l*bits(), (), vec![0usize; l]).unwrap();
            bm.set(b, 1);
            quickcheck::TestResult::from_bool(bm.first_set() == Some(b))
        }
        quickcheck(f as fn(usize, usize) -> quickcheck::TestResult);
    }
}