arcis-compiler 0.9.7

use crate::{ArcisField, STATISTICAL_SECURITY_FACTOR};
use ff::PrimeField;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum DataSize {
    /// Size is a power of two.
    PowerOfTwo(u8),
    /// Size is the same as a container. Container size does not have to be a power of two.
    Full, // Needs to be after `PowerOfTwo` for `PartialOrd` derivation.
}

impl DataSize {
    pub fn size(&self, full_size: Size) -> Size {
        match self {
            DataSize::PowerOfTwo(pow) => (1 as Size) << pow,
            DataSize::Full => full_size,
        }
    }
}

type Location = usize;
type Size = u16;

#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct PackLocation {
    /// The index of the container.
    pub index: Location,
    /// The offset in bits of the location of the data in the container.
    pub bit_offset: Size,
}

impl PackLocation {
    fn new(index: Location, bit_offset: Size) -> Self {
        PackLocation { index, bit_offset }
    }
}

#[derive(Debug, Default)]
struct CurrentState {
    /// The size of a container, in bits. Only relevant for `DataSize::PowerOfTwo` items.
    container_size: Size,
    /// Number of containers containing `DataSize::Full` items.
    n_full_containers: Location,
    /// Loads of the containers. MUST be decreasing.
    loads: Vec<Size>,
    /// Last place where something was inserted.
    last_insert: Option<Location>,
}
impl CurrentState {
    pub fn new(container_size: Size) -> CurrentState {
        CurrentState {
            container_size,
            ..Default::default()
        }
    }
    /// Load of the context to the left of index.
    fn left_load(&self, index: Location) -> Size {
        if index == 0 {
            self.container_size
        } else {
            self.loads[index - 1]
        }
    }
    /// Load of the container at index.
    fn location_load(&self, index: Location) -> Size {
        self.loads.get(index).copied().unwrap_or(0)
    }
    /// Inserts an item of size `item_size` at location `index`,
    /// creating a new container if necessary.
    fn insert_at(&mut self, index: Location, item_size: Size) -> Size {
        let res = self.location_load(index);
        let new_load = res + item_size;
        assert!(new_load <= self.left_load(index));
        assert!(new_load <= self.container_size);
        if index >= self.loads.len() {
            assert_eq!(index, self.loads.len());
            self.loads.resize(index + 1, 0);
        }
        self.loads[index] = new_load;
        self.last_insert = Some(index);
        res
    }
    /// Is there enough room to insert an item of size `item_size` at `index`.
    fn can_insert_at(&mut self, index: Location, item_size: Size) -> bool {
        let at_size = self.location_load(index);
        let new_size = at_size + item_size;
        new_size <= self.container_size
    }
    /// Inserts an item of size `data_size`, returning its location.
    /// Note: `data_size` must be given in decreasing order for the function to work.
    pub fn insert(&mut self, data_size: DataSize) -> PackLocation {
        match data_size {
            DataSize::PowerOfTwo(pow) => {
                let item_size = (1 as Size) << pow;
                let last = self.last_insert.unwrap_or(0);
                // Because everything is a power of two, everything has the same load in the
                // interval [self.n_full_containers, last[. So we only need to check these 2 places.
                let insert_loc = [self.n_full_containers, last]
                    .into_iter()
                    .find(|index| self.can_insert_at(*index, item_size))
                    .unwrap_or(last + 1);
                let old_size = self.location_load(insert_loc);
                self.insert_at(insert_loc, item_size);
                PackLocation::new(insert_loc, old_size)
            }
            DataSize::Full => {
                let res = PackLocation::new(self.n_full_containers, 0);
                assert_eq!(self.n_full_containers, self.loads.len());
                self.n_full_containers += 1;
                self.loads.push(self.container_size);
                res
            }
        }
    }
    pub fn state(self) -> Vec<Option<Size>> {
        let len = self.loads.len();
        (0..len)
            .map(|index| {
                if index < self.n_full_containers {
                    None
                } else {
                    Some(self.loads[index])
                }
            })
            .collect()
    }
}
impl DataSize {
    /// Returns in the same order as data, for each item,
    /// its container_id and where in the container it starts.
    /// The algorithm is First-fit-decreasing (FFD).
    /// See <https://en.wikipedia.org/wiki/First-fit-decreasing_bin_packing>
    /// "A special case of divisible item sizes occurs in memory allocation in computer systems,
    /// where the item sizes are all powers of 2. In this case it always finds the optimal packing."
    /// Returns:
    /// * Where each data is put.
    /// * The final state of each container,
    ///   * `None` if filled by `DataSize::Full`
    ///   * `Some(size)` if loaded up to `size`.
    fn pack(data: Vec<DataSize>, container_size: Size) -> (Vec<PackLocation>, Vec<Option<Size>>) {
        let mut to_sort: Vec<_> = data
            .into_iter()
            .enumerate()
            .map(|(i, size)| (size, i))
            .collect();
        // sort according to reverse order of DataSize
        to_sort.sort_by(|a, b| b.0.cmp(&a.0));
        let mut current_state = CurrentState::new(container_size);
        let mut res = vec![PackLocation::new(0, 0); to_sort.len()];
        for (size, index) in to_sort {
            res[index] = current_state.insert(size);
        }
        (res, current_state.state())
    }
    pub fn pack_arcis(data: Vec<DataSize>) -> (Vec<PackLocation>, Vec<Option<Size>>) {
        Self::pack(data, Self::PACKING_SIZE)
    }
    pub const PACKING_SIZE: Size =
        ArcisField::CAPACITY as Size - (STATISTICAL_SECURITY_FACTOR as Size);
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::Rng;
    fn test_pack(sizes: Vec<DataSize>, container_size: Size, expected_len: Option<usize>) {
        let computed_expected = {
            let max_exp = sizes.iter().fold(0, |acc, x| {
                if let DataSize::PowerOfTwo(pow) = x {
                    acc.max(*pow)
                } else {
                    acc
                }
            });
            let mut pow_count = vec![0usize; 1 + max_exp as usize];
            let mut n_full = 0usize;
            for size in &sizes {
                match size {
                    DataSize::PowerOfTwo(pow) => {
                        pow_count[*pow as usize] += 1;
                    }
                    DataSize::Full => {
                        n_full += 1;
                    }
                }
            }
            let mut free_spaces = 0usize;
            let mut n_used_by_pows = 0usize;
            for (exp, count) in pow_count.iter().enumerate().rev() {
                let current_pow = 1 << exp;
                if container_size & current_pow != 0 {
                    free_spaces += n_used_by_pows * (current_pow as usize);
                }
                let count_spaces = count * current_pow as usize;
                if count_spaces <= free_spaces {
                    // We pack in all the free spaces.
                    free_spaces -= count_spaces;
                    continue;
                }
                let remaining_to_pack = (count_spaces - free_spaces) / current_pow as usize;
                let n_packable_per_pack = container_size >> exp;
                let ceil = remaining_to_pack.div_ceil(n_packable_per_pack as usize);
                n_used_by_pows += ceil;
                free_spaces = (ceil * n_packable_per_pack as usize - remaining_to_pack)
                    * current_pow as usize;
            }
            n_full + n_used_by_pows
        };
        if let Some(expected_len) = expected_len {
            assert_eq!(computed_expected, expected_len);
        }
        let (chosen_pack, v) = DataSize::pack(sizes.clone(), container_size);
        assert_eq!(v.len(), computed_expected);
        let mut packs = vec![false; computed_expected * container_size as usize];
        for (size, location) in std::iter::zip(sizes, chosen_pack) {
            let size_begin = location.bit_offset;
            let location = location.index;
            assert!(location < computed_expected);
            let data_size = size.size(container_size);
            let size_end = size_begin + data_size;
            assert!(size_end <= container_size);
            for index in size_begin..size_end {
                let location = location * container_size as usize + index as usize;
                assert!(!packs[location]);
                packs[location] = true;
            }
        }
    }
    #[test]
    fn powers_of_two() {
        let rng = &mut crate::utils::test_rng::get();
        for exp in [1u8, 2, 3, 4, 5, 6, 7, 8] {
            let pow_2 = (1 as Size) << exp;
            for is_size_exact_pow_2 in [true, false, false, false] {
                let offset = if is_size_exact_pow_2 {
                    0
                } else {
                    rng.gen_range(0..pow_2)
                };
                let container_size = pow_2 + offset;
                for n_items in [1, 4, 16, 64, 256, 1024, 4096] {
                    let mut sizes = Vec::new();
                    for _ in 0..n_items {
                        if rng.gen_bool(0.5) {
                            sizes.push(DataSize::PowerOfTwo(rng.gen_range(0..=exp)));
                        } else {
                            sizes.push(DataSize::Full);
                        }
                    }
                    let optimal = if is_size_exact_pow_2 {
                        Some(
                            sizes
                                .iter()
                                .map(|size| size.size(container_size) as usize)
                                .sum::<usize>()
                                .div_ceil(container_size as usize),
                        )
                    } else {
                        None
                    };
                    test_pack(sizes, container_size, optimal);
                }
            }
        }
    }
    #[test]
    fn all_full() {
        let container_size = 7;
        let n_data = 63;
        let sizes = vec![DataSize::Full; n_data];
        test_pack(sizes, container_size, Some(n_data));
    }
}