<commit_before>/*
   This Source Code Form is subject to the terms of the Mozilla Public
   License, v. 2.0. If a copy of the MPL was not distributed with this
   file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

use net::packet::size::*;

#[derive(Debug, PartialEq, Eq)]
struct PacketDescriptor(header::PacketType, IdSizeType);

mod header {

    use std::convert::{From, Into};
    //use net::packet::size::*;

    /// internal representation of PacketHeader. Use the accessors
    #[derive(Debug, PartialEq, Eq)]
    struct PacketHeader(u64);
    /// internal repsentation. Use the methods
    #[derive(Debug, PartialEq, Eq)]
    struct AckBitField(u64);

    #[derive(Debug, PartialEq, Eq)]
    pub enum PacketType {
        HeartBeat, Data
    }
    const MASK_TYPE : u8 = 0b1000_0000;

    impl<'a> From<&'a u8> for PacketType {
        fn from(other: &'a u8) -> PacketType {
            match other & MASK_TYPE {
                MASK_TYPE => PacketType::Data,
                _ => PacketType::HeartBeat
            }
        }
    }
    impl Into<PacketType> for u8 {
        fn into(self) -> PacketType {
            match self & MASK_TYPE {
                MASK_TYPE => PacketType::Data,
                _ => PacketType::HeartBeat
            }
        }
    }

    #[test]
    fn equivalence_packet_type() {
        let vals = [
            0b1100_0000, 0b1100_0001, 0b1101_0000, 0b1111_1111,
            0b0000_0000, 0b0000_1111, 0b1011_1111, 0b1010_1010
        ];

        for v in vals.iter() {
            let t1 = PacketType::from(v);
            let t2 : PacketType = v.into();
            assert_eq!(t1, t2);
        }
    }

    impl PacketHeader {
        pub fn header(&self) -> super::PacketDescriptor {
            let first_byte : u8 = (self.0 >> 56) as u8;
            super::PacketDescriptor(first_byte.into(), first_byte.into())
        }

        pub fn ack_bitfield(&self) -> AckBitField {
            AckBitField((self.0 << 8) & 0xFFFF_FFFF_FFFF_FF00)
        }
    }

    #[test]
    fn header() {
        use net::packet::size::IdSizeType;
        use super::PacketDescriptor;

        let tmp = PacketHeader(0xC000_0000_0000_0000);
        assert_eq!(PacketDescriptor(PacketType::Data, IdSizeType::Short), tmp.header());
        let tmp = PacketHeader(0x8000_0000_0000_0000);
        assert_eq!(PacketDescriptor(PacketType::Data, IdSizeType::Long), tmp.header());
        let tmp = PacketHeader(0x4000_0000_0000_0000);
        assert_eq!(PacketDescriptor(PacketType::HeartBeat, IdSizeType::Short), tmp.header());
        let tmp = PacketHeader(0);
        assert_eq!(PacketDescriptor(PacketType::HeartBeat, IdSizeType::Long), tmp.header());
    }

    #[test]
    fn ack_bitfield() {
        let tmp = PacketHeader(0xdead_beef_b00b_cafe);
        assert_eq!(AckBitField(0xad_beef_b00b_cafe_00), tmp.ack_bitfield());
    }

}

mod body {
    use net::packet::size::*;

    type RawData = [u8];

    #[derive(Debug, PartialEq, Eq)]
    struct DataChunkHeader<Size : IdSize> {
        size: u16,
        number: u16,
        head: Size,
        data: RawData
    }

    #[derive(Debug)]
    enum PacketBody<Size : IdSize> {
        HeartBeat,
        Data(RawData),
        DataChunk(DataChunkHeader<Size>)
    }

}
<commit_msg> reorg net::packet::packet<commit_after>/*
   This Source Code Form is subject to the terms of the Mozilla Public
   License, v. 2.0. If a copy of the MPL was not distributed with this
   file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

use std::convert::{From, Into};
use net::packet::size::*;

mod header {

    use std::convert::{From, Into};
    use net::packet::size::*;

    /// a packet could have data or be an heartbeat
    #[derive(Debug, PartialEq, Eq)]
    pub enum PacketType {
        HeartBeat, Data
    }
    const MASK_TYPE : u8 = 0b1000_0000;

    /// hold the packet attributes
    #[derive(Debug, PartialEq, Eq)]
    struct PacketDescriptor(PacketType, IdSizeType);

    /// internal binary representation of PacketFlags. Use the accessors
    #[derive(Debug, PartialEq, Eq)]
    struct PacketFlags(u64);

    /// internal binary reprentation. Use the methods
    #[derive(Debug, PartialEq, Eq)]
    struct AckBitField(u64);

    #[derive(Debug, PartialEq, Eq)]
    struct PacketHeader<Size : IdSize> {
        flags: PacketFlags,
        acks:  AckBitField,
        packet_id: Size,
        last_ack:  Size
    }

    impl<'a> From<&'a u8> for PacketDescriptor {
        fn from(other: &'a u8) -> PacketDescriptor {
            PacketDescriptor(PacketType::from(other), IdSizeType::from(other))
        }
    }
    impl From<u8> for PacketDescriptor {
        fn from(other: u8) -> PacketDescriptor { PacketDescriptor::from(&other) }
    }

    impl<'a> From<&'a u8> for PacketType {
        fn from(other: &'a u8) -> PacketType {
            match other & MASK_TYPE {
                MASK_TYPE => PacketType::Data,
                _ => PacketType::HeartBeat
            }
        }
    }
    impl Into<PacketType> for u8 {
        fn into(self) -> PacketType {
            match self & MASK_TYPE {
                MASK_TYPE => PacketType::Data,
                _ => PacketType::HeartBeat
            }
        }
    }

    #[test]
    fn equivalence_packet_type() {
        let vals = [
            0b1100_0000, 0b1100_0001, 0b1101_0000, 0b1111_1111,
            0b0000_0000, 0b0000_1111, 0b1011_1111, 0b1010_1010
        ];

        for v in vals.iter() {
            let t1 = PacketType::from(v);
            let t2 : PacketType = v.into();
            assert_eq!(t1, t2);
        }
    }

    impl PacketFlags {
        pub fn header(&self) -> PacketDescriptor {
            let first_byte : u8 = (self.0 >> 56) as u8;
            PacketDescriptor::from(first_byte)
        }

        pub fn ack_bitfield(&self) -> AckBitField {
            AckBitField((self.0 << 8) & 0xFFFF_FFFF_FFFF_FF00)
        }
    }

    #[test]
    fn header() {
        use net::packet::size::IdSizeType;
        use super::PacketDescriptor;

        let tmp = PacketFlags(0xC000_0000_0000_0000);
        assert_eq!(PacketDescriptor(PacketType::Data, IdSizeType::Short), tmp.header());
        let tmp = PacketFlags(0x8000_0000_0000_0000);
        assert_eq!(PacketDescriptor(PacketType::Data, IdSizeType::Long), tmp.header());
        let tmp = PacketFlags(0x4000_0000_0000_0000);
        assert_eq!(PacketDescriptor(PacketType::HeartBeat, IdSizeType::Short), tmp.header());
        let tmp = PacketFlags(0);
        assert_eq!(PacketDescriptor(PacketType::HeartBeat, IdSizeType::Long), tmp.header());
    }

    #[test]
    fn ack_bitfield() {
        let tmp = PacketFlags(0xdead_beef_b00b_cafe);
        assert_eq!(AckBitField(0xad_beef_b00b_cafe_00), tmp.ack_bitfield());
    }

}

mod body {
    use net::packet::size::*;

    type RawData = [u8];

    #[derive(Debug, PartialEq, Eq)]
    struct DataChunkHeader<Size : IdSize> {
        size: u16,
        number: u16,
        head: Size,
        data: RawData
    }

    #[derive(Debug)]
    enum PacketBody<Size : IdSize> {
        HeartBeat,
        Data(RawData),
        DataChunk(DataChunkHeader<Size>)
    }

}
<|endoftext|>


<commit_before>//! FBX binary parser.

use std::io;
use std::io::Read;

pub use self::error::{Result, Error, Warning};
pub use self::event::{Event, FbxHeader, FbxFooter, StartNode};
pub use self::event::{Attributes, Attribute, SpecialAttributeType};
pub use self::event::{PrimitiveAttribute, ArrayAttribute, SpecialAttribute};
pub use self::event::ArrayAttributeReader;
use self::event::{EventBuilder, NodeHeader, StartNodeBuilder};
use self::event::read_fbx_header;
pub use self::reader::{ParserSource, BasicSource, SeekableSource, LimitedSeekReader};

mod error;
mod event;
mod reader;
pub mod utils;

/// Warnings store.
#[derive(Default, Debug, Clone)]
pub struct Warnings(Vec<Warning>);

impl Warnings {
    /// Creates a new `Warnings`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Add a warning.
    pub fn warn(&mut self, warning: Warning) {
        warn!("FBX binary parser warning: {}", warning);
        self.0.push(warning);
    }

    /// Returns the inner vector.
    pub fn inner(self) -> Vec<Warning> {
        self.0
    }
}

impl ::std::ops::Deref for Warnings {
    type Target = [Warning];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Parser of a FBX binary node.
pub trait Parser<R: ParserSource> {
    /// Returns the root parser.
    fn root_parser(&self) -> &RootParser<R>;
    /// Parses FBX from the given stream and returns the next event.
    fn next_event(&mut self) -> Result<Event<R>>;
    /// Skips to the end of the current node.
    ///
    /// Returns `Ok(true)` if the current node is skipped and closed,
    /// `Ok(false)` if no nodes are open (i.e. the parser is reading under implicit root node),
    /// `Err(err)` if error happened.
    fn skip_current_node(&mut self) -> Result<bool>;
}

/// Parser state without error.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
    /// Reading file header.
    Header,
    /// A node started.
    ///
    /// This status means:
    ///
    /// - there might be some attributes remain unread, or
    /// - the most recent opened node might close without no null node header.
    NodeStarted,
    /// A node ended.
    ///
    /// This status means:
    ///
    /// - if the next event is `NodeEnd`, there must be a null node header, and
    /// - if the parser got an extra null header, it indicates end of implicit root node.
    NodeEnded,
}

/// Information about opened (but not yet closed) node.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct OpenNode {
    /// Start offset of the node attribute.
    ///
    /// Note that this doesn't mean start offset of node header.
    begin: u64,
    /// End offset of the node.
    end: u64,
    /// End offset of attributes of the node.
    attributes_end: u64,
}

/// Pull parser for whole FBX with binary format.
#[derive(Debug)]
pub struct RootParser<R> {
    /// Source reader.
    source: R,
    /// Parser state.
    ///
    /// `Ok(State)` if the parser is working without critical error,
    /// `Err(Error)` if the parsing failed and cannot be continued.
    state: Result<State>,
    /// Warnings.
    warnings: Warnings,
    /// FBX version.
    fbx_version: Option<u32>,
    /// Open nodes stack.
    open_nodes: Vec<OpenNode>,
    /// Node name of the recent opened node.
    recent_node_name: String,
}

impl<R: Read> RootParser<BasicSource<R>> {
    /// Creates a new binary parser.
    pub fn new(source: R) -> Self {
        RootParser {
            source: BasicSource::new(source),
            state: Ok(State::Header),
            warnings: Warnings::new(),
            fbx_version: None,
            open_nodes: Vec::new(),
            recent_node_name: String::new(),
        }
    }
}

impl<R: Read + io::Seek> RootParser<SeekableSource<R>> {
    /// Creates a new binary parser.
    pub fn from_seekable(source: R) -> Self {
        RootParser {
            source: SeekableSource::new(source),
            state: Ok(State::Header),
            warnings: Warnings::new(),
            fbx_version: None,
            open_nodes: Vec::new(),
            recent_node_name: String::new(),
        }
    }
}

impl<R: ParserSource> RootParser<R> {
    /// Returns FBX version of the reading input.
    ///
    /// Returns `None` if unknown yet.
    pub fn fbx_version(&self) -> Option<u32> {
        self.fbx_version
    }

    /// Returns the parser error if available.
    pub fn error(&self) -> Option<&Error> {
        self.state.as_ref().err()
    }

    /// Returns reference to the warnings.
    pub fn warnings(&self) -> &[Warning] {
        &self.warnings
    }

    /// Returns the number of the opened (and not closed) node.
    pub fn num_open_nodes(&self) -> usize {
        self.open_nodes.len()
    }

    /// Returns the node name of the recent opened node.
    pub fn recent_node_name(&self) -> &str {
        &self.recent_node_name
    }

    /// Creates subtree parser for the current node.
    pub fn subtree_parser(&mut self) -> SubtreeParser<R> {
        SubtreeParser::new(self)
    }

    /// Set the parser state as finished parsing.
    fn set_finish(&mut self) {
        self.state = Err(Error::Finished);
    }

    /// Set the parser state as error.
    fn set_error(&mut self, err: &Error) {
        error!("FBX binary parser error: {}", err);
        debug!("Parser: {:#?}", self);
        self.state = Err(err.clone());
    }

    /// Add warning.
    fn warn(&mut self, warning: Warning) {
        self.warnings.warn(warning);
        debug!("Parser: {:#?}", self);
    }

    /// Reads FBX header.
    fn read_fbx_header(&mut self) -> Result<EventBuilder> {
        let header = read_fbx_header(self)?;
        self.fbx_version = Some(header.version);
        self.state = Ok(State::NodeEnded);
        Ok(header.into())
    }

    /// Gets event after node start.
    fn read_after_node_start(&mut self) -> Result<EventBuilder> {
        // Attributes of recent opened node might remain partially unread.
        // They should be skipped before getting a next node event.
        self.skip_attributes()?;

        // Most recent opened node might ends here without a null node header.
        if let Some(end) = self.open_nodes.last().map(|v| v.end) {
            if self.source.position() == end {
                // Most recent opened node ends here (without a null node header).
                self.state = Ok(State::NodeEnded);
                self.open_nodes.pop();
                return Ok(EventBuilder::EndNode);
            }
        }

        let builder = self.read_node_event()?;
        Ok(builder)
    }

    /// Gets event after node end.
    fn read_after_node_end(&mut self) -> Result<EventBuilder> {
        self.read_node_event()
    }

    /// Reads a next node-related event from the source.
    ///
    /// This always returns `Ok(EventBuilder::StartNode)`, `Ok(EventBuilder::EndNode)`,
    /// `Ok(EventBuilder::EndFbx)` or `Err(_)`.
    fn read_node_event(&mut self) -> Result<EventBuilder> {
        let header = NodeHeader::read_from_parser(self)?;
        if header.is_node_end() {
            if let Some(last_node) = self.open_nodes.pop() {
                // There is open nodes, so this is not end of the FBX.
                let current_pos = self.source.position();
                if current_pos != last_node.end {
                    // Invalid node header.
                    return Err(Error::WrongNodeEndOffset {
                        begin: last_node.begin,
                        expected_end: last_node.end,
                        real_end: current_pos,
                    });
                }
            } else {
                assert_eq!(self.state.as_ref().ok(),
                           Some(&State::NodeEnded),
                           "End of implicit root node is read with unexpected parser state {:?}",
                           self.state);
                // No open nodes, so this `EndNode` event indicates the end of
                // the implicit root node.
                // FBX file has no more nodes.
                let footer = self.read_fbx_footer();
                return Ok(footer.into());
            }

            self.state = Ok(State::NodeEnded);
            Ok(EventBuilder::EndNode)
        } else {
            // Reuse node name buffer.
            self.recent_node_name = {
                // Take node name buffer without consuming it and take inner `Vec<u8>` buffer.
                // Note that empty `String` doesn't allocate.
                // `std::mem::uninitialized()` can be used instead of empty `String` but it is `unsafe`.
                let mut vecbuf = ::std::mem::replace(&mut self.recent_node_name, String::new()).into_bytes();
                // Resize buffer.
                // This reallocates only if the buffer is too small.
                vecbuf.resize(header.bytelen_name as usize, 0);
                // Read the node name into the buffer
                self.source.read_exact(&mut vecbuf)?;
                // Covert the name into `String`.
                // If conversion failed, the buffer can be left empty.
                // This is because no more nodes would be loaded and
                // the buffer would no longer be used.
                String::from_utf8(vecbuf).map_err(Error::node_name_invalid_utf8)?
            };

            let current_pos = self.source.position();
            self.open_nodes.push(OpenNode {
                begin: current_pos,
                end: header.end_offset,
                attributes_end: current_pos + header.bytelen_attributes,
            });

            // Zero or more attributes come after node start.
            self.state = Ok(State::NodeStarted);
            Ok(StartNodeBuilder {
                header: header,
            }
            .into())
        }
    }

    /// Reads an FBX footer.
    fn read_fbx_footer(&mut self) -> Result<FbxFooter> {
        self.set_finish();
        FbxFooter::read_from_parser(self)
    }

    /// Skip attributes of the most recent opened node.
    fn skip_attributes(&mut self) -> io::Result<()> {
        let attributes_end = self.open_nodes
            .last()
            .expect("`RootParser::skip_attributes()` is called but no nodes are open")
            .attributes_end;
        self.source.skip_to(attributes_end)
    }
}

impl<R: ParserSource> Parser<R> for RootParser<R> {
    fn root_parser(&self) -> &RootParser<R> {
        self
    }

    fn next_event(&mut self) -> Result<Event<R>> {
        let builder = match self.state.clone()? {
            State::Header => self.read_fbx_header(),
            State::NodeStarted => self.read_after_node_start(),
            State::NodeEnded => self.read_after_node_end(),
        };
        if let Err(ref err) = builder {
            self.set_error(err);
        }
        Ok(builder?.build(self))
    }

    fn skip_current_node(&mut self) -> Result<bool> {
        if let Some(end) = self.open_nodes.pop().map(|v| v.end) {
            self.source.skip_to(end)?;
            self.state = Ok(State::NodeEnded);
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

/// Pull parser for a subtree of the FBX binary.
pub struct SubtreeParser<'a, R: 'a> {
    /// Root parser.
    root_parser: &'a mut RootParser<R>,
    /// Initial depth.
    ///
    /// Depth of the implicit root node is `0`.
    initial_depth: usize,
}

impl<'a, R: 'a + ParserSource> SubtreeParser<'a, R> {
    /// Creates a new `SubtreeParser`.
    pub fn new(root_parser: &'a mut RootParser<R>) -> Self {
        let initial_depth = root_parser.num_open_nodes();
        SubtreeParser {
            root_parser: root_parser,
            initial_depth: initial_depth,
        }
    }

    /// Checks if the subtree parser can emit more events.
    ///
    /// Returns `Ok(())` if more events can be read,
    /// `Err(Error::Finished)` if the subtree is all read,
    /// `Err(_)` if error happened.
    fn check_finished(&self) -> Result<()> {
        if let Some(err) = self.root_parser.error() {
            return Err(err.clone());
        }
        if self.root_parser.num_open_nodes() < self.initial_depth {
            return Err(Error::Finished);
        }
        Ok(())
    }
}

impl<'a, R: 'a + ParserSource> Parser<R> for SubtreeParser<'a, R> {
    fn root_parser(&self) -> &RootParser<R> {
        self.root_parser
    }

    fn next_event(&mut self) -> Result<Event<R>> {
        self.check_finished()?;
        self.root_parser.next_event()
    }

    fn skip_current_node(&mut self) -> Result<bool> {
        self.check_finished()?;
        self.root_parser.skip_current_node()
    }
}
<commit_msg>Add a method to get the recent node name with ownership<commit_after>//! FBX binary parser.

use std::io;
use std::io::Read;

pub use self::error::{Result, Error, Warning};
pub use self::event::{Event, FbxHeader, FbxFooter, StartNode};
pub use self::event::{Attributes, Attribute, SpecialAttributeType};
pub use self::event::{PrimitiveAttribute, ArrayAttribute, SpecialAttribute};
pub use self::event::ArrayAttributeReader;
use self::event::{EventBuilder, NodeHeader, StartNodeBuilder};
use self::event::read_fbx_header;
pub use self::reader::{ParserSource, BasicSource, SeekableSource, LimitedSeekReader};

mod error;
mod event;
mod reader;
pub mod utils;

/// Warnings store.
#[derive(Default, Debug, Clone)]
pub struct Warnings(Vec<Warning>);

impl Warnings {
    /// Creates a new `Warnings`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Add a warning.
    pub fn warn(&mut self, warning: Warning) {
        warn!("FBX binary parser warning: {}", warning);
        self.0.push(warning);
    }

    /// Returns the inner vector.
    pub fn inner(self) -> Vec<Warning> {
        self.0
    }
}

impl ::std::ops::Deref for Warnings {
    type Target = [Warning];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Parser of a FBX binary node.
pub trait Parser<R: ParserSource> {
    /// Returns the root parser.
    fn root_parser(&self) -> &RootParser<R>;
    /// Parses FBX from the given stream and returns the next event.
    fn next_event(&mut self) -> Result<Event<R>>;
    /// Skips to the end of the current node.
    ///
    /// Returns `Ok(true)` if the current node is skipped and closed,
    /// `Ok(false)` if no nodes are open (i.e. the parser is reading under implicit root node),
    /// `Err(err)` if error happened.
    fn skip_current_node(&mut self) -> Result<bool>;
}

/// Parser state without error.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
    /// Reading file header.
    Header,
    /// A node started.
    ///
    /// This status means:
    ///
    /// - there might be some attributes remain unread, or
    /// - the most recent opened node might close without no null node header.
    NodeStarted,
    /// A node ended.
    ///
    /// This status means:
    ///
    /// - if the next event is `NodeEnd`, there must be a null node header, and
    /// - if the parser got an extra null header, it indicates end of implicit root node.
    NodeEnded,
}

/// Information about opened (but not yet closed) node.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct OpenNode {
    /// Start offset of the node attribute.
    ///
    /// Note that this doesn't mean start offset of node header.
    begin: u64,
    /// End offset of the node.
    end: u64,
    /// End offset of attributes of the node.
    attributes_end: u64,
}

/// Pull parser for whole FBX with binary format.
#[derive(Debug)]
pub struct RootParser<R> {
    /// Source reader.
    source: R,
    /// Parser state.
    ///
    /// `Ok(State)` if the parser is working without critical error,
    /// `Err(Error)` if the parsing failed and cannot be continued.
    state: Result<State>,
    /// Warnings.
    warnings: Warnings,
    /// FBX version.
    fbx_version: Option<u32>,
    /// Open nodes stack.
    open_nodes: Vec<OpenNode>,
    /// Node name of the recent opened node.
    recent_node_name: String,
}

impl<R: Read> RootParser<BasicSource<R>> {
    /// Creates a new binary parser.
    pub fn new(source: R) -> Self {
        RootParser {
            source: BasicSource::new(source),
            state: Ok(State::Header),
            warnings: Warnings::new(),
            fbx_version: None,
            open_nodes: Vec::new(),
            recent_node_name: String::new(),
        }
    }
}

impl<R: Read + io::Seek> RootParser<SeekableSource<R>> {
    /// Creates a new binary parser.
    pub fn from_seekable(source: R) -> Self {
        RootParser {
            source: SeekableSource::new(source),
            state: Ok(State::Header),
            warnings: Warnings::new(),
            fbx_version: None,
            open_nodes: Vec::new(),
            recent_node_name: String::new(),
        }
    }
}

impl<R: ParserSource> RootParser<R> {
    /// Returns FBX version of the reading input.
    ///
    /// Returns `None` if unknown yet.
    pub fn fbx_version(&self) -> Option<u32> {
        self.fbx_version
    }

    /// Returns the parser error if available.
    pub fn error(&self) -> Option<&Error> {
        self.state.as_ref().err()
    }

    /// Returns reference to the warnings.
    pub fn warnings(&self) -> &[Warning] {
        &self.warnings
    }

    /// Returns the number of the opened (and not closed) node.
    pub fn num_open_nodes(&self) -> usize {
        self.open_nodes.len()
    }

    /// Returns the node name of the recent opened node.
    pub fn recent_node_name(&self) -> &str {
        &self.recent_node_name
    }

    /// Returns the node name of the recent opened node with ownership.
    pub fn take_recent_node_name(&mut self) -> String {
        ::std::mem::replace(&mut self.recent_node_name, String::new())
    }

    /// Creates subtree parser for the current node.
    pub fn subtree_parser(&mut self) -> SubtreeParser<R> {
        SubtreeParser::new(self)
    }

    /// Set the parser state as finished parsing.
    fn set_finish(&mut self) {
        self.state = Err(Error::Finished);
    }

    /// Set the parser state as error.
    fn set_error(&mut self, err: &Error) {
        error!("FBX binary parser error: {}", err);
        debug!("Parser: {:#?}", self);
        self.state = Err(err.clone());
    }

    /// Add warning.
    fn warn(&mut self, warning: Warning) {
        self.warnings.warn(warning);
        debug!("Parser: {:#?}", self);
    }

    /// Reads FBX header.
    fn read_fbx_header(&mut self) -> Result<EventBuilder> {
        let header = read_fbx_header(self)?;
        self.fbx_version = Some(header.version);
        self.state = Ok(State::NodeEnded);
        Ok(header.into())
    }

    /// Gets event after node start.
    fn read_after_node_start(&mut self) -> Result<EventBuilder> {
        // Attributes of recent opened node might remain partially unread.
        // They should be skipped before getting a next node event.
        self.skip_attributes()?;

        // Most recent opened node might ends here without a null node header.
        if let Some(end) = self.open_nodes.last().map(|v| v.end) {
            if self.source.position() == end {
                // Most recent opened node ends here (without a null node header).
                self.state = Ok(State::NodeEnded);
                self.open_nodes.pop();
                return Ok(EventBuilder::EndNode);
            }
        }

        let builder = self.read_node_event()?;
        Ok(builder)
    }

    /// Gets event after node end.
    fn read_after_node_end(&mut self) -> Result<EventBuilder> {
        self.read_node_event()
    }

    /// Reads a next node-related event from the source.
    ///
    /// This always returns `Ok(EventBuilder::StartNode)`, `Ok(EventBuilder::EndNode)`,
    /// `Ok(EventBuilder::EndFbx)` or `Err(_)`.
    fn read_node_event(&mut self) -> Result<EventBuilder> {
        let header = NodeHeader::read_from_parser(self)?;
        if header.is_node_end() {
            if let Some(last_node) = self.open_nodes.pop() {
                // There is open nodes, so this is not end of the FBX.
                let current_pos = self.source.position();
                if current_pos != last_node.end {
                    // Invalid node header.
                    return Err(Error::WrongNodeEndOffset {
                        begin: last_node.begin,
                        expected_end: last_node.end,
                        real_end: current_pos,
                    });
                }
            } else {
                assert_eq!(self.state.as_ref().ok(),
                           Some(&State::NodeEnded),
                           "End of implicit root node is read with unexpected parser state {:?}",
                           self.state);
                // No open nodes, so this `EndNode` event indicates the end of
                // the implicit root node.
                // FBX file has no more nodes.
                let footer = self.read_fbx_footer();
                return Ok(footer.into());
            }

            self.state = Ok(State::NodeEnded);
            Ok(EventBuilder::EndNode)
        } else {
            // Reuse node name buffer.
            self.recent_node_name = {
                // Take node name buffer without consuming it and take inner `Vec<u8>` buffer.
                // Note that empty `String` doesn't allocate.
                // `std::mem::uninitialized()` can be used instead of empty `String` but it is `unsafe`.
                let mut vecbuf = ::std::mem::replace(&mut self.recent_node_name, String::new()).into_bytes();
                // Resize buffer.
                // This reallocates only if the buffer is too small.
                vecbuf.resize(header.bytelen_name as usize, 0);
                // Read the node name into the buffer
                self.source.read_exact(&mut vecbuf)?;
                // Covert the name into `String`.
                // If conversion failed, the buffer can be left empty.
                // This is because no more nodes would be loaded and
                // the buffer would no longer be used.
                String::from_utf8(vecbuf).map_err(Error::node_name_invalid_utf8)?
            };

            let current_pos = self.source.position();
            self.open_nodes.push(OpenNode {
                begin: current_pos,
                end: header.end_offset,
                attributes_end: current_pos + header.bytelen_attributes,
            });

            // Zero or more attributes come after node start.
            self.state = Ok(State::NodeStarted);
            Ok(StartNodeBuilder {
                header: header,
            }
            .into())
        }
    }

    /// Reads an FBX footer.
    fn read_fbx_footer(&mut self) -> Result<FbxFooter> {
        self.set_finish();
        FbxFooter::read_from_parser(self)
    }

    /// Skip attributes of the most recent opened node.
    fn skip_attributes(&mut self) -> io::Result<()> {
        let attributes_end = self.open_nodes
            .last()
            .expect("`RootParser::skip_attributes()` is called but no nodes are open")
            .attributes_end;
        self.source.skip_to(attributes_end)
    }
}

impl<R: ParserSource> Parser<R> for RootParser<R> {
    fn root_parser(&self) -> &RootParser<R> {
        self
    }

    fn next_event(&mut self) -> Result<Event<R>> {
        let builder = match self.state.clone()? {
            State::Header => self.read_fbx_header(),
            State::NodeStarted => self.read_after_node_start(),
            State::NodeEnded => self.read_after_node_end(),
        };
        if let Err(ref err) = builder {
            self.set_error(err);
        }
        Ok(builder?.build(self))
    }

    fn skip_current_node(&mut self) -> Result<bool> {
        if let Some(end) = self.open_nodes.pop().map(|v| v.end) {
            self.source.skip_to(end)?;
            self.state = Ok(State::NodeEnded);
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

/// Pull parser for a subtree of the FBX binary.
pub struct SubtreeParser<'a, R: 'a> {
    /// Root parser.
    root_parser: &'a mut RootParser<R>,
    /// Initial depth.
    ///
    /// Depth of the implicit root node is `0`.
    initial_depth: usize,
}

impl<'a, R: 'a + ParserSource> SubtreeParser<'a, R> {
    /// Creates a new `SubtreeParser`.
    pub fn new(root_parser: &'a mut RootParser<R>) -> Self {
        let initial_depth = root_parser.num_open_nodes();
        SubtreeParser {
            root_parser: root_parser,
            initial_depth: initial_depth,
        }
    }

    /// Checks if the subtree parser can emit more events.
    ///
    /// Returns `Ok(())` if more events can be read,
    /// `Err(Error::Finished)` if the subtree is all read,
    /// `Err(_)` if error happened.
    fn check_finished(&self) -> Result<()> {
        if let Some(err) = self.root_parser.error() {
            return Err(err.clone());
        }
        if self.root_parser.num_open_nodes() < self.initial_depth {
            return Err(Error::Finished);
        }
        Ok(())
    }
}

impl<'a, R: 'a + ParserSource> Parser<R> for SubtreeParser<'a, R> {
    fn root_parser(&self) -> &RootParser<R> {
        self.root_parser
    }

    fn next_event(&mut self) -> Result<Event<R>> {
        self.check_finished()?;
        self.root_parser.next_event()
    }

    fn skip_current_node(&mut self) -> Result<bool> {
        self.check_finished()?;
        self.root_parser.skip_current_node()
    }
}
<|endoftext|>


<commit_before>use prolog::ast::*;
use prolog::fixtures::*;

use std::cell::Cell;
use std::cmp::{min, max};
use std::vec::Vec;

pub struct ArithExprIterator<'a> {
    state_stack: Vec<TermIterState<'a>>
}

pub type ArithCont = (Code, Option<ArithmeticTerm>);

impl<'a> ArithExprIterator<'a> {
    fn push_subterm(&mut self, lvl: Level, term: &'a Term) {
        self.state_stack.push(TermIterState::to_state(lvl, term));
    }

    fn new(term: &'a Term) -> Result<Self, ArithmeticError> {
        let state = match term {
            &Term::AnonVar =>
                return Err(ArithmeticError::InvalidTerm),
            &Term::Clause(_, _, ref terms, _) =>
                TermIterState::Clause(0, ClauseType::Root, terms),
            &Term::Constant(ref cell, ref cons) =>
                TermIterState::Constant(Level::Shallow, cell, cons),
            &Term::Cons(_, _, _) =>
                return Err(ArithmeticError::InvalidTerm),
            &Term::Var(ref cell, ref var) =>
                TermIterState::Var(Level::Shallow, cell, var)
        };

        Ok(ArithExprIterator { state_stack: vec![state] })
    }
}

impl Term {
    pub fn arith_expr_iter<'a>(&'a self) -> Result<ArithExprIterator<'a>, ArithmeticError> {
        ArithExprIterator::new(self)
    }
}

pub enum ArithTermRef<'a> {
    Constant(&'a Constant),
    Op(ClauseType<'a>, &'a Vec<Box<Term>>),
    Var(&'a Cell<VarReg>, &'a Var)
}

impl<'a> Iterator for ArithExprIterator<'a> {
    type Item = Result<ArithTermRef<'a>, ArithmeticError>;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(iter_state) = self.state_stack.pop() {
            match iter_state {
                TermIterState::AnonVar(_) =>
                    return Some(Err(ArithmeticError::UninstantiatedVar)),
                TermIterState::Clause(child_num, ct, child_terms) => {
                    if child_num == child_terms.len() {
                        return Some(Ok(ArithTermRef::Op(ct, child_terms)));
                    } else {
                        self.state_stack.push(TermIterState::Clause(child_num + 1, ct, child_terms));
                        self.push_subterm(ct.level_of_subterms(), child_terms[child_num].as_ref());
                    }
                },
                TermIterState::Constant(_, _, c) =>
                    return Some(Ok(ArithTermRef::Constant(c))),
                TermIterState::Var(_, cell, var) =>
                    return Some(Ok(ArithTermRef::Var(cell, var))),
                _ =>
                    return Some(Err(ArithmeticError::InvalidTerm))
            };
        }

        None
    }
}

pub struct ArithmeticEvaluator<'a> {
    bindings: &'a AllocVarDict<'a>,
    interm: Vec<ArithmeticTerm>,
    interm_c: usize
}

impl<'a> ArithmeticEvaluator<'a> {
    pub fn new(bindings: &'a AllocVarDict<'a>, target_int: usize) -> Self {
        ArithmeticEvaluator { bindings, interm: Vec::new(), interm_c: target_int }
    }

    fn get_unary_instr(name: &Atom, a1: ArithmeticTerm, t: usize)
                       -> Result<ArithmeticInstruction, ArithmeticError>
    {
        match name.as_str() {
            "-" => Ok(ArithmeticInstruction::Neg(a1, t)),
             _  => Err(ArithmeticError::InvalidOp)
        }
    }

    fn gen_binary_instr(name: &Atom, a1: ArithmeticTerm, a2: ArithmeticTerm, t: usize)
                        -> Result<ArithmeticInstruction, ArithmeticError>
    {
        match name.as_str() {
            "+"    => Ok(ArithmeticInstruction::Add(a1, a2, t)),
            "-"    => Ok(ArithmeticInstruction::Sub(a1, a2, t)),
            "/"    => Ok(ArithmeticInstruction::Div(a1, a2, t)),
            "//"   => Ok(ArithmeticInstruction::IDiv(a1, a2, t)),
            "div"  => Ok(ArithmeticInstruction::FIDiv(a1, a2, t)),
            "rdiv" => Ok(ArithmeticInstruction::RDiv(a1, a2, t)),
            "*"    => Ok(ArithmeticInstruction::Mul(a1, a2, t)),
            ">>"   => Ok(ArithmeticInstruction::Shr(a1, a2, t)),
            "<<"   => Ok(ArithmeticInstruction::Shl(a1, a2, t)),
            "/\\"  => Ok(ArithmeticInstruction::And(a1, a2, t)),
            "\\/"  => Ok(ArithmeticInstruction::Or(a1, a2, t)),
            "xor"  => Ok(ArithmeticInstruction::Xor(a1, a2, t)),
            "mod"  => Ok(ArithmeticInstruction::Mod(a1, a2, t)),
            "rem"  => Ok(ArithmeticInstruction::Rem(a1, a2, t)),
             _     => Err(ArithmeticError::InvalidOp)
        }
    }

    fn incr_interm(&mut self) -> usize {
        let temp = self.interm_c;

        self.interm.push(ArithmeticTerm::Interm(temp));
        self.interm_c += 1;

        temp
    }

    fn instr_from_clause(&mut self, name: &Atom, terms: &Vec<Box<Term>>)
                         -> Result<ArithmeticInstruction, ArithmeticError>
    {
        match terms.len() {
            1 => {
                let a1 = self.interm.pop().unwrap();

                let ninterm = if a1.interm_or(0) == 0 {
                    self.incr_interm()
                } else {
                    self.interm.push(a1.clone());
                    a1.interm_or(0)
                };

                Self::get_unary_instr(name, a1, ninterm)
            },
            2 => {
                let a2 = self.interm.pop().unwrap();
                let a1 = self.interm.pop().unwrap();

                let min_interm = min(a1.interm_or(0), a2.interm_or(0));

                let ninterm = if min_interm == 0 {
                    let max_interm = max(a1.interm_or(0), a2.interm_or(0));

                    if max_interm == 0 {
                        self.incr_interm()
                    } else {
                        self.interm.push(ArithmeticTerm::Interm(max_interm));
                        self.interm_c = max_interm + 1;
                        max_interm
                    }
                } else {
                    self.interm.push(ArithmeticTerm::Interm(min_interm));
                    self.interm_c = min_interm + 1;
                    min_interm
                };

                Self::gen_binary_instr(name, a1, a2, ninterm)
            },
            _ => Err(ArithmeticError::InvalidOp)
        }
    }

    fn push_constant(&mut self, c: &Constant) -> Result<(), ArithmeticError> {
        match c {
            &Constant::Number(ref n) =>
                self.interm.push(ArithmeticTerm::Number(n.clone())),
            _ =>
                return Err(ArithmeticError::InvalidAtom),
        }

        Ok(())
    }

    pub fn eval(&mut self, term: &Term) -> Result<ArithCont, ArithmeticError> {
        let mut code = Vec::new();

        for term_ref in term.arith_expr_iter()? {
            match term_ref? {
                ArithTermRef::Constant(c) => self.push_constant(c)?,
                ArithTermRef::Var(cell, name) => {
                    let r = if cell.get().norm().reg_num() == 0 {
                        match self.bindings.get(name) {
                            Some(&VarData::Temp(_, t, _)) if t != 0 => RegType::Temp(t),
                            Some(&VarData::Perm(p)) if p != 0 => RegType::Perm(p),
                            _ => return Err(ArithmeticError::UninstantiatedVar)
                        }
                    } else {
                        cell.get().norm()
                    };

                    self.interm.push(ArithmeticTerm::Reg(r));
                },
                ArithTermRef::Op(ClauseType::Deep(_, _, name, _), terms) => {
                    code.push(Line::Arithmetic(self.instr_from_clause(&*name, terms)?));
                },
                ArithTermRef::Op(ClauseType::Root, terms) => {
                    let name = term.name().unwrap();
                    code.push(Line::Arithmetic(self.instr_from_clause(&*name, terms)?));
                },
                _ =>
                    return Err(ArithmeticError::InvalidTerm)
            }
        }

        Ok((code, self.interm.pop()))
    }
}
<commit_msg>abstract out iterator for arithmetic terms.<commit_after>use prolog::ast::*;
use prolog::fixtures::*;
use prolog::tabled_rc::*;

use std::cell::Cell;
use std::cmp::{min, max};
use std::vec::Vec;

pub struct ArithInstructionIterator<'a> {
    state_stack: Vec<TermIterState<'a>>
}

pub type ArithCont = (Code, Option<ArithmeticTerm>);

impl<'a> ArithInstructionIterator<'a> {
    fn push_subterm(&mut self, lvl: Level, term: &'a Term) {
        self.state_stack.push(TermIterState::to_state(lvl, term));
    }

    fn new(term: &'a Term) -> Result<Self, ArithmeticError> {
        let state = match term {
            &Term::AnonVar =>
                return Err(ArithmeticError::InvalidTerm),
            &Term::Clause(_, _, ref terms, _) =>
                TermIterState::Clause(0, ClauseType::Root, terms),
            &Term::Constant(ref cell, ref cons) =>
                TermIterState::Constant(Level::Shallow, cell, cons),
            &Term::Cons(_, _, _) =>
                return Err(ArithmeticError::InvalidTerm),
            &Term::Var(ref cell, ref var) =>
                TermIterState::Var(Level::Shallow, cell, var)
        };

        Ok(ArithInstructionIterator { state_stack: vec![state] })
    }
}

pub enum ArithTermRef<'a> {
    Constant(&'a Constant),
    Op(ClauseType<'a>, &'a Vec<Box<Term>>),
    Var(&'a Cell<VarReg>, &'a Var)
}

impl<'a> Iterator for ArithInstructionIterator<'a> {
    type Item = Result<ArithTermRef<'a>, ArithmeticError>;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(iter_state) = self.state_stack.pop() {
            match iter_state {
                TermIterState::AnonVar(_) =>
                    return Some(Err(ArithmeticError::UninstantiatedVar)),
                TermIterState::Clause(child_num, ct, child_terms) => {
                    if child_num == child_terms.len() {
                        return Some(Ok(ArithTermRef::Op(ct, child_terms)));
                    } else {
                        self.state_stack.push(TermIterState::Clause(child_num + 1, ct, child_terms));
                        self.push_subterm(ct.level_of_subterms(), child_terms[child_num].as_ref());
                    }
                },
                TermIterState::Constant(_, _, c) =>
                    return Some(Ok(ArithTermRef::Constant(c))),
                TermIterState::Var(_, cell, var) =>
                    return Some(Ok(ArithTermRef::Var(cell, var))),
                _ =>
                    return Some(Err(ArithmeticError::InvalidTerm))
            };
        }

        None
    }
}

pub struct ArithmeticEvaluator<'a> {
    bindings: &'a AllocVarDict<'a>,
    interm: Vec<ArithmeticTerm>,
    interm_c: usize
}

pub trait ArithmeticTermIter<'a> {
    type Iter : Iterator<Item=Result<ArithTermRef<'a>, ArithmeticError>>;

    fn iter(&self) -> Result<Self::Iter, ArithmeticError>;
    fn root_name(&self) -> Result<TabledRc<Atom>, ArithmeticError>;
}

impl<'a> ArithmeticTermIter<'a> for &'a Term {
    type Iter = ArithInstructionIterator<'a>;

    fn iter(&self) -> Result<Self::Iter, ArithmeticError> {
        ArithInstructionIterator::new(self)
    }

    fn root_name(&self) -> Result<TabledRc<Atom>, ArithmeticError> {
        match self {
            &&Term::Clause(_, ref name, _, _) => Ok(name.clone()),
            _ => Err(ArithmeticError::InvalidTerm)
        }
    }
}

impl<'a> ArithmeticEvaluator<'a>
{
    pub fn new(bindings: &'a AllocVarDict<'a>, target_int: usize) -> Self {
        ArithmeticEvaluator { bindings, interm: Vec::new(), interm_c: target_int }
    }

    fn get_unary_instr(name: &Atom, a1: ArithmeticTerm, t: usize)
                       -> Result<ArithmeticInstruction, ArithmeticError>
    {
        match name.as_str() {
            "-" => Ok(ArithmeticInstruction::Neg(a1, t)),
             _  => Err(ArithmeticError::InvalidOp)
        }
    }

    fn get_binary_instr(name: &Atom, a1: ArithmeticTerm, a2: ArithmeticTerm, t: usize)
                        -> Result<ArithmeticInstruction, ArithmeticError>
    {
        match name.as_str() {
            "+"    => Ok(ArithmeticInstruction::Add(a1, a2, t)),
            "-"    => Ok(ArithmeticInstruction::Sub(a1, a2, t)),
            "/"    => Ok(ArithmeticInstruction::Div(a1, a2, t)),
            "//"   => Ok(ArithmeticInstruction::IDiv(a1, a2, t)),
            "div"  => Ok(ArithmeticInstruction::FIDiv(a1, a2, t)),
            "rdiv" => Ok(ArithmeticInstruction::RDiv(a1, a2, t)),
            "*"    => Ok(ArithmeticInstruction::Mul(a1, a2, t)),
            ">>"   => Ok(ArithmeticInstruction::Shr(a1, a2, t)),
            "<<"   => Ok(ArithmeticInstruction::Shl(a1, a2, t)),
            "/\\"  => Ok(ArithmeticInstruction::And(a1, a2, t)),
            "\\/"  => Ok(ArithmeticInstruction::Or(a1, a2, t)),
            "xor"  => Ok(ArithmeticInstruction::Xor(a1, a2, t)),
            "mod"  => Ok(ArithmeticInstruction::Mod(a1, a2, t)),
            "rem"  => Ok(ArithmeticInstruction::Rem(a1, a2, t)),
             _     => Err(ArithmeticError::InvalidOp)
        }
    }

    fn incr_interm(&mut self) -> usize {
        let temp = self.interm_c;

        self.interm.push(ArithmeticTerm::Interm(temp));
        self.interm_c += 1;

        temp
    }

    fn instr_from_clause(&mut self, name: &Atom, terms: &Vec<Box<Term>>)
                         -> Result<ArithmeticInstruction, ArithmeticError>
    {
        match terms.len() {
            1 => {
                let a1 = self.interm.pop().unwrap();

                let ninterm = if a1.interm_or(0) == 0 {
                    self.incr_interm()
                } else {
                    self.interm.push(a1.clone());
                    a1.interm_or(0)
                };

                Self::get_unary_instr(name, a1, ninterm)
            },
            2 => {
                let a2 = self.interm.pop().unwrap();
                let a1 = self.interm.pop().unwrap();

                let min_interm = min(a1.interm_or(0), a2.interm_or(0));

                let ninterm = if min_interm == 0 {
                    let max_interm = max(a1.interm_or(0), a2.interm_or(0));

                    if max_interm == 0 {
                        self.incr_interm()
                    } else {
                        self.interm.push(ArithmeticTerm::Interm(max_interm));
                        self.interm_c = max_interm + 1;
                        max_interm
                    }
                } else {
                    self.interm.push(ArithmeticTerm::Interm(min_interm));
                    self.interm_c = min_interm + 1;
                    min_interm
                };

                Self::get_binary_instr(name, a1, a2, ninterm)
            },
            _ => Err(ArithmeticError::InvalidOp)
        }
    }

    fn push_constant(&mut self, c: &Constant) -> Result<(), ArithmeticError> {
        match c {
            &Constant::Number(ref n) =>
                self.interm.push(ArithmeticTerm::Number(n.clone())),
            _ =>
                return Err(ArithmeticError::InvalidAtom),
        }

        Ok(())
    }

    pub fn eval<Iter>(&mut self, src: Iter) -> Result<ArithCont, ArithmeticError>
        where Iter: ArithmeticTermIter<'a>
    {
        let mut code = vec![];

        for term_ref in src.iter()?
        {
            match term_ref? {
                ArithTermRef::Constant(c) => self.push_constant(c)?,
                ArithTermRef::Var(cell, name) => {
                    let r = if cell.get().norm().reg_num() == 0 {
                        match self.bindings.get(name) {
                            Some(&VarData::Temp(_, t, _)) if t != 0 => RegType::Temp(t),
                            Some(&VarData::Perm(p)) if p != 0 => RegType::Perm(p),
                            _ => return Err(ArithmeticError::UninstantiatedVar)
                        }
                    } else {
                        cell.get().norm()
                    };

                    self.interm.push(ArithmeticTerm::Reg(r));
                },
                ArithTermRef::Op(ClauseType::Deep(_, _, name, _), terms) => {
                    code.push(Line::Arithmetic(self.instr_from_clause(&*name, terms)?));
                },
                ArithTermRef::Op(ClauseType::Root, terms) => {
                    let name = src.root_name()?;
                    code.push(Line::Arithmetic(self.instr_from_clause(&*name, terms)?));
                },
                _ =>
                    return Err(ArithmeticError::InvalidTerm)
            }
        }

        Ok((code, self.interm.pop()))
    }
}
<|endoftext|>


<commit_before>use std::{mem, sync::Arc};

use bifrost::{conshash::ConsistentHashing, raft::client::RaftClient, rpc::RPCError};
use dovahkiin::{expr::serde::Expr, types::OwnedValue};
use futures::stream::{FuturesUnordered, StreamExt};
use itertools::Itertools;

use crate::{
    client::{client_by_server_id, client_by_server_name},
    index::{
        ranged::{client::cursor::ClientCursor, lsm::btree::Ordering},
        EntryKey, IndexerClients, SCHEMA_SCAN_PATT_SIZE,
    },
    ram::cell::{OwnedCell, ReadError},
};

const SCAN_BUFFER_SIZE: u16 = 64;

pub struct IndexedDataClient {
    conshash: Arc<ConsistentHashing>,
    index_clients: Arc<IndexerClients>,
}

pub struct DataCursor<'a> {
    index_cursor: Option<ClientCursor>,
    buffer: Vec<OwnedCell>,
    projection: Vec<u64>,
    selection: Expr,
    proc: Expr,
    client: &'a IndexedDataClient,
    pos: usize,
}

impl IndexedDataClient {
    pub fn new(conshash: &Arc<ConsistentHashing>, raft_client: &Arc<RaftClient>) -> Self {
        Self {
            conshash: conshash.clone(),
            index_clients: Arc::new(IndexerClients::new(conshash, raft_client)),
        }
    }
    pub async fn range_index_scan<'a>(
        &self,
        schema: u32,
        field: u64,
        key: u64,
        selection: OwnedValue,
        projection: OwnedValue,
    ) -> DataCursor<'a> {
        unimplemented!()
    }
    pub async fn scan_all<'a>(
        &'a self,
        schema: u32,
        projection: Vec<u64>, // Column array
        selection: Expr,      // Checker expression
        proc: Expr,
        ordering: Ordering,
    ) -> Result<DataCursor<'a>, RPCError> {
        let key = EntryKey::for_schema(schema);
        let index_cursor = self
            .index_clients
            .range_seek(
                &key,
                ordering,
                SCAN_BUFFER_SIZE,
                Some(SCHEMA_SCAN_PATT_SIZE),
            )
            .await?;
        let mut cursor = DataCursor {
            index_cursor,
            projection,
            selection,
            proc,
            client: self,
            buffer: vec![],
            pos: 0,
        };
        cursor.refresh_batch().await;
        Ok(cursor)
    }
}

impl<'a> DataCursor<'a> {
    pub async fn next(&mut self) -> Result<Option<OwnedCell>, RPCError> {
        if self.buffer.len() <= self.pos {
            if self.next_block().await? {
                if !self.refresh_batch().await {
                    return Ok(None);
                }
            } else {
                return Ok(None);
            }
        }
        let cell = mem::take(&mut self.buffer[self.pos]);
        self.pos += 1;
        return Ok(Some(cell));
    }

    pub async fn next_block(&mut self) -> Result<bool, RPCError> {
        if let Some(cursor) = &mut self.index_cursor {
            if cursor.next_block().await? {
                return Ok(true);
            }
        }
        // If cannot get next block, set the index cursor to none
        // We are done with this cursor
        self.index_cursor = None;
        self.buffer = vec![];
        self.pos = 0;
        return Ok(false);
    }

    pub async fn refresh_batch(&mut self) -> bool {
        if let Some(cursor) = &self.index_cursor {
            let all_ids = cursor.current_block();
            let mut tasks = all_ids
                .iter()
                .enumerate()
                .filter_map(|(i, id)| {
                    self.client
                        .conshash
                        .get_server_id_by(id)
                        .map(|sid| (i, sid, id))
                })
                .group_by(|(_i, sid, _id)| *sid)
                .into_iter()
                .map(|(sid, pairs)| {
                    let mut ids = vec![];
                    let mut idx = vec![];
                    for (i, _, id) in pairs {
                        idx.push(i);
                        ids.push(*id);
                    }
                    let projection = self.projection.clone();
                    let selection = self.selection.clone();
                    let proc = self.proc.clone();
                    let server_name = self.client.conshash.to_server_name(sid);
                    async move {
                        match client_by_server_name(sid, server_name).await {
                            Ok(client) => {
                                let read_res = client
                                    .read_all_cells_proced(ids, projection, selection, proc)
                                    .await
                                    .map(|v| v.into_iter().zip(idx).collect_vec());
                                match read_res {
                                    Ok(cells) => Ok(cells
                                        .into_iter()
                                        .filter_map(|(c, i)| c.ok().map(|c| (c, i)))
                                        .collect_vec()),
                                    Err(e) => return Err(e),
                                }
                            }
                            Err(e) => return Err(e),
                        }
                    }
                })
                .collect::<FuturesUnordered<_>>();
            let mut all_cells = vec![];
            while let Some(task_res) = tasks.next().await {
                if let Ok(mut cells) = task_res {
                    all_cells.append(&mut cells);
                }
            }
            all_cells.sort_by(|(_, i1), (_, i2)| i1.cmp(i2));
            self.buffer = all_cells.into_iter().map(|(c, _)| c).collect_vec();
            self.pos = 0;
            true
        } else {
            self.buffer = vec![];
            self.pos = 0;
            false
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        index::ranged::lsm::btree::Ordering,
        ram::{
            cell::OwnedCell,
            schema::{Field, IndexType, Schema},
        },
        server::*,
    };
    use bifrost_hasher::hash_str;
    use dovahkiin::{expr::serde::Expr, types::*};

    #[tokio::test(flavor = "multi_thread")]
    async fn scan_all() {
        const DATA_1: &'static str = "DATA_1";
        const DATA_2: &'static str = "DATA_2";
        let _ = env_logger::try_init();
        let server_addr = String::from("127.0.0.1:6701");
        let server_group = String::from("indexed_scan_all_test");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 8,
                memory_size: 512 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
                index_enabled: true,
                services: vec![Service::Cell, Service::Query],
            },
            &server_addr,
            &server_group,
        )
        .await;
        // Require schema to be scannable to insert special scan key to the range indexer
        let fields = Field::new(
            "*",
            Type::Map,
            false,
            false,
            Some(vec![
                Field::new(
                    DATA_1,
                    Type::U64,
                    false,
                    false,
                    None,
                    vec![IndexType::Ranged],
                ),
                Field::new(DATA_2, Type::U32, false, false, None, vec![]),
            ]),
            vec![],
        );
        let schema_id_1 = 123;
        let schema_id_2 = 234;
        let schema_1 = Schema::new_with_id(
            schema_id_1,
            &String::from("schema_1"),
            None,
            fields.clone(),
            false,
            true, // Scannable
        );
        let schema_2 = Schema::new_with_id(
            schema_id_2,
            &String::from("schema_2"),
            None,
            fields,
            false,
            true, // Scannable
        );
        let client = server.data_client(&vec![server_addr]).await.unwrap();
        client.new_schema_with_id(schema_1).await.unwrap().unwrap();
        client.new_schema_with_id(schema_2).await.unwrap().unwrap();
        let num = 1024;
        for i in 0..num {
            let id = Id::new(1, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 2) as u32);
            let cell = OwnedCell::new_with_id(schema_id_1, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let idx_data_client = server.indexed_data_client();
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        for i in 0..num {
            let id = Id::new(2, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 3) as u32);
            let cell = OwnedCell::new_with_id(schema_id_2, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_2,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(2, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 3) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
        }
        // Testing projection
        let f1_id = hash_str(DATA_1);
        let mut cursor = idx_data_client
        .scan_all(
            schema_id_1,
            vec![f1_id],
            Expr::nothing(),
            Expr::nothing(),
            Ordering::Forward,
        )
        .await
        .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert!(cell[DATA_2].u32().is_none());
            debug!("Checked cell id {:?} from index", id);
        }
    }
}
<commit_msg>range scan<commit_after>use std::{mem, sync::Arc};

use bifrost::{conshash::ConsistentHashing, raft::client::RaftClient, rpc::RPCError};
use dovahkiin::{expr::serde::Expr, types::OwnedValue};
use futures::stream::{FuturesUnordered, StreamExt};
use itertools::Itertools;

use crate::{
    client::{client_by_server_id, client_by_server_name},
    index::{
        ranged::{client::cursor::ClientCursor, lsm::btree::Ordering},
        EntryKey, IndexerClients, SCHEMA_SCAN_PATT_SIZE,
    },
    ram::cell::{OwnedCell, ReadError},
};

const SCAN_BUFFER_SIZE: u16 = 64;

pub struct IndexedDataClient {
    conshash: Arc<ConsistentHashing>,
    index_clients: Arc<IndexerClients>,
}

pub struct DataCursor<'a> {
    index_cursor: Option<ClientCursor>,
    buffer: Vec<OwnedCell>,
    projection: Vec<u64>,
    selection: Expr,
    proc: Expr,
    client: &'a IndexedDataClient,
    pos: usize,
}

impl IndexedDataClient {
    pub fn new(conshash: &Arc<ConsistentHashing>, raft_client: &Arc<RaftClient>) -> Self {
        Self {
            conshash: conshash.clone(),
            index_clients: Arc::new(IndexerClients::new(conshash, raft_client)),
        }
    }
    pub async fn range_index_scan<'a>(
        &self,
        schema: u32,
        field: u64,
        key: u64,
        selection: OwnedValue,
        projection: OwnedValue,
    ) -> DataCursor<'a> {
        unimplemented!()
    }
    pub async fn scan_all<'a>(
        &'a self,
        schema: u32,
        projection: Vec<u64>, // Column array
        selection: Expr,      // Checker expression
        proc: Expr,
        ordering: Ordering,
    ) -> Result<DataCursor<'a>, RPCError> {
        let key = EntryKey::for_schema(schema);
        let index_cursor = self
            .index_clients
            .range_seek(
                &key,
                ordering,
                SCAN_BUFFER_SIZE,
                Some(SCHEMA_SCAN_PATT_SIZE),
            )
            .await?;
        let mut cursor = DataCursor {
            index_cursor,
            projection,
            selection,
            proc,
            client: self,
            buffer: vec![],
            pos: 0,
        };
        cursor.refresh_batch().await;
        Ok(cursor)
    }
}

impl<'a> DataCursor<'a> {
    pub async fn next(&mut self) -> Result<Option<OwnedCell>, RPCError> {
        if self.buffer.len() <= self.pos {
            if self.next_block().await? {
                if !self.refresh_batch().await {
                    return Ok(None);
                }
            } else {
                return Ok(None);
            }
        }
        if self.buffer.is_empty() {
            return Ok(None);
        }
        let cell = mem::take(&mut self.buffer[self.pos]);
        self.pos += 1;
        return Ok(Some(cell));
    }

    pub async fn next_block(&mut self) -> Result<bool, RPCError> {
        if let Some(cursor) = &mut self.index_cursor {
            if cursor.next_block().await? {
                return Ok(true);
            }
        }
        // If cannot get next block, set the index cursor to none
        // We are done with this cursor
        self.index_cursor = None;
        self.buffer = vec![];
        self.pos = 0;
        return Ok(false);
    }

    pub async fn refresh_batch(&mut self) -> bool {
        if let Some(cursor) = &self.index_cursor {
            let all_ids = cursor.current_block();
            let mut tasks = all_ids
                .iter()
                .enumerate()
                .filter_map(|(i, id)| {
                    self.client
                        .conshash
                        .get_server_id_by(id)
                        .map(|sid| (i, sid, id))
                })
                .group_by(|(_i, sid, _id)| *sid)
                .into_iter()
                .map(|(sid, pairs)| {
                    let mut ids = vec![];
                    let mut idx = vec![];
                    for (i, _, id) in pairs {
                        idx.push(i);
                        ids.push(*id);
                    }
                    let projection = self.projection.clone();
                    let selection = self.selection.clone();
                    let proc = self.proc.clone();
                    let server_name = self.client.conshash.to_server_name(sid);
                    async move {
                        match client_by_server_name(sid, server_name).await {
                            Ok(client) => {
                                let read_res = client
                                    .read_all_cells_proced(ids, projection, selection, proc)
                                    .await
                                    .map(|v| v.into_iter().zip(idx).collect_vec());
                                match read_res {
                                    Ok(cells) => Ok(cells
                                        .into_iter()
                                        .filter_map(|(c, i)| c.ok().map(|c| (c, i)))
                                        .collect_vec()),
                                    Err(e) => return Err(e),
                                }
                            }
                            Err(e) => return Err(e),
                        }
                    }
                })
                .collect::<FuturesUnordered<_>>();
            let mut all_cells = vec![];
            while let Some(task_res) = tasks.next().await {
                if let Ok(mut cells) = task_res {
                    all_cells.append(&mut cells);
                }
            }
            all_cells.sort_by(|(_, i1), (_, i2)| i1.cmp(i2));
            self.buffer = all_cells.into_iter().map(|(c, _)| c).collect_vec();
            self.pos = 0;
            true
        } else {
            self.buffer = vec![];
            self.pos = 0;
            false
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        index::ranged::lsm::btree::Ordering,
        ram::{
            cell::OwnedCell,
            schema::{Field, IndexType, Schema},
        },
        server::*,
    };
    use bifrost_hasher::hash_str;
    use dovahkiin::{expr::serde::Expr, types::*, integrated::lisp::*};

    #[tokio::test(flavor = "multi_thread")]
    async fn scan_all() {
        const DATA_1: &'static str = "DATA_1";
        const DATA_2: &'static str = "DATA_2";
        let _ = env_logger::try_init();
        let server_addr = String::from("127.0.0.1:6701");
        let server_group = String::from("indexed_scan_all_test");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 8,
                memory_size: 512 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
                index_enabled: true,
                services: vec![Service::Cell, Service::Query],
            },
            &server_addr,
            &server_group,
        )
        .await;
        // Require schema to be scannable to insert special scan key to the range indexer
        let fields = Field::new(
            "*",
            Type::Map,
            false,
            false,
            Some(vec![
                Field::new(
                    DATA_1,
                    Type::U64,
                    false,
                    false,
                    None,
                    vec![IndexType::Ranged],
                ),
                Field::new(DATA_2, Type::U32, false, false, None, vec![]),
            ]),
            vec![],
        );
        let schema_id_1 = 123;
        let schema_id_2 = 234;
        let schema_1 = Schema::new_with_id(
            schema_id_1,
            &String::from("schema_1"),
            None,
            fields.clone(),
            false,
            true, // Scannable
        );
        let schema_2 = Schema::new_with_id(
            schema_id_2,
            &String::from("schema_2"),
            None,
            fields,
            false,
            true, // Scannable
        );
        let client = server.data_client(&vec![server_addr]).await.unwrap();
        client.new_schema_with_id(schema_1).await.unwrap().unwrap();
        client.new_schema_with_id(schema_2).await.unwrap().unwrap();
        let num = 1024;
        for i in 0..num {
            let id = Id::new(1, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 2) as u32);
            let cell = OwnedCell::new_with_id(schema_id_1, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let idx_data_client = server.indexed_data_client();
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        for i in 0..num {
            let id = Id::new(2, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 3) as u32);
            let cell = OwnedCell::new_with_id(schema_id_2, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_2,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(2, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 3) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
        }
        // Testing selection
        let select_expr = parse_to_serde_expr("(and (>= DATA_1 10u64) (< DATA_1 100u64))").unwrap()[0].clone();
        let mut cursor = idx_data_client
        .scan_all(
            schema_id_1,
            vec![],
            select_expr,
            Expr::nothing(),
            Ordering::Forward,
        )
        .await
        .unwrap();
        // Start from 10 to 100 due to the selection expression
        for i in 10..100 {
            let id = Id::new(1, i);
            let cell_res = match cursor.next().await {
                Ok(r) => r,
                Err(e) => {
                    panic!("Error next for {}, {:?}", i, e);
                }
            };
            let cell = match cell_res {
                Some(c) => c,
                None => {
                    panic!("Have none for {}", i);
                }
            };
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! ncurses-compatible database discovery
//!
//! Does not support hashed database, only filesystem!

use std::env;
use std::fs;
use std::path::PathBuf;

/// Return path to database entry for `term`
pub fn get_dbpath_for_term(term: &str) -> Option<PathBuf> {
    if term.len() == 0 {
        return None;
    }

    let mut dirs_to_search = Vec::new();
    let first_char = term.chars().next().unwrap();

    // Find search directory
    match env::var_os("TERMINFO") {
        Some(dir) => dirs_to_search.push(PathBuf::from(&dir)),
        None => {
            if let Some(mut homedir) = env::home_dir() {
                // ncurses compatibility;
                homedir.push(".terminfo");
                dirs_to_search.push(homedir)
            }
            match env::var("TERMINFO_DIRS") {
                Ok(dirs) => for i in dirs.split(':') {
                    if i == "" {
                        dirs_to_search.push(PathBuf::from("/usr/share/terminfo"));
                    } else {
                        dirs_to_search.push(PathBuf::from(i));
                    }
                },
                // Found nothing in TERMINFO_DIRS, use the default paths:
                // According to  /etc/terminfo/README, after looking at
                // ~/.terminfo, ncurses will search /etc/terminfo, then
                // /lib/terminfo, and eventually /usr/share/terminfo.
                Err(..) => {
                    dirs_to_search.push(PathBuf::from("/etc/terminfo"));
                    dirs_to_search.push(PathBuf::from("/lib/terminfo"));
                    dirs_to_search.push(PathBuf::from("/usr/share/terminfo"));
                }
            }
        }
    };

    // Look for the terminal in all of the search directories
    for mut p in dirs_to_search {
        if fs::metadata(&p).is_ok() {
            p.push(&first_char.to_string());
            p.push(&term);
            if fs::metadata(&p).is_ok() {
                return Some(p);
            }
            p.pop();
            p.pop();

            // on some installations the dir is named after the hex of the char
            // (e.g. OS X)
            p.push(&format!("{:x}", first_char as usize));
            p.push(term);
            if fs::metadata(&p).is_ok() {
                return Some(p);
            }
        }
    }
    None
}
<commit_msg>Avoid unnecessary allocation<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! ncurses-compatible database discovery
//!
//! Does not support hashed database, only filesystem!

use std::env;
use std::fs;
use std::path::PathBuf;

/// Return path to database entry for `term`
pub fn get_dbpath_for_term(term: &str) -> Option<PathBuf> {
    if term.len() == 0 {
        return None;
    }

    let mut dirs_to_search = Vec::new();
    let first_char = term.chars().next().unwrap();

    // Find search directory
    match env::var_os("TERMINFO") {
        Some(dir) => dirs_to_search.push(PathBuf::from(dir)),
        None => {
            if let Some(mut homedir) = env::home_dir() {
                // ncurses compatibility;
                homedir.push(".terminfo");
                dirs_to_search.push(homedir)
            }
            match env::var("TERMINFO_DIRS") {
                Ok(dirs) => for i in dirs.split(':') {
                    if i == "" {
                        dirs_to_search.push(PathBuf::from("/usr/share/terminfo"));
                    } else {
                        dirs_to_search.push(PathBuf::from(i));
                    }
                },
                // Found nothing in TERMINFO_DIRS, use the default paths:
                // According to  /etc/terminfo/README, after looking at
                // ~/.terminfo, ncurses will search /etc/terminfo, then
                // /lib/terminfo, and eventually /usr/share/terminfo.
                Err(..) => {
                    dirs_to_search.push(PathBuf::from("/etc/terminfo"));
                    dirs_to_search.push(PathBuf::from("/lib/terminfo"));
                    dirs_to_search.push(PathBuf::from("/usr/share/terminfo"));
                }
            }
        }
    };

    // Look for the terminal in all of the search directories
    for mut p in dirs_to_search {
        if fs::metadata(&p).is_ok() {
            p.push(&first_char.to_string());
            p.push(&term);
            if fs::metadata(&p).is_ok() {
                return Some(p);
            }
            p.pop();
            p.pop();

            // on some installations the dir is named after the hex of the char
            // (e.g. OS X)
            p.push(&format!("{:x}", first_char as usize));
            p.push(term);
            if fs::metadata(&p).is_ok() {
                return Some(p);
            }
        }
    }
    None
}
<|endoftext|>