<commit_before>use crate::{
    ebml::{self, custom_error, ebml_header, EBMLHeader},
    elements::{
        segment, segment_element, simple_block, Cluster, Info, SeekHead, SegmentElement,
        TrackEntry, Tracks,
    },
};
use av_data::{packet::Packet, params::*, rational::Rational64, timeinfo::TimeInfo};
use av_format::{
    buffer::Buffered,
    common::GlobalInfo,
    demuxer::{Demuxer, Descr, Descriptor, Event},
    error::*,
    stream::Stream,
};
use log::{debug, error, trace};
use nom::{self, Err, IResult, Needed, Offset};
use std::{collections::VecDeque, io::SeekFrom};

#[derive(Debug, Clone, Default)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
    pub params: Option<DemuxerParams>,
}

#[derive(Debug, Clone, Default)]
pub struct DemuxerParams {
    pub track_numbers: Option<Vec<u64>>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
            params: None,
        }
    }

    pub fn with_params(params: DemuxerParams) -> MkvDemuxer {
        MkvDemuxer {
            params: Some(params),
            ..Default::default()
        }
    }

    pub fn parse_until_tracks<'a>(
        &mut self,
        original_input: &'a [u8],
    ) -> IResult<&'a [u8], (), ebml::Error<'a>> {
        let (i1, header) = ebml_header(original_input)?;

        self.header = Some(header);

        let (mut input, _) = segment(i1)?;

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            // println!("offset: {}", original_input.offset(input));

            let (i3, element) = segment_element(input)?;

            match element {
                SegmentElement::SeekHead(s) => {
                    trace!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return Err(Err::Error(custom_error(input, 1)));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    trace!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return Err(Err::Error(custom_error(input, 1)));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    trace!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return Err(Err::Error(custom_error(input, 1)));
                    } else {
                        let mut t = t;

                        // Only keep tracks we're interested in
                        if let Some(params) = &self.params {
                            if let Some(track_numbers) = &params.track_numbers {
                                t.tracks = t
                                    .tracks
                                    .into_iter()
                                    .filter(|tr| track_numbers.contains(&tr.track_number))
                                    .collect();
                            }
                        };

                        self.tracks = Some(t);
                    }
                }
                el => {
                    debug!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<dyn Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self
                    .info
                    .as_ref()
                    .and_then(|info| info.duration)
                    .map(|d| d as u64);
                if let Some(ref mut t) = self.tracks {
                    for tr in t.tracks.iter_mut() {
                        info.add_stream(track_to_stream(self.info.as_ref().unwrap(), tr));
                        tr.stream_index = info.streams.last().unwrap().index;
                    }
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => buf.data().len() + usize::from(size),
                    _ => 1024,
                };
                Err(Error::MoreDataNeeded(sz))
            }
            e => {
                error!("{:?}", e);
                Err(Error::InvalidData)
            }
        }
    }

    fn read_event(&mut self, buf: &Box<dyn Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            // println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            debug!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets(self.tracks.as_ref().unwrap());
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        _el => {
                            // println!("got element: {:#?}", el);
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                }
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(usize::from(size)))
                }
                e => {
                    error!("{:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1" => Some("av1".to_owned()),
        "V_VP8" => Some("vp8".to_owned()),
        "V_VP9" => Some("vp9".to_owned()),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    use av_data::audiosample::*;
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let map = if audio.channel_positions.is_none() {
            Some(ChannelMap::default_map(audio.channels as usize))
        } else {
            unimplemented!("Convert matroska map to rust-av map")
        };
        // Allow for consistency reasons
        #[allow(clippy::redundant_field_names)]
        let a = AudioInfo {
            rate: rate as usize,
            map: map,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

// TODO: make sure the timecode_scale isn't 0
pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = if let Some(ts) = t.track_timecode_scale {
        (ts * info.timecode_scale as f64) as i64
    } else {
        info.timecode_scale as i64
    };

    Stream {
        id: t.track_uid as isize,
        index: t.stream_index,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None,
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self, tracks: &Tracks) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                debug!("parsing simple block: {:?}", block);
                if let Some(index) = tracks.lookup(block.track_number) {
                    let packet = Packet {
                        data: i.into(),
                        t: TimeInfo {
                            pts: Some(i64::from(block.timecode)),
                            dts: None,
                            duration: None,
                            timebase: None,
                            user_private: None,
                        },
                        pos: None,
                        stream_index: index as isize,
                        is_key: block.keyframe,
                        is_corrupted: false,
                    };

                    v.push(Event::NewPacket(packet));
                }
            } else {
                error!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<dyn Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe(&self) -> &Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &dyn Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    },
};

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
    use super::*;
    use av_format::{buffer::*, demuxer::Context};
    use log::info;
    use nom::Offset;
    use std::io::Cursor;

    const webm: &[u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        info!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                info!("offset: {} bytes", webm.offset(i));
            }
            e => {
                info!("could not parse: {:?}", e);
            }
        }

        info!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => info!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => info!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(
            Box::new(MkvDemuxer::new()),
            Box::new(AccReader::new(Cursor::new(webm))),
        );
        info!("read headers: {:?}", context.read_headers().unwrap());
        info!("streams: {:?}", context.info.streams);
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
    }
}
<commit_msg>Remove dead code from demuxer<commit_after>use crate::{
    ebml::{self, custom_error, ebml_header, EBMLHeader},
    elements::{
        segment, segment_element, simple_block, Cluster, Info, SeekHead, SegmentElement,
        TrackEntry, Tracks,
    },
};
use av_data::{packet::Packet, params::*, rational::Rational64, timeinfo::TimeInfo};
use av_format::{
    buffer::Buffered,
    common::GlobalInfo,
    demuxer::{Demuxer, Descr, Descriptor, Event},
    error::*,
    stream::Stream,
};
use log::{debug, error, trace};
use nom::{self, Err, IResult, Needed, Offset};
use std::{collections::VecDeque, io::SeekFrom};

#[derive(Debug, Clone, Default)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
    pub params: Option<DemuxerParams>,
}

#[derive(Debug, Clone, Default)]
pub struct DemuxerParams {
    pub track_numbers: Option<Vec<u64>>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
            params: None,
        }
    }

    pub fn with_params(params: DemuxerParams) -> MkvDemuxer {
        MkvDemuxer {
            params: Some(params),
            ..Default::default()
        }
    }

    pub fn parse_until_tracks<'a>(
        &mut self,
        original_input: &'a [u8],
    ) -> IResult<&'a [u8], (), ebml::Error<'a>> {
        let (i1, header) = ebml_header(original_input)?;

        self.header = Some(header);

        let (mut input, _) = segment(i1)?;

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            let (i3, element) = segment_element(input)?;

            match element {
                SegmentElement::SeekHead(s) => {
                    trace!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return Err(Err::Error(custom_error(input, 1)));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    trace!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return Err(Err::Error(custom_error(input, 1)));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    trace!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return Err(Err::Error(custom_error(input, 1)));
                    } else {
                        let mut t = t;

                        // Only keep tracks we're interested in
                        if let Some(params) = &self.params {
                            if let Some(track_numbers) = &params.track_numbers {
                                t.tracks = t
                                    .tracks
                                    .into_iter()
                                    .filter(|tr| track_numbers.contains(&tr.track_number))
                                    .collect();
                            }
                        };

                        self.tracks = Some(t);
                    }
                }
                el => {
                    debug!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<dyn Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self
                    .info
                    .as_ref()
                    .and_then(|info| info.duration)
                    .map(|d| d as u64);
                if let Some(ref mut t) = self.tracks {
                    for tr in t.tracks.iter_mut() {
                        info.add_stream(track_to_stream(self.info.as_ref().unwrap(), tr));
                        tr.stream_index = info.streams.last().unwrap().index;
                    }
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => buf.data().len() + usize::from(size),
                    _ => 1024,
                };
                Err(Error::MoreDataNeeded(sz))
            }
            e => {
                error!("{:?}", e);
                Err(Error::InvalidData)
            }
        }
    }

    fn read_event(&mut self, buf: &Box<dyn Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            debug!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets(self.tracks.as_ref().unwrap());
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        _el => {}
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                }
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(usize::from(size)))
                }
                e => {
                    error!("{:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1" => Some("av1".to_owned()),
        "V_VP8" => Some("vp8".to_owned()),
        "V_VP9" => Some("vp9".to_owned()),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    use av_data::audiosample::*;
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let map = if audio.channel_positions.is_none() {
            Some(ChannelMap::default_map(audio.channels as usize))
        } else {
            unimplemented!("Convert matroska map to rust-av map")
        };
        // Allow for consistency reasons
        #[allow(clippy::redundant_field_names)]
        let a = AudioInfo {
            rate: rate as usize,
            map: map,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

// TODO: make sure the timecode_scale isn't 0
pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = if let Some(ts) = t.track_timecode_scale {
        (ts * info.timecode_scale as f64) as i64
    } else {
        info.timecode_scale as i64
    };

    Stream {
        id: t.track_uid as isize,
        index: t.stream_index,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None,
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self, tracks: &Tracks) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                debug!("parsing simple block: {:?}", block);
                if let Some(index) = tracks.lookup(block.track_number) {
                    let packet = Packet {
                        data: i.into(),
                        t: TimeInfo {
                            pts: Some(i64::from(block.timecode)),
                            dts: None,
                            duration: None,
                            timebase: None,
                            user_private: None,
                        },
                        pos: None,
                        stream_index: index as isize,
                        is_key: block.keyframe,
                        is_corrupted: false,
                    };

                    v.push(Event::NewPacket(packet));
                }
            } else {
                error!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<dyn Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe(&self) -> &Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &dyn Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    },
};

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
    use super::*;
    use av_format::{buffer::*, demuxer::Context};
    use log::info;
    use nom::Offset;
    use std::io::Cursor;

    const webm: &[u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        info!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                info!("offset: {} bytes", webm.offset(i));
            }
            e => {
                info!("could not parse: {:?}", e);
            }
        }

        info!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => info!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => info!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(
            Box::new(MkvDemuxer::new()),
            Box::new(AccReader::new(Cursor::new(webm))),
        );
        info!("read headers: {:?}", context.read_headers().unwrap());
        info!("streams: {:?}", context.info.streams);
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::collections::BTreeMap;

use rustc_serialize::json::{self, Json};
use url::form_urlencoded;
use kite::document::DocRef;
use kite::query::Query;
use kite::collectors::top_score::TopScoreCollector;
use kite::collectors::total_count::TotalCountCollector;

use query_parser::{QueryParseContext, parse as parse_query};

use api::persistent;
use api::iron::prelude::*;
use api::iron::status;
use api::router::Router;
use api::utils::json_response;

pub fn view_count(req: &mut Request) -> IronResult<Response> {
    let ref system = get_system!(req);
    let ref index_name = read_path_parameter!(req, "index").unwrap_or("");

    // Lock index array
    let indices = system.indices.read().unwrap();

    // Get index
    let index = get_index_or_404!(indices, *index_name);
    let index_reader = index.store.reader();

    let count = match json_from_request_body!(req) {
        Some(query_json) => {
            // Parse query
            let query = parse_query(&QueryParseContext::new().set_mappings(&index.mappings).no_score(), query_json.as_object().unwrap().get("query").unwrap());
            debug!("{:#?}", query);

            match query {
                Ok(query) => {
                    let mut collector = TotalCountCollector::new();
                    index_reader.search(&mut collector, &query.build(&index_reader.schema())).unwrap();
                    collector.get_total_count()
                }
                Err(_) => {
                    // TODO: What specifically is bad about the Query?
                    let mut response = Response::with((status::BadRequest,
                                                       "{\"message\": \"Query error\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            }
        }
        None => {
            let mut collector = TotalCountCollector::new();
            index_reader.search(&mut collector, &Query::new_match_all()).unwrap();
            collector.get_total_count()
        }
    };

    return Ok(json_response(status::Ok, format!("{{\"count\": {}}}", count)));
}

pub fn view_search(req: &mut Request) -> IronResult<Response> {
    let ref system = get_system!(req);
    let ref index_name = read_path_parameter!(req, "index").unwrap_or("");

    // Lock index array
    let indices = system.indices.read().unwrap();

    // Get index
    let index = get_index_or_404!(indices, *index_name);
    let index_reader = index.store.reader();

    match json_from_request_body!(req) {
        Some(query_json) => {
            // Parse query
            let query = parse_query(&QueryParseContext::new().set_mappings(&index.mappings), query_json.as_object().unwrap().get("query").unwrap());
            debug!("{:#?}", query);

            match query {
                Ok(query) => {
                    let mut from = 0;
                    let mut size = 10;
                    let mut fields = Vec::new();

                    // TODO: Rewrite this
                    if let Some(ref url_query) = req.url.query() {
                        for (key, value) in form_urlencoded::parse(url_query.as_bytes()) {
                            match key.as_ref() {
                                "from" => {
                                    from = value.as_ref().parse().expect("need a number");
                                }
                                "size" => {
                                    size = value.as_ref().parse().expect("need a number");
                                }
                                "fields" => {
                                    for field_name in value.split(",") {
                                        let field_ref = match index_reader.schema().get_field_by_name(field_name) {
                                            Some(field_ref) => field_ref,
                                            None => {
                                                warn!("unknown field {:?}", field_name);
                                                continue;
                                            }
                                        };

                                        fields.push((field_name.to_owned(), field_ref));
                                    }
                                }
                                // terminate_after
                                // explain
                                // version
                                // timeout
                                // fielddata_fields
                                // track_scores
                                // stats
                                // suggest_field
                                _ => warn!("unrecognised GET parameter {:?}", key),
                            }
                        }
                    }

                    // Do the search
                    let mut collector = TopScoreCollector::new(from + size);
                    index_reader.search(&mut collector, &query.build(&index_reader.schema())).unwrap();

                    // TODO: {"took":5,"timed_out":false,"_shards":{"total":5,"successful":5,"failed":0},"hits":{"total":4,"max_score":1.0,"hits":[{"_index":"wagtail","_type":"searchtests_searchtest_searchtests_searchtestchild","_id":"searchtests_searchtest:5380","_score":1.0,"fields":{"pk":["5380"]}},{"_index":"wagtail","_type":"searchtests_searchtest","_id":"searchtests_searchtest:5379","_score":1.0,"fields":{"pk":["5379"]}}]}}
                    Ok(json_response(status::Ok,
                                     format!("{{\"hits\": {{\"total\": {}, \"hits\": {}}}}}",
                                             0, // TODO
                                             json::encode(&Json::Array(collector.into_sorted_vec().iter()
                                                .skip(from)
                                                .map(|doc_match| {
                                                    let mut field_values = BTreeMap::new();

                                                    for &(ref field_name, field_ref) in fields.iter() {
                                                        let value = match index_reader.read_stored_field(field_ref, DocRef::from_u64(doc_match.doc_id())) {
                                                            Ok(Some(value)) => {
                                                                value.as_json()
                                                            }
                                                            Ok(None) => Json::Array(vec![]),
                                                            Err(_) => Json::Array(vec![]),
                                                        };

                                                        field_values.insert(field_name.clone(), Json::Array(vec![value]));
                                                    }

                                                    let mut hit = BTreeMap::new();
                                                    hit.insert("_score".to_owned(), Json::F64(doc_match.score().unwrap()));
                                                    hit.insert("fields".to_owned(), Json::Object(field_values));
                                                    Json::Object(hit)
                                                })
                                                .collect())
                                            ).unwrap())))
                }
                Err(_) => {
                    // TODO: What specifically is bad about the Query?
                    let mut response = Response::with((status::BadRequest,
                                                       "{\"message\": \"Query error\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    Ok(response)
                }
            }
        }
        None => Ok(json_response(status::BadRequest, "{\"message\": \"Missing query\"}")),
    }
}
<commit_msg>Fixed output when stored field has no value<commit_after>use std::io::Read;
use std::collections::BTreeMap;

use rustc_serialize::json::{self, Json};
use url::form_urlencoded;
use kite::document::DocRef;
use kite::query::Query;
use kite::collectors::top_score::TopScoreCollector;
use kite::collectors::total_count::TotalCountCollector;

use query_parser::{QueryParseContext, parse as parse_query};

use api::persistent;
use api::iron::prelude::*;
use api::iron::status;
use api::router::Router;
use api::utils::json_response;

pub fn view_count(req: &mut Request) -> IronResult<Response> {
    let ref system = get_system!(req);
    let ref index_name = read_path_parameter!(req, "index").unwrap_or("");

    // Lock index array
    let indices = system.indices.read().unwrap();

    // Get index
    let index = get_index_or_404!(indices, *index_name);
    let index_reader = index.store.reader();

    let count = match json_from_request_body!(req) {
        Some(query_json) => {
            // Parse query
            let query = parse_query(&QueryParseContext::new().set_mappings(&index.mappings).no_score(), query_json.as_object().unwrap().get("query").unwrap());
            debug!("{:#?}", query);

            match query {
                Ok(query) => {
                    let mut collector = TotalCountCollector::new();
                    index_reader.search(&mut collector, &query.build(&index_reader.schema())).unwrap();
                    collector.get_total_count()
                }
                Err(_) => {
                    // TODO: What specifically is bad about the Query?
                    let mut response = Response::with((status::BadRequest,
                                                       "{\"message\": \"Query error\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            }
        }
        None => {
            let mut collector = TotalCountCollector::new();
            index_reader.search(&mut collector, &Query::new_match_all()).unwrap();
            collector.get_total_count()
        }
    };

    return Ok(json_response(status::Ok, format!("{{\"count\": {}}}", count)));
}

pub fn view_search(req: &mut Request) -> IronResult<Response> {
    let ref system = get_system!(req);
    let ref index_name = read_path_parameter!(req, "index").unwrap_or("");

    // Lock index array
    let indices = system.indices.read().unwrap();

    // Get index
    let index = get_index_or_404!(indices, *index_name);
    let index_reader = index.store.reader();

    match json_from_request_body!(req) {
        Some(query_json) => {
            // Parse query
            let query = parse_query(&QueryParseContext::new().set_mappings(&index.mappings), query_json.as_object().unwrap().get("query").unwrap());
            debug!("{:#?}", query);

            match query {
                Ok(query) => {
                    let mut from = 0;
                    let mut size = 10;
                    let mut fields = Vec::new();

                    // TODO: Rewrite this
                    if let Some(ref url_query) = req.url.query() {
                        for (key, value) in form_urlencoded::parse(url_query.as_bytes()) {
                            match key.as_ref() {
                                "from" => {
                                    from = value.as_ref().parse().expect("need a number");
                                }
                                "size" => {
                                    size = value.as_ref().parse().expect("need a number");
                                }
                                "fields" => {
                                    for field_name in value.split(",") {
                                        let field_ref = match index_reader.schema().get_field_by_name(field_name) {
                                            Some(field_ref) => field_ref,
                                            None => {
                                                warn!("unknown field {:?}", field_name);
                                                continue;
                                            }
                                        };

                                        fields.push((field_name.to_owned(), field_ref));
                                    }
                                }
                                // terminate_after
                                // explain
                                // version
                                // timeout
                                // fielddata_fields
                                // track_scores
                                // stats
                                // suggest_field
                                _ => warn!("unrecognised GET parameter {:?}", key),
                            }
                        }
                    }

                    // Do the search
                    let mut collector = TopScoreCollector::new(from + size);
                    index_reader.search(&mut collector, &query.build(&index_reader.schema())).unwrap();

                    // TODO: {"took":5,"timed_out":false,"_shards":{"total":5,"successful":5,"failed":0},"hits":{"total":4,"max_score":1.0,"hits":[{"_index":"wagtail","_type":"searchtests_searchtest_searchtests_searchtestchild","_id":"searchtests_searchtest:5380","_score":1.0,"fields":{"pk":["5380"]}},{"_index":"wagtail","_type":"searchtests_searchtest","_id":"searchtests_searchtest:5379","_score":1.0,"fields":{"pk":["5379"]}}]}}
                    Ok(json_response(status::Ok,
                                     format!("{{\"hits\": {{\"total\": {}, \"hits\": {}}}}}",
                                             0, // TODO
                                             json::encode(&Json::Array(collector.into_sorted_vec().iter()
                                                .skip(from)
                                                .map(|doc_match| {
                                                    let mut field_values = BTreeMap::new();

                                                    for &(ref field_name, field_ref) in fields.iter() {
                                                        let value = match index_reader.read_stored_field(field_ref, DocRef::from_u64(doc_match.doc_id())) {
                                                            Ok(Some(value)) => {
                                                                Json::Array(vec![value.as_json()])
                                                            }
                                                            Ok(None) => Json::Array(vec![]),
                                                            Err(_) => Json::Array(vec![]),
                                                        };

                                                        field_values.insert(field_name.clone(), value);
                                                    }

                                                    let mut hit = BTreeMap::new();
                                                    hit.insert("_score".to_owned(), Json::F64(doc_match.score().unwrap()));
                                                    hit.insert("fields".to_owned(), Json::Object(field_values));
                                                    Json::Object(hit)
                                                })
                                                .collect())
                                            ).unwrap())))
                }
                Err(_) => {
                    // TODO: What specifically is bad about the Query?
                    let mut response = Response::with((status::BadRequest,
                                                       "{\"message\": \"Query error\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    Ok(response)
                }
            }
        }
        None => Ok(json_response(status::BadRequest, "{\"message\": \"Missing query\"}")),
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

use gdk::{RGBA, SELECTION_PRIMARY};
use gtk;
use gtk::{
    BoxExt,
    Clipboard,
    CssProvider,
    CssProviderExt,
    EditableExt,
    EditableSignals,
    EntryExt,
    LabelExt,
    OrientableExt,
    PackType,
    StateFlags,
    StyleContextExt,
    WidgetExt,
    STYLE_PROVIDER_PRIORITY_APPLICATION,
};
use gtk::Orientation::Horizontal;
use pango::EllipsizeMode;
use relm::{Relm, Widget};
use relm_derive::widget;

use self::Msg::*;
use self::ItemMsg::{Color, Text};

#[derive(Msg)]
pub enum Msg {
    BarVisible(bool),
    Copy,
    Cut,
    DeleteNextChar,
    DeleteNextWord,
    DeletePreviousWord,
    End,
    EntryActivate(String),
    EntryChanged(String),
    EntryText(String),
    EntryShown(bool),
    Identifier(String),
    NextChar,
    NextWord,
    Paste,
    PasteSelection,
    PreviousChar,
    PreviousWord,
    ShowIdentifier,
    SmartHome,
}

pub struct Model {
    identifier_label: &'static str,
    identifier_visible: bool,
    relm: Relm<StatusBar>,
    visible: bool,
}

#[widget]
impl Widget for StatusBar {
    fn init_view(&mut self) {
        // Adjust the look of the entry.
        let style_context = self.command_entry.get_style_context();
        // TODO: remove the next line when relm supports css.
        let style = include_bytes!("../../style/command-input.css");
        let provider = CssProvider::new();
        provider.load_from_data(style).unwrap();
        style_context.add_provider(&provider, STYLE_PROVIDER_PRIORITY_APPLICATION);
    }

    fn model(relm: &Relm<Self>, _: ()) -> Model {
        Model {
            identifier_label: ":",
            identifier_visible: false,
            relm: relm.clone(),
            visible: true,
        }
    }

    /// Set whether the entry is visible or not.
    fn set_entry_shown(&mut self, visible: bool) {
        // TODO: document this use of lock.
        let _lock = self.model.relm.stream().lock();
        self.command_entry.set_text("");
        self.model.identifier_visible = visible;
        self.command_entry.set_visible(visible);

        if visible {
            self.command_entry.grab_focus();
        }
    }

    /// Show the identifier.
    fn show_identifier(&mut self) {
        self.model.identifier_visible = true;
    }

    fn update(&mut self, msg: Msg) {
        match msg {
            BarVisible(visible) => self.model.visible = visible,
            Copy => self.copy(),
            Cut => self.cut(),
            DeleteNextChar => self.delete_next_char(),
            DeleteNextWord => self.delete_next_word(),
            DeletePreviousWord => self.delete_previous_word(),
            End => self.end(),
            EntryActivate(_) | EntryChanged(_) => (), // NOTE: to be listened by the user.
            EntryShown(visible) => self.set_entry_shown(visible),
            EntryText(input) => self.set_input(&input),
            Identifier(identifier) => self.set_identifier(&identifier),
            NextChar => self.next_char(),
            NextWord => self.next_word(),
            Paste => self.paste(),
            PasteSelection => self.paste_selection(),
            PreviousChar => self.previous_char(),
            PreviousWord => self.previous_word(),
            ShowIdentifier => self.show_identifier(),
            SmartHome => self.smart_home(),
        }
    }

    view! {
        #[container]
        gtk::Box {
            property_height_request: 20, // TODO: is this still useful?
            orientation: Horizontal,
            visible: self.model.visible,
            #[name="identifier_label"]
            gtk::Label {
                text: self.model.identifier_label,
                visible: self.model.identifier_visible,
            },
            #[name="command_entry"]
            gtk::Entry {
                activate(entry) => EntryActivate(entry.get_text().to_string()),
                changed(entry) => EntryChanged(entry.get_text().to_string()),
                has_frame: false,
                hexpand: true,
                widget_name: "mg-input-command",
            },
        }
    }
}

impl StatusBar {
    /// Copy the selection to the clipboard.
    fn copy(&self) {
        self.command_entry.copy_clipboard();
    }

    /// Cut the selection to the clipboard.
    fn cut(&self) {
        self.command_entry.cut_clipboard();
    }

    /// Delete the character after the cursor.
    fn delete_next_char(&self) {
        if !self.delete_selection() {
            let text = self.get_command();
            if !text.is_empty() {
                let pos = self.command_entry.get_position();
                let len = text.chars().count();
                if pos < len as i32 {
                    // NOTE: Lock to avoid moving the cursor when updating the text entry.
                    let _lock = self.model.relm.stream().lock();
                    self.command_entry.delete_text(pos, pos + 1);
                }
            }
        }
        self.emit_entry_changed();
    }

    /// Delete the word after the cursor.
    fn delete_next_word(&self) {
        if !self.delete_selection() {
            let text = self.get_command();
            if !text.is_empty() {
                let pos = self.command_entry.get_position();
                let end = text.chars().enumerate()
                    .skip(pos as usize)
                    .skip_while(|&(_, c)| !c.is_alphanumeric())
                    .skip_while(|&(_, c)| c.is_alphanumeric())
                    .map(|(index, _)| index)
                    .next()
                    .unwrap_or_else(|| text.chars().count());
                // NOTE: Lock to avoid moving the cursor when updating the text entry.
                let _lock = self.model.relm.stream().lock();
                self.command_entry.delete_text(pos, end as i32);
            }
        }
        self.emit_entry_changed();
    }

    /// Delete the word before the cursor.
    fn delete_previous_word(&self) {
        if !self.delete_selection() {
            let text = self.get_command();
            if !text.is_empty() {
                let pos = self.command_entry.get_position();
                let len = text.chars().count();
                let start = text.chars().rev().enumerate()
                    .skip(len - pos as usize)
                    .skip_while(|&(_, c)| !c.is_alphanumeric())
                    .skip_while(|&(_, c)| c.is_alphanumeric())
                    .map(|(index, _)| len - index)
                    .next()
                    .unwrap_or_default();
                // NOTE: Lock to avoid moving the cursor when updating the text entry.
                let _lock = self.model.relm.stream().lock();
                self.command_entry.delete_text(start as i32, pos);
            }
        }
        self.emit_entry_changed();
    }

    /// Delete the selected text.
    fn delete_selection(&self) -> bool {
        if self.command_entry.get_selection_bounds().is_some() {
            // NOTE: Lock to avoid moving the cursor when updating the text entry.
            let _lock = self.model.relm.stream().lock();
            self.command_entry.delete_selection();
            true
        }
        else {
            false
        }
    }

    /// Emit the EntryChanged event with the current text.
    fn emit_entry_changed(&self) {
        self.model.relm.stream().emit(EntryChanged(self.command_entry.get_text().to_string()));
    }

    /// Go to the end of the command entry.
    fn end(&self) {
        let text = self.get_command();
        self.command_entry.set_position(text.chars().count() as i32);
    }

    /// Get the text of the command entry.
    fn get_command(&self) -> String {
        self.command_entry.get_text().to_string()
    }

    /// Go forward one character in the command entry.
    fn next_char(&self) {
        let pos = self.command_entry.get_position();
        let text = self.get_command();
        if pos < text.chars().count() as i32 {
            self.command_entry.set_position(pos + 1);
        }
    }

    /// Go forward one word in the command entry.
    fn next_word(&self) {
        let pos = self.command_entry.get_position();
        let text = self.get_command();
        let position = text.chars().enumerate()
            .skip(pos as usize)
            .skip_while(|&(_, c)| !c.is_alphanumeric())
            .skip_while(|&(_, c)| c.is_alphanumeric())
            .next()
            .map(|(index, _)| index)
            .unwrap_or_else(|| text.chars().count());
        self.command_entry.set_position(position as i32);
    }

    /// Paste the clipboard to the selection or the cursor position.
    fn paste(&self) {
        self.command_entry.paste_clipboard();
    }

    /// Paste the selection clipboard to the selection or the cursor position.
    fn paste_selection(&self) {
        self.delete_selection();
        let clipboard = Clipboard::get(&SELECTION_PRIMARY);
        if let Some(text) = clipboard.wait_for_text() {
            let mut position = self.command_entry.get_position();
            self.command_entry.insert_text(&text, &mut position);
            self.command_entry.set_position(position);
        }
    }

    /// Go back one character in the command entry.
    fn previous_char(&self) {
        let pos = self.command_entry.get_position();
        if pos > 0 {
            self.command_entry.set_position(pos - 1);
        }
    }

    /// Go back one word in the command entry.
    fn previous_word(&self) {
        let pos = self.command_entry.get_position();
        let text = self.get_command();
        let len = text.chars().count();
        let position = text.chars().rev().enumerate()
            .skip(len - pos as usize)
            .skip_while(|&(_, c)| !c.is_alphanumeric())
            .skip_while(|&(_, c)| c.is_alphanumeric())
            .next()
            .map(|(index, _)| len - index)
            .unwrap_or_default();
        self.command_entry.set_position(position as i32);
    }

    /// Seth the prefix identifier shown at the left of the command entry.
    fn set_identifier(&self, identifier: &str) {
        self.identifier_label.set_text(identifier);
    }

    /// Set the text of the input entry and move the cursor at the end.
    fn set_input(&self, command: &str) {
        // NOTE: Prevent updating the completions when the user selects a completion entry.
        let _lock = self.model.relm.stream().lock();
        self.command_entry.set_text(command);
        self.command_entry.set_position(command.chars().count() as i32);
    }

    /// Go to the beginning of the command entry.
    /// If the cursor is already at the beginning, go after the spaces after the command name.
    fn smart_home(&self) {
        let pos = self.command_entry.get_position();
        if pos == 0 {
            let text = self.get_command();
            let position = text.chars().enumerate()
                .skip_while(|&(_, c)| c.is_whitespace())
                .skip_while(|&(_, c)| !c.is_whitespace())
                .skip_while(|&(_, c)| c.is_whitespace())
                .map(|(index, _)| index)
                .next()
                .unwrap_or_default();
            self.command_entry.set_position(position as i32);
        }
        else {
            self.command_entry.set_position(0);
        }
    }
}

#[derive(Msg)]
pub enum ItemMsg {
    // Set the color of the status bar item.
    Color(Option<RGBA>),
    // Set the text of the status bar item.
    Text(String),
}

/// A status bar text item.
#[widget]
impl Widget for StatusBarItem {
    fn model() -> () {
        ()
    }

    fn update(&mut self, msg: ItemMsg) {
        match msg {
            Color(color) => self.label.override_color(StateFlags::NORMAL, color.as_ref()),
            Text(text) => self.label.set_text(&text),
        }
    }

    view! {
        #[parent="status-bar-item"]
        #[name="label"]
        gtk::Label {
            ellipsize: EllipsizeMode::End,
            child: {
                pack_type: PackType::End,
                padding: 3,
            },
        }
    }
}
<commit_msg>Restore doc comments<commit_after>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

use gdk::{RGBA, SELECTION_PRIMARY};
use gtk;
use gtk::{
    BoxExt,
    Clipboard,
    CssProvider,
    CssProviderExt,
    EditableExt,
    EditableSignals,
    EntryExt,
    LabelExt,
    OrientableExt,
    PackType,
    StateFlags,
    StyleContextExt,
    WidgetExt,
    STYLE_PROVIDER_PRIORITY_APPLICATION,
};
use gtk::Orientation::Horizontal;
use pango::EllipsizeMode;
use relm::{Relm, Widget};
use relm_derive::widget;

use self::Msg::*;
use self::ItemMsg::{Color, Text};

#[derive(Msg)]
pub enum Msg {
    BarVisible(bool),
    Copy,
    Cut,
    DeleteNextChar,
    DeleteNextWord,
    DeletePreviousWord,
    End,
    EntryActivate(String),
    EntryChanged(String),
    EntryText(String),
    EntryShown(bool),
    Identifier(String),
    NextChar,
    NextWord,
    Paste,
    PasteSelection,
    PreviousChar,
    PreviousWord,
    ShowIdentifier,
    SmartHome,
}

pub struct Model {
    identifier_label: &'static str,
    identifier_visible: bool,
    relm: Relm<StatusBar>,
    visible: bool,
}

#[widget]
impl Widget for StatusBar {
    fn init_view(&mut self) {
        // Adjust the look of the entry.
        let style_context = self.command_entry.get_style_context();
        // TODO: remove the next line when relm supports css.
        let style = include_bytes!("../../style/command-input.css");
        let provider = CssProvider::new();
        provider.load_from_data(style).unwrap();
        style_context.add_provider(&provider, STYLE_PROVIDER_PRIORITY_APPLICATION);
    }

    fn model(relm: &Relm<Self>, _: ()) -> Model {
        Model {
            identifier_label: ":",
            identifier_visible: false,
            relm: relm.clone(),
            visible: true,
        }
    }

    /// Set whether the entry is visible or not.
    fn set_entry_shown(&mut self, visible: bool) {
        // TODO: document this use of lock.
        let _lock = self.model.relm.stream().lock();
        self.command_entry.set_text("");
        self.model.identifier_visible = visible;
        self.command_entry.set_visible(visible);

        if visible {
            self.command_entry.grab_focus();
        }
    }

    /// Show the identifier.
    fn show_identifier(&mut self) {
        self.model.identifier_visible = true;
    }

    fn update(&mut self, msg: Msg) {
        match msg {
            BarVisible(visible) => self.model.visible = visible,
            Copy => self.copy(),
            Cut => self.cut(),
            DeleteNextChar => self.delete_next_char(),
            DeleteNextWord => self.delete_next_word(),
            DeletePreviousWord => self.delete_previous_word(),
            End => self.end(),
            EntryActivate(_) | EntryChanged(_) => (), // NOTE: to be listened by the user.
            EntryShown(visible) => self.set_entry_shown(visible),
            EntryText(input) => self.set_input(&input),
            Identifier(identifier) => self.set_identifier(&identifier),
            NextChar => self.next_char(),
            NextWord => self.next_word(),
            Paste => self.paste(),
            PasteSelection => self.paste_selection(),
            PreviousChar => self.previous_char(),
            PreviousWord => self.previous_word(),
            ShowIdentifier => self.show_identifier(),
            SmartHome => self.smart_home(),
        }
    }

    view! {
        #[container]
        gtk::Box {
            property_height_request: 20, // TODO: is this still useful?
            orientation: Horizontal,
            visible: self.model.visible,
            #[name="identifier_label"]
            gtk::Label {
                text: self.model.identifier_label,
                visible: self.model.identifier_visible,
            },
            #[name="command_entry"]
            gtk::Entry {
                activate(entry) => EntryActivate(entry.get_text().to_string()),
                changed(entry) => EntryChanged(entry.get_text().to_string()),
                has_frame: false,
                hexpand: true,
                widget_name: "mg-input-command",
            },
        }
    }
}

impl StatusBar {
    /// Copy the selection to the clipboard.
    fn copy(&self) {
        self.command_entry.copy_clipboard();
    }

    /// Cut the selection to the clipboard.
    fn cut(&self) {
        self.command_entry.cut_clipboard();
    }

    /// Delete the character after the cursor.
    fn delete_next_char(&self) {
        if !self.delete_selection() {
            let text = self.get_command();
            if !text.is_empty() {
                let pos = self.command_entry.get_position();
                let len = text.chars().count();
                if pos < len as i32 {
                    // NOTE: Lock to avoid moving the cursor when updating the text entry.
                    let _lock = self.model.relm.stream().lock();
                    self.command_entry.delete_text(pos, pos + 1);
                }
            }
        }
        self.emit_entry_changed();
    }

    /// Delete the word after the cursor.
    fn delete_next_word(&self) {
        if !self.delete_selection() {
            let text = self.get_command();
            if !text.is_empty() {
                let pos = self.command_entry.get_position();
                let end = text.chars().enumerate()
                    .skip(pos as usize)
                    .skip_while(|&(_, c)| !c.is_alphanumeric())
                    .skip_while(|&(_, c)| c.is_alphanumeric())
                    .map(|(index, _)| index)
                    .next()
                    .unwrap_or_else(|| text.chars().count());
                // NOTE: Lock to avoid moving the cursor when updating the text entry.
                let _lock = self.model.relm.stream().lock();
                self.command_entry.delete_text(pos, end as i32);
            }
        }
        self.emit_entry_changed();
    }

    /// Delete the word before the cursor.
    fn delete_previous_word(&self) {
        if !self.delete_selection() {
            let text = self.get_command();
            if !text.is_empty() {
                let pos = self.command_entry.get_position();
                let len = text.chars().count();
                let start = text.chars().rev().enumerate()
                    .skip(len - pos as usize)
                    .skip_while(|&(_, c)| !c.is_alphanumeric())
                    .skip_while(|&(_, c)| c.is_alphanumeric())
                    .map(|(index, _)| len - index)
                    .next()
                    .unwrap_or_default();
                // NOTE: Lock to avoid moving the cursor when updating the text entry.
                let _lock = self.model.relm.stream().lock();
                self.command_entry.delete_text(start as i32, pos);
            }
        }
        self.emit_entry_changed();
    }

    /// Delete the selected text.
    fn delete_selection(&self) -> bool {
        if self.command_entry.get_selection_bounds().is_some() {
            // NOTE: Lock to avoid moving the cursor when updating the text entry.
            let _lock = self.model.relm.stream().lock();
            self.command_entry.delete_selection();
            true
        }
        else {
            false
        }
    }

    /// Emit the EntryChanged event with the current text.
    fn emit_entry_changed(&self) {
        self.model.relm.stream().emit(EntryChanged(self.command_entry.get_text().to_string()));
    }

    /// Go to the end of the command entry.
    fn end(&self) {
        let text = self.get_command();
        self.command_entry.set_position(text.chars().count() as i32);
    }

    /// Get the text of the command entry.
    fn get_command(&self) -> String {
        self.command_entry.get_text().to_string()
    }

    /// Go forward one character in the command entry.
    fn next_char(&self) {
        let pos = self.command_entry.get_position();
        let text = self.get_command();
        if pos < text.chars().count() as i32 {
            self.command_entry.set_position(pos + 1);
        }
    }

    /// Go forward one word in the command entry.
    fn next_word(&self) {
        let pos = self.command_entry.get_position();
        let text = self.get_command();
        let position = text.chars().enumerate()
            .skip(pos as usize)
            .skip_while(|&(_, c)| !c.is_alphanumeric())
            .skip_while(|&(_, c)| c.is_alphanumeric())
            .next()
            .map(|(index, _)| index)
            .unwrap_or_else(|| text.chars().count());
        self.command_entry.set_position(position as i32);
    }

    /// Paste the clipboard to the selection or the cursor position.
    fn paste(&self) {
        self.command_entry.paste_clipboard();
    }

    /// Paste the selection clipboard to the selection or the cursor position.
    fn paste_selection(&self) {
        self.delete_selection();
        let clipboard = Clipboard::get(&SELECTION_PRIMARY);
        if let Some(text) = clipboard.wait_for_text() {
            let mut position = self.command_entry.get_position();
            self.command_entry.insert_text(&text, &mut position);
            self.command_entry.set_position(position);
        }
    }

    /// Go back one character in the command entry.
    fn previous_char(&self) {
        let pos = self.command_entry.get_position();
        if pos > 0 {
            self.command_entry.set_position(pos - 1);
        }
    }

    /// Go back one word in the command entry.
    fn previous_word(&self) {
        let pos = self.command_entry.get_position();
        let text = self.get_command();
        let len = text.chars().count();
        let position = text.chars().rev().enumerate()
            .skip(len - pos as usize)
            .skip_while(|&(_, c)| !c.is_alphanumeric())
            .skip_while(|&(_, c)| c.is_alphanumeric())
            .next()
            .map(|(index, _)| len - index)
            .unwrap_or_default();
        self.command_entry.set_position(position as i32);
    }

    /// Seth the prefix identifier shown at the left of the command entry.
    fn set_identifier(&self, identifier: &str) {
        self.identifier_label.set_text(identifier);
    }

    /// Set the text of the input entry and move the cursor at the end.
    fn set_input(&self, command: &str) {
        // NOTE: Prevent updating the completions when the user selects a completion entry.
        let _lock = self.model.relm.stream().lock();
        self.command_entry.set_text(command);
        self.command_entry.set_position(command.chars().count() as i32);
    }

    /// Go to the beginning of the command entry.
    /// If the cursor is already at the beginning, go after the spaces after the command name.
    fn smart_home(&self) {
        let pos = self.command_entry.get_position();
        if pos == 0 {
            let text = self.get_command();
            let position = text.chars().enumerate()
                .skip_while(|&(_, c)| c.is_whitespace())
                .skip_while(|&(_, c)| !c.is_whitespace())
                .skip_while(|&(_, c)| c.is_whitespace())
                .map(|(index, _)| index)
                .next()
                .unwrap_or_default();
            self.command_entry.set_position(position as i32);
        }
        else {
            self.command_entry.set_position(0);
        }
    }
}
#[derive(Msg)]
#[allow(unused_doc_comments)]
pub enum ItemMsg {
    /// Set the color of the status bar item.
    Color(Option<RGBA>),
    /// Set the text of the status bar item.
    Text(String),
}

/// A status bar text item.
#[widget]
impl Widget for StatusBarItem {
    fn model() -> () {
        ()
    }

    fn update(&mut self, msg: ItemMsg) {
        match msg {
            Color(color) => self.label.override_color(StateFlags::NORMAL, color.as_ref()),
            Text(text) => self.label.set_text(&text),
        }
    }

    view! {
        #[parent="status-bar-item"]
        #[name="label"]
        gtk::Label {
            ellipsize: EllipsizeMode::End,
            child: {
                pack_type: PackType::End,
                padding: 3,
            },
        }
    }
}
<|endoftext|>


<commit_before>use std::iter::Peekable;

#[derive(PartialEq, Debug)]
pub enum TokenType {
    NOT,
    EQ,
    LT,
    LE,

    ASSIGN,

    PLUS,
    MINUS,
    ASTER,
    SLASH,

    Comma,
    SemiColon,

    LParenth,
    RParenth,
    LBracket,
    RBracket,
    LBrace,
    RBrace,

    KwVoid,
    KwBool,
    KwInt,
    KwFloat,

    KwTrue,
    KwFalse,

    KwIf,
    KwElse,
    KwWhile,
    KwReturn,

    Identifier(String),

    INum(i32),
    FNum(f32),

    Str(String),
}

pub struct Lexer<I: Iterator<Item = char>> {
    input: Peekable<I>,
    line: i32,
    col: i32,
}

impl<I: Iterator<Item = char>> Lexer<I> {
    fn new(input: I) -> Lexer<I> {
        Lexer {
            input: input.peekable(),
            line: 0,
            col: 0,
        }
    }

    fn lex_word(&mut self, c: char) -> TokenType {
        let mut ident = self.input
            .by_ref()
            .take_while(|&c| c.is_alphabetic() || c.is_numeric() || c == '_')
            .collect::<String>();
        ident.insert(0, c);

        match ident.as_ref() {
            // handle type keywords
            "void" => TokenType::KwVoid,
            "bool" => TokenType::KwBool,
            "int" => TokenType::KwInt,
            "float" => TokenType::KwFloat,

            // handle bool literals
            "true" => TokenType::KwTrue,
            "false" => TokenType::KwFalse,

            // handle keywords
            "if" => TokenType::KwIf,
            "else" => TokenType::KwElse,
            "while" => TokenType::KwWhile,
            "return" => TokenType::KwReturn,

            // handle identifier
            _ => TokenType::Identifier(ident),
        }
    }

    fn lex_number(&mut self, c: char) -> TokenType {
        let mut num = self.input
            .by_ref()
            .take_while(|&c| c.is_numeric() || c == '.')
            .collect::<String>();
        num.insert(0, c);

        if let Ok(n) = num.parse::<i32>() {
            TokenType::INum(n)
        } else if let Ok(n) = num.parse::<f32>() {
            TokenType::FNum(n)
        } else {
            panic!("wrong number format: {}", num);
        }
    }

    fn lex_string(&mut self) -> TokenType {
        TokenType::Str(self.input.by_ref().take_while(|&c| c != '\"').collect())
    }
}

impl<I: Iterator<Item = char>> Iterator for Lexer<I> {
    type Item = TokenType;

    fn next(&mut self) -> Option<TokenType> {
        loop {
            // handle EOF
            if self.input.peek().is_none() {
                return None;
            }

            return Some(match self.input.next().unwrap() {
                // handle WS
                ' ' | '\t' | '\n' | '\r' => continue,

                // handle Syntax
                ',' => TokenType::Comma,
                ';' => TokenType::SemiColon,
                '(' => TokenType::LParenth,
                ')' => TokenType::RParenth,
                '[' => TokenType::LBracket,
                ']' => TokenType::RBracket,
                '{' => TokenType::LBrace,
                '}' => TokenType::RBrace,

                // handle Ops (simple)
                '!' => TokenType::NOT,
                '+' => TokenType::PLUS,
                '-' => TokenType::MINUS,
                '*' => TokenType::ASTER,
                '/' => TokenType::SLASH,

                // handle Ops (complex)
                '=' => {
                    if let Some(&'=') = self.input.peek() {
                        self.input.next();
                        TokenType::EQ
                    } else {
                        TokenType::ASSIGN
                    }
                }
                '<' => {
                    if let Some(&'=') = self.input.peek() {
                        self.input.next();
                        TokenType::LE
                    } else {
                        TokenType::LT
                    }
                }

                // handle keywords, types and identifier
                c @ 'A'...'z' => self.lex_word(c),

                // handle number
                c @ '0'...'9' => self.lex_number(c),

                // handle strings
                '\"' => self.lex_string(),

                x @ _ => panic!("unexpected input: `{}`", x),
            });
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn syntax() {
        let mut l = Lexer::new("! == < <= = + - * / , ; ( [ { } ] )".chars());
        assert_eq!(l.next(), Some(TokenType::NOT));
        assert_eq!(l.next(), Some(TokenType::EQ));
        assert_eq!(l.next(), Some(TokenType::LT));
        assert_eq!(l.next(), Some(TokenType::LE));
        assert_eq!(l.next(), Some(TokenType::ASSIGN));
        assert_eq!(l.next(), Some(TokenType::PLUS));
        assert_eq!(l.next(), Some(TokenType::MINUS));
        assert_eq!(l.next(), Some(TokenType::ASTER));
        assert_eq!(l.next(), Some(TokenType::SLASH));
        assert_eq!(l.next(), Some(TokenType::Comma));
        assert_eq!(l.next(), Some(TokenType::SemiColon));
        assert_eq!(l.next(), Some(TokenType::LParenth));
        assert_eq!(l.next(), Some(TokenType::LBracket));
        assert_eq!(l.next(), Some(TokenType::LBrace));
        assert_eq!(l.next(), Some(TokenType::RBrace));
        assert_eq!(l.next(), Some(TokenType::RBracket));
        assert_eq!(l.next(), Some(TokenType::RParenth));
        assert_eq!(l.next(), None);
    }

    #[test]
    fn keywords() {
        let mut l = Lexer::new("void bool int float true false if else while return".chars());
        assert_eq!(l.next(), Some(TokenType::KwVoid));
        assert_eq!(l.next(), Some(TokenType::KwBool));
        assert_eq!(l.next(), Some(TokenType::KwInt));
        assert_eq!(l.next(), Some(TokenType::KwFloat));
        assert_eq!(l.next(), Some(TokenType::KwTrue));
        assert_eq!(l.next(), Some(TokenType::KwFalse));
        assert_eq!(l.next(), Some(TokenType::KwIf));
        assert_eq!(l.next(), Some(TokenType::KwElse));
        assert_eq!(l.next(), Some(TokenType::KwWhile));
        assert_eq!(l.next(), Some(TokenType::KwReturn));
        assert_eq!(l.next(), None);
    }

    #[test]
    fn identifiers() {
        let mut l = Lexer::new("foo b4r ba_".chars());
        assert_eq!(l.next(), Some(TokenType::Identifier("foo".to_string())));
        assert_eq!(l.next(), Some(TokenType::Identifier("b4r".to_string())));
        assert_eq!(l.next(), Some(TokenType::Identifier("ba_".to_string())));
        assert_eq!(l.next(), None);
    }

    #[test]
    fn numbers() {
        let mut l = Lexer::new("42 3.14".chars());
        assert_eq!(l.next(), Some(TokenType::INum(42)));
        assert_eq!(l.next(), Some(TokenType::FNum(3.14)));
        assert_eq!(l.next(), None);
    }

    #[test]
    fn strings() {
        let mut l = Lexer::new("\"foo\" \"bar\"".chars());
        assert_eq!(l.next(), Some(TokenType::Str("foo".to_string())));
        assert_eq!(l.next(), Some(TokenType::Str("bar".to_string())));
        assert_eq!(l.next(), None);
    }
}
<commit_msg>Front: Improve lexer whitespace handling and position tracking<commit_after>use std::iter::Peekable;

use front::Position;

#[derive(PartialEq, Debug)]
pub enum TokenType {
    NOT,
    EQ,
    LT,
    LE,

    ASSIGN,

    PLUS,
    MINUS,
    ASTER,
    SLASH,

    Comma,
    SemiColon,

    LParenth,
    RParenth,
    LBracket,
    RBracket,
    LBrace,
    RBrace,

    KwVoid,
    KwBool,
    KwInt,
    KwFloat,

    KwTrue,
    KwFalse,

    KwIf,
    KwElse,
    KwWhile,
    KwReturn,

    Identifier(String),

    INum(i32),
    FNum(f32),

    Str(String),
}

#[derive(PartialEq, Debug)]
pub struct Token {
    tok: TokenType,
    pos: Position,
}

pub struct Lexer<I: Iterator<Item = char>> {
    input: Peekable<I>,
    pos: Position,
}

impl<I: Iterator<Item = char>> Lexer<I> {
    fn new(input: I) -> Lexer<I> {
        Lexer {
            input: input.peekable(),
            pos: Position { line: 1, col: 1 },
        }
    }

    fn get_char(&mut self) -> Option<char> {
        self.input.next().map(|c| {
            self.update_position(c);
            c
        })
    }

    fn get_char_if<P>(&mut self, predicate: P) -> Option<char>
        where P: Fn(&char) -> bool
    {
        if self.input.peek().map(predicate).unwrap_or(false) {
            self.get_char()
        } else {
            None
        }
    }

    fn update_position(&mut self, c: char) {
        match c {
            '\n' => {
                self.pos.line += 1;
                self.pos.col = 1;
            }
            '\t' => self.pos.col += 7,
            _ => self.pos.col += 1,
        }
    }

    fn lex_word(&mut self, c: char) -> TokenType {
        let mut s = c.to_string();
        while let Some(c) = self.get_char_if(|&c| c.is_alphabetic() || c.is_numeric() || c == '_') {
            s.push(c);
        }

        match s.as_ref() {
            // handle type keywords
            "void" => TokenType::KwVoid,
            "bool" => TokenType::KwBool,
            "int" => TokenType::KwInt,
            "float" => TokenType::KwFloat,

            // handle bool literals
            "true" => TokenType::KwTrue,
            "false" => TokenType::KwFalse,

            // handle keywords
            "if" => TokenType::KwIf,
            "else" => TokenType::KwElse,
            "while" => TokenType::KwWhile,
            "return" => TokenType::KwReturn,

            // handle identifier
            _ => TokenType::Identifier(s),
        }
    }

    fn lex_number(&mut self, c: char) -> TokenType {
        let mut num = c.to_string();
        while let Some(n) = self.get_char_if(|&c| c.is_numeric() || c == '.') {
            num.push(n);
        }

        if let Ok(n) = num.parse::<i32>() {
            TokenType::INum(n)
        } else if let Ok(n) = num.parse::<f32>() {
            TokenType::FNum(n)
        } else {
            panic!("wrong number format: {}", num);
        }
    }

    fn lex_string(&mut self) -> TokenType {
        let mut s = "".to_string();
        while let Some(c) = self.get_char_if(|&c| c != '\"') {
            s.push(c);
        }

        // pop off \"
        if self.input.peek().is_some() {
            self.get_char();
        }

        TokenType::Str(s)
    }
}

impl<I: Iterator<Item = char>> Iterator for Lexer<I> {
    type Item = Token;

    fn next(&mut self) -> Option<Token> {
        loop {
            // handle EOF
            if self.input.peek().is_none() {
                return None;
            }

            let pos = self.pos;

            let tok = match self.get_char().unwrap() {
                ' ' | '\t' | '\n' | '\r' => continue,

                // handle Syntax
                ',' => TokenType::Comma,
                ';' => TokenType::SemiColon,
                '(' => TokenType::LParenth,
                ')' => TokenType::RParenth,
                '[' => TokenType::LBracket,
                ']' => TokenType::RBracket,
                '{' => TokenType::LBrace,
                '}' => TokenType::RBrace,

                // handle Ops (simple)
                '!' => TokenType::NOT,
                '+' => TokenType::PLUS,
                '-' => TokenType::MINUS,
                '*' => TokenType::ASTER,
                '/' => TokenType::SLASH,

                // handle Ops (complex)
                '=' => {
                    if self.get_char_if(|&c| c == '=').is_some() {
                        TokenType::EQ
                    } else {
                        TokenType::ASSIGN
                    }
                }
                '<' => {
                    if self.get_char_if(|&c| c == '=').is_some() {
                        TokenType::LE
                    } else {
                        TokenType::LT
                    }
                }

                // handle keywords, types and identifier
                c @ 'A'...'z' => self.lex_word(c),

                // handle number
                c @ '0'...'9' => self.lex_number(c),

                // handle strings
                '\"' => self.lex_string(),

                x @ _ => panic!("unexpected input: `{}`", x),
            };

            return Some(Token {
                tok: tok,
                pos: pos,
            });
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn syntax() {
        let mut l = Lexer::new("! == < <= = + - * / , ; ([{}])".chars());
        assert_eq!(TokenType::NOT, l.next().unwrap().tok);
        assert_eq!(TokenType::EQ, l.next().unwrap().tok);
        assert_eq!(TokenType::LT, l.next().unwrap().tok);
        assert_eq!(TokenType::LE, l.next().unwrap().tok);
        assert_eq!(TokenType::ASSIGN, l.next().unwrap().tok);
        assert_eq!(TokenType::PLUS, l.next().unwrap().tok);
        assert_eq!(TokenType::MINUS, l.next().unwrap().tok);
        assert_eq!(TokenType::ASTER, l.next().unwrap().tok);
        assert_eq!(TokenType::SLASH, l.next().unwrap().tok);
        assert_eq!(TokenType::Comma, l.next().unwrap().tok);
        assert_eq!(TokenType::SemiColon, l.next().unwrap().tok);
        assert_eq!(TokenType::LParenth, l.next().unwrap().tok);
        assert_eq!(TokenType::LBracket, l.next().unwrap().tok);
        assert_eq!(TokenType::LBrace, l.next().unwrap().tok);
        assert_eq!(TokenType::RBrace, l.next().unwrap().tok);
        assert_eq!(TokenType::RBracket, l.next().unwrap().tok);
        assert_eq!(TokenType::RParenth, l.next().unwrap().tok);
        assert_eq!(None, l.next());
    }

    #[test]
    fn keywords() {
        let mut l = Lexer::new("void bool int float true false if else while,return".chars());
        assert_eq!(TokenType::KwVoid, l.next().unwrap().tok);
        assert_eq!(TokenType::KwBool, l.next().unwrap().tok);
        assert_eq!(TokenType::KwInt, l.next().unwrap().tok);
        assert_eq!(TokenType::KwFloat, l.next().unwrap().tok);
        assert_eq!(TokenType::KwTrue, l.next().unwrap().tok);
        assert_eq!(TokenType::KwFalse, l.next().unwrap().tok);
        assert_eq!(TokenType::KwIf, l.next().unwrap().tok);
        assert_eq!(TokenType::KwElse, l.next().unwrap().tok);
        assert_eq!(TokenType::KwWhile, l.next().unwrap().tok);
        assert_eq!(TokenType::Comma, l.next().unwrap().tok);
        assert_eq!(TokenType::KwReturn, l.next().unwrap().tok);
        assert_eq!(None, l.next());
    }

    #[test]
    fn identifiers() {
        let mut l = Lexer::new("foo b4r ba_".chars());
        assert_eq!(TokenType::Identifier("foo".to_string()),
                   l.next().unwrap().tok);
        assert_eq!(TokenType::Identifier("b4r".to_string()),
                   l.next().unwrap().tok);
        assert_eq!(TokenType::Identifier("ba_".to_string()),
                   l.next().unwrap().tok);
        assert_eq!(None, l.next());
    }

    #[test]
    fn numbers() {
        let mut l = Lexer::new("42 3.14".chars());
        assert_eq!(TokenType::INum(42), l.next().unwrap().tok);
        assert_eq!(TokenType::FNum(3.14), l.next().unwrap().tok);
        assert_eq!(None, l.next());
    }

    #[test]
    fn strings() {
        let mut l = Lexer::new("\"foo\" \"bar\"".chars());
        assert_eq!(TokenType::Str("foo".to_string()), l.next().unwrap().tok);
        assert_eq!(TokenType::Str("bar".to_string()), l.next().unwrap().tok);
        assert_eq!(None, l.next());
    }

    #[test]
    fn position() {
        // operator
        let mut l = Lexer::new("+ == -".chars());
        assert_eq!(1, l.next().unwrap().pos.col);
        assert_eq!(3, l.next().unwrap().pos.col);
        assert_eq!(6, l.next().unwrap().pos.col);

        // number
        l = Lexer::new("42 3.14".chars());
        assert_eq!(1, l.next().unwrap().pos.col);
        assert_eq!(4, l.next().unwrap().pos.col);

        // keyword
        l = Lexer::new("if while true".chars());
        assert_eq!(1, l.next().unwrap().pos.col);
        assert_eq!(4, l.next().unwrap().pos.col);
        assert_eq!(10, l.next().unwrap().pos.col);

        // strings
        l = Lexer::new("\"foo\" \"bar\"".chars());
        assert_eq!(1, l.next().unwrap().pos.col);
        assert_eq!(7, l.next().unwrap().pos.col);

        // lines
        l = Lexer::new("foo\n\"bar\nbaz\"\nbas".chars());
        assert_eq!((1, 1), l.next().unwrap().pos.to_pair());
        assert_eq!((2, 1), l.next().unwrap().pos.to_pair());
        assert_eq!((4, 1), l.next().unwrap().pos.to_pair());
    }
}
<|endoftext|>


<commit_before>pub mod header;
pub mod program_header;
pub mod dyn;
pub mod rela;
pub mod sym;
pub mod strtab;
pub mod gnu_hash;

use std::path::Path;
use std::fs::File;
use std::io;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom::Start;

#[derive(Debug)]
pub struct Elf {
    pub header: header::Header,
    pub program_headers: Vec<program_header::ProgramHeader>,
    pub dynamic: Option<Vec<dyn::Dyn>>,
    pub symtab: Vec<sym::Sym>,
    pub rela: Vec<rela::Rela>,
    pub pltrela: Vec<rela::Rela>,
    pub soname: Option<String>,
    pub interpreter: Option<String>,
    pub libraries: Vec<String>,
    pub is_lib: bool,
    pub size: usize,
    pub entry: usize,
}

impl Elf {
    pub fn from_path<'a>(path: &Path) -> io::Result<Elf> {
        let mut fd = try!(File::open(&path));
        let metadata = fd.metadata().unwrap();
        if metadata.len() < header::EHDR_SIZE as u64 {
            let error = io::Error::new(io::ErrorKind::Other,
                                       format!("Error: {:?} size is smaller than an ELF header",
                                               path.as_os_str()));
            Err(error)
        } else {
            let mut elf_header = [0; header::EHDR_SIZE];
            try!(fd.read(&mut elf_header));

            let header = header::Header::from_bytes(&elf_header);
            let entry = header.e_entry as usize;
            let is_lib = header.e_type == header::ET_DYN;

            let mut bytes = vec![0u8; ((header.e_phnum * header.e_phentsize) as usize)];
            try!(fd.seek(Start(header.e_phoff)));
            try!(fd.read(&mut bytes));

            let program_headers =
                program_header::ProgramHeader::from_bytes(bytes, header.e_phnum as usize);

            let dynamic = try!(dyn::from_fd(&mut fd, &program_headers));

            let mut bias: usize = 0;
            for ph in &program_headers {
                if ph.p_type == program_header::PT_LOAD {
                    bias = ((::std::u64::MAX - ph.p_vaddr).wrapping_add(1)) as usize; // my name's David
                    break;
                }
            }

            let mut interpreter = None;
            for ph in &program_headers {
                if ph.p_type == program_header::PT_INTERP {
                    let mut bytes = vec![0u8; (ph.p_filesz - 1) as usize];
                    try!(fd.seek(Start(ph.p_offset)));
                    try!(fd.read(&mut bytes));
                    interpreter = Some(String::from_utf8(bytes).unwrap())
                }
            }

            //            let mut strtab = None;
            //            let mut symtab = None;
            let mut soname = None;
            let mut libraries = vec![];
            let mut symtab = vec![];
            let mut rela = vec![];
            let mut pltrela = vec![];
            if let Some(ref dynamic) = dynamic {
                let link_info = dyn::LinkInfo::new(&dynamic, bias); // we explicitly overflow the values here with our bias
                let strtab = try!(strtab::Strtab::from_fd(&mut fd,
                                                          link_info.strtab,
                                                          link_info.strsz));
                let soname_ = strtab.get(link_info.soname).to_owned();
                if soname_ != "" {
                    soname = Some(soname_)
                }
                let needed = dyn::get_needed(dynamic, &strtab, link_info.needed_count);
                libraries = Vec::with_capacity(link_info.needed_count);
                for lib in needed {
                    libraries.push(lib.to_owned());
                }

                let num_syms = (link_info.strtab - link_info.symtab) / link_info.syment; // old caveat about how this is probably not safe but rdr has been doing it with tons of binaries and never any problems
                symtab = try!(sym::from_fd(&mut fd, link_info.symtab, num_syms));

                rela = try!(rela::from_fd(&mut fd, link_info.rela, link_info.relasz));
                pltrela = try!(rela::from_fd(&mut fd, link_info.jmprel, link_info.pltrelsz));

            }

            let elf = Elf {
                header: header,
                program_headers: program_headers,
                dynamic: dynamic,
                symtab: symtab,
                rela: rela,
                pltrela: pltrela,
                soname: soname,
                interpreter: interpreter,
                libraries: libraries,
                is_lib: is_lib,
                size: metadata.len() as usize,
                entry: entry,
            };

            Ok(elf)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;
    use elf;
    #[test]
    fn read_ls() {
        let _ = elf::Elf::from_path(Path::new("/bin/ls"));
    }
}
<commit_msg>style: removed commented code<commit_after>pub mod header;
pub mod program_header;
pub mod dyn;
pub mod rela;
pub mod sym;
pub mod strtab;
pub mod gnu_hash;

use std::path::Path;
use std::fs::File;
use std::io;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom::Start;

#[derive(Debug)]
pub struct Elf {
    pub header: header::Header,
    pub program_headers: Vec<program_header::ProgramHeader>,
    pub dynamic: Option<Vec<dyn::Dyn>>,
    pub symtab: Vec<sym::Sym>,
    pub rela: Vec<rela::Rela>,
    pub pltrela: Vec<rela::Rela>,
    pub soname: Option<String>,
    pub interpreter: Option<String>,
    pub libraries: Vec<String>,
    pub is_lib: bool,
    pub size: usize,
    pub entry: usize,
}

impl Elf {
    pub fn from_path<'a>(path: &Path) -> io::Result<Elf> {
        let mut fd = try!(File::open(&path));
        let metadata = fd.metadata().unwrap();
        if metadata.len() < header::EHDR_SIZE as u64 {
            let error = io::Error::new(io::ErrorKind::Other,
                                       format!("Error: {:?} size is smaller than an ELF header",
                                               path.as_os_str()));
            Err(error)
        } else {
            let mut elf_header = [0; header::EHDR_SIZE];
            try!(fd.read(&mut elf_header));

            let header = header::Header::from_bytes(&elf_header);
            let entry = header.e_entry as usize;
            let is_lib = header.e_type == header::ET_DYN;

            let mut bytes = vec![0u8; ((header.e_phnum * header.e_phentsize) as usize)];
            try!(fd.seek(Start(header.e_phoff)));
            try!(fd.read(&mut bytes));

            let program_headers =
                program_header::ProgramHeader::from_bytes(bytes, header.e_phnum as usize);

            let dynamic = try!(dyn::from_fd(&mut fd, &program_headers));

            let mut bias: usize = 0;
            for ph in &program_headers {
                if ph.p_type == program_header::PT_LOAD {
                    bias = ((::std::u64::MAX - ph.p_vaddr).wrapping_add(1)) as usize; // my name's David
                    break;
                }
            }

            let mut interpreter = None;
            for ph in &program_headers {
                if ph.p_type == program_header::PT_INTERP {
                    let mut bytes = vec![0u8; (ph.p_filesz - 1) as usize];
                    try!(fd.seek(Start(ph.p_offset)));
                    try!(fd.read(&mut bytes));
                    interpreter = Some(String::from_utf8(bytes).unwrap())
                }
            }

            let mut soname = None;
            let mut libraries = vec![];
            let mut symtab = vec![];
            let mut rela = vec![];
            let mut pltrela = vec![];
            if let Some(ref dynamic) = dynamic {
                let link_info = dyn::LinkInfo::new(&dynamic, bias); // we explicitly overflow the values here with our bias
                let strtab = try!(strtab::Strtab::from_fd(&mut fd,
                                                          link_info.strtab,
                                                          link_info.strsz));
                let soname_ = strtab.get(link_info.soname).to_owned();
                if soname_ != "" {
                    soname = Some(soname_)
                }
                let needed = dyn::get_needed(dynamic, &strtab, link_info.needed_count);
                libraries = Vec::with_capacity(link_info.needed_count);
                for lib in needed {
                    libraries.push(lib.to_owned());
                }

                let num_syms = (link_info.strtab - link_info.symtab) / link_info.syment; // old caveat about how this is probably not safe but rdr has been doing it with tons of binaries and never any problems
                symtab = try!(sym::from_fd(&mut fd, link_info.symtab, num_syms));

                rela = try!(rela::from_fd(&mut fd, link_info.rela, link_info.relasz));
                pltrela = try!(rela::from_fd(&mut fd, link_info.jmprel, link_info.pltrelsz));

            }

            let elf = Elf {
                header: header,
                program_headers: program_headers,
                dynamic: dynamic,
                symtab: symtab,
                rela: rela,
                pltrela: pltrela,
                soname: soname,
                interpreter: interpreter,
                libraries: libraries,
                is_lib: is_lib,
                size: metadata.len() as usize,
                entry: entry,
            };

            Ok(elf)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;
    use elf;
    #[test]
    fn read_ls() {
        let _ = elf::Elf::from_path(Path::new("/bin/ls"));
    }
}
<|endoftext|>