// Copyright © 2021 piet-gpu 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. This file may not be copied, modified, or distributed
// except according to those

//! A simple best-fit allocator.

use std::collections::{BTreeMap, BTreeSet};

/// An allocator that tracks free ranges and returns best fit.
pub struct BestFit {
    // map offset to size of free block
    free_by_ix: BTreeMap<u32, u32>,
    // size and offset
    free_by_size: BTreeSet<(u32, u32)>,
}

impl BestFit {
    pub fn new(size: u32) -> BestFit {
        let mut free_by_ix = BTreeMap::new();
        free_by_ix.insert(0, size);
        let mut free_by_size = BTreeSet::new();
        free_by_size.insert((size, 0));
        BestFit {
            free_by_ix,
            free_by_size,
        }
    }

    pub fn alloc(&mut self, size: u32) -> Option<u32> {
        let block = *self.free_by_size.range((size, 0)..).next()?;
        let ix = block.1;
        self.free_by_ix.remove(&ix);
        self.free_by_size.remove(&block);
        let fragment_size = block.0 - size;
        if fragment_size > 0 {
            let fragment_ix = ix + size;
            self.free_by_ix.insert(fragment_ix, fragment_size);
            self.free_by_size.insert((fragment_size, fragment_ix));
        }
        Some(ix)
    }

    pub fn free(&mut self, ix: u32, size: u32) {
        let next_ix = size + ix;
        if let Some((&prev_ix, &prev_size)) = self.free_by_ix.range(..ix).rev().next() {
            if prev_ix + prev_size == ix {
                self.free_by_size.remove(&(prev_size, prev_ix));
                if let Some(&next_size) = self.free_by_ix.get(&next_ix) {
                    // consolidate with prev and next
                    let new_size = prev_size + size + next_size;
                    *self.free_by_ix.get_mut(&prev_ix).unwrap() = new_size;
                    self.free_by_ix.remove(&next_ix);
                    self.free_by_size.remove(&(next_size, next_ix));
                    self.free_by_size.insert((new_size, prev_ix));
                } else {
                    // consolidate with prev
                    let new_size = prev_size + size;
                    *self.free_by_ix.get_mut(&prev_ix).unwrap() = new_size;
                    self.free_by_size.insert((new_size, prev_ix));
                }
                return;
            }
        }
        if let Some(&next_size) = self.free_by_ix.get(&next_ix) {
            // consolidate with next
            let new_size = size + next_size;
            self.free_by_ix.remove(&next_ix);
            self.free_by_ix.insert(ix, new_size);
            self.free_by_size.remove(&(next_size, next_ix));
            self.free_by_size.insert((new_size, ix));
        } else {
            // new isolated free block
            self.free_by_ix.insert(ix, size);
            self.free_by_size.insert((size, ix));
        }
    }
}