<commit_before>use crate::error::IndentexError;
use std::path::{Path, PathBuf};
use std::vec::Vec;

const INDENTEX_GLOB: &str = "*.inden.tex";

pub fn walk_indentex_files<T: AsRef<Path>>(rootdir: T) -> Result<Vec<PathBuf>, IndentexError> {
    use ignore::types::TypesBuilder;
    use ignore::WalkBuilder;

    // Create a type matcher
    let mut tb = TypesBuilder::new();
    tb.add("indentex", INDENTEX_GLOB).unwrap();
    tb.select("indentex");
    let matcher = tb.build().unwrap();

    // Walk the path
    let mut files: Vec<PathBuf> = Vec::new();
    for res in WalkBuilder::new(rootdir.as_ref()).types(matcher).build() {
        let item = res?;
        if item.file_type().unwrap().is_file() {
            files.push(item.path().to_path_buf());
        }
    }

    Ok(files)
}

#[inline]
fn is_indentex_file<T: AsRef<Path>>(filepath: T) -> bool {
    use globset::Glob;

    let glob = Glob::new(INDENTEX_GLOB).unwrap().compile_matcher();
    glob.is_match(filepath.as_ref())
}

/// Rename an `*.inden.tex` file into `*_inden.tex`
pub fn rename_indentex_file<T: AsRef<Path>>(old_path: T) -> Result<PathBuf, IndentexError> {
    if !is_indentex_file(old_path.as_ref()) {
        return Err(IndentexError::InvalidExtension);
    }

    let mut new_pathbuf = old_path.as_ref().to_path_buf();
    // Strip both extensions, first `.tex` and then `.inden`
    // We assume that the file `old_path` is a valid indentex file,
    // i.e. it has the
    new_pathbuf.set_extension("");
    new_pathbuf.set_extension("");
    // Get the full filename (i.e. with all dots etc. if there are any)
    let mut new_name = new_pathbuf.file_name().unwrap().to_os_string();
    new_name.push("_indentex.tex");
    new_pathbuf.pop();
    new_pathbuf.push(new_name);

    Ok(new_pathbuf)
}

/// Read a file line by line, trim the ends of lines and _copy_ them into a vec of strings
pub fn read_and_trim_lines<T: AsRef<Path>>(path: T) -> Result<Vec<String>, IndentexError> {
    use std::fs::File;
    use std::io::{BufRead, BufReader};

    if !is_indentex_file(path.as_ref()) {
        return Err(IndentexError::InvalidExtension);
    }

    let file = File::open(path.as_ref())?;
    let buf = BufReader::new(file);

    buf.lines()
        .map(|r| Ok(r?.trim_end().to_string()))
        .collect()
}

pub fn write_to_file<T, U>(path: T, data: U) -> Result<(), IndentexError>
where
    T: AsRef<Path>,
    U: AsRef<str>,
{
    use std::fs::File;
    use std::io::{BufWriter, Write};

    let file = File::create(path.as_ref())?;
    let mut buf = BufWriter::new(file);
    buf.write_all(data.as_ref().as_bytes())?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::path::{Path, PathBuf};

    #[test]
    fn is_indentex_file() {
        use super::is_indentex_file;

        assert!(!is_indentex_file(Path::new("foo")));
        assert!(!is_indentex_file(Path::new("foo.tex")));
        assert!(!is_indentex_file(Path::new("foo_inden.tex")));

        assert!(is_indentex_file(Path::new("foo.inden.tex")));
        assert!(is_indentex_file(Path::new("foo.bar.inden.tex")));
    }

    #[test]
    fn rename_indentex_file() {
        use super::rename_indentex_file;

        assert_eq!(
            rename_indentex_file(Path::new("./foo.inden.tex")).unwrap(),
            PathBuf::from("./foo_indentex.tex")
        );
        assert_eq!(
            rename_indentex_file(Path::new("./foo.bar.inden.tex")).unwrap(),
            PathBuf::from("./foo.bar_indentex.tex")
        );
        assert_eq!(
            rename_indentex_file(Path::new("./.foo.bar.inden.tex")).unwrap(),
            PathBuf::from("./.foo.bar_indentex.tex")
        );
        assert_eq!(
            rename_indentex_file(Path::new("foo.inden.tex")).unwrap(),
            PathBuf::from("foo_indentex.tex")
        );
        assert!(rename_indentex_file(Path::new("foo.bar.tex")).is_err())
    }
}
<commit_msg>Run rustfmt<commit_after>use crate::error::IndentexError;
use std::path::{Path, PathBuf};
use std::vec::Vec;

const INDENTEX_GLOB: &str = "*.inden.tex";

pub fn walk_indentex_files<T: AsRef<Path>>(rootdir: T) -> Result<Vec<PathBuf>, IndentexError> {
    use ignore::types::TypesBuilder;
    use ignore::WalkBuilder;

    // Create a type matcher
    let mut tb = TypesBuilder::new();
    tb.add("indentex", INDENTEX_GLOB).unwrap();
    tb.select("indentex");
    let matcher = tb.build().unwrap();

    // Walk the path
    let mut files: Vec<PathBuf> = Vec::new();
    for res in WalkBuilder::new(rootdir.as_ref()).types(matcher).build() {
        let item = res?;
        if item.file_type().unwrap().is_file() {
            files.push(item.path().to_path_buf());
        }
    }

    Ok(files)
}

#[inline]
fn is_indentex_file<T: AsRef<Path>>(filepath: T) -> bool {
    use globset::Glob;

    let glob = Glob::new(INDENTEX_GLOB).unwrap().compile_matcher();
    glob.is_match(filepath.as_ref())
}

/// Rename an `*.inden.tex` file into `*_inden.tex`
pub fn rename_indentex_file<T: AsRef<Path>>(old_path: T) -> Result<PathBuf, IndentexError> {
    if !is_indentex_file(old_path.as_ref()) {
        return Err(IndentexError::InvalidExtension);
    }

    let mut new_pathbuf = old_path.as_ref().to_path_buf();
    // Strip both extensions, first `.tex` and then `.inden`
    // We assume that the file `old_path` is a valid indentex file,
    // i.e. it has the
    new_pathbuf.set_extension("");
    new_pathbuf.set_extension("");
    // Get the full filename (i.e. with all dots etc. if there are any)
    let mut new_name = new_pathbuf.file_name().unwrap().to_os_string();
    new_name.push("_indentex.tex");
    new_pathbuf.pop();
    new_pathbuf.push(new_name);

    Ok(new_pathbuf)
}

/// Read a file line by line, trim the ends of lines and _copy_ them into a vec of strings
pub fn read_and_trim_lines<T: AsRef<Path>>(path: T) -> Result<Vec<String>, IndentexError> {
    use std::fs::File;
    use std::io::{BufRead, BufReader};

    if !is_indentex_file(path.as_ref()) {
        return Err(IndentexError::InvalidExtension);
    }

    let file = File::open(path.as_ref())?;
    let buf = BufReader::new(file);

    buf.lines().map(|r| Ok(r?.trim_end().to_string())).collect()
}

pub fn write_to_file<T, U>(path: T, data: U) -> Result<(), IndentexError>
where
    T: AsRef<Path>,
    U: AsRef<str>,
{
    use std::fs::File;
    use std::io::{BufWriter, Write};

    let file = File::create(path.as_ref())?;
    let mut buf = BufWriter::new(file);
    buf.write_all(data.as_ref().as_bytes())?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::path::{Path, PathBuf};

    #[test]
    fn is_indentex_file() {
        use super::is_indentex_file;

        assert!(!is_indentex_file(Path::new("foo")));
        assert!(!is_indentex_file(Path::new("foo.tex")));
        assert!(!is_indentex_file(Path::new("foo_inden.tex")));

        assert!(is_indentex_file(Path::new("foo.inden.tex")));
        assert!(is_indentex_file(Path::new("foo.bar.inden.tex")));
    }

    #[test]
    fn rename_indentex_file() {
        use super::rename_indentex_file;

        assert_eq!(
            rename_indentex_file(Path::new("./foo.inden.tex")).unwrap(),
            PathBuf::from("./foo_indentex.tex")
        );
        assert_eq!(
            rename_indentex_file(Path::new("./foo.bar.inden.tex")).unwrap(),
            PathBuf::from("./foo.bar_indentex.tex")
        );
        assert_eq!(
            rename_indentex_file(Path::new("./.foo.bar.inden.tex")).unwrap(),
            PathBuf::from("./.foo.bar_indentex.tex")
        );
        assert_eq!(
            rename_indentex_file(Path::new("foo.inden.tex")).unwrap(),
            PathBuf::from("foo_indentex.tex")
        );
        assert!(rename_indentex_file(Path::new("foo.bar.tex")).is_err())
    }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};
use std::collections::HashMap;

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
        }
    }
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile,
    indices: HashMap<u32, IndexFile>
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexFile {
    id: u32,
    file: File
}

#[derive(Debug,Clone)]
pub struct IndexEntry {
    index: u32,
    id: u32,
    size: u32,
    offset: u64
}

impl IndexEntry {
    pub fn index(&self) -> u32 {
        self.index
    }

    /// Gets the id of this block.
    pub fn id(&self) -> u32 {
        self.id
    }

    /// Gets the absolute size of the entry data, not counting the 8-10
    /// byte header in the blocks.
    pub fn size(&self) -> u32 {
        self.size
    }

    /// Gets the absolute offset of the very first 520-byte block of this
    /// entry in the main data file.
    pub fn offset(&self) -> u64 {
        self.offset
    }
}

impl IndexFile {
    pub fn last_entry(&self) -> u64 {
         self.file.metadata().unwrap().len() / 6u64
    }

    pub fn entry(&mut self, id: u32) -> Option<IndexEntry> {
        let ref mut file = self.file;
        let mut tmp: [u8; 6] = [0; 6];

        // Seek to the proper position and read into the temp buffer
        // TODO this must be done safer.. what if it doesn't exist?
        file.seek(SeekFrom::Start(id as u64 * 6u64));
        file.read(&mut tmp);

        // Decode the size and offset from the temp buffer
        let size: u32 = ((tmp[0] as u32) << 16) | ((tmp[1] as u32) << 8) | (tmp[2] as u32);
        let offset: u64 = ((tmp[3] as u64) << 16) | ((tmp[4] as u64) << 8) | (tmp[5] as u64);

        Some(IndexEntry {index: self.id, id: id, size: size, offset: offset * 520u64})
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        // Find all valid index files
        let mut indices: HashMap<u32, IndexFile> = HashMap::new();
        let entries = fs::read_dir(&path).unwrap();
        for entry in entries {
            let e = entry.unwrap();
            let fname = e.file_name().into_string().unwrap();

            // Is this an index?
            if fname.starts_with("main_file_cache.idx") {
                // Parse the index id into an integer
                let idx = fname[19..].parse::<u32>().unwrap();

                // Add the index file to our map with indices
                indices.insert(idx, IndexFile {id: idx, file: File::open(e.path()).unwrap()});
            }
        }

        // Create the filesystem object and return it
        let file = File::open(mainfile_path).ok();
        let mainfile = MainFile{file: file};

        Ok(FileSystem {path: path, mainfile: mainfile, indices: indices})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }

    /// Gets an index with a specific id if it exists. The index can only exist if the file exists
    /// on the file system.
    pub fn index(&mut self, index: u32) -> Option<&mut IndexFile> {
        self.indices.get_mut(&index)
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64));
        file.read(&mut data);

        return Some(data);
    }
}
<commit_msg>Add stub for entry reading<commit_after>use std::path::PathBuf;
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};
use std::collections::HashMap;

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
        }
    }
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile,
    indices: HashMap<u32, IndexFile>
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexFile {
    id: u32,
    file: File
}

#[derive(Debug,Clone)]
pub struct IndexEntry {
    index: u32,
    id: u32,
    size: u32,
    offset: u64
}

impl IndexEntry {
    pub fn index(&self) -> u32 {
        self.index
    }

    /// Gets the id of this block.
    pub fn id(&self) -> u32 {
        self.id
    }

    /// Gets the absolute size of the entry data, not counting the 8-10
    /// byte header in the blocks.
    pub fn size(&self) -> u32 {
        self.size
    }

    /// Gets the absolute offset of the very first 520-byte block of this
    /// entry in the main data file.
    pub fn offset(&self) -> u64 {
        self.offset
    }
}

impl IndexFile {
    pub fn last_entry(&self) -> u64 {
         self.file.metadata().unwrap().len() / 6u64
    }

    pub fn entry(&mut self, id: u32) -> Option<IndexEntry> {
        let ref mut file = self.file;
        let mut tmp: [u8; 6] = [0; 6];

        // Seek to the proper position and read into the temp buffer
        // TODO this must be done safer.. what if it doesn't exist?
        file.seek(SeekFrom::Start(id as u64 * 6u64));
        file.read(&mut tmp);

        // Decode the size and offset from the temp buffer
        let size: u32 = ((tmp[0] as u32) << 16) | ((tmp[1] as u32) << 8) | (tmp[2] as u32);
        let offset: u64 = ((tmp[3] as u64) << 16) | ((tmp[4] as u64) << 8) | (tmp[5] as u64);

        Some(IndexEntry {index: self.id, id: id, size: size, offset: offset * 520u64})
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        // Find all valid index files
        let mut indices: HashMap<u32, IndexFile> = HashMap::new();
        let entries = fs::read_dir(&path).unwrap();
        for entry in entries {
            let e = entry.unwrap();
            let fname = e.file_name().into_string().unwrap();

            // Is this an index?
            if fname.starts_with("main_file_cache.idx") {
                // Parse the index id into an integer
                let idx = fname[19..].parse::<u32>().unwrap();

                // Add the index file to our map with indices
                indices.insert(idx, IndexFile {id: idx, file: File::open(e.path()).unwrap()});
            }
        }

        // Create the filesystem object and return it
        let file = File::open(mainfile_path).ok();
        let mainfile = MainFile{file: file};

        Ok(FileSystem {path: path, mainfile: mainfile, indices: indices})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }

    /// Gets an index with a specific id if it exists. The index can only exist if the file exists
    /// on the file system.
    pub fn index(&mut self, index: u32) -> Option<&mut IndexFile> {
        self.indices.get_mut(&index)
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64));
        file.read(&mut data);

        return Some(data);
    }

    pub fn read_entry(&mut self, entry: IndexEntry) -> Option<&[u8]> {
        // Create a vec with what we assume is the size. If not, the vec will
        // perfectly resize itself, so it's only an estimation to help us speed up.
        let mut data: Vec<i8> = Vec::with_capacity(entry.size() as usize);

        let file = self.file().unwrap();

        None
    }

}
<|endoftext|>


<commit_before>use crate::size_hint;
use std::{fmt, iter::FusedIterator};

pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    FlattenOk {
        iter,
        inner_front: None,
        inner_back: None,
    }
}

/// An iterator adaptor that flattens `Result::Ok` values and
/// allows `Result::Err` values through unchanged.
///
/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    iter: I,
    inner_front: Option<T::IntoIter>,
    inner_back: Option<T::IntoIter>,
}

impl<I, T, E> Iterator for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    type Item = Result<T::Item, E>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Handle the front inner iterator.
            if let Some(inner) = &mut self.inner_front {
                if let Some(item) = inner.next() {
                    return Some(Ok(item));
                } else {
                    // This is necessary for the iterator to implement `FusedIterator`
                    // with only the orginal iterator being fused.
                    self.inner_front = None;
                }
            }

            match self.iter.next() {
                Some(Ok(ok)) => self.inner_front = Some(ok.into_iter()),
                Some(Err(e)) => return Some(Err(e)),
                None => {
                    // Handle the back inner iterator.
                    if let Some(inner) = &mut self.inner_back {
                        if let Some(item) = inner.next() {
                            return Some(Ok(item));
                        } else {
                            // This is necessary for the iterator to implement `FusedIterator`
                            // with only the orginal iterator being fused.
                            self.inner_back = None;
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let inner_hint = |inner: &Option<T::IntoIter>| {
            inner
                .as_ref()
                .map(Iterator::size_hint)
                .unwrap_or((0, Some(0)))
        };
        let inner_front = inner_hint(&self.inner_front);
        let inner_back = inner_hint(&self.inner_back);
        // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
        let outer = match self.iter.size_hint() {
            (0, Some(0)) => (0, Some(0)),
            _ => (0, None),
        };

        size_hint::add(size_hint::add(inner_front, inner_back), outer)
    }
}

impl<I, T, E> Clone for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + Clone,
    T: IntoIterator,
    T::IntoIter: Clone,
{
    #[inline]
    clone_fields!(iter, inner_front, inner_back);
}

impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + fmt::Debug,
    T: IntoIterator,
    T::IntoIter: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FlattenOk")
            .field("iter", &self.iter)
            .field("inner_front", &self.inner_front)
            .field("inner_back", &self.inner_back)
            .finish()
    }
}

/// Only the iterator being flattened needs to implement [`FusedIterator`].
impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
where
    I: FusedIterator<Item = Result<T, E>>,
    T: IntoIterator,
{
}
<commit_msg>add impl of DoubleEndedIterator to FlattenOk<commit_after>use crate::size_hint;
use std::{
    fmt,
    iter::{DoubleEndedIterator, FusedIterator},
};

pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    FlattenOk {
        iter,
        inner_front: None,
        inner_back: None,
    }
}

/// An iterator adaptor that flattens `Result::Ok` values and
/// allows `Result::Err` values through unchanged.
///
/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    iter: I,
    inner_front: Option<T::IntoIter>,
    inner_back: Option<T::IntoIter>,
}

impl<I, T, E> Iterator for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    type Item = Result<T::Item, E>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Handle the front inner iterator.
            if let Some(inner) = &mut self.inner_front {
                if let Some(item) = inner.next() {
                    return Some(Ok(item));
                } else {
                    // This is necessary for the iterator to implement `FusedIterator`
                    // with only the orginal iterator being fused.
                    self.inner_front = None;
                }
            }

            match self.iter.next() {
                Some(Ok(ok)) => self.inner_front = Some(ok.into_iter()),
                Some(Err(e)) => return Some(Err(e)),
                None => {
                    // Handle the back inner iterator.
                    if let Some(inner) = &mut self.inner_back {
                        if let Some(item) = inner.next() {
                            return Some(Ok(item));
                        } else {
                            // This is necessary for the iterator to implement `FusedIterator`
                            // with only the orginal iterator being fused.
                            self.inner_back = None;
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let inner_hint = |inner: &Option<T::IntoIter>| {
            inner
                .as_ref()
                .map(Iterator::size_hint)
                .unwrap_or((0, Some(0)))
        };
        let inner_front = inner_hint(&self.inner_front);
        let inner_back = inner_hint(&self.inner_back);
        // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
        let outer = match self.iter.size_hint() {
            (0, Some(0)) => (0, Some(0)),
            _ => (0, None),
        };

        size_hint::add(size_hint::add(inner_front, inner_back), outer)
    }
}

impl<I, T, E> DoubleEndedIterator for FlattenOk<I, T, E>
where
    I: DoubleEndedIterator<Item = Result<T, E>>,
    T: IntoIterator,
    T::IntoIter: DoubleEndedIterator,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            // Handle the back inner iterator.
            if let Some(inner) = &mut self.inner_back {
                if let Some(item) = inner.next_back() {
                    return Some(Ok(item));
                } else {
                    // This is necessary for the iterator to implement `FusedIterator`
                    // with only the orginal iterator being fused.
                    self.inner_back = None;
                }
            }

            match self.iter.next_back() {
                Some(Ok(ok)) => self.inner_back = Some(ok.into_iter()),
                Some(Err(e)) => return Some(Err(e)),
                None => {
                    // Handle the front inner iterator.
                    if let Some(inner) = &mut self.inner_front {
                        if let Some(item) = inner.next_back() {
                            return Some(Ok(item));
                        } else {
                            // This is necessary for the iterator to implement `FusedIterator`
                            // with only the orginal iterator being fused.
                            self.inner_front = None;
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }
}

impl<I, T, E> Clone for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + Clone,
    T: IntoIterator,
    T::IntoIter: Clone,
{
    #[inline]
    clone_fields!(iter, inner_front, inner_back);
}

impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + fmt::Debug,
    T: IntoIterator,
    T::IntoIter: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FlattenOk")
            .field("iter", &self.iter)
            .field("inner_front", &self.inner_front)
            .field("inner_back", &self.inner_back)
            .finish()
    }
}

/// Only the iterator being flattened needs to implement [`FusedIterator`].
impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
where
    I: FusedIterator<Item = Result<T, E>>,
    T: IntoIterator,
{
}
<|endoftext|>


<commit_before>use std::{fmt, iter::FusedIterator};

pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    FlattenOk { iter, inner: None }
}

/// An iterator adaptor that flattens `Result::Ok` values and
/// allows `Result::Err` values through unchanged.
///
/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    iter: I,
    inner: Option<T::IntoIter>,
}

impl<I, T, E> Iterator for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    type Item = Result<T::Item, E>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(inner) = &mut self.inner {
                if let Some(item) = inner.next() {
                    return Some(Ok(item));
                } else {
                    // This is necessary for the iterator to implement `FusedIterator`
                    // with only the orginal iterator being fused.
                    self.inner = None;
                }
            }

            match self.iter.next() {
                Some(Ok(ok)) => self.inner = Some(ok.into_iter()),
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
    }
}

impl<I, T, E> Clone for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + Clone,
    T: IntoIterator,
    T::IntoIter: Clone,
{
    #[inline]
    clone_fields!(iter, inner);
}

impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + fmt::Debug,
    T: IntoIterator,
    T::IntoIter: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FlattenOk")
            .field("iter", &self.iter)
            .field("inner", &self.inner)
            .finish()
    }
}

/// Only the iterator being flattened needs to implement [`FusedIterator`].
impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
where
    I: FusedIterator<Item = Result<T, E>>,
    T: IntoIterator,
{
}
<commit_msg>add size_hint impl for FlattenOk<commit_after>use std::{fmt, iter::FusedIterator};

pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    FlattenOk { iter, inner: None }
}

/// An iterator adaptor that flattens `Result::Ok` values and
/// allows `Result::Err` values through unchanged.
///
/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    iter: I,
    inner: Option<T::IntoIter>,
}

impl<I, T, E> Iterator for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>>,
    T: IntoIterator,
{
    type Item = Result<T::Item, E>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(inner) = &mut self.inner {
                if let Some(item) = inner.next() {
                    return Some(Ok(item));
                } else {
                    // This is necessary for the iterator to implement `FusedIterator`
                    // with only the orginal iterator being fused.
                    self.inner = None;
                }
            }

            match self.iter.next() {
                Some(Ok(ok)) => self.inner = Some(ok.into_iter()),
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if let Some(inner) = &self.inner {
            // If we have an inner iterator, then its lower bound is our lower bound,
            // but we still don't know the upper bound.
            let (inner_lower, inner_higher) = inner.size_hint();
            let (_, outer_higher) = self.iter.size_hint();
            if outer_higher == Some(0) {
                // If there is nothing remaining in the outer iterator, we know the upper bound.
                (inner_lower, inner_higher)
            } else {
                // However, if the outer iterator could have more items in it, we don't
                // know the upper bound.
                (inner_lower, None)
            }
        } else if self.iter.size_hint() == (0, Some(0)) {
            // If the outer iterator is empty, we have no items.
            (0, Some(0))
        } else {
            // Otherwise we do not know anything about the number of items.
            (0, None)
        }
    }
}

impl<I, T, E> Clone for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + Clone,
    T: IntoIterator,
    T::IntoIter: Clone,
{
    #[inline]
    clone_fields!(iter, inner);
}

impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
where
    I: Iterator<Item = Result<T, E>> + fmt::Debug,
    T: IntoIterator,
    T::IntoIter: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FlattenOk")
            .field("iter", &self.iter)
            .field("inner", &self.inner)
            .finish()
    }
}

/// Only the iterator being flattened needs to implement [`FusedIterator`].
impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
where
    I: FusedIterator<Item = Result<T, E>>,
    T: IntoIterator,
{
}
<|endoftext|>


<commit_before>//! Various http headers
// This is mostly copy of [hyper](https://github.com/hyperium/hyper/tree/master/src/header)

use std::fmt;
use std::str::FromStr;

use bytes::{Bytes, BytesMut};
use mime::Mime;
use modhttp::header::GetAll;
use modhttp::Error as HttpError;

pub use modhttp::header::*;

use error::ParseError;
use httpmessage::HttpMessage;

mod common;
mod shared;
#[doc(hidden)]
pub use self::common::*;
#[doc(hidden)]
pub use self::shared::*;

#[doc(hidden)]
/// A trait for any object that will represent a header field and value.
pub trait Header
where
    Self: IntoHeaderValue,
{
    /// Returns the name of the header field
    fn name() -> HeaderName;

    /// Parse a header
    fn parse<T: HttpMessage>(msg: &T) -> Result<Self, ParseError>;
}

#[doc(hidden)]
/// A trait for any object that can be Converted to a `HeaderValue`
pub trait IntoHeaderValue: Sized {
    /// The type returned in the event of a conversion error.
    type Error: Into<HttpError>;

    /// Cast from PyObject to a concrete Python object type.
    fn try_into(self) -> Result<HeaderValue, Self::Error>;
}

impl IntoHeaderValue for HeaderValue {
    type Error = InvalidHeaderValue;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        Ok(self)
    }
}

impl<'a> IntoHeaderValue for &'a str {
    type Error = InvalidHeaderValue;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        self.parse()
    }
}

impl<'a> IntoHeaderValue for &'a [u8] {
    type Error = InvalidHeaderValue;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_bytes(self)
    }
}

impl IntoHeaderValue for Bytes {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(self)
    }
}

impl IntoHeaderValue for Vec<u8> {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(Bytes::from(self))
    }
}

impl IntoHeaderValue for String {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(Bytes::from(self))
    }
}

impl IntoHeaderValue for Mime {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(Bytes::from(format!("{}", self)))
    }
}

/// Represents supported types of content encodings
#[derive(Copy, Clone, PartialEq, Debug)]
pub enum ContentEncoding {
    /// Automatically select encoding based on encoding negotiation
    Auto,
    /// A format using the Brotli algorithm
    #[cfg(feature = "brotli")]
    Br,
    /// A format using the zlib structure with deflate algorithm
    #[cfg(feature = "flate2")]
    Deflate,
    /// Gzip algorithm
    #[cfg(feature = "flate2")]
    Gzip,
    /// Indicates the identity function (i.e. no compression, nor modification)
    Identity,
}

impl ContentEncoding {
    #[inline]
    pub fn is_compression(&self) -> bool {
        match *self {
            ContentEncoding::Identity | ContentEncoding::Auto => false,
            _ => true,
        }
    }

    #[inline]
    pub fn as_str(&self) -> &'static str {
        match *self {
            #[cfg(feature = "brotli")]
            ContentEncoding::Br => "br",
            #[cfg(feature = "flate2")]
            ContentEncoding::Gzip => "gzip",
            #[cfg(feature = "flate2")]
            ContentEncoding::Deflate => "deflate",
            ContentEncoding::Identity | ContentEncoding::Auto => "identity",
        }
    }

    #[inline]
    /// default quality value
    pub fn quality(&self) -> f64 {
        match *self {
            #[cfg(feature = "brotli")]
            ContentEncoding::Br => 1.1,
            #[cfg(feature = "flate2")]
            ContentEncoding::Gzip => 1.0,
            #[cfg(feature = "flate2")]
            ContentEncoding::Deflate => 0.9,
            ContentEncoding::Identity | ContentEncoding::Auto => 0.1,
        }
    }
}

// TODO: remove memory allocation
impl<'a> From<&'a str> for ContentEncoding {
    fn from(s: &'a str) -> ContentEncoding {
        match AsRef::<str>::as_ref(&s.trim().to_lowercase()) {
            #[cfg(feature = "brotli")]
            "br" => ContentEncoding::Br,
            #[cfg(feature = "flate2")]
            "gzip" => ContentEncoding::Gzip,
            #[cfg(feature = "flate2")]
            "deflate" => ContentEncoding::Deflate,
            _ => ContentEncoding::Identity,
        }
    }
}

#[doc(hidden)]
pub(crate) struct Writer {
    buf: BytesMut,
}

impl Writer {
    fn new() -> Writer {
        Writer {
            buf: BytesMut::new(),
        }
    }
    fn take(&mut self) -> Bytes {
        self.buf.take().freeze()
    }
}

impl fmt::Write for Writer {
    #[inline]
    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.buf.extend_from_slice(s.as_bytes());
        Ok(())
    }

    #[inline]
    fn write_fmt(&mut self, args: fmt::Arguments) -> fmt::Result {
        fmt::write(self, args)
    }
}

#[inline]
#[doc(hidden)]
/// Reads a comma-delimited raw header into a Vec.
pub fn from_comma_delimited<T: FromStr>(
    all: GetAll<HeaderValue>,
) -> Result<Vec<T>, ParseError> {
    let mut result = Vec::new();
    for h in all {
        let s = h.to_str().map_err(|_| ParseError::Header)?;
        result.extend(
            s.split(',')
                .filter_map(|x| match x.trim() {
                    "" => None,
                    y => Some(y),
                })
                .filter_map(|x| x.trim().parse().ok()),
        )
    }
    Ok(result)
}

#[inline]
#[doc(hidden)]
/// Reads a single string when parsing a header.
pub fn from_one_raw_str<T: FromStr>(val: Option<&HeaderValue>) -> Result<T, ParseError> {
    if let Some(line) = val {
        let line = line.to_str().map_err(|_| ParseError::Header)?;
        if !line.is_empty() {
            return T::from_str(line).or(Err(ParseError::Header));
        }
    }
    Err(ParseError::Header)
}

#[inline]
#[doc(hidden)]
/// Format an array into a comma-delimited string.
pub fn fmt_comma_delimited<T>(f: &mut fmt::Formatter, parts: &[T]) -> fmt::Result
where
    T: fmt::Display,
{
    let mut iter = parts.iter();
    if let Some(part) = iter.next() {
        fmt::Display::fmt(part, f)?;
    }
    for part in iter {
        f.write_str(", ")?;
        fmt::Display::fmt(part, f)?;
    }
    Ok(())
}
<commit_msg>fix comment<commit_after>//! Various http headers
// This is mostly copy of [hyper](https://github.com/hyperium/hyper/tree/master/src/header)

use std::fmt;
use std::str::FromStr;

use bytes::{Bytes, BytesMut};
use mime::Mime;
use modhttp::header::GetAll;
use modhttp::Error as HttpError;

pub use modhttp::header::*;

use error::ParseError;
use httpmessage::HttpMessage;

mod common;
mod shared;
#[doc(hidden)]
pub use self::common::*;
#[doc(hidden)]
pub use self::shared::*;

#[doc(hidden)]
/// A trait for any object that will represent a header field and value.
pub trait Header
where
    Self: IntoHeaderValue,
{
    /// Returns the name of the header field
    fn name() -> HeaderName;

    /// Parse a header
    fn parse<T: HttpMessage>(msg: &T) -> Result<Self, ParseError>;
}

#[doc(hidden)]
/// A trait for any object that can be Converted to a `HeaderValue`
pub trait IntoHeaderValue: Sized {
    /// The type returned in the event of a conversion error.
    type Error: Into<HttpError>;

    /// Try to convert value to a Header value.
    fn try_into(self) -> Result<HeaderValue, Self::Error>;
}

impl IntoHeaderValue for HeaderValue {
    type Error = InvalidHeaderValue;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        Ok(self)
    }
}

impl<'a> IntoHeaderValue for &'a str {
    type Error = InvalidHeaderValue;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        self.parse()
    }
}

impl<'a> IntoHeaderValue for &'a [u8] {
    type Error = InvalidHeaderValue;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_bytes(self)
    }
}

impl IntoHeaderValue for Bytes {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(self)
    }
}

impl IntoHeaderValue for Vec<u8> {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(Bytes::from(self))
    }
}

impl IntoHeaderValue for String {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(Bytes::from(self))
    }
}

impl IntoHeaderValue for Mime {
    type Error = InvalidHeaderValueBytes;

    #[inline]
    fn try_into(self) -> Result<HeaderValue, Self::Error> {
        HeaderValue::from_shared(Bytes::from(format!("{}", self)))
    }
}

/// Represents supported types of content encodings
#[derive(Copy, Clone, PartialEq, Debug)]
pub enum ContentEncoding {
    /// Automatically select encoding based on encoding negotiation
    Auto,
    /// A format using the Brotli algorithm
    #[cfg(feature = "brotli")]
    Br,
    /// A format using the zlib structure with deflate algorithm
    #[cfg(feature = "flate2")]
    Deflate,
    /// Gzip algorithm
    #[cfg(feature = "flate2")]
    Gzip,
    /// Indicates the identity function (i.e. no compression, nor modification)
    Identity,
}

impl ContentEncoding {
    #[inline]
    pub fn is_compression(&self) -> bool {
        match *self {
            ContentEncoding::Identity | ContentEncoding::Auto => false,
            _ => true,
        }
    }

    #[inline]
    pub fn as_str(&self) -> &'static str {
        match *self {
            #[cfg(feature = "brotli")]
            ContentEncoding::Br => "br",
            #[cfg(feature = "flate2")]
            ContentEncoding::Gzip => "gzip",
            #[cfg(feature = "flate2")]
            ContentEncoding::Deflate => "deflate",
            ContentEncoding::Identity | ContentEncoding::Auto => "identity",
        }
    }

    #[inline]
    /// default quality value
    pub fn quality(&self) -> f64 {
        match *self {
            #[cfg(feature = "brotli")]
            ContentEncoding::Br => 1.1,
            #[cfg(feature = "flate2")]
            ContentEncoding::Gzip => 1.0,
            #[cfg(feature = "flate2")]
            ContentEncoding::Deflate => 0.9,
            ContentEncoding::Identity | ContentEncoding::Auto => 0.1,
        }
    }
}

// TODO: remove memory allocation
impl<'a> From<&'a str> for ContentEncoding {
    fn from(s: &'a str) -> ContentEncoding {
        match AsRef::<str>::as_ref(&s.trim().to_lowercase()) {
            #[cfg(feature = "brotli")]
            "br" => ContentEncoding::Br,
            #[cfg(feature = "flate2")]
            "gzip" => ContentEncoding::Gzip,
            #[cfg(feature = "flate2")]
            "deflate" => ContentEncoding::Deflate,
            _ => ContentEncoding::Identity,
        }
    }
}

#[doc(hidden)]
pub(crate) struct Writer {
    buf: BytesMut,
}

impl Writer {
    fn new() -> Writer {
        Writer {
            buf: BytesMut::new(),
        }
    }
    fn take(&mut self) -> Bytes {
        self.buf.take().freeze()
    }
}

impl fmt::Write for Writer {
    #[inline]
    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.buf.extend_from_slice(s.as_bytes());
        Ok(())
    }

    #[inline]
    fn write_fmt(&mut self, args: fmt::Arguments) -> fmt::Result {
        fmt::write(self, args)
    }
}

#[inline]
#[doc(hidden)]
/// Reads a comma-delimited raw header into a Vec.
pub fn from_comma_delimited<T: FromStr>(
    all: GetAll<HeaderValue>,
) -> Result<Vec<T>, ParseError> {
    let mut result = Vec::new();
    for h in all {
        let s = h.to_str().map_err(|_| ParseError::Header)?;
        result.extend(
            s.split(',')
                .filter_map(|x| match x.trim() {
                    "" => None,
                    y => Some(y),
                })
                .filter_map(|x| x.trim().parse().ok()),
        )
    }
    Ok(result)
}

#[inline]
#[doc(hidden)]
/// Reads a single string when parsing a header.
pub fn from_one_raw_str<T: FromStr>(val: Option<&HeaderValue>) -> Result<T, ParseError> {
    if let Some(line) = val {
        let line = line.to_str().map_err(|_| ParseError::Header)?;
        if !line.is_empty() {
            return T::from_str(line).or(Err(ParseError::Header));
        }
    }
    Err(ParseError::Header)
}

#[inline]
#[doc(hidden)]
/// Format an array into a comma-delimited string.
pub fn fmt_comma_delimited<T>(f: &mut fmt::Formatter, parts: &[T]) -> fmt::Result
where
    T: fmt::Display,
{
    let mut iter = parts.iter();
    if let Some(part) = iter.next() {
        fmt::Display::fmt(part, f)?;
    }
    for part in iter {
        f.write_str(", ")?;
        fmt::Display::fmt(part, f)?;
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::fmt::{self, Display, Formatter};
use std::cmp::Ordering;

#[derive(Debug)]
pub struct ParsedFile {
    pub path: PathBuf,
    pub method: IncludeMethod,
    pub mappings: Vec<FileMapping>,
    pub included_files: Vec<ParsedFile>,
}

#[derive(Debug)]
pub enum IncludeMethod {
    DTS,
    CPP,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FileMapping {
    pub parent_start: usize,
    pub child_start: usize,
    pub len: usize,
}

impl PartialOrd for FileMapping {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for FileMapping {
    fn cmp(&self, other: &Self) -> Ordering {
        self.parent_start.cmp(&other.parent_start)
    }
}

impl FileMapping {
    pub fn parent_end(&self) -> usize {
        self.parent_start + self.len
    }

    pub fn child_end(&self) -> usize {
        self.parent_start + self.len
    }
}

impl ParsedFile {
    pub fn new(path: PathBuf, method: IncludeMethod) -> Self {
        ParsedFile {
            path: path,
            method: method,
            mappings: vec![],
            included_files: vec![],
        }
    }

    // TODO: this is a bad idea as a file can be included twice
    pub fn assign_mapping(&mut self, path: &Path, mapping: FileMapping) -> Result<(), String> {
        if self.path == path {
            self.mappings.push(mapping);
            self.mappings.sort();
            // TODO: sanity checks?
            Ok(())
        } else {
            for file in &mut self.included_files {
                if file.assign_mapping(path, mapping).is_ok() {
                    return Ok(());
                }
            }
            Err(format!("Did not find file: {:?}", path))
        }
    }

    pub fn add_include(&mut self, other: ParsedFile) {
        self.included_files.push(other);
        if self.included_files.iter().all(|f| f.mappings.len() >= 1) {
            self.included_files.sort_by_key(|f| f.mappings[0]);
        }
        // TODO: sanity checks?
    }

    pub fn bounds_of_tree(&self) -> Result<(usize, usize), String> {
        let first_map = self.mappings.first().ok_or("No mappings")?;

        let mut start = first_map.parent_start;
        let mut end = first_map.parent_end();

        if let Some(last_map) = self.mappings.last() {
            let e = last_map.parent_end();
            if e > end {
                end = e;
            }
        }

        if let Some(file) = self.included_files.first() {
            let (s, e) = file.bounds_of_tree()?;
            if s < start {
                start = s;
            }
            if e > end {
                end = e;
            }
        }

        if let Some(file) = self.included_files.last() {
            let (_, e) = file.bounds_of_tree()?;
            if e > end {
                end = e;
            }
        }

        Ok((start, end))
    }

    pub fn file_from_offset_mut(&mut self, offset: usize) -> Result<&mut ParsedFile, String> {
        if !self.mappings.is_empty() &&
           self.mappings
            .iter()
            .any(|m| offset >= m.parent_start && offset < m.parent_start + m.len) {
            return Ok(self);
        }

        for f in &mut self.included_files {
            if let Ok(f) = f.file_from_offset_mut(offset) {
                return Ok(f);
            }
        }

        Err(format!("Byte offset is not within parsed files: {}", offset))
    }

    pub fn offset_after_location(&mut self, start: usize, offset: isize) {
        for map in &mut self.mappings {
            if map.parent_start > start {
                if offset.is_positive() {
                    map.parent_start += offset as usize;
                } else {
                    map.parent_start -= offset.abs() as usize;
                }
            }
        }

        for file in &mut self.included_files {
            file.offset_after_location(start, offset);
        }
    }

    // Takes start of region to split around, end of region,
    // and length to remove due to DTS include statement
    pub fn split_mappings(&mut self, start: usize, end: usize, offset: usize) {
        {
            let mut remainders: Vec<FileMapping> = Vec::new();

            for map in &mut self.mappings {
                if map.parent_start < start && map.parent_end() > start {
                    // parent_start -- start -- parent_end
                    let remainder = map.parent_end() - start;

                    if map.parent_end() > start + offset {
                        // parent_start -- start -- end -- parent_end
                        remainders.push(FileMapping {
                            parent_start: end,
                            // TODO: this aint right
                            child_start: map.child_start + start - map.parent_start + offset,
                            len: remainder - offset,
                        });
                    } //TODO: check for split across two includes

                    map.len = start - map.parent_start;
                } else if map.parent_start == start {
                    // split is at begining of the mapping
                    // offset the start
                    {
                        let offset = end as isize - start as isize;
                        if offset.is_positive() {
                            map.parent_start += offset as usize;
                        } else {
                            map.parent_start -= offset.abs() as usize;
                        }
                    }
                    // shrink the len by the offset
                    map.len -= offset;
                }
            }

            self.mappings.extend_from_slice(&remainders);
            self.mappings.sort();
        }

        for file in &mut self.included_files {
            file.split_mappings(start, end, offset);
        }
    }

    fn write(&self, f: &mut Formatter, prefix: &str) -> fmt::Result {
        let mut next_prefix = prefix.to_string();
        next_prefix.push_str(" |-");
        writeln!(f, "{} {:?}: {}", prefix, self.path, self.method).and(self.included_files
            .iter()
            .fold(Ok(()), |res, x| res.and(x.write(f, &next_prefix))))
    }
}

impl Display for ParsedFile {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.write(f, "")
    }
}

impl Display for IncludeMethod {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            IncludeMethod::DTS => write!(f, "DTS"),
            IncludeMethod::CPP => write!(f, "CPP"),
        }
    }
}
<commit_msg>impl Display for ParsedFile is now readable<commit_after>use std::path::{Path, PathBuf};
use std::fmt::{self, Display, Formatter};
use std::cmp::Ordering;

#[derive(Debug)]
pub struct ParsedFile {
    pub path: PathBuf,
    pub method: IncludeMethod,
    pub mappings: Vec<FileMapping>,
    pub included_files: Vec<ParsedFile>,
}

#[derive(Debug)]
pub enum IncludeMethod {
    DTS,
    CPP,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FileMapping {
    pub parent_start: usize,
    pub child_start: usize,
    pub len: usize,
}

impl PartialOrd for FileMapping {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for FileMapping {
    fn cmp(&self, other: &Self) -> Ordering {
        self.parent_start.cmp(&other.parent_start)
    }
}

impl FileMapping {
    pub fn parent_end(&self) -> usize {
        self.parent_start + self.len
    }

    pub fn child_end(&self) -> usize {
        self.parent_start + self.len
    }
}

impl ParsedFile {
    pub fn new(path: PathBuf, method: IncludeMethod) -> Self {
        ParsedFile {
            path: path,
            method: method,
            mappings: vec![],
            included_files: vec![],
        }
    }

    // TODO: this is a bad idea as a file can be included twice
    pub fn assign_mapping(&mut self, path: &Path, mapping: FileMapping) -> Result<(), String> {
        if self.path == path {
            self.mappings.push(mapping);
            self.mappings.sort();
            // TODO: sanity checks?
            Ok(())
        } else {
            for file in &mut self.included_files {
                if file.assign_mapping(path, mapping).is_ok() {
                    return Ok(());
                }
            }
            Err(format!("Did not find file: {:?}", path))
        }
    }

    pub fn add_include(&mut self, other: ParsedFile) {
        self.included_files.push(other);
        if self.included_files.iter().all(|f| f.mappings.len() >= 1) {
            self.included_files.sort_by_key(|f| f.mappings[0]);
        }
        // TODO: sanity checks?
    }

    pub fn bounds_of_tree(&self) -> Result<(usize, usize), String> {
        let first_map = self.mappings.first().ok_or("No mappings")?;

        let mut start = first_map.parent_start;
        let mut end = first_map.parent_end();

        if let Some(last_map) = self.mappings.last() {
            let e = last_map.parent_end();
            if e > end {
                end = e;
            }
        }

        if let Some(file) = self.included_files.first() {
            let (s, e) = file.bounds_of_tree()?;
            if s < start {
                start = s;
            }
            if e > end {
                end = e;
            }
        }

        if let Some(file) = self.included_files.last() {
            let (_, e) = file.bounds_of_tree()?;
            if e > end {
                end = e;
            }
        }

        Ok((start, end))
    }

    pub fn file_from_offset_mut(&mut self, offset: usize) -> Result<&mut ParsedFile, String> {
        if !self.mappings.is_empty() &&
           self.mappings
            .iter()
            .any(|m| offset >= m.parent_start && offset < m.parent_start + m.len) {
            return Ok(self);
        }

        for f in &mut self.included_files {
            if let Ok(f) = f.file_from_offset_mut(offset) {
                return Ok(f);
            }
        }

        Err(format!("Byte offset is not within parsed files: {}", offset))
    }

    pub fn offset_after_location(&mut self, start: usize, offset: isize) {
        for map in &mut self.mappings {
            if map.parent_start > start {
                if offset.is_positive() {
                    map.parent_start += offset as usize;
                } else {
                    map.parent_start -= offset.abs() as usize;
                }
            }
        }

        for file in &mut self.included_files {
            file.offset_after_location(start, offset);
        }
    }

    // Takes start of region to split around, end of region,
    // and length to remove due to DTS include statement
    pub fn split_mappings(&mut self, start: usize, end: usize, offset: usize) {
        {
            let mut remainders: Vec<FileMapping> = Vec::new();

            for map in &mut self.mappings {
                if map.parent_start < start && map.parent_end() > start {
                    // parent_start -- start -- parent_end
                    let remainder = map.parent_end() - start;

                    if map.parent_end() > start + offset {
                        // parent_start -- start -- end -- parent_end
                        remainders.push(FileMapping {
                            parent_start: end,
                            // TODO: this aint right
                            child_start: map.child_start + start - map.parent_start + offset,
                            len: remainder - offset,
                        });
                    } //TODO: check for split across two includes

                    map.len = start - map.parent_start;
                } else if map.parent_start == start {
                    // split is at begining of the mapping
                    // offset the start
                    {
                        let offset = end as isize - start as isize;
                        if offset.is_positive() {
                            map.parent_start += offset as usize;
                        } else {
                            map.parent_start -= offset.abs() as usize;
                        }
                    }
                    // shrink the len by the offset
                    map.len -= offset;
                }
            }

            self.mappings.extend_from_slice(&remainders);
            self.mappings.sort();
        }

        for file in &mut self.included_files {
            file.split_mappings(start, end, offset);
        }
    }

    fn write(&self, f: &mut Formatter, prefix: &str) -> fmt::Result {
        let mut next_prefix = prefix.to_string();
        next_prefix.push_str(" |-");
        writeln!(f, "{} {:?}: {}", prefix, self.path, self.method)?;
        for file in &self.included_files {
            file.write(f, &next_prefix)?;
		}

		Ok(())
    }
}

impl Display for ParsedFile {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.write(f, "")
    }
}

impl Display for IncludeMethod {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            IncludeMethod::DTS => write!(f, "DTS"),
            IncludeMethod::CPP => write!(f, "CPP"),
        }
    }
}
<|endoftext|>


<commit_before>use l2::ast::*;
use l2::iterators::{FactIterator, QueryIterator};

use std::cell::Cell;
use std::cmp::max;
use std::collections::HashMap;
use std::fmt;
use std::vec::Vec;

impl fmt::Display for FactInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &FactInstruction::GetStructure(Level::Deep, ref name, ref arity, ref r) =>
                write!(f, "get_structure {}/{}, {}", name, arity, r),
            &FactInstruction::GetStructure(Level::Shallow, ref name, ref arity, ref r) =>
                write!(f, "get_structure {}/{}, A{}", name, arity, r.reg_num()),
            &FactInstruction::GetValue(ref x, ref a) =>
                write!(f, "get_value {}, A{}", x, a),
            &FactInstruction::GetVariable(ref x, ref a) =>
                write!(f, "get_variable {}, A{}", x, a),
            &FactInstruction::UnifyVariable(ref r) =>
                write!(f, "unify_variable {}", r),
            &FactInstruction::UnifyValue(ref r) =>
                write!(f, "unify_value {}", r)
        }
    }
}

impl fmt::Display for QueryInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &QueryInstruction::PutStructure(Level::Deep, ref name, ref arity, ref r) =>
                write!(f, "put_structure {}/{}, A{}", name, arity, r.reg_num()),
            &QueryInstruction::PutStructure(Level::Shallow, ref name, ref arity, ref r) =>
                write!(f, "put_structure {}/{}, {}", name, arity, r),
            &QueryInstruction::PutValue(ref x, ref a) =>
                write!(f, "put_value {}, A{}", x, a),
            &QueryInstruction::PutVariable(ref x, ref a) =>
                write!(f, "put_variable {}, A{}", x, a),
            &QueryInstruction::SetVariable(ref r) =>
                write!(f, "set_variable {}", r),
            &QueryInstruction::SetValue(ref r) =>
                write!(f, "set_value {}", r),
        }
    }
}

impl fmt::Display for ControlInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &ControlInstruction::Allocate(num_cells) =>
                write!(f, "allocate {}", num_cells),
            &ControlInstruction::Call(ref name, ref arity) =>
                write!(f, "call {}/{}", name, arity),
            &ControlInstruction::Deallocate =>
                write!(f, "deallocate"),
            &ControlInstruction::Proceed =>
                write!(f, "proceed")
        }
    }
}

trait CompilationTarget<'a> {
    type Iterator : Iterator<Item=TermRef<'a>>;

    fn iter(&'a Term) -> Self::Iterator;

    fn to_structure(Level, Atom, usize, RegType) -> Self;

    fn argument_to_variable(RegType, usize) -> Self;
    fn argument_to_value(RegType, usize) -> Self;
    fn subterm_to_variable(RegType) -> Self;
    fn subterm_to_value(RegType) -> Self;

    fn clause_arg_to_instr(RegType) -> Self;
}

impl<'a> CompilationTarget<'a> for FactInstruction {
    type Iterator = FactIterator<'a>;

    fn iter(term: &'a Term) -> Self::Iterator {
        term.breadth_first_iter()
    }

    fn to_structure(lvl: Level, atom: Atom, arity: usize, reg: RegType) -> Self {
        FactInstruction::GetStructure(lvl, atom, arity, reg)
    }

    fn argument_to_variable(arg: RegType, val: usize) -> Self {
        FactInstruction::GetVariable(arg, val)
    }

    fn argument_to_value(arg: RegType, val: usize) -> Self {
        FactInstruction::GetValue(arg, val)
    }

    fn subterm_to_variable(val: RegType) -> Self {
        FactInstruction::UnifyVariable(val)
    }

    fn subterm_to_value(val: RegType) -> Self {
        FactInstruction::UnifyValue(val)
    }

    fn clause_arg_to_instr(val: RegType) -> Self {
        FactInstruction::UnifyVariable(val)
    }
}

impl<'a> CompilationTarget<'a> for QueryInstruction {
    type Iterator = QueryIterator<'a>;

    fn iter(term: &'a Term) -> Self::Iterator {
        term.post_order_iter()
    }

    fn to_structure(lvl: Level, atom: Atom, arity: usize, reg: RegType) -> Self {
        QueryInstruction::PutStructure(lvl, atom, arity, reg)
    }

    fn argument_to_variable(arg: RegType, val: usize) -> Self {
        QueryInstruction::PutVariable(arg, val)
    }

    fn argument_to_value(arg: RegType, val: usize) -> Self {
        QueryInstruction::PutValue(arg, val)
    }

    fn subterm_to_variable(val: RegType) -> Self {
        QueryInstruction::SetVariable(val)
    }

    fn subterm_to_value(val: RegType) -> Self {
        QueryInstruction::SetValue(val)
    }

    fn clause_arg_to_instr(val: RegType) -> Self {
        QueryInstruction::SetValue(val)
    }
}

struct TermMarker<'a> {
    bindings: HashMap<&'a Var, VarReg>,
    arg_c:    usize,
    perm_c:   usize,
    temp_c:   usize
}

impl<'a> TermMarker<'a> {
    fn new() -> TermMarker<'a> {
        TermMarker { bindings: HashMap::new(),
                     arg_c:    1,
                     perm_c:   1,
                     temp_c:   1 }
    }

    fn contains_var(&self, var: &'a Var) -> bool {
        self.bindings.contains_key(var)
    }

    fn get(&self, var: &'a Var) -> VarReg {
        *self.bindings.get(var).unwrap()
    }

    fn insert(&mut self, var: &'a Var, r: VarReg) {
        self.bindings.insert(var, r);
    }

    fn mark_non_var(&mut self, lvl: Level, cell: &Cell<RegType>) {
        let reg_type = cell.get();

        if reg_type.reg_num() == 0 {
            match lvl {
                Level::Deep if reg_type.is_perm() => {
                    let perm = self.perm_c;
                    self.perm_c += 1;
                    cell.set(RegType::Perm(perm));
                },
                Level::Deep => {
                    let temp = self.temp_c;
                    self.temp_c += 1;
                    cell.set(RegType::Temp(temp));
                },
                Level::Shallow if reg_type.is_perm() => {
                    let arg = self.arg_c;
                    self.arg_c += 1;
                    cell.set(RegType::Perm(arg));
                },
                Level::Shallow => {
                    let arg = self.arg_c;
                    self.arg_c += 1;
                    cell.set(RegType::Temp(arg));
                }
            };
        }
    }

    fn mark_old_var(&mut self, lvl: Level, var: &'a Var) -> VarReg
    {
        let reg = self.get(var);

        match lvl {
            Level::Deep => VarReg::Norm(reg.norm()),
            Level::Shallow => {
                let reg = VarReg::ArgAndNorm(reg.norm(), self.arg_c);

                self.arg_c += 1;
                self.insert(var, reg);

                reg
            }
        }
    }

    fn mark_new_var(&mut self, lvl: Level, var: &'a Var, reg: RegType) -> VarReg
    {
        let inner_reg = if reg.is_perm() {
            let perm = self.perm_c;
            self.perm_c += 1;
            RegType::Perm(perm)
        } else {
            let temp = self.temp_c;
            self.temp_c += 1;
            RegType::Temp(temp)
        };

        let reg = match lvl {
            Level::Deep => VarReg::Norm(inner_reg),
            Level::Shallow => {
                let reg = VarReg::ArgAndNorm(inner_reg, self.arg_c);
                self.arg_c += 1;
                reg
            }
        };

        self.insert(var, reg);
        reg
    }

    fn advance_at_header(&mut self, term: &'a Term) {
        self.arg_c = 1;
        self.temp_c = max(term.subterms(), self.temp_c) + 1;
    }

    fn advance(&mut self, term: &'a Term) {
        self.arg_c = 1;
        self.temp_c = term.subterms() + 1;
    }
}

#[derive(Copy, Clone)]
enum TermStatus {
    New, Old, Recurrent
}

pub struct CodeGenerator<'a> {
    marker: TermMarker<'a>
}

type VariableFixture<'a>  = (TermStatus, Vec<&'a Cell<VarReg>>);
type VariableFixtures<'a> = HashMap<&'a Var, VariableFixture<'a>>;

impl<'a> CodeGenerator<'a> {
    pub fn new() -> Self {
        CodeGenerator { marker: TermMarker::new() }
    }

    pub fn vars(&self) -> &HashMap<&Var, VarReg> {
        &self.marker.bindings
    }

    fn to_structure<Target>(&mut self,
                            lvl: Level,
                            name: &'a Atom,
                            cell: &'a Cell<RegType>,
                            arity: usize)
                            -> Target
        where Target: CompilationTarget<'a>
    {
        self.marker.mark_non_var(lvl, cell);
        Target::to_structure(lvl, name.clone(), arity, cell.get())
    }

    fn var_term<Target>(&mut self,
                        lvl: Level,
                        cell: &'a Cell<VarReg>,
                        var: &'a Var)
                        -> Target
        where Target: CompilationTarget<'a>
    {
        if !self.marker.contains_var(var) {
            let reg = self.marker.mark_new_var(lvl, var, cell.get().norm());
            cell.set(reg);

            match reg {
                VarReg::ArgAndNorm(arg, norm) =>
                    Target::argument_to_variable(arg, norm),
                VarReg::Norm(norm) =>
                    Target::subterm_to_variable(norm)
            }
        } else {
            let reg = self.marker.mark_old_var(lvl, var);
            cell.set(reg);

            match reg {
                VarReg::ArgAndNorm(arg, norm) =>
                    Target::argument_to_value(arg, norm),
                VarReg::Norm(norm) =>
                    Target::subterm_to_value(norm)
            }
        }
    }

    fn non_var_subterm<Target>(&mut self, cell: &'a Cell<RegType>) -> Target
        where Target: CompilationTarget<'a>
    {
        self.marker.mark_non_var(Level::Deep, cell);
        Target::clause_arg_to_instr(cell.get())
    }

    fn subterm_to_instr<Target>(&mut self, subterm: &'a Term) -> Target
        where Target: CompilationTarget<'a>
    {
        match subterm {
            &Term::Atom(ref cell, _) | &Term::Clause(ref cell, _, _) =>
                self.non_var_subterm(cell),
            &Term::Var(ref cell, ref var) =>
                self.var_term(Level::Deep, cell, var)
        }
    }

    fn compile_target<Target>(&mut self, term: &'a Term) -> Vec<Target>
        where Target: CompilationTarget<'a>
    {
        let iter       = Target::iter(term);
        let mut target = Vec::new();

        for term in iter {
            match term {
                TermRef::Atom(lvl, term, atom) =>
                    target.push(self.to_structure(lvl, atom, term, 0)),
                TermRef::Clause(lvl, term, atom, terms) => {
                    target.push(self.to_structure(lvl, atom, term, terms.len()));

                    for subterm in terms {
                        target.push(self.subterm_to_instr(subterm.as_ref()));
                    }
                },
                TermRef::Var(lvl @ Level::Shallow, ref cell, ref var) =>
                    target.push(self.var_term(lvl, cell, var)),
                _ => {}
            };
        }

        target
    }

    fn mark_vars_in_term<Iter>(iter: Iter, vs: &mut VariableFixtures<'a>)
        where Iter : Iterator<Item=TermRef<'a>>
    {
        for term in iter {
            if let TermRef::Var(_, reg_cell, var) = term {
                let mut status =
                    vs.entry(var)
                      .or_insert((TermStatus::New, Vec::new()));

                status.1.push(reg_cell);

                match status.0 {
                    TermStatus::Old | TermStatus::Recurrent =>
                        status.0 = TermStatus::Recurrent,
                    _ => {}
                };
            }
        }

        for &mut (ref mut term_status, ref mut cb) in vs.values_mut() {
            match *term_status {
                TermStatus::New => *term_status = TermStatus::Old,
                TermStatus::Recurrent => {
                    for cell_reg in cb.drain(0..) {
                        cell_reg.set(VarReg::Norm(RegType::Perm(0)));
                    }
                },
                _ => {}
            }
        }
    }

    fn mark_perm_vars(rule: &'a Rule) -> VariableFixtures {
        let &Rule { head: (ref p0, ref p1), ref clauses } = rule;
        let mut vfs = HashMap::new();

        let iter = p0.breadth_first_iter().chain(p1.breadth_first_iter());

        Self::mark_vars_in_term(iter, &mut vfs);

        for term in clauses {
            Self::mark_vars_in_term(term.breadth_first_iter(), &mut vfs);
        }

        vfs
    }

    fn add_conditional_call(compiled_query: &mut Code, term: &Term) {
        match term {
            &Term::Atom(_, ref atom) => {
                let call = ControlInstruction::Call(atom.clone(), 0);
                compiled_query.push(Line::Control(call));
            },
            &Term::Clause(_, ref atom, ref terms) => {
                let call = ControlInstruction::Call(atom.clone(), terms.len());
                compiled_query.push(Line::Control(call));
            },
            _ => {}
        }
    }

    pub fn compile_rule(&mut self, rule: &'a Rule) -> Code {
        let vfs = Self::mark_perm_vars(&rule);
        let &Rule { head: (ref p0, ref p1), ref clauses } = rule;
        let mut perm_vars = 0;

        for &(term_status, _) in vfs.values() {
            if let TermStatus::Recurrent = term_status {
                perm_vars += 1;
            }
        }

        let mut body = Vec::new();

        body.push(Line::Control(ControlInstruction::Allocate(perm_vars)));

        self.marker.advance(p0);
        body.push(Line::Fact(self.compile_target(p0)));

        self.marker.advance_at_header(p1);
        body.push(Line::Query(self.compile_target(p1)));
        Self::add_conditional_call(&mut body, p1);

        body = clauses.iter()
            .map(|ref term| self.compile_query(term))
            .fold(body, |mut body, ref mut cqs| {
                body.append(cqs);
                body
            });

        body.push(Line::Control(ControlInstruction::Deallocate));

        body
    }

    pub fn compile_fact(&mut self, term: &'a Term) -> Code {
        self.marker.advance(term);

        let mut compiled_fact = vec![Line::Fact(self.compile_target(term))];
        let proceed = Line::Control(ControlInstruction::Proceed);

        compiled_fact.push(proceed);
        compiled_fact
    }

    pub fn compile_query(&mut self, term: &'a Term) -> Code {
        self.marker.advance(term);

        let mut compiled_query = vec![Line::Query(self.compile_target(term))];
        Self::add_conditional_call(&mut compiled_query, term);

        compiled_query
    }
}
<commit_msg>minor tweak<commit_after>use l2::ast::*;
use l2::iterators::{FactIterator, QueryIterator};

use std::cell::Cell;
use std::cmp::max;
use std::collections::HashMap;
use std::fmt;
use std::vec::Vec;

impl fmt::Display for FactInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &FactInstruction::GetStructure(Level::Deep, ref name, ref arity, ref r) =>
                write!(f, "get_structure {}/{}, {}", name, arity, r),
            &FactInstruction::GetStructure(Level::Shallow, ref name, ref arity, ref r) =>
                write!(f, "get_structure {}/{}, A{}", name, arity, r.reg_num()),
            &FactInstruction::GetValue(ref x, ref a) =>
                write!(f, "get_value {}, A{}", x, a),
            &FactInstruction::GetVariable(ref x, ref a) =>
                write!(f, "get_variable {}, A{}", x, a),
            &FactInstruction::UnifyVariable(ref r) =>
                write!(f, "unify_variable {}", r),
            &FactInstruction::UnifyValue(ref r) =>
                write!(f, "unify_value {}", r)
        }
    }
}

impl fmt::Display for QueryInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &QueryInstruction::PutStructure(Level::Deep, ref name, ref arity, ref r) =>
                write!(f, "put_structure {}/{}, A{}", name, arity, r.reg_num()),
            &QueryInstruction::PutStructure(Level::Shallow, ref name, ref arity, ref r) =>
                write!(f, "put_structure {}/{}, {}", name, arity, r),
            &QueryInstruction::PutValue(ref x, ref a) =>
                write!(f, "put_value {}, A{}", x, a),
            &QueryInstruction::PutVariable(ref x, ref a) =>
                write!(f, "put_variable {}, A{}", x, a),
            &QueryInstruction::SetVariable(ref r) =>
                write!(f, "set_variable {}", r),
            &QueryInstruction::SetValue(ref r) =>
                write!(f, "set_value {}", r),
        }
    }
}

impl fmt::Display for ControlInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &ControlInstruction::Allocate(num_cells) =>
                write!(f, "allocate {}", num_cells),
            &ControlInstruction::Call(ref name, ref arity) =>
                write!(f, "call {}/{}", name, arity),
            &ControlInstruction::Deallocate =>
                write!(f, "deallocate"),
            &ControlInstruction::Proceed =>
                write!(f, "proceed")
        }
    }
}

trait CompilationTarget<'a> {
    type Iterator : Iterator<Item=TermRef<'a>>;

    fn iter(&'a Term) -> Self::Iterator;

    fn to_structure(Level, Atom, usize, RegType) -> Self;

    fn argument_to_variable(RegType, usize) -> Self;
    fn argument_to_value(RegType, usize) -> Self;
    fn subterm_to_variable(RegType) -> Self;
    fn subterm_to_value(RegType) -> Self;

    fn clause_arg_to_instr(RegType) -> Self;
}

impl<'a> CompilationTarget<'a> for FactInstruction {
    type Iterator = FactIterator<'a>;

    fn iter(term: &'a Term) -> Self::Iterator {
        term.breadth_first_iter()
    }

    fn to_structure(lvl: Level, atom: Atom, arity: usize, reg: RegType) -> Self {
        FactInstruction::GetStructure(lvl, atom, arity, reg)
    }

    fn argument_to_variable(arg: RegType, val: usize) -> Self {
        FactInstruction::GetVariable(arg, val)
    }

    fn argument_to_value(arg: RegType, val: usize) -> Self {
        FactInstruction::GetValue(arg, val)
    }

    fn subterm_to_variable(val: RegType) -> Self {
        FactInstruction::UnifyVariable(val)
    }

    fn subterm_to_value(val: RegType) -> Self {
        FactInstruction::UnifyValue(val)
    }

    fn clause_arg_to_instr(val: RegType) -> Self {
        FactInstruction::UnifyVariable(val)
    }
}

impl<'a> CompilationTarget<'a> for QueryInstruction {
    type Iterator = QueryIterator<'a>;

    fn iter(term: &'a Term) -> Self::Iterator {
        term.post_order_iter()
    }

    fn to_structure(lvl: Level, atom: Atom, arity: usize, reg: RegType) -> Self {
        QueryInstruction::PutStructure(lvl, atom, arity, reg)
    }

    fn argument_to_variable(arg: RegType, val: usize) -> Self {
        QueryInstruction::PutVariable(arg, val)
    }

    fn argument_to_value(arg: RegType, val: usize) -> Self {
        QueryInstruction::PutValue(arg, val)
    }

    fn subterm_to_variable(val: RegType) -> Self {
        QueryInstruction::SetVariable(val)
    }

    fn subterm_to_value(val: RegType) -> Self {
        QueryInstruction::SetValue(val)
    }

    fn clause_arg_to_instr(val: RegType) -> Self {
        QueryInstruction::SetValue(val)
    }
}

struct TermMarker<'a> {
    bindings: HashMap<&'a Var, VarReg>,
    arg_c:    usize,
    perm_c:   usize,
    temp_c:   usize
}

impl<'a> TermMarker<'a> {
    fn new() -> TermMarker<'a> {
        TermMarker { bindings: HashMap::new(),
                     arg_c:    1,
                     perm_c:   1,
                     temp_c:   1 }
    }

    fn contains_var(&self, var: &'a Var) -> bool {
        self.bindings.contains_key(var)
    }

    fn get(&self, var: &'a Var) -> VarReg {
        *self.bindings.get(var).unwrap()
    }

    fn insert(&mut self, var: &'a Var, r: VarReg) {
        self.bindings.insert(var, r);
    }

    fn mark_non_var(&mut self, lvl: Level, cell: &Cell<RegType>) {
        let reg_type = cell.get();

        if reg_type.reg_num() == 0 {
            match lvl {
                Level::Deep if reg_type.is_perm() => {
                    let perm = self.perm_c;
                    self.perm_c += 1;
                    cell.set(RegType::Perm(perm));
                },
                Level::Deep => {
                    let temp = self.temp_c;
                    self.temp_c += 1;
                    cell.set(RegType::Temp(temp));
                },
                Level::Shallow if reg_type.is_perm() => {
                    let arg = self.arg_c;
                    self.arg_c += 1;
                    cell.set(RegType::Perm(arg));
                },
                Level::Shallow => {
                    let arg = self.arg_c;
                    self.arg_c += 1;
                    cell.set(RegType::Temp(arg));
                }
            };
        }
    }

    fn mark_old_var(&mut self, lvl: Level, var: &'a Var) -> VarReg
    {
        let reg = self.get(var);

        match lvl {
            Level::Deep => VarReg::Norm(reg.norm()),
            Level::Shallow => {
                let reg = VarReg::ArgAndNorm(reg.norm(), self.arg_c);

                self.arg_c += 1;
                self.insert(var, reg);

                reg
            }
        }
    }

    fn mark_new_var(&mut self, lvl: Level, var: &'a Var, reg: RegType) -> VarReg
    {
        let inner_reg = if reg.is_perm() {
            let perm = self.perm_c;
            self.perm_c += 1;
            RegType::Perm(perm)
        } else {
            let temp = self.temp_c;
            self.temp_c += 1;
            RegType::Temp(temp)
        };

        let reg = match lvl {
            Level::Deep => VarReg::Norm(inner_reg),
            Level::Shallow => {
                let reg = VarReg::ArgAndNorm(inner_reg, self.arg_c);
                self.arg_c += 1;
                reg
            }
        };

        self.insert(var, reg);
        reg
    }

    fn advance_at_header(&mut self, term: &'a Term) {
        self.arg_c = 1;
        self.temp_c = max(term.subterms(), self.temp_c) + 1;
    }

    fn advance(&mut self, term: &'a Term) {
        self.arg_c = 1;
        self.temp_c = term.subterms() + 1;
    }
}

#[derive(Copy, Clone)]
enum TermStatus {
    New, Old, Recurrent
}

pub struct CodeGenerator<'a> {
    marker: TermMarker<'a>
}

type VariableFixture<'a>  = (TermStatus, Vec<&'a Cell<VarReg>>);
type VariableFixtures<'a> = HashMap<&'a Var, VariableFixture<'a>>;

impl<'a> CodeGenerator<'a> {
    pub fn new() -> Self {
        CodeGenerator { marker: TermMarker::new() }
    }

    pub fn vars(&self) -> &HashMap<&Var, VarReg> {
        &self.marker.bindings
    }

    fn to_structure<Target>(&mut self,
                            lvl: Level,
                            name: &'a Atom,
                            cell: &'a Cell<RegType>,
                            arity: usize)
                            -> Target
        where Target: CompilationTarget<'a>
    {
        self.marker.mark_non_var(lvl, cell);
        Target::to_structure(lvl, name.clone(), arity, cell.get())
    }

    fn var_term<Target>(&mut self,
                        lvl: Level,
                        cell: &'a Cell<VarReg>,
                        var: &'a Var)
                        -> Target
        where Target: CompilationTarget<'a>
    {
        if !self.marker.contains_var(var) {
            let reg = self.marker.mark_new_var(lvl, var, cell.get().norm());
            cell.set(reg);

            match reg {
                VarReg::ArgAndNorm(arg, norm) =>
                    Target::argument_to_variable(arg, norm),
                VarReg::Norm(norm) =>
                    Target::subterm_to_variable(norm)
            }
        } else {
            let reg = self.marker.mark_old_var(lvl, var);
            cell.set(reg);

            match reg {
                VarReg::ArgAndNorm(arg, norm) =>
                    Target::argument_to_value(arg, norm),
                VarReg::Norm(norm) =>
                    Target::subterm_to_value(norm)
            }
        }
    }

    fn non_var_subterm<Target>(&mut self, cell: &'a Cell<RegType>) -> Target
        where Target: CompilationTarget<'a>
    {
        self.marker.mark_non_var(Level::Deep, cell);
        Target::clause_arg_to_instr(cell.get())
    }

    fn subterm_to_instr<Target>(&mut self, subterm: &'a Term) -> Target
        where Target: CompilationTarget<'a>
    {
        match subterm {
            &Term::Atom(ref cell, _) | &Term::Clause(ref cell, _, _) =>
                self.non_var_subterm(cell),
            &Term::Var(ref cell, ref var) =>
                self.var_term(Level::Deep, cell, var)
        }
    }

    fn compile_target<Target>(&mut self, term: &'a Term) -> Vec<Target>
        where Target: CompilationTarget<'a>
    {
        let iter       = Target::iter(term);
        let mut target = Vec::new();

        for term in iter {
            match term {
                TermRef::Atom(lvl, term, atom) =>
                    target.push(self.to_structure(lvl, atom, term, 0)),
                TermRef::Clause(lvl, term, atom, terms) => {
                    target.push(self.to_structure(lvl, atom, term, terms.len()));

                    for subterm in terms {
                        target.push(self.subterm_to_instr(subterm.as_ref()));
                    }
                },
                TermRef::Var(lvl @ Level::Shallow, ref cell, ref var) =>
                    target.push(self.var_term(lvl, cell, var)),
                _ => {}
            };
        }

        target
    }

    fn mark_vars_in_term<Iter>(iter: Iter, vs: &mut VariableFixtures<'a>)
        where Iter : Iterator<Item=TermRef<'a>>
    {
        for term in iter {
            if let TermRef::Var(_, reg_cell, var) = term {
                let mut status = vs.entry(var)
                                   .or_insert((TermStatus::New, Vec::new()));

                status.1.push(reg_cell);

                match status.0 {
                    TermStatus::Old => status.0 = TermStatus::Recurrent,
                    _ => {}
                };
            }
        }

        for &mut (ref mut term_status, ref mut cb) in vs.values_mut() {
            match *term_status {
                TermStatus::New => *term_status = TermStatus::Old,
                TermStatus::Recurrent => {
                    for cell_reg in cb.drain(0..) {
                        cell_reg.set(VarReg::Norm(RegType::Perm(0)));
                    }
                },
                _ => {}
            }
        }
    }

    fn mark_perm_vars(rule: &'a Rule) -> VariableFixtures {
        let &Rule { head: (ref p0, ref p1), ref clauses } = rule;
        let mut vfs = HashMap::new();

        let iter = p0.breadth_first_iter().chain(p1.breadth_first_iter());

        Self::mark_vars_in_term(iter, &mut vfs);

        for term in clauses {
            Self::mark_vars_in_term(term.breadth_first_iter(), &mut vfs);
        }

        vfs
    }

    fn add_conditional_call(compiled_query: &mut Code, term: &Term) {
        match term {
            &Term::Atom(_, ref atom) => {
                let call = ControlInstruction::Call(atom.clone(), 0);
                compiled_query.push(Line::Control(call));
            },
            &Term::Clause(_, ref atom, ref terms) => {
                let call = ControlInstruction::Call(atom.clone(), terms.len());
                compiled_query.push(Line::Control(call));
            },
            _ => {}
        }
    }

    pub fn compile_rule(&mut self, rule: &'a Rule) -> Code {
        let vfs = Self::mark_perm_vars(&rule);
        let &Rule { head: (ref p0, ref p1), ref clauses } = rule;
        let mut perm_vars = 0;

        for &(term_status, _) in vfs.values() {
            if let TermStatus::Recurrent = term_status {
                perm_vars += 1;
            }
        }

        let mut body = Vec::new();

        body.push(Line::Control(ControlInstruction::Allocate(perm_vars)));

        self.marker.advance(p0);
        body.push(Line::Fact(self.compile_target(p0)));

        self.marker.advance_at_header(p1);
        body.push(Line::Query(self.compile_target(p1)));
        Self::add_conditional_call(&mut body, p1);

        body = clauses.iter()
            .map(|ref term| self.compile_query(term))
            .fold(body, |mut body, ref mut cqs| {
                body.append(cqs);
                body
            });

        body.push(Line::Control(ControlInstruction::Deallocate));

        body
    }

    pub fn compile_fact(&mut self, term: &'a Term) -> Code {
        self.marker.advance(term);

        let mut compiled_fact = vec![Line::Fact(self.compile_target(term))];
        let proceed = Line::Control(ControlInstruction::Proceed);

        compiled_fact.push(proceed);
        compiled_fact
    }

    pub fn compile_query(&mut self, term: &'a Term) -> Code {
        self.marker.advance(term);

        let mut compiled_query = vec![Line::Query(self.compile_target(term))];
        Self::add_conditional_call(&mut compiled_query, term);

        compiled_query
    }
}
<|endoftext|>