libreda-logic 0.0.3

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Compact data-structure for multiple linked lists with the same type of elements.

#![allow(unused)] // TODO remove once stabilized, or move this module to a separate crate

use crate::slab_alloc::{SlabAlloc, SlabIndex, ZeroBitCounter};

type Idx = SlabIndex<usize, ZeroBitCounter>;

/// Pointer to the first element of a linked list.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum ListIndex {
    /// An empty list.
    Empty,
    /// Pointer to a list element.
    Element(Idx),
}

impl Default for ListIndex {
    fn default() -> Self {
        Self::Empty
    }
}

impl ListIndex {
    /// Get the index of the first element in the list.
    pub fn first(&self) -> Option<ElementIndex> {
        match self {
            ListIndex::Empty => None,
            ListIndex::Element(idx) => Some(ElementIndex { idx: *idx }),
        }
    }
}

impl From<Idx> for ListIndex {
    fn from(idx: Idx) -> Self {
        ListIndex::Element(idx)
    }
}

impl From<ListIndex> for Option<Idx> {
    fn from(idx: ListIndex) -> Self {
        match idx {
            ListIndex::Empty => None,
            ListIndex::Element(i) => Some(i),
        }
    }
}

/// Pointer to an element in a linked list.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct ElementIndex {
    idx: Idx,
}

impl From<Idx> for ElementIndex {
    fn from(idx: Idx) -> Self {
        Self { idx }
    }
}

impl From<ElementIndex> for Idx {
    fn from(idx: ElementIndex) -> Self {
        idx.idx
    }
}

/// Container for many double-linked lists.
#[derive(Clone)]
pub struct LinkedLists<T> {
    storage: SlabAlloc<LinkedListEntry<T>, usize, ZeroBitCounter>,
}

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

impl<T> LinkedLists<T> {
    /// Create an empty container.
    pub fn new() -> Self {
        Self {
            storage: SlabAlloc::new(),
        }
    }

    /// Create a new empty linked list.
    /// Returns the index to the list.
    /// The list effectively is created only once an element is inserted.
    pub fn empty_list(&self) -> ListIndex {
        ListIndex::Empty
    }

    /// Prepend an element to a list.
    /// Modifies the `list` and returns an index to the added element.
    pub fn prepend(&mut self, list: &mut ListIndex, value: T) -> ElementIndex {
        match list {
            ListIndex::Empty => {
                let element_index = self.one_element(value);
                *list = ListIndex::Element(element_index.into());
                element_index
            }
            ListIndex::Element(idx) => {
                let element_index = ElementIndex { idx: *idx };
                *idx = self.insert_before(element_index, value).into();
                ElementIndex { idx: *idx }
            }
        }
    }

    pub fn one_element(&mut self, value: T) -> ElementIndex {
        let entry = LinkedListEntry {
            value,
            prev_entry: None,
            next_entry: None,
        };
        self.storage.insert(entry).into()
    }

    fn link(&mut self, i1: Option<Idx>, i2: Option<Idx>) {
        if let Some(i1) = i1 {
            if let Some(e1) = self.storage.get_mut(i1) {
                e1.next_entry = i2
            }
        }
        if let Some(i2) = i2 {
            if let Some(e2) = self.storage.get_mut(i2) {
                e2.prev_entry = i1
            }
        }
    }

    /// Insert an element after the given element.
    pub fn insert_after(&mut self, idx: ElementIndex, value: T) -> ElementIndex {
        let entry = LinkedListEntry {
            value,
            prev_entry: None,
            next_entry: None,
        };

        let prev_idx = idx.into();

        let next_idx = self.storage.get(prev_idx).and_then(|prev| prev.next_entry);

        let new_idx = self.storage.insert(entry);

        self.link(Some(prev_idx), Some(new_idx));
        self.link(Some(new_idx), next_idx);

        new_idx.into()
    }

    /// Insert an element before the given element.
    pub fn insert_before(&mut self, idx: ElementIndex, value: T) -> ElementIndex {
        let entry = LinkedListEntry {
            value,
            prev_entry: None,
            next_entry: None,
        };

        let next_idx = idx.into();

        let prev_idx = self.storage.get(next_idx).and_then(|next| next.prev_entry);

        let new_idx = self.storage.insert(entry);

        self.link(prev_idx, Some(new_idx));
        self.link(Some(new_idx), Some(next_idx));

        new_idx.into()
    }

    /// Get a reference to the element (if it exists).
    pub fn get(&self, idx: ElementIndex) -> Option<&T> {
        self.storage.get(idx.into()).map(|entry| &entry.value)
    }

    /// Get a mutable reference to the element (if it exists).
    pub fn get_mut(&mut self, idx: ElementIndex) -> Option<&mut T> {
        self.storage
            .get_mut(idx.into())
            .map(|entry| &mut entry.value)
    }

    /// Get the index of the next element in the list.
    pub fn next(&self, idx: ElementIndex) -> Option<ElementIndex> {
        self.storage
            .get(idx.into())
            .and_then(|element| element.next_entry)
            .map(|i| i.into())
    }

    /// Get the index of the previous element in the list.
    pub fn prev(&self, idx: ElementIndex) -> Option<ElementIndex> {
        self.storage
            .get(idx.into())
            .and_then(|element| element.prev_entry)
            .map(|i| i.into())
    }

    /// Remove an element from a list.
    /// Panics if the element does not exist.
    fn remove_element(&mut self, idx: ElementIndex) -> T {
        let idx = idx.into();

        let LinkedListEntry {
            value,
            prev_entry,
            next_entry,
        } = self.storage.remove(idx).expect("index not present");

        // Link the new neighbours.
        {
            if let Some(p) = prev_entry {
                if let Some(prev_entry) = self.storage.get_mut(p) {
                    prev_entry.next_entry = next_entry;
                }
            }
            if let Some(n) = next_entry {
                if let Some(next_entry) = self.storage.get_mut(n) {
                    next_entry.prev_entry = prev_entry;
                }
            }
        }

        value
    }

    /// Remove an element from a list.
    /// The given element must be contained in the given list. Otherwise the function does not work correctly.
    pub fn remove(&mut self, list: &mut ListIndex, idx: ElementIndex) -> T {
        if list == &ListIndex::Element(idx.into()) {
            // We are about to remove the head of the list.
            let next_idx = self.next(idx);
            match next_idx {
                Some(next_idx) => *list = ListIndex::Element(next_idx.into()),
                None => *list = ListIndex::Empty,
            }
        }

        self.remove_element(idx)
    }

    /// Get an iterator over the indices of the elements in the given list.
    fn iter_idx(&self, list: &ListIndex) -> LinkedListIdxIter<T> {
        LinkedListIdxIter {
            container: self,
            current_index: list.first(),
            iterate_forward: true,
        }
    }

    /// Get an iterator over the elements of the given list.
    fn iter(&self, list: &ListIndex) -> LinkedListIter<T> {
        self.iter_idx(list).into()
    }

    /// Iterate through the indices of the elements starting at `idx`.
    fn iter_idx_from(&self, idx: ElementIndex) -> LinkedListIdxIter<T> {
        LinkedListIdxIter {
            container: self,
            current_index: Some(idx),
            iterate_forward: true,
        }
    }

    /// Iterate towards the end of the list starting with the element at `idx`.
    fn iter_from(&self, idx: ElementIndex) -> LinkedListIter<T> {
        self.iter_idx_from(idx).into()
    }
}

/// Element of a linked list.
#[derive(Clone)]
struct LinkedListEntry<T> {
    value: T,
    /// Pointer to the previous entry.
    prev_entry: Option<Idx>,
    /// Pointer to the next entry.
    next_entry: Option<Idx>,
}

/// Iterator over a linked list.
pub struct LinkedListIter<'a, T> {
    idx_iter: LinkedListIdxIter<'a, T>,
}

impl<'a, T> From<LinkedListIdxIter<'a, T>> for LinkedListIter<'a, T> {
    fn from(it: LinkedListIdxIter<'a, T>) -> Self {
        Self { idx_iter: it }
    }
}

impl<'a, T> Iterator for LinkedListIter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.idx_iter
            .next()
            .and_then(|idx| self.idx_iter.container.get(idx))
    }
}

/// Iterator over a indices of the elements in a linked list.
pub struct LinkedListIdxIter<'a, T> {
    container: &'a LinkedLists<T>,
    current_index: Option<ElementIndex>,
    iterate_forward: bool,
}

impl<'a, T> Iterator for LinkedListIdxIter<'a, T> {
    type Item = ElementIndex;

    fn next(&mut self) -> Option<Self::Item> {
        let output = self.current_index;

        // Find the index of the next element in the list.
        let next = self
            .current_index
            .and_then(|idx| match self.iterate_forward {
                false => self.container.prev(idx),
                true => self.container.next(idx),
            });
        self.current_index = next;

        output
    }
}

#[test]
fn test_create_linked_list() {
    let mut lists = LinkedLists::new();

    let e1 = lists.one_element(1);
    let e2 = lists.insert_after(e1, 2);

    assert_eq!(lists.get(e1), Some(&1));
    assert_eq!(lists.get(e2), Some(&2));

    assert_eq!(lists.prev(e1), None);
    assert_eq!(lists.prev(e2), Some(e1));
    assert_eq!(lists.next(e1), Some(e2));
    assert_eq!(lists.next(e2), None);
}

#[test]
fn test_iter_empty_list() {
    let lists: LinkedLists<()> = LinkedLists::new();
    let empty = lists.empty_list();
    assert_eq!(lists.iter(&empty).count(), 0);
}

#[test]
fn test_linked_list_iter() {
    let mut lists = LinkedLists::new();

    let mut l = lists.empty_list();
    lists.prepend(&mut l, 3);
    lists.prepend(&mut l, 2);
    lists.prepend(&mut l, 1);

    let v: Vec<_> = lists.iter(&l).copied().collect();

    assert_eq!(v, vec![1, 2, 3]);
}

#[test]
fn test_prepend() {
    let mut lists = LinkedLists::new();

    let mut l = lists.empty_list();

    let e2 = lists.prepend(&mut l, 2);
    assert_eq!(lists.get(e2), Some(&2));
    let e1 = lists.prepend(&mut l, 1);

    assert_eq!(lists.get(e1), Some(&1));
    assert_eq!(lists.get(e2), Some(&2));
}

#[test]
fn test_remove_element() {
    let mut lists = LinkedLists::new();

    let mut l = lists.empty_list();
    let e3 = lists.prepend(&mut l, 3);
    let e2 = lists.prepend(&mut l, 2);
    let e1 = lists.prepend(&mut l, 1);

    assert_eq!(lists.remove(&mut l, e2), 2);
    assert_eq!(lists.iter(&l).count(), 2);

    assert_eq!(lists.remove(&mut l, e1), 1);
    assert_eq!(lists.iter(&l).count(), 1);

    assert_eq!(lists.remove(&mut l, e3), 3);
    assert_eq!(lists.iter(&l).count(), 0);
}