path: root/forest/src/bytes.rs

//! This module defines functions to turn a forest of forest labels
//! into a sequence of bytes.
//!
//! To be more specific, a forest consists of two parts: a vector of
//! node labels and a graph specifying the relations between nodes.
//!
//! # Number format (endianness)
//!
//! Every number mentionned in this format is an unsigned integer of
//! 64 bits, or 8 bytes.  Every such number is specified in the big
//! endian format.
//!
//! # Parts of the sequence of bytes
//!
//! The sequence of bytes has three parts:
//!
//! ## Header
//!
//! The header specifies metadata of this forest.
//!
//! ### Special mark
//!
//! The first three bytes form the string "rep", as a special mark.
//!
//! ### Number of nodes
//!
//! The next 8 bytes specifies the number of nodes of this forest.
//!
//! ## Graph
//!
//! Next comes the underlying graph for this forest.  This part
//! consists of a vector of vectors of numbers.  Each vector of
//! numbers represents the list of children of a node.  So the number
//! of vectors of numbers is equal to the number of nodes.
//!
//! ### Vector of vectors of numbers
//!
//! The vectors are not simply concatenated one after another, as that
//! way one cannot read a random node in a constant amount of time.
//!
//! Instead, we first specify the number of children of each node
//! first, along with the offset for that node of the vector of
//! children.
//!
//! As an example, if a graph has three nodes, represented as the
//! adjacency list: `[[1, 2], [2], []]`, then its representation as a
//! sequence of bytes is as follows:
//! ```text
//! 2, x, 1, y, 0, z, 1, 2, 2,
//!                   ^     ^  ^
//!                   x     y  z
//! ```
//!
//! This has the advantage that we can read the children of the `n`-th
//! node in constant time.
//!
//! ## Vector of labels
//!
//! Each label occupies a fixed number of bytes, so we simply put the
//! labels one after another.  The only thing to note here is the
//! format of the labels.
//!
//! ### Labels
//!
//! Each label has 3 parts:
//!
//! 1. Status: either Packed, Cloned, or Plain.  If the node is
//!    cloned, it has a clone index.  So in total this part occupies 1
//!    byte for the status and 8 bytes for the clone index.
//! 2. Start and end: the range in the input sentence.  We just store
//!    two numbers here.  Hence this part occupies 16 bytes.
//! 3. Grammar label: either a terminal, a non-terminal, or a rule.
//!    Each variant needs a number to speify its index.  So in total
//!    this part occupies 1 byte for the variant discriminant and 8
//!    bytes for the number.
//!
//! To sum up, each label occupies 34 bytes.

use super::{Forest, ForestLabel, ForestLabelType};
use grammar::{GrammarLabel, GrammarLabelType, TNT};
use graph::error::Error as GError;

use std::array::TryFromSliceError;

/// Convert any type that implements the `Forest` trait into a
/// sequence of bytes.
pub fn forest_to_bytes<F: Forest<GrammarLabel>>(forest: &F) -> Vec<u8> {
    // First calculate the total number of bytes.
    let nodes_len = forest.nodes_len();

    let degrees: Vec<_> = forest
        .nodes()
        .map(|node| forest.degree(node).unwrap_or(0))
        .collect();

    let total_degree: usize = degrees.iter().copied().sum();

    let label_offset: usize = 8 // total number of bytes = len
        + 3                     // special mark
        + 8                     // number of nodes
        + 8                     // offset of labels
        + 16 * nodes_len        // degree & offset for each node
        + 8 * total_degree      // children of each node
        ;

    let len: usize = label_offset // see that variable for details.
        + 34 * nodes_len          // labels
        ;

    // Then fill in the bytes.

    let mut bytes: Vec<u8> = Vec::with_capacity(len);

    // First the headers

    bytes.extend(len.to_be_bytes());
    bytes.extend([114, 101, 112]); // rep
    bytes.extend(nodes_len.to_be_bytes());
    bytes.extend(label_offset.to_be_bytes());

    let mut accumulated: usize = 8 + 3 + 8 + 8 + 16 * nodes_len;

    // Then degrees and offsets

    for degree in degrees.iter().copied() {
        bytes.extend(degree.to_be_bytes());
        bytes.extend(accumulated.to_be_bytes());
        // bytes.extend((8 + 3 + 8 + 8 + 16 * nodes_len + 8 * accumulated).to_be_bytes());

        accumulated += 8 * degree;
    }

    // Then the children

    bytes.extend(forest.nodes().flat_map(|node| {
        forest
            .children_of(node)
            .unwrap()
            .flat_map(|child| child.to_be_bytes())
    }));

    // Finally the labels

    'nodes_loop: for node in forest.nodes() {
        let label = match forest.vertex_label(node) {
            Ok(Some(label)) => label,
            _ => continue 'nodes_loop,
        };

        if label.is_plain() {
            bytes.extend(std::iter::repeat(0).take(9));
        } else if label.is_packed() {
            bytes.extend(std::iter::once(1).chain(std::iter::repeat(0).take(8)));
        } else if let Some(index) = label.clone_index() {
            bytes.extend(std::iter::once(2).chain(index.to_be_bytes()));
        }

        let label = label.label();

        bytes.extend(label.start().to_be_bytes());
        bytes.extend(label.end().unwrap_or(0).to_be_bytes());

        let label = label.label();

        match label {
            GrammarLabelType::TNT(TNT::Ter(t)) => {
                bytes.extend(std::iter::once(0).chain(t.to_be_bytes()));
            }
            GrammarLabelType::TNT(TNT::Non(n)) => {
                bytes.extend(std::iter::once(1).chain(n.to_be_bytes()));
            }
            GrammarLabelType::Rule(r) => {
                bytes.extend(std::iter::once(2).chain(r.to_be_bytes()));
            }
        }
    }

    if bytes.len() != len {
        dbg!();
    }

    bytes
}

/// The type of a forest that is represented as a sequence of bytes.
#[derive(Debug, Clone, Default)]
pub struct BytesForest {
    /// The sequence of bytes.
    bytes: Vec<u8>,
    /// The number of nodes.
    nodes_len: usize,
    /// The offset of labels
    offset: usize,
}

/// The errors encountered when parsing a sequence of bytes into a
/// forest.
#[derive(Debug)]
pub enum BytesError {
    /// File size and forest length do not match.
    LengthNotMatched(usize, usize),
    /// The underlying sequence has the wrong length.
    WrongBytesLen(usize, usize),
    /// The file is too small.
    TooSmall(usize, usize),
    /// The special mark at the beginning is corrupted.
    SpecialCorrupted(u8, u8, u8),
    /// An error occurs during conversion of slices.
    WrongSliceConversion(TryFromSliceError),
    /// The data indicates a wrong offset.
    InvalidOffset(usize, usize),
    /// Wrong node index.
    WrongNodeIndex(usize, usize),
}

impl std::fmt::Display for BytesError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::WrongBytesLen(wrong, correct) => write!(
                f,
                "the sequence has length {wrong} but should have length {correct}"
            ),
            Self::LengthNotMatched(file, forest) => write!(
                f,
                "the file size is {file}, but the file indicates that the size should be {forest}"
            ),
            Self::TooSmall(len, threshold) => {
                write!(f, "the file size {len} is less than {threshold} bytes")
            }
            Self::SpecialCorrupted(a, b, c) => {
                write!(f, "the special mark is [{a}, {b}, {c}] which is wrong")
            }
            Self::WrongSliceConversion(e) => {
                write!(f, "error during slice conversion: {e}")
            }
            Self::InvalidOffset(offset, len) => write!(
                f,
                "Data says offset = {offset} but the total length is {len}"
            ),
            Self::WrongNodeIndex(index, len) => write!(
                f,
                "node index {index} is greater than or equal to the length {len} of nodes"
            ),
        }
    }
}

impl From<TryFromSliceError> for BytesError {
    fn from(value: TryFromSliceError) -> Self {
        Self::WrongSliceConversion(value)
    }
}

impl std::error::Error for BytesError {}

impl TryFrom<Vec<u8>> for BytesForest {
    type Error = BytesError;

    fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
        let forest_len = bytes.len();

        if forest_len < 27 {
            return Err(BytesError::TooSmall(forest_len, 27));
        }

        let indicated_len = usize::from_be_bytes((&bytes[..8]).try_into()?);

        if forest_len != indicated_len {
            return Err(BytesError::LengthNotMatched(forest_len, indicated_len));
        }

        let nodes_len = usize::from_be_bytes((&bytes[11..19]).try_into()?);

        let special_marks = &bytes[8..11];

        if special_marks != [114, 101, 112] {
            return Err(BytesError::SpecialCorrupted(
                special_marks[0],
                special_marks[1],
                special_marks[2],
            ));
        }

        let offset = usize::from_be_bytes((&bytes[19..27]).try_into()?);

        let threshold = std::cmp::max(offset + 34 * nodes_len, 27 + 16 * nodes_len);

        if forest_len < threshold {
            return Err(BytesError::TooSmall(forest_len, threshold));
        }

        let mut total_degree = 0usize;

        for node in 0..nodes_len {
            let degree =
                usize::from_be_bytes((&bytes[(27 + 16 * node)..(35 + 16 * node)]).try_into()?);

            total_degree += degree;
        }

        let total_degree = total_degree;

        let correct_len: usize = 27 + 50 * nodes_len + 8 * total_degree;

        if forest_len != correct_len {
            return Err(BytesError::WrongBytesLen(forest_len, correct_len));
        }

        for node in 0..nodes_len {
            let degree =
                usize::from_be_bytes((&bytes[(27 + 16 * node)..(35 + 16 * node)]).try_into()?);

            let node_offset =
                usize::from_be_bytes((&bytes[(35 + 16 * node)..(43 + 16 * node)]).try_into()?);

            if forest_len <= node_offset + 8 * degree {
                return Err(BytesError::InvalidOffset(node_offset, forest_len));
            }

            for n in 0..degree {
                let target = usize::from_be_bytes(
                    (&bytes[(node_offset + 8 * n)..(node_offset + 8 * n + 8)]).try_into()?,
                );

                if nodes_len <= target {
                    return Err(BytesError::WrongNodeIndex(target, nodes_len));
                }
            }
        }

        Ok(Self {
            bytes,
            nodes_len,
            offset,
        })
    }
}

/// The delegation of iterators for children of nodes.
#[derive(Debug, Clone)]
pub struct BytesChildrenIter<'a> {
    vec: &'a [u8],
    offset: usize,
    label_offset: usize,
    degree: usize,
    current: usize,
}

/// The delegation of iterators for children and labels of nodes.
#[derive(Debug, Clone)]
pub struct BytesLabelChildrenIter<'a>(BytesChildrenIter<'a>);

impl<'a> Iterator for BytesLabelChildrenIter<'a> {
    type Item = (ForestLabel<GrammarLabel>, usize);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let iter = &self.0;

        if iter.current >= iter.degree {
            return None;
        }

        let len = iter.vec.len();

        if iter.offset + 8 * iter.degree >= len {
            dbg!();
            return None;
        }

        if iter.label_offset + 34 * iter.degree > len {
            dbg!();
            return None;
        }

        let result_index = usize::from_be_bytes(
            (&iter.vec[(iter.offset + 8 * iter.current)..(iter.offset + 8 * iter.current + 8)])
                .try_into()
                .ok()?,
        );

        let result_label = read_label(
            &iter.vec[(iter.label_offset + iter.current * 34)
                ..(iter.label_offset + iter.current * 34 + 34)],
        )
        .ok()?;

        self.0.current += 1;

        Some((result_label, result_index))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.0.degree, Some(self.0.degree))
    }
}

impl<'a> ExactSizeIterator for BytesLabelChildrenIter<'a> {
    #[inline]
    fn len(&self) -> usize {
        self.0.degree
    }
}

impl<'a> Iterator for BytesChildrenIter<'a> {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.current >= self.degree {
            return None;
        }

        let len = self.vec.len();

        if self.offset + 8 * self.degree >= len {
            dbg!();
            return None;
        }

        let result = usize::from_be_bytes(
            (&self.vec[(self.offset + 8 * self.current)..(self.offset + 8 * self.current + 8)])
                .try_into()
                .ok()?,
        );

        self.current += 1;

        Some(result)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.degree, Some(self.degree))
    }
}

impl<'a> ExactSizeIterator for BytesChildrenIter<'a> {
    #[inline]
    fn len(&self) -> usize {
        self.degree
    }
}

fn read_label(data: &[u8]) -> Result<ForestLabel<GrammarLabel>, GError> {
    let len = data.len();

    if len < 34 {
        return Err(GError::IndexOutOfBounds(34, len));
    }

    let pack_status: ForestLabelType = match data.first().copied().unwrap() {
        0 => ForestLabelType::Plain,
        1 => ForestLabelType::Packed,
        _ => ForestLabelType::Cloned(usize::from_be_bytes(
            data[1..9].try_into().expect("should not happen"),
        )),
    };

    let start = usize::from_be_bytes(data[9..17].try_into().expect("should not happen"));
    let end = usize::from_be_bytes(data[17..25].try_into().expect("should not happen"));

    let grammarlabeltype = match data.get(25).unwrap() {
        0 => GrammarLabelType::TNT(TNT::Ter(usize::from_be_bytes(
            data[26..34].try_into().expect("should not happen"),
        ))),
        1 => GrammarLabelType::TNT(TNT::Non(usize::from_be_bytes(
            data[26..34].try_into().expect("should not happen"),
        ))),
        _ => GrammarLabelType::Rule(usize::from_be_bytes(
            data[26..34].try_into().expect("should not happen"),
        )),
    };

    let label = if end == 0 {
        GrammarLabel::new(grammarlabeltype, start)
    } else {
        GrammarLabel::new_closed(grammarlabeltype, start, end)
    };

    Ok(ForestLabel::new(label, pack_status))
}

impl graph::Graph for BytesForest {
    type Iter<'a> = BytesChildrenIter<'a>
    where
        Self: 'a;

    fn is_empty(&self) -> bool {
        self.bytes.is_empty()
    }

    fn nodes_len(&self) -> usize {
        self.nodes_len
    }

    fn children_of(&self, node_id: usize) -> Result<Self::Iter<'_>, GError> {
        if node_id >= self.nodes_len {
            return Err(GError::IndexOutOfBounds(node_id, self.nodes_len));
        }

        let degree = usize::from_be_bytes(
            (&self.bytes[(27 + 16 * node_id)..(35 + 16 * node_id)])
                .try_into()
                .unwrap(),
        );

        let offset = usize::from_be_bytes(
            (&self.bytes[(35 + 16 * node_id)..(43 + 16 * node_id)])
                .try_into()
                .unwrap(),
        );

        let label_offset = usize::from_be_bytes((&self.bytes[19..27]).try_into().unwrap());

        let current = 0;

        let vec = &self.bytes;

        Ok(BytesChildrenIter {
            vec,
            offset,
            label_offset,
            degree,
            current,
        })
    }

    fn degree(&self, node_id: usize) -> Result<usize, GError> {
        if node_id >= self.nodes_len {
            return Err(GError::IndexOutOfBounds(node_id, self.nodes_len));
        }

        let degree = usize::from_be_bytes(
            (&self.bytes[(27 + 16 * node_id)..(35 + 16 * node_id)])
                .try_into()
                .unwrap(),
        );

        Ok(degree)
    }

    fn is_empty_node(&self, node_id: usize) -> Result<bool, GError> {
        Ok(self.degree(node_id)? == 0)
    }

    fn has_edge(&self, source: usize, target: usize) -> Result<bool, GError> {
        for child in self.children_of(source)? {
            if child == target {
                return Ok(true);
            }
        }

        Ok(false)
    }

    fn replace_by_builder(&mut self, builder: impl graph::Builder<Result = Self>) {
        let result = builder.build();

        self.nodes_len = result.nodes_len;
        self.offset = result.offset;
        self.bytes = result.bytes;
    }
}

impl graph::VLabelGraph<ForestLabel<GrammarLabel>> for BytesForest {
    type TupleIter<'a> = BytesLabelChildrenIter<'a>
    where
        Self: 'a,
        ForestLabel<GrammarLabel>: 'a;

    fn verlabren_of(&self, node_id: usize) -> Result<Self::TupleIter<'_>, GError> {
        if node_id >= self.nodes_len {
            return Err(GError::IndexOutOfBounds(node_id, self.nodes_len));
        }

        let degree = usize::from_be_bytes(
            (&self.bytes[(27 + 16 * node_id)..(35 + 16 * node_id)])
                .try_into()
                .unwrap(),
        );

        let offset = usize::from_be_bytes(
            (&self.bytes[(35 + 16 * node_id)..(43 + 16 * node_id)])
                .try_into()
                .unwrap(),
        );

        let label_offset = usize::from_be_bytes((&self.bytes[19..27]).try_into().unwrap());

        let current = 0;

        let vec = &self.bytes;

        Ok(BytesLabelChildrenIter(BytesChildrenIter {
            vec,
            offset,
            label_offset,
            degree,
            current,
        }))
    }
}

impl graph::LabelGraph<ForestLabel<GrammarLabel>> for BytesForest {
    fn query_label(&self, label: ForestLabel<GrammarLabel>) -> Option<usize> {
        use graph::Graph;
        let nodes_len = self.nodes_len();

        for index in 0..nodes_len {
            let index_label = read_label(
                &self.bytes[(self.offset + index * 34)..(self.offset + index * 34 + 34)],
            )
            .ok()?;

            if index_label == label {
                return Some(index);
            }
        }

        None
    }

    fn vertex_label(&self, node_id: usize) -> Result<Option<ForestLabel<GrammarLabel>>, GError> {
        read_label(&self.bytes[(self.offset + node_id * 34)..(self.offset + node_id * 34 + 34)])
            .map(Some)
    }

    type Iter<'a> = std::iter::Empty<usize>
    where
        Self: 'a;

    type TupleIter<'a>=std::iter::Empty<(ForestLabel<GrammarLabel>, usize)>
    where
        Self: 'a,
        ForestLabel<GrammarLabel>: 'a;

    type LabelIter<'a> = std::iter::Empty<(&'a ForestLabel<GrammarLabel>, std::iter::Empty<usize>)>
    where
        Self: 'a,
        ForestLabel<GrammarLabel>: 'a;

    type EdgeLabelIter<'a> = std::iter::Empty<ForestLabel<GrammarLabel>>
    where
        Self: 'a,
        ForestLabel<GrammarLabel>: 'a;

    fn edge_label(
        &self,
        _source: usize,
        _target: usize,
    ) -> Result<Self::EdgeLabelIter<'_>, GError> {
        unimplemented!("bytes forests do not have edge labels")
    }

    fn find_children_with_label(
        &self,
        _node_id: usize,
        _label: &ForestLabel<GrammarLabel>,
    ) -> Result<<Self as graph::LabelGraph<ForestLabel<GrammarLabel>>>::Iter<'_>, GError> {
        unimplemented!("bytes forests do not have edge labels")
    }

    fn labels_of(&self, _node_id: usize) -> Result<Self::LabelIter<'_>, GError> {
        unimplemented!("bytes forests do not have edge labels")
    }

    fn tuples_of(&self, _node_id: usize) -> Result<Self::TupleIter<'_>, GError> {
        unimplemented!("bytes forests do not have edge labels")
    }

    fn has_edge_label(
        &self,
        _node_id: usize,
        _label: &ForestLabel<GrammarLabel>,
        _target: usize,
    ) -> Result<bool, GError> {
        unimplemented!("bytes forests do not have edge labels")
    }
}