<commit_before>#[macro_use]
extern crate glium;

use glium::Surface;
use glium::glutin;

mod support;

#[derive(Copy, Clone, Debug)]
struct PerInstance {
    pub id: u32,
    pub w_position: (f32, f32, f32),
    pub color: (f32, f32, f32),
}
implement_vertex!(PerInstance, id, w_position, color);

fn main() {
    use glium::DisplayBuild;

    // building the display, ie. the main object
    let display = glutin::WindowBuilder::new()
        .with_depth_buffer(24)
        .build_glium()
        .unwrap();

    // building the vertex and index buffers
    let vertex_buffer = support::load_wavefront(&display, include_bytes!("support/teapot.obj"));

    // the program
    let program = program!(&display,
        140 => {
            vertex: "
                #version 140

                uniform mat4 persp_matrix;
                uniform mat4 view_matrix;

                in uint id;
                in vec3 w_position;
                in vec3 color;
                in vec3 position;
                in vec3 normal;
                out vec3 v_normal;
                out vec3 v_color;

                void main() {
                    v_normal = normal;
                    v_color = color;
                    gl_Position = persp_matrix * view_matrix * vec4(position * 0.005 + w_position, 1.0);
                }
            ",

            fragment: "
                #version 140

                in vec3 v_normal;
                in vec3 v_color;
                out vec4 f_color;

                const vec3 LIGHT = vec3(-0.2, 0.8, 0.1);

                void main() {
                    float lum = max(dot(normalize(v_normal), normalize(LIGHT)), 0.0);
                    vec3 color = (0.3 + 0.7 * lum) * v_color;
                    f_color = vec4(color, 1.0);
                }
            ",
        },
    ).unwrap();

    let picking_program = program!(&display,
        140 => {
            vertex: "
                #version 140

                uniform mat4 persp_matrix;
                uniform mat4 view_matrix;

                in uint id;
                in vec3 w_position;
                in vec3 color;
                in vec3 position;
                in vec3 normal;
                flat out uint v_id;

                void main() {
                    v_id = id;
                    gl_Position = persp_matrix * view_matrix * vec4(position * 0.005 + w_position, 1.0);
                }
            ",

            fragment: "
                #version 140

                flat in uint v_id;
                out uint f_id;

                void main() {
                    f_id = v_id;
                }
            ",
        },
    ).unwrap();

    //
    let mut camera = support::camera::CameraState::new();
    camera.set_position((0.0, 0.0, 1.5));
    camera.set_direction((0.0, 0.0, 1.0));

    let mut per_instance = vec![
        PerInstance { id: 1, w_position: (-1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
        PerInstance { id: 2, w_position: ( 0.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
        PerInstance { id: 3, w_position: ( 1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
    ];

    let mut picking_texture: Option<glium::texture::UnsignedTexture2d> = None;
    let mut picking_pbo: glium::texture::pixel_buffer::PixelBuffer<u32>
        = glium::texture::pixel_buffer::PixelBuffer::new_empty(&display, 1);

    let mut cursor_position: Option<(i32, i32)> = None;

    // the main loop
    support::start_loop(|| {
        camera.update();

        // determing which object has been picked at the previous frame
        let picked_object = {
            let data = picking_pbo.read().map(|d| d[0]).unwrap_or(0);
            if data != 0 {
                Some(data - 1)//for this to work the id of a PerInstance has to be the index of it +1
            } else {
                None
            }
        };

        if let Some(index) = picked_object {
            per_instance[index as usize] = PerInstance { id: index + 1, w_position: per_instance[index as usize].w_position, color: (0.0, 1.0, 0.0) };
        } else {
            per_instance = vec![
                PerInstance { id: 1, w_position: (-1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
                PerInstance { id: 2, w_position: ( 0.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
                PerInstance { id: 3, w_position: ( 1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
            ];
        }

        // drawing the models and pass the picking texture
        // building the uniforms
        let uniforms = uniform! {
            persp_matrix: camera.get_perspective(),
            view_matrix: camera.get_view(),
        };

        // draw parameters
        let params = glium::DrawParameters {
            depth: glium::Depth {
                test: glium::DepthTest::IfLess,
                write: true,
                .. Default::default()
            },
            .. Default::default()
        };

        let per_instance_buffer = glium::vertex::VertexBuffer::new(&display, &per_instance).unwrap();

        // drawing a frame
        let mut target = display.draw();
        target.clear_color_and_depth((0.0, 0.0, 0.0, 0.0), 1.0);

        //update picking texture
        if picking_texture.is_none()
        || (picking_texture.as_ref().unwrap().get_width(), picking_texture.as_ref().unwrap().get_height().unwrap()) != target.get_dimensions() {
            println!("new pick_tex: {:?}", target.get_dimensions());
            picking_texture = Some(glium::texture::UnsignedTexture2d::empty_with_format(
                &display,
                glium::texture::UncompressedUintFormat::U32,
                glium::texture::MipmapsOption::NoMipmap,
                target.get_dimensions().0, target.get_dimensions().1
            ).unwrap())
        }

        if let Some(picking_tex) = picking_texture.as_ref() {
            let depth_buffer = glium::framebuffer::DepthRenderBuffer::new(
                &display,
                glium::texture::DepthFormat::F32,
                picking_texture.as_ref().unwrap().get_width(),
                picking_texture.as_ref().unwrap().get_height().unwrap()
            ).unwrap();
            let mut picking_target = glium::framebuffer::SimpleFrameBuffer::with_depth_buffer(&display, picking_tex, &depth_buffer).unwrap();
            picking_target.clear_color_and_depth((0.0, 0.0, 0.0, 1.0), 1.0);
            picking_target.draw((&vertex_buffer, per_instance_buffer.per_instance().unwrap()),
                        &glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList),
                        &picking_program, &uniforms, &params).unwrap();
        }
        target.draw((&vertex_buffer, per_instance_buffer.per_instance().unwrap()),
                    &glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList),
                    &program, &uniforms, &params).unwrap();
        target.finish().unwrap();

        // committing into the picking pbo
        if let (Some(cursor), Some(ref picking_texture)) = (cursor_position, picking_texture.as_ref()) {
            let read_target = glium::Rect {
                //left: ((cursor.0 as f32 + 1.0) * 0.5 * picking_texture.get_width() as f32) as u32,
                //bottom: ((cursor.1 as f32 + 1.0) * 0.5 * picking_texture.get_height().unwrap() as f32) as u32,
                left: cursor.0 as u32,
                bottom: cursor.1 as u32,
                width: 1,
                height: 1,
            };

            if read_target.left >= 0 && read_target.left < picking_texture.get_width() &&
               read_target.bottom >= 0 && read_target.bottom < picking_texture.get_height().unwrap() {
                //println!("writing to pixel_buffer");
                picking_texture.main_level()
                    .first_layer()
                    .into_image(None).unwrap()
                    .raw_read_to_pixel_buffer(&read_target, &picking_pbo);
            } else {
                picking_pbo.write(&[0]);
            }
        } else {
            picking_pbo.write(&[0]);
        }

        // polling and handling the events received by the window
        for event in display.poll_events() {
            match event {
                glutin::Event::Closed => return support::Action::Stop,
                glutin::Event::MouseMoved(m) => cursor_position = Some(m),
                ev => camera.process_input(&ev),
            }
        }

        support::Action::Continue
    });
}
<commit_msg>fixed inverted y-axis<commit_after>#[macro_use]
extern crate glium;

use glium::Surface;
use glium::glutin;

mod support;

#[derive(Copy, Clone, Debug)]
struct PerInstance {
    pub id: u32,
    pub w_position: (f32, f32, f32),
    pub color: (f32, f32, f32),
}
implement_vertex!(PerInstance, id, w_position, color);

fn main() {
    use glium::DisplayBuild;

    // building the display, ie. the main object
    let display = glutin::WindowBuilder::new()
        .with_depth_buffer(24)
        .build_glium()
        .unwrap();

    // building the vertex and index buffers
    let vertex_buffer = support::load_wavefront(&display, include_bytes!("support/teapot.obj"));

    // the program
    let program = program!(&display,
        140 => {
            vertex: "
                #version 140

                uniform mat4 persp_matrix;
                uniform mat4 view_matrix;

                in uint id;
                in vec3 w_position;
                in vec3 color;
                in vec3 position;
                in vec3 normal;
                out vec3 v_normal;
                out vec3 v_color;

                void main() {
                    v_normal = normal;
                    v_color = color;
                    gl_Position = persp_matrix * view_matrix * vec4(position * 0.005 + w_position, 1.0);
                }
            ",

            fragment: "
                #version 140

                in vec3 v_normal;
                in vec3 v_color;
                out vec4 f_color;

                const vec3 LIGHT = vec3(-0.2, 0.8, 0.1);

                void main() {
                    float lum = max(dot(normalize(v_normal), normalize(LIGHT)), 0.0);
                    vec3 color = (0.3 + 0.7 * lum) * v_color;
                    f_color = vec4(color, 1.0);
                }
            ",
        },
    ).unwrap();

    let picking_program = program!(&display,
        140 => {
            vertex: "
                #version 140

                uniform mat4 persp_matrix;
                uniform mat4 view_matrix;

                in uint id;
                in vec3 w_position;
                in vec3 color;
                in vec3 position;
                in vec3 normal;
                flat out uint v_id;

                void main() {
                    v_id = id;
                    gl_Position = persp_matrix * view_matrix * vec4(position * 0.005 + w_position, 1.0);
                }
            ",

            fragment: "
                #version 140

                flat in uint v_id;
                out uint f_id;

                void main() {
                    f_id = v_id;
                }
            ",
        },
    ).unwrap();

    //
    let mut camera = support::camera::CameraState::new();
    camera.set_position((0.0, 0.0, 1.5));
    camera.set_direction((0.0, 0.0, 1.0));

    let mut per_instance = vec![
        PerInstance { id: 1, w_position: (-1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
        PerInstance { id: 2, w_position: ( 0.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
        PerInstance { id: 3, w_position: ( 1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
    ];

    let mut picking_texture: Option<glium::texture::UnsignedTexture2d> = None;
    let mut picking_pbo: glium::texture::pixel_buffer::PixelBuffer<u32>
        = glium::texture::pixel_buffer::PixelBuffer::new_empty(&display, 1);

    let mut cursor_position: Option<(i32, i32)> = None;

    // the main loop
    support::start_loop(|| {
        camera.update();

        // determing which object has been picked at the previous frame
        let picked_object = {
            let data = picking_pbo.read().map(|d| d[0]).unwrap_or(0);
            if data != 0 {
                Some(data - 1)//for this to work the id of a PerInstance has to be the index of it +1
            } else {
                None
            }
        };

        if let Some(index) = picked_object {
            per_instance[index as usize] = PerInstance { id: index + 1, w_position: per_instance[index as usize].w_position, color: (0.0, 1.0, 0.0) };
        } else {
            per_instance = vec![
                PerInstance { id: 1, w_position: (-1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
                PerInstance { id: 2, w_position: ( 0.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
                PerInstance { id: 3, w_position: ( 1.0, 0.0, 0.0), color: (1.0, 0.0, 0.0)},
            ];
        }

        // building the uniforms
        let uniforms = uniform! {
            persp_matrix: camera.get_perspective(),
            view_matrix: camera.get_view(),
        };

        // draw parameters
        let params = glium::DrawParameters {
            depth: glium::Depth {
                test: glium::DepthTest::IfLess,
                write: true,
                .. Default::default()
            },
            .. Default::default()
        };

        let per_instance_buffer = glium::vertex::VertexBuffer::new(&display, &per_instance).unwrap();

        // drawing a frame
        let mut target = display.draw();
        target.clear_color_and_depth((0.0, 0.0, 0.0, 0.0), 1.0);

        //update picking texture
        if picking_texture.is_none()
        || (picking_texture.as_ref().unwrap().get_width(), picking_texture.as_ref().unwrap().get_height().unwrap()) != target.get_dimensions() {
            println!("new pick_tex: {:?}", target.get_dimensions());
            picking_texture = Some(glium::texture::UnsignedTexture2d::empty_with_format(
                &display,
                glium::texture::UncompressedUintFormat::U32,
                glium::texture::MipmapsOption::NoMipmap,
                target.get_dimensions().0, target.get_dimensions().1
            ).unwrap())
        }

        // drawing the models and pass the picking texture
        if let Some(picking_tex) = picking_texture.as_ref() {
            let depth_buffer = glium::framebuffer::DepthRenderBuffer::new(
                &display,
                glium::texture::DepthFormat::F32,
                picking_texture.as_ref().unwrap().get_width(),
                picking_texture.as_ref().unwrap().get_height().unwrap()
            ).unwrap();
            let mut picking_target = glium::framebuffer::SimpleFrameBuffer::with_depth_buffer(&display, picking_tex, &depth_buffer).unwrap();
            picking_target.clear_color_and_depth((0.0, 0.0, 0.0, 1.0), 1.0);
            picking_target.draw((&vertex_buffer, per_instance_buffer.per_instance().unwrap()),
                        &glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList),
                        &picking_program, &uniforms, &params).unwrap();
        }
        target.draw((&vertex_buffer, per_instance_buffer.per_instance().unwrap()),
                    &glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList),
                    &program, &uniforms, &params).unwrap();
        target.finish().unwrap();

        // committing into the picking pbo
        if let (Some(cursor), Some(ref picking_texture)) = (cursor_position, picking_texture.as_ref()) {
            let read_target = glium::Rect {
                //left: ((cursor.0 as f32 + 1.0) * 0.5 * picking_texture.get_width() as f32) as u32,
                //bottom: ((cursor.1 as f32 + 1.0) * 0.5 * picking_texture.get_height().unwrap() as f32) as u32,
                left: cursor.0 as u32,
                bottom: picking_texture.get_height().unwrap() - cursor.1 as u32,
                width: 1,
                height: 1,
            };

            if read_target.left >= 0 && read_target.left < picking_texture.get_width() &&
               read_target.bottom >= 0 && read_target.bottom < picking_texture.get_height().unwrap() {
                //println!("writing to pixel_buffer");
                picking_texture.main_level()
                    .first_layer()
                    .into_image(None).unwrap()
                    .raw_read_to_pixel_buffer(&read_target, &picking_pbo);
            } else {
                picking_pbo.write(&[0]);
            }
        } else {
            picking_pbo.write(&[0]);
        }

        // polling and handling the events received by the window
        for event in display.poll_events() {
            match event {
                glutin::Event::Closed => return support::Action::Stop,
                glutin::Event::MouseMoved(m) => cursor_position = Some(m),
                ev => camera.process_input(&ev),
            }
        }

        support::Action::Continue
    });
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate wlroots;

use std::cell::{Cell, RefCell};
use std::rc::Rc;
use wlroots::compositor::Compositor;
use wlroots::cursor::{Cursor, XCursor, XCursorTheme};
use wlroots::device::Device;
use wlroots::key_event::KeyEvent;
use wlroots::manager::{InputManagerHandler, OutputManagerHandler};
use wlroots::output::{Output, OutputLayout};
use wlroots::pointer;
use wlroots::wlroots_sys::gl;
use wlroots::wlroots_sys::wlr_button_state::{WLR_BUTTON_PRESSED, WLR_BUTTON_RELEASED};
use wlroots::xkbcommon::xkb::keysyms::KEY_Escape;

struct OutputHandler {
    color: Rc<Cell<[f32; 4]>>,
    cursor: Rc<RefCell<Cursor>>
}

struct InputHandler {
    color: Rc<Cell<[f32; 4]>>,
    default_color: [f32; 4],
    cursor: Rc<RefCell<Cursor>>
}

impl OutputManagerHandler for OutputHandler {
    fn output_added(&mut self, mut output: Output) {
        let mut cursor = self.cursor.borrow_mut();
        {
            let xcursor = cursor.xcursor().expect("XCursor was not set!");
            let image = &xcursor.images()[0];
            // TODO use output config if present instead of auto
            let mut layout = cursor
                .output_layout()
                .as_ref()
                .expect("Could not get output layout");
            layout.borrow_mut().add_auto(&mut output);
            if output.set_cursor(image).is_err() {
                wlr_log!(L_DEBUG, "Failed to set hardware cursor");
                return;
            }
        }
        let (x, y) = cursor.coords();
        // https://en.wikipedia.org/wiki/Mouse_warping
        cursor.warp(None, x, y);
    }
}

impl InputManagerHandler for InputHandler {
    fn button(&mut self, event: pointer::ButtonEvent) {
        if event.state() == WLR_BUTTON_RELEASED {
            self.color.set(self.default_color.clone())
        } else {
            let mut red: [f32; 4] = [0.25, 0.25, 0.25, 1.0];
            red[event.button() as usize % 3] = 1.0;
            self.color.set(red);
        }
    }
}

fn managers(mut cursor: Cursor) -> (OutputHandler, InputHandler) {
    let mut layout = Rc::new(RefCell::new(OutputLayout::new()));
    // TODO Ensure this can be safe...
    // e.g what's stopping me from simply dropping layout now that I gave it to
    // cursor?
    cursor.attach_output_layout(layout);
    let cursor = Rc::new(RefCell::new(cursor));
    let color = Rc::new(Cell::new([0.25, 0.25, 0.25, 1.0]));
    (OutputHandler {
         color: color.clone(),
         cursor: cursor.clone()
     },
     InputHandler {
         color: color.clone(),
         default_color: color.get(),
         cursor: cursor.clone()
     })
}

fn main() {
    let mut cursor = Cursor::new().expect("Could not create cursor");
    let xcursor_theme = XCursorTheme::load_theme(None, 16).expect("Could not load theme");
    let mut xcursor = xcursor_theme
        .get_cursor("left_ptr".into())
        .expect("Could not load cursor from theme");
    cursor.set_xcursor(Some(xcursor));

    let (mut output_manager, input_manager) = managers(cursor);
    let compositor = Compositor::new(Box::new(input_manager), Box::new(output_manager));
    compositor.run();
}
<commit_msg>Draw colors when you press the mouse buttons<commit_after>#[macro_use]
extern crate wlroots;

use std::cell::{Cell, RefCell};
use std::rc::Rc;
use wlroots::compositor::Compositor;
use wlroots::cursor::{Cursor, XCursor, XCursorTheme};
use wlroots::device::Device;
use wlroots::key_event::KeyEvent;
use wlroots::manager::{InputManagerHandler, OutputManagerHandler};
use wlroots::output::{Output, OutputLayout};
use wlroots::pointer;
use wlroots::wlroots_sys::gl;
use wlroots::wlroots_sys::wlr_button_state::{WLR_BUTTON_PRESSED, WLR_BUTTON_RELEASED};
use wlroots::xkbcommon::xkb::keysyms::KEY_Escape;

struct OutputHandler {
    color: Rc<Cell<[f32; 4]>>,
    cursor: Rc<RefCell<Cursor>>
}

struct InputHandler {
    color: Rc<Cell<[f32; 4]>>,
    default_color: [f32; 4],
    cursor: Rc<RefCell<Cursor>>
}

impl OutputManagerHandler for OutputHandler {
    fn output_added(&mut self, mut output: Output) {
        let mut cursor = self.cursor.borrow_mut();
        {
            let xcursor = cursor.xcursor().expect("XCursor was not set!");
            let image = &xcursor.images()[0];
            // TODO use output config if present instead of auto
            let mut layout = cursor
                .output_layout()
                .as_ref()
                .expect("Could not get output layout");
            layout.borrow_mut().add_auto(&mut output);
            if output.set_cursor(image).is_err() {
                wlr_log!(L_DEBUG, "Failed to set hardware cursor");
                return;
            }
        }
        let (x, y) = cursor.coords();
        // https://en.wikipedia.org/wiki/Mouse_warping
        cursor.warp(None, x, y);
    }

    fn output_frame(&mut self, output: Output) {
        output.make_current();
        unsafe {
            gl::ClearColor(self.color.get()[0], self.color.get()[1], self.color.get()[2], 1.0);
            gl::Clear(gl::COLOR_BUFFER_BIT);
        }
        output.swap_buffers();
    }
}

impl InputManagerHandler for InputHandler {
    fn button(&mut self, event: pointer::ButtonEvent) {
        if event.state() == WLR_BUTTON_RELEASED {
            self.color.set(self.default_color.clone())
        } else {
            let mut red: [f32; 4] = [0.25, 0.25, 0.25, 1.0];
            red[event.button() as usize % 3] = 1.0;
            self.color.set(red);
        }
    }
}

fn managers(mut cursor: Cursor) -> (OutputHandler, InputHandler) {
    let mut layout = Rc::new(RefCell::new(OutputLayout::new()));
    // TODO Ensure this can be safe...
    // e.g what's stopping me from simply dropping layout now that I gave it to
    // cursor?
    cursor.attach_output_layout(layout);
    let cursor = Rc::new(RefCell::new(cursor));
    let color = Rc::new(Cell::new([0.25, 0.25, 0.25, 1.0]));
    (OutputHandler {
         color: color.clone(),
         cursor: cursor.clone()
     },
     InputHandler {
         color: color.clone(),
         default_color: color.get(),
         cursor: cursor.clone()
     })
}

fn main() {
    let mut cursor = Cursor::new().expect("Could not create cursor");
    let xcursor_theme = XCursorTheme::load_theme(None, 16).expect("Could not load theme");
    let mut xcursor = xcursor_theme
        .get_cursor("left_ptr".into())
        .expect("Could not load cursor from theme");
    cursor.set_xcursor(Some(xcursor));

    let (mut output_manager, input_manager) = managers(cursor);
    let compositor = Compositor::new(Box::new(input_manager), Box::new(output_manager));
    compositor.run();
}
<|endoftext|>


<commit_before>use crate::{IdentFragment, ToTokens, TokenStreamExt};
use std::fmt;
use std::ops::BitOr;

pub use proc_macro2::*;

pub struct HasIterator; // True
pub struct ThereIsNoIteratorInRepetition; // False

impl BitOr<ThereIsNoIteratorInRepetition> for ThereIsNoIteratorInRepetition {
    type Output = ThereIsNoIteratorInRepetition;
    fn bitor(self, _rhs: ThereIsNoIteratorInRepetition) -> ThereIsNoIteratorInRepetition {
        ThereIsNoIteratorInRepetition
    }
}

impl BitOr<ThereIsNoIteratorInRepetition> for HasIterator {
    type Output = HasIterator;
    fn bitor(self, _rhs: ThereIsNoIteratorInRepetition) -> HasIterator {
        HasIterator
    }
}

impl BitOr<HasIterator> for ThereIsNoIteratorInRepetition {
    type Output = HasIterator;
    fn bitor(self, _rhs: HasIterator) -> HasIterator {
        HasIterator
    }
}

impl BitOr<HasIterator> for HasIterator {
    type Output = HasIterator;
    fn bitor(self, _rhs: HasIterator) -> HasIterator {
        HasIterator
    }
}

/// Extension traits used by the implementation of `quote!`. These are defined
/// in separate traits, rather than as a single trait due to ambiguity issues.
///
/// These traits expose a `quote_into_iter` method which should allow calling
/// whichever impl happens to be applicable. Calling that method repeatedly on
/// the returned value should be idempotent.
pub mod ext {
    use super::RepInterp;
    use super::{HasIterator as HasIter, ThereIsNoIteratorInRepetition as DoesNotHaveIter};
    use crate::ToTokens;
    use std::slice;

    /// Extension trait providing the `quote_into_iter` method on iterators.
    pub trait RepIteratorExt: Iterator + Sized {
        fn quote_into_iter(self) -> (Self, HasIter) {
            (self, HasIter)
        }
    }

    impl<T: Iterator> RepIteratorExt for T {}

    /// Extension trait providing the `quote_into_iter` method for
    /// non-iterable types. These types interpolate the same value in each
    /// iteration of the repetition.
    pub trait RepToTokensExt {
        /// Pretend to be an iterator for the purposes of `quote_into_iter`.
        /// This allows repeated calls to `quote_into_iter` to continue
        /// correctly returning DoesNotHaveIter.
        fn next(&self) -> Option<&Self> {
            Some(self)
        }

        fn quote_into_iter(&self) -> (&Self, DoesNotHaveIter) {
            (self, DoesNotHaveIter)
        }
    }

    impl<T: ToTokens + ?Sized> RepToTokensExt for T {}

    /// Extension trait providing the `quote_into_iter` method for types that
    /// can be referenced as an iterator.
    pub trait RepAsIteratorExt<'q> {
        type Iter: Iterator;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter);
    }

    impl<'q, 'a, T: RepAsIteratorExt<'q> + ?Sized> RepAsIteratorExt<'q> for &'a T {
        type Iter = T::Iter;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            <T as RepAsIteratorExt>::quote_into_iter(*self)
        }
    }

    impl<'q, 'a, T: RepAsIteratorExt<'q> + ?Sized> RepAsIteratorExt<'q> for &'a mut T {
        type Iter = T::Iter;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            <T as RepAsIteratorExt>::quote_into_iter(*self)
        }
    }

    impl<'q, T: 'q> RepAsIteratorExt<'q> for [T] {
        type Iter = slice::Iter<'q, T>;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            (self.iter(), HasIter)
        }
    }

    impl<'q, T: 'q> RepAsIteratorExt<'q> for Vec<T> {
        type Iter = slice::Iter<'q, T>;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            (self.iter(), HasIter)
        }
    }

    macro_rules! array_rep_slice {
        ($($l:tt)*) => {
            $(
                impl<'q, T: 'q> RepAsIteratorExt<'q> for [T; $l] {
                    type Iter = slice::Iter<'q, T>;

                    fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
                        (self.iter(), HasIter)
                    }
                }
            )*
        }
    }

    array_rep_slice!(
        0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
        17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32
    );

    impl<'q, T: RepAsIteratorExt<'q>> RepAsIteratorExt<'q> for RepInterp<T> {
        type Iter = T::Iter;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            self.0.quote_into_iter()
        }
    }
}

// Helper type used within interpolations to allow for repeated binding names.
// Implements the relevant traits, and exports a dummy `next()` method.
#[derive(Copy, Clone)]
pub struct RepInterp<T>(pub T);

impl<T> RepInterp<T> {
    // This method is intended to look like `Iterator::next`, and is called when
    // a name is bound multiple times, as the previous binding will shadow the
    // original `Iterator` object. This allows us to avoid advancing the
    // iterator multiple times per iteration.
    pub fn next(self) -> Option<T> {
        Some(self.0)
    }
}

impl<T: Iterator> Iterator for RepInterp<T> {
    type Item = T::Item;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl<T: ToTokens> ToTokens for RepInterp<T> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.0.to_tokens(tokens);
    }
}

fn is_ident_start(c: u8) -> bool {
    (b'a' <= c && c <= b'z') || (b'A' <= c && c <= b'Z') || c == b'_'
}

fn is_ident_continue(c: u8) -> bool {
    (b'a' <= c && c <= b'z') || (b'A' <= c && c <= b'Z') || c == b'_' || (b'0' <= c && c <= b'9')
}

fn is_ident(token: &str) -> bool {
    let mut iter = token.bytes();
    let first_ok = iter.next().map(is_ident_start).unwrap_or(false);

    first_ok && iter.all(is_ident_continue)
}

pub fn parse(tokens: &mut TokenStream, span: Span, s: &str) {
    if is_ident(s) {
        // Fast path, since idents are the most common token.
        tokens.append(Ident::new(s, span));
    } else {
        let s: TokenStream = s.parse().expect("invalid token stream");
        tokens.extend(s.into_iter().map(|mut t| {
            t.set_span(span);
            t
        }));
    }
}

macro_rules! push_punct {
    ($name:ident $char1:tt) => {
        pub fn $name(tokens: &mut TokenStream, span: Span) {
            let mut punct = Punct::new($char1, Spacing::Alone);
            punct.set_span(span);
            tokens.append(punct);
        }
    };
    ($name:ident $char1:tt $char2:tt) => {
        pub fn $name(tokens: &mut TokenStream, span: Span) {
            let mut punct = Punct::new($char1, Spacing::Joint);
            punct.set_span(span);
            tokens.append(punct);
            let mut punct = Punct::new($char2, Spacing::Alone);
            punct.set_span(span);
            tokens.append(punct);
        }
    };
    ($name:ident $char1:tt $char2:tt $char3:tt) => {
        pub fn $name(tokens: &mut TokenStream, span: Span) {
            let mut punct = Punct::new($char1, Spacing::Joint);
            punct.set_span(span);
            tokens.append(punct);
            let mut punct = Punct::new($char2, Spacing::Joint);
            punct.set_span(span);
            tokens.append(punct);
            let mut punct = Punct::new($char3, Spacing::Alone);
            punct.set_span(span);
            tokens.append(punct);
        }
    };
}

push_punct!(push_add '+');
push_punct!(push_add_eq '+' '=');
push_punct!(push_and '&');
push_punct!(push_and_and '&' '&');
push_punct!(push_and_eq '&' '=');
push_punct!(push_at '@');
push_punct!(push_bang '!');
push_punct!(push_caret '^');
push_punct!(push_caret_eq '^' '=');
push_punct!(push_colon ':');
push_punct!(push_colon2 ':' ':');
push_punct!(push_comma ',');
push_punct!(push_div '/');
push_punct!(push_div_eq '/' '=');
push_punct!(push_dot '.');
push_punct!(push_dot2 '.' '.');
push_punct!(push_dot3 '.' '.' '.');
push_punct!(push_dot_dot_eq '.' '.' '=');
push_punct!(push_eq '=');
push_punct!(push_eq_eq '=' '=');
push_punct!(push_ge '>' '=');
push_punct!(push_gt '>');
push_punct!(push_le '<' '=');
push_punct!(push_lt '<');
push_punct!(push_mul_eq '*' '=');
push_punct!(push_ne '!' '=');
push_punct!(push_or '|');
push_punct!(push_or_eq '|' '=');
push_punct!(push_or_or '|' '|');
push_punct!(push_pound '#');
push_punct!(push_question '?');
push_punct!(push_rarrow '-' '>');
push_punct!(push_larrow '<' '-');
push_punct!(push_rem '%');
push_punct!(push_rem_eq '%' '=');
push_punct!(push_fat_arrow '=' '>');
push_punct!(push_semi ';');
push_punct!(push_shl '<' '<');
push_punct!(push_shl_eq '<' '<' '=');
push_punct!(push_shr '>' '>');
push_punct!(push_shr_eq '>' '>' '=');
push_punct!(push_star '*');
push_punct!(push_sub '-');
push_punct!(push_sub_eq '-' '=');

// Helper method for constructing identifiers from the `format_ident!` macro,
// handling `r#` prefixes.
//
// Directly parsing the input string may produce a valid identifier,
// although the input string was invalid, due to ignored characters such as
// whitespace and comments. Instead, we always create a non-raw identifier
// to validate that the string is OK, and only parse again if needed.
//
// The `is_ident` method defined above is insufficient for validation, as it
// will reject non-ASCII identifiers.
pub fn mk_ident(id: &str, span: Option<Span>) -> Ident {
    let span = span.unwrap_or_else(Span::call_site);

    let is_raw = id.starts_with("r#");
    let unraw = Ident::new(if is_raw { &id[2..] } else { id }, span);
    if !is_raw {
        return unraw;
    }

    // At this point, the identifier is raw, and the unraw-ed version of it was
    // successfully converted into an identifier. Try to produce a valid raw
    // identifier by running the `TokenStream` parser, and unwrapping the first
    // token as an `Ident`.
    //
    // FIXME: When `Ident::new_raw` becomes stable, this method should be
    // updated to call it when available.
    match id.parse::<TokenStream>() {
        Ok(ts) => {
            let mut iter = ts.into_iter();
            match (iter.next(), iter.next()) {
                (Some(TokenTree::Ident(mut id)), None) => {
                    id.set_span(span);
                    id
                }
                _ => unreachable!("valid raw ident fails to parse"),
            }
        }
        Err(_) => unreachable!("valid raw ident fails to parse"),
    }
}

// Adapts from `IdentFragment` to `fmt::Display` for use by the `format_ident!`
// macro, and exposes span information from these fragments.
//
// This struct also has forwarding implementations of the formatting traits
// `Octal`, `LowerHex`, `UpperHex`, and `Binary` to allow for their use within
// `format_ident!`.
#[derive(Copy, Clone)]
pub struct IdentFragmentAdapter<T: IdentFragment>(pub T);

impl<T: IdentFragment> IdentFragmentAdapter<T> {
    pub fn span(&self) -> Option<Span> {
        self.0.span()
    }
}

impl<T: IdentFragment> fmt::Display for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        IdentFragment::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::Octal> fmt::Octal for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Octal::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::LowerHex> fmt::LowerHex for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::LowerHex::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::UpperHex> fmt::UpperHex for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::UpperHex::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::Binary> fmt::Binary for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Binary::fmt(&self.0, f)
    }
}
<commit_msg>Repetition for BTreeSet<commit_after>use crate::{IdentFragment, ToTokens, TokenStreamExt};
use std::fmt;
use std::ops::BitOr;

pub use proc_macro2::*;

pub struct HasIterator; // True
pub struct ThereIsNoIteratorInRepetition; // False

impl BitOr<ThereIsNoIteratorInRepetition> for ThereIsNoIteratorInRepetition {
    type Output = ThereIsNoIteratorInRepetition;
    fn bitor(self, _rhs: ThereIsNoIteratorInRepetition) -> ThereIsNoIteratorInRepetition {
        ThereIsNoIteratorInRepetition
    }
}

impl BitOr<ThereIsNoIteratorInRepetition> for HasIterator {
    type Output = HasIterator;
    fn bitor(self, _rhs: ThereIsNoIteratorInRepetition) -> HasIterator {
        HasIterator
    }
}

impl BitOr<HasIterator> for ThereIsNoIteratorInRepetition {
    type Output = HasIterator;
    fn bitor(self, _rhs: HasIterator) -> HasIterator {
        HasIterator
    }
}

impl BitOr<HasIterator> for HasIterator {
    type Output = HasIterator;
    fn bitor(self, _rhs: HasIterator) -> HasIterator {
        HasIterator
    }
}

/// Extension traits used by the implementation of `quote!`. These are defined
/// in separate traits, rather than as a single trait due to ambiguity issues.
///
/// These traits expose a `quote_into_iter` method which should allow calling
/// whichever impl happens to be applicable. Calling that method repeatedly on
/// the returned value should be idempotent.
pub mod ext {
    use super::RepInterp;
    use super::{HasIterator as HasIter, ThereIsNoIteratorInRepetition as DoesNotHaveIter};
    use crate::ToTokens;
    use std::collections::btree_set::{self, BTreeSet};
    use std::slice;

    /// Extension trait providing the `quote_into_iter` method on iterators.
    pub trait RepIteratorExt: Iterator + Sized {
        fn quote_into_iter(self) -> (Self, HasIter) {
            (self, HasIter)
        }
    }

    impl<T: Iterator> RepIteratorExt for T {}

    /// Extension trait providing the `quote_into_iter` method for
    /// non-iterable types. These types interpolate the same value in each
    /// iteration of the repetition.
    pub trait RepToTokensExt {
        /// Pretend to be an iterator for the purposes of `quote_into_iter`.
        /// This allows repeated calls to `quote_into_iter` to continue
        /// correctly returning DoesNotHaveIter.
        fn next(&self) -> Option<&Self> {
            Some(self)
        }

        fn quote_into_iter(&self) -> (&Self, DoesNotHaveIter) {
            (self, DoesNotHaveIter)
        }
    }

    impl<T: ToTokens + ?Sized> RepToTokensExt for T {}

    /// Extension trait providing the `quote_into_iter` method for types that
    /// can be referenced as an iterator.
    pub trait RepAsIteratorExt<'q> {
        type Iter: Iterator;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter);
    }

    impl<'q, 'a, T: RepAsIteratorExt<'q> + ?Sized> RepAsIteratorExt<'q> for &'a T {
        type Iter = T::Iter;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            <T as RepAsIteratorExt>::quote_into_iter(*self)
        }
    }

    impl<'q, 'a, T: RepAsIteratorExt<'q> + ?Sized> RepAsIteratorExt<'q> for &'a mut T {
        type Iter = T::Iter;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            <T as RepAsIteratorExt>::quote_into_iter(*self)
        }
    }

    impl<'q, T: 'q> RepAsIteratorExt<'q> for [T] {
        type Iter = slice::Iter<'q, T>;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            (self.iter(), HasIter)
        }
    }

    impl<'q, T: 'q> RepAsIteratorExt<'q> for Vec<T> {
        type Iter = slice::Iter<'q, T>;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            (self.iter(), HasIter)
        }
    }

    impl<'q, T: 'q> RepAsIteratorExt<'q> for BTreeSet<T> {
        type Iter = btree_set::Iter<'q, T>;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            (self.iter(), HasIter)
        }
    }

    macro_rules! array_rep_slice {
        ($($l:tt)*) => {
            $(
                impl<'q, T: 'q> RepAsIteratorExt<'q> for [T; $l] {
                    type Iter = slice::Iter<'q, T>;

                    fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
                        (self.iter(), HasIter)
                    }
                }
            )*
        }
    }

    array_rep_slice!(
        0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
        17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32
    );

    impl<'q, T: RepAsIteratorExt<'q>> RepAsIteratorExt<'q> for RepInterp<T> {
        type Iter = T::Iter;

        fn quote_into_iter(&'q self) -> (Self::Iter, HasIter) {
            self.0.quote_into_iter()
        }
    }
}

// Helper type used within interpolations to allow for repeated binding names.
// Implements the relevant traits, and exports a dummy `next()` method.
#[derive(Copy, Clone)]
pub struct RepInterp<T>(pub T);

impl<T> RepInterp<T> {
    // This method is intended to look like `Iterator::next`, and is called when
    // a name is bound multiple times, as the previous binding will shadow the
    // original `Iterator` object. This allows us to avoid advancing the
    // iterator multiple times per iteration.
    pub fn next(self) -> Option<T> {
        Some(self.0)
    }
}

impl<T: Iterator> Iterator for RepInterp<T> {
    type Item = T::Item;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl<T: ToTokens> ToTokens for RepInterp<T> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.0.to_tokens(tokens);
    }
}

fn is_ident_start(c: u8) -> bool {
    (b'a' <= c && c <= b'z') || (b'A' <= c && c <= b'Z') || c == b'_'
}

fn is_ident_continue(c: u8) -> bool {
    (b'a' <= c && c <= b'z') || (b'A' <= c && c <= b'Z') || c == b'_' || (b'0' <= c && c <= b'9')
}

fn is_ident(token: &str) -> bool {
    let mut iter = token.bytes();
    let first_ok = iter.next().map(is_ident_start).unwrap_or(false);

    first_ok && iter.all(is_ident_continue)
}

pub fn parse(tokens: &mut TokenStream, span: Span, s: &str) {
    if is_ident(s) {
        // Fast path, since idents are the most common token.
        tokens.append(Ident::new(s, span));
    } else {
        let s: TokenStream = s.parse().expect("invalid token stream");
        tokens.extend(s.into_iter().map(|mut t| {
            t.set_span(span);
            t
        }));
    }
}

macro_rules! push_punct {
    ($name:ident $char1:tt) => {
        pub fn $name(tokens: &mut TokenStream, span: Span) {
            let mut punct = Punct::new($char1, Spacing::Alone);
            punct.set_span(span);
            tokens.append(punct);
        }
    };
    ($name:ident $char1:tt $char2:tt) => {
        pub fn $name(tokens: &mut TokenStream, span: Span) {
            let mut punct = Punct::new($char1, Spacing::Joint);
            punct.set_span(span);
            tokens.append(punct);
            let mut punct = Punct::new($char2, Spacing::Alone);
            punct.set_span(span);
            tokens.append(punct);
        }
    };
    ($name:ident $char1:tt $char2:tt $char3:tt) => {
        pub fn $name(tokens: &mut TokenStream, span: Span) {
            let mut punct = Punct::new($char1, Spacing::Joint);
            punct.set_span(span);
            tokens.append(punct);
            let mut punct = Punct::new($char2, Spacing::Joint);
            punct.set_span(span);
            tokens.append(punct);
            let mut punct = Punct::new($char3, Spacing::Alone);
            punct.set_span(span);
            tokens.append(punct);
        }
    };
}

push_punct!(push_add '+');
push_punct!(push_add_eq '+' '=');
push_punct!(push_and '&');
push_punct!(push_and_and '&' '&');
push_punct!(push_and_eq '&' '=');
push_punct!(push_at '@');
push_punct!(push_bang '!');
push_punct!(push_caret '^');
push_punct!(push_caret_eq '^' '=');
push_punct!(push_colon ':');
push_punct!(push_colon2 ':' ':');
push_punct!(push_comma ',');
push_punct!(push_div '/');
push_punct!(push_div_eq '/' '=');
push_punct!(push_dot '.');
push_punct!(push_dot2 '.' '.');
push_punct!(push_dot3 '.' '.' '.');
push_punct!(push_dot_dot_eq '.' '.' '=');
push_punct!(push_eq '=');
push_punct!(push_eq_eq '=' '=');
push_punct!(push_ge '>' '=');
push_punct!(push_gt '>');
push_punct!(push_le '<' '=');
push_punct!(push_lt '<');
push_punct!(push_mul_eq '*' '=');
push_punct!(push_ne '!' '=');
push_punct!(push_or '|');
push_punct!(push_or_eq '|' '=');
push_punct!(push_or_or '|' '|');
push_punct!(push_pound '#');
push_punct!(push_question '?');
push_punct!(push_rarrow '-' '>');
push_punct!(push_larrow '<' '-');
push_punct!(push_rem '%');
push_punct!(push_rem_eq '%' '=');
push_punct!(push_fat_arrow '=' '>');
push_punct!(push_semi ';');
push_punct!(push_shl '<' '<');
push_punct!(push_shl_eq '<' '<' '=');
push_punct!(push_shr '>' '>');
push_punct!(push_shr_eq '>' '>' '=');
push_punct!(push_star '*');
push_punct!(push_sub '-');
push_punct!(push_sub_eq '-' '=');

// Helper method for constructing identifiers from the `format_ident!` macro,
// handling `r#` prefixes.
//
// Directly parsing the input string may produce a valid identifier,
// although the input string was invalid, due to ignored characters such as
// whitespace and comments. Instead, we always create a non-raw identifier
// to validate that the string is OK, and only parse again if needed.
//
// The `is_ident` method defined above is insufficient for validation, as it
// will reject non-ASCII identifiers.
pub fn mk_ident(id: &str, span: Option<Span>) -> Ident {
    let span = span.unwrap_or_else(Span::call_site);

    let is_raw = id.starts_with("r#");
    let unraw = Ident::new(if is_raw { &id[2..] } else { id }, span);
    if !is_raw {
        return unraw;
    }

    // At this point, the identifier is raw, and the unraw-ed version of it was
    // successfully converted into an identifier. Try to produce a valid raw
    // identifier by running the `TokenStream` parser, and unwrapping the first
    // token as an `Ident`.
    //
    // FIXME: When `Ident::new_raw` becomes stable, this method should be
    // updated to call it when available.
    match id.parse::<TokenStream>() {
        Ok(ts) => {
            let mut iter = ts.into_iter();
            match (iter.next(), iter.next()) {
                (Some(TokenTree::Ident(mut id)), None) => {
                    id.set_span(span);
                    id
                }
                _ => unreachable!("valid raw ident fails to parse"),
            }
        }
        Err(_) => unreachable!("valid raw ident fails to parse"),
    }
}

// Adapts from `IdentFragment` to `fmt::Display` for use by the `format_ident!`
// macro, and exposes span information from these fragments.
//
// This struct also has forwarding implementations of the formatting traits
// `Octal`, `LowerHex`, `UpperHex`, and `Binary` to allow for their use within
// `format_ident!`.
#[derive(Copy, Clone)]
pub struct IdentFragmentAdapter<T: IdentFragment>(pub T);

impl<T: IdentFragment> IdentFragmentAdapter<T> {
    pub fn span(&self) -> Option<Span> {
        self.0.span()
    }
}

impl<T: IdentFragment> fmt::Display for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        IdentFragment::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::Octal> fmt::Octal for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Octal::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::LowerHex> fmt::LowerHex for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::LowerHex::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::UpperHex> fmt::UpperHex for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::UpperHex::fmt(&self.0, f)
    }
}

impl<T: IdentFragment + fmt::Binary> fmt::Binary for IdentFragmentAdapter<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Binary::fmt(&self.0, f)
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The SAX-RS Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[link(name = "sax",
       vers = "0.1",
       author = "Brendan Zabarauskas",
       url = "https://github.com/bjz/sax-rs")];

#[comment = "Wrapper for libxml2's SAX parser."];
#[crate_type = "lib"];

#[feature(globs)];

extern mod extra;

use std::cast;
use std::comm::{Port, stream};
use std::libc::{c_char, c_int};
use std::str;
// use std::task;

use error::ErrorData;

pub mod error;
pub mod ffi;
mod extfn;

/// Events to be sent by the parser.
#[deriving(Eq, Clone)]
pub enum ParseEvent {
    /// The document has begun to be processed.
    StartDocument,
    /// The document processing has finished.
    EndDocument,
    /// An opening tag has was parsed.
    StartElement(~str, Attributes),
    /// A closing tag was parsed.
    EndElement(~str),
    /// Some characters between tags have been recived.
    Characters(~str),
    /// A comment tag was parsed.
    Comment(~str),
    /// A `CDATA` block was parsed.
    CdataBlock(~str),
}

impl ToStr for ParseEvent {
    fn to_str(&self) -> ~str {
        match *self {
            StartDocument => ~"START DOCUMENT",
            EndDocument => ~"END DOCUMENT",
            StartElement(ref name, ref atts) => format!("<{}{}>", *name, atts.to_str()),
            EndElement(ref name) => format!("</{}>", *name),
            Characters(ref ch) => ch.clone(),
            Comment(ref value) => format!("", *value),
            CdataBlock(ref value) => format!("<![CDATA[{}]]>", *value),
        }
    }
}

#[deriving(Eq, Clone)]
pub struct Attribute {
    name: ~str,
    value: ~str,
}

/// A list of attributes
#[deriving(Eq, Clone)]
pub struct Attributes(~[Attribute]);

impl Attributes {
    unsafe fn from_buf(atts: **ffi::xmlChar) -> Attributes {
        let mut ret = Attributes(~[]);
        let mut ptr = atts as **c_char;
        while !ptr.is_null() && !(*ptr).is_null() {
            ret.push(
                Attribute {
                    name: str::raw::from_c_str(*ptr),
                    value: str::raw::from_c_str(*ptr.offset(1)),
                }
            );
            ptr = ptr.offset(2);
        }
        ret
    }

    pub fn find<'a>(&'a self, name: &str) -> Option<&'a str> {
        self.iter()
            .find(|att| name == att.name)
            .map(|att| att.value.as_slice())
    }

    pub fn get<'a>(&'a self, name: &str) -> &'a str {
        self.find(name).expect(format!("Could not find an attribute with the name \"{}\"", name))
    }

    pub fn find_clone(&self, name: &str) -> Option<~str> {
        self.find(name).map(|v| v.to_owned())
    }

    pub fn get_clone(&self, name: &str) -> ~str {
        self.get(name).to_owned()
    }
}

impl ToStr for Attributes {
    fn to_str(&self) -> ~str {
        self.iter().map(|att| {
            format!(" {}=\"{}\"", att.name, att.value)
        }).to_owned_vec().concat()
    }
}

/// Either a parse event wrapped in `Ok` or some Error data wrapped in `Err`.
pub type ParseResult = Result<ParseEvent, ErrorData>;

/// A port to recieve `ParseResult`s from the parser.
pub struct SaxPort {
    priv port: Port<ParseResult>,
}

impl GenericPort<ParseResult> for SaxPort {
    /// Recives a new parse message.
    ///
    /// # Failure
    ///
    /// Fails if the method is called again after the final `Ok(EndDocument)`
    /// parse result has been recived.
    fn recv(&self) -> ParseResult {
        self.port.try_recv().expect(
            "Could not get a new parse result, the parser has already finished!"
        )
    }

    /// Receives a parse result wrapped in `Some`, or `None` if the parser has
    /// finished.
    fn try_recv(&self) -> Option<ParseResult> {
        self.port.try_recv()
    }
}

/// Parses the entire XML string.
///
/// # Returns
///
/// A port that recieves parse results.
///
/// # Example
///
/// ~~~rust
/// let port = parse_xml("<hullo><world /></hullo>");
/// loop {
///     if port.recv() == Ok(EndDocument) { break }
/// }
/// ~~~
#[inline(never)]
pub fn parse_xml(src: &str) -> SaxPort {
    let len = src.len() as c_int;
    src.to_c_str().with_ref(|c_str| {
        let (port, chan) = stream();
        // do task::spawn {
            unsafe {
                ffi::xmlSAXUserParseMemory(&extfn::new_handler(),
                                           cast::transmute(&chan),
                                           c_str, len);
                ffi::xmlCleanupParser();
            }
        // }
        SaxPort { port: port }
    })
}

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

    fn get_mock_atts() -> Attributes {
        Attributes(~[
            Attribute { name: ~"foo", value: ~"0" },
            Attribute { name: ~"bar", value: ~"1" },
            Attribute { name: ~"baz", value: ~"2" }
        ])
    }

    #[test]
    fn test_att_get() {
        let atts = get_mock_atts();
        assert_eq!(atts.get("foo"), "0");
        assert_eq!(atts.get("bar"), "1");
        assert_eq!(atts.get("baz"), "2");
    }

    #[test]
    fn test_att_find() {
        let atts = get_mock_atts();
        assert_eq!(atts.find("foo"), Some("0"));
        assert_eq!(atts.find("bar"), Some("1"));
        assert_eq!(atts.find("baz"), Some("2"));
    }

    #[test]
    fn test() {
        let sax = parse_xml(
            "<hello><this />test</hello>"
        );
        loop {
            match sax.recv() {
                Ok(StartDocument) => (),
                Ok(EndDocument) => break,
                Ok(event) => println(event.to_str()),
                Err(err) => println(err.to_str()),
            }
        }
    }
}
<commit_msg>Update for pkgid<commit_after>// Copyright 2013 The SAX-RS Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[comment = "Wrapper for libxml2's SAX parser."];
#[pkgid = "github.com/bjz/sax-rs#sax:0.1"];

#[crate_type = "lib"];

#[feature(globs)];

extern mod extra;

use std::cast;
use std::comm::{Port, stream};
use std::libc::{c_char, c_int};
use std::str;
// use std::task;

use error::ErrorData;

pub mod error;
pub mod ffi;
mod extfn;

/// Events to be sent by the parser.
#[deriving(Eq, Clone)]
pub enum ParseEvent {
    /// The document has begun to be processed.
    StartDocument,
    /// The document processing has finished.
    EndDocument,
    /// An opening tag has was parsed.
    StartElement(~str, Attributes),
    /// A closing tag was parsed.
    EndElement(~str),
    /// Some characters between tags have been recived.
    Characters(~str),
    /// A comment tag was parsed.
    Comment(~str),
    /// A `CDATA` block was parsed.
    CdataBlock(~str),
}

impl ToStr for ParseEvent {
    fn to_str(&self) -> ~str {
        match *self {
            StartDocument => ~"START DOCUMENT",
            EndDocument => ~"END DOCUMENT",
            StartElement(ref name, ref atts) => format!("<{}{}>", *name, atts.to_str()),
            EndElement(ref name) => format!("</{}>", *name),
            Characters(ref ch) => ch.clone(),
            Comment(ref value) => format!("", *value),
            CdataBlock(ref value) => format!("<![CDATA[{}]]>", *value),
        }
    }
}

#[deriving(Eq, Clone)]
pub struct Attribute {
    name: ~str,
    value: ~str,
}

/// A list of attributes
#[deriving(Eq, Clone)]
pub struct Attributes(~[Attribute]);

impl Attributes {
    unsafe fn from_buf(atts: **ffi::xmlChar) -> Attributes {
        let mut ret = Attributes(~[]);
        let mut ptr = atts as **c_char;
        while !ptr.is_null() && !(*ptr).is_null() {
            ret.push(
                Attribute {
                    name: str::raw::from_c_str(*ptr),
                    value: str::raw::from_c_str(*ptr.offset(1)),
                }
            );
            ptr = ptr.offset(2);
        }
        ret
    }

    pub fn find<'a>(&'a self, name: &str) -> Option<&'a str> {
        self.iter()
            .find(|att| name == att.name)
            .map(|att| att.value.as_slice())
    }

    pub fn get<'a>(&'a self, name: &str) -> &'a str {
        self.find(name).expect(format!("Could not find an attribute with the name \"{}\"", name))
    }

    pub fn find_clone(&self, name: &str) -> Option<~str> {
        self.find(name).map(|v| v.to_owned())
    }

    pub fn get_clone(&self, name: &str) -> ~str {
        self.get(name).to_owned()
    }
}

impl ToStr for Attributes {
    fn to_str(&self) -> ~str {
        self.iter().map(|att| {
            format!(" {}=\"{}\"", att.name, att.value)
        }).to_owned_vec().concat()
    }
}

/// Either a parse event wrapped in `Ok` or some Error data wrapped in `Err`.
pub type ParseResult = Result<ParseEvent, ErrorData>;

/// A port to recieve `ParseResult`s from the parser.
pub struct SaxPort {
    priv port: Port<ParseResult>,
}

impl GenericPort<ParseResult> for SaxPort {
    /// Recives a new parse message.
    ///
    /// # Failure
    ///
    /// Fails if the method is called again after the final `Ok(EndDocument)`
    /// parse result has been recived.
    fn recv(&self) -> ParseResult {
        self.port.try_recv().expect(
            "Could not get a new parse result, the parser has already finished!"
        )
    }

    /// Receives a parse result wrapped in `Some`, or `None` if the parser has
    /// finished.
    fn try_recv(&self) -> Option<ParseResult> {
        self.port.try_recv()
    }
}

/// Parses the entire XML string.
///
/// # Returns
///
/// A port that recieves parse results.
///
/// # Example
///
/// ~~~rust
/// let port = parse_xml("<hullo><world /></hullo>");
/// loop {
///     if port.recv() == Ok(EndDocument) { break }
/// }
/// ~~~
#[inline(never)]
pub fn parse_xml(src: &str) -> SaxPort {
    let len = src.len() as c_int;
    src.to_c_str().with_ref(|c_str| {
        let (port, chan) = stream();
        // do task::spawn {
            unsafe {
                ffi::xmlSAXUserParseMemory(&extfn::new_handler(),
                                           cast::transmute(&chan),
                                           c_str, len);
                ffi::xmlCleanupParser();
            }
        // }
        SaxPort { port: port }
    })
}

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

    fn get_mock_atts() -> Attributes {
        Attributes(~[
            Attribute { name: ~"foo", value: ~"0" },
            Attribute { name: ~"bar", value: ~"1" },
            Attribute { name: ~"baz", value: ~"2" }
        ])
    }

    #[test]
    fn test_att_get() {
        let atts = get_mock_atts();
        assert_eq!(atts.get("foo"), "0");
        assert_eq!(atts.get("bar"), "1");
        assert_eq!(atts.get("baz"), "2");
    }

    #[test]
    fn test_att_find() {
        let atts = get_mock_atts();
        assert_eq!(atts.find("foo"), Some("0"));
        assert_eq!(atts.find("bar"), Some("1"));
        assert_eq!(atts.find("baz"), Some("2"));
    }

    #[test]
    fn test() {
        let sax = parse_xml(
            "<hello><this />test</hello>"
        );
        loop {
            match sax.recv() {
                Ok(StartDocument) => (),
                Ok(EndDocument) => break,
                Ok(event) => println(event.to_str()),
                Err(err) => println(err.to_str()),
            }
        }
    }
}
<|endoftext|>


<commit_before>use tf;
use libc::{c_char, c_int};
use std::ffi::CString;
use std::marker;
use std::path::Path;
use std::ptr;
use super::AnyTensor;
use super::{Buffer, BufferTrait};
use super::Code;
use super::DataType;
use super::Graph;
use super::GraphTrait;
use super::Operation;
use super::OperationTrait;
use super::Result;
use super::SessionOptions;
use super::Status;
use super::Tensor;
use super::TensorType;

/// Aggregation type for a saved model bundle.
#[derive(Debug)]
pub struct SavedModelBundle {
    /// The loaded session.
    pub session: Session,
    /// A meta graph definition as raw protocol buffer.
    pub meta_graph_def: Vec<u8>,
}

impl SavedModelBundle {

    /// Loads a session from an exported model, creating a bundle
    pub fn load<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
        (options: &SessionOptions,
         tags: Tags,
         graph: &mut Graph,
         export_dir: P)
         -> Result<SavedModelBundle> {
        let mut status = Status::new();

        let export_dir_cstr =
            export_dir.as_ref()
                     .to_str()
                     .and_then(|s| CString::new(s.as_bytes()).ok())
                     .ok_or_else(|| invalid_arg!("Invalid export directory path"))?;

        let tags_cstr: Vec<_> = tags.into_iter()
                                    .map(|t| CString::new(t.as_ref()))
                                    .collect::<::std::result::Result<_, _>>()
                                    .map_err(|_| invalid_arg!("Invalid tag name"))?;
        let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();

        // The empty TF_Buffer will be filled by LoadSessionFromSavedModel
        let mut meta = unsafe { Buffer::<u8>::from_ptr(ptr::null_mut(), 0) };

        let inner = unsafe {
            tf::TF_LoadSessionFromSavedModel(options.inner,
                                             ptr::null(),
                                             export_dir_cstr.as_ptr(),
                                             tags_ptr.as_ptr(),
                                             tags_ptr.len() as c_int,
                                             graph.inner(),
                                             meta.inner_mut(),
                                             status.inner())
        };
        if inner.is_null() {
            Err(status)
        } else {
            let session = Session { inner: inner };
            Ok(SavedModelBundle {
                session: session,
                meta_graph_def: Vec::from(meta.as_ref())
            })
        }
    }

}

/// Manages a single graph and execution.
#[derive(Debug)]
pub struct Session {
    inner: *mut tf::TF_Session,
}

impl Session {
    /// Creates a session.
    pub fn new(options: &SessionOptions, graph: &Graph) -> Result<Self> {
        let mut status = Status::new();
        let inner = unsafe { tf::TF_NewSession(graph.inner(), options.inner, status.inner()) };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Loads a session from an exported model.
    pub fn from_saved_model<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
        (options: &SessionOptions,
         tags: Tags,
         graph: &mut Graph,
         export_dir: P)
         -> Result<Self> {
        let mut status = Status::new();

        let export_dir_cstr = export_dir.as_ref()
                     .to_str()
                     .and_then(|s| CString::new(s.as_bytes()).ok())
                     .ok_or_else(|| invalid_arg!("Invalid export directory path"))?;

        let tags_cstr: Vec<_> = tags.into_iter()
                                    .map(|t| CString::new(t.as_ref()))
                                    .collect::<::std::result::Result<_, _>>()
                                    .map_err(|_| invalid_arg!("Invalid tag name"))?;
        // keeping tags_cstr to retain strings in memory
        let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();

        let inner = unsafe {
            tf::TF_LoadSessionFromSavedModel(options.inner,
                                             ptr::null(),
                                             export_dir_cstr.as_ptr(),
                                             tags_ptr.as_ptr(),
                                             tags_ptr.len() as c_int,
                                             graph.inner(),
                                             ptr::null_mut(),
                                             status.inner())
        };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Closes the session.
    pub fn close(&mut self) -> Result<()> {
        let mut status = Status::new();
        unsafe {
            tf::TF_CloseSession(self.inner, status.inner());
        }
        status.into_result()
    }

    /// Runs the graph, feeding the inputs and then fetching the outputs requested in the step.
    pub fn run(&mut self, step: &mut StepWithGraph) -> Result<()> {
        // In case we're running it a second time and not all outputs were taken out.
        step.drop_output_tensors();

        let mut status = Status::new();
        let maybe_tensors: Result<_> = step.input_tensors.iter().map(|t| t.inner()).collect();
        let input_tensors: Vec<_> = maybe_tensors?;
        unsafe {
            tf::TF_SessionRun(self.inner,
                              ptr::null(),
                              step.input_ports.as_ptr(),
                              input_tensors.as_ptr() as *const *const tf::TF_Tensor,
                              input_tensors.len() as c_int,
                              step.output_ports.as_ptr(),
                              step.output_tensors.as_mut_ptr(),
                              step.output_tensors.len() as c_int,
                              step.target_operations.as_mut_ptr(),
                              step.target_operations.len() as c_int,
                              ptr::null_mut(),
                              status.inner());
        };
        status.into_result()
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        let mut status = Status::new();
        unsafe {
            tf::TF_DeleteSession(self.inner, status.inner());
        }
        // TODO: What do we do with the status?
    }
}

////////////////////////

/// An opaque token for retrieving an output from a computation.
#[derive(Copy,Clone,Debug)]
pub struct OutputToken {
    index: usize,
}

/// Manages the inputs and outputs for a single execution of a graph.
///
/// Typical usage involves creating an instance of this struct,
/// adding some inputs to it, requesting some outputs, passing it to `Session::run`
/// and then taking the outputs out of it.
///
/// This will be renamed to Step once the old API goes away.
#[derive(Debug)]
pub struct StepWithGraph<'l> {
    input_ports: Vec<tf::TF_Output>,
    input_tensors: Vec<&'l AnyTensor>,

    output_ports: Vec<tf::TF_Output>,
    output_tensors: Vec<*mut tf::TF_Tensor>,

    target_operations: Vec<*const tf::TF_Operation>,

    phantom: marker::PhantomData<&'l ()>,
}

impl<'l> StepWithGraph<'l> {
    /// Creates a StepWithGraph.
    pub fn new() -> Self {
        StepWithGraph {
            input_ports: vec![],
            input_tensors: vec![],

            output_ports: vec![],
            output_tensors: vec![],

            target_operations: vec![],

            phantom: marker::PhantomData,
        }
    }

    /// Adds an input to be fed to the graph.
    pub fn add_input<T: TensorType>(&mut self,
                                    operation: &Operation,
                                    index: c_int,
                                    tensor: &'l Tensor<T>) {
        self.input_ports.push(tf::TF_Output {
                                  oper: operation.inner(),
                                  index: index,
                              });
        self.input_tensors.push(tensor);
    }

    /// Requests that an output is fetched from the graph after running this step.
    /// Returns an index that you can then use to fetch this output from the step after running it.
    pub fn request_output(&mut self, operation: &Operation, index: c_int) -> OutputToken {
        self.output_ports.push(tf::TF_Output {
                                   oper: operation.inner(),
                                   index: index,
                               });
        self.output_tensors.push(ptr::null_mut());
        OutputToken { index: self.output_tensors.len() - 1 }
    }

    /// Extracts a tensor output given an index. A given index can only be extracted
    /// once per `Session::run`.
    /// Returns an error if output_idx is out of range, output is unavailable or the
    /// requested type does not match the type of the actual tensor.
    pub fn take_output<T: TensorType>(&mut self, token: OutputToken) -> Result<Tensor<T>> {
        let output_idx = token.index;
        if output_idx >= self.output_tensors.len() {
            return Err(Status::new_set(Code::OutOfRange,
                                       &format!("Requested output index is out of range: {} vs \
                                                 {}",
                                                output_idx,
                                                self.output_tensors.len()))
                               .unwrap());
        }
        if self.output_tensors[output_idx].is_null() {
            return Err(Status::new_set(Code::Unavailable,
                                       "Output not available. Either it was already taken, or \
                                        this step has not been sucessfully run yet.")
                               .unwrap());
        }
        let actual_data_type = self.output_data_type(output_idx).unwrap();
        if actual_data_type != T::data_type() {
            return Err(invalid_arg!("Requested tensor type does not match actual tensor type: \
                                     {} vs {}",
                                    actual_data_type,
                                    T::data_type()));
        }
        let tensor = unsafe { Tensor::from_tf_tensor(self.output_tensors[output_idx]).unwrap() };
        self.output_tensors[output_idx] = ptr::null_mut();
        Ok(tensor)
    }

    /// Adds a target operation to be executed when running the graph.
    pub fn add_target(&mut self, operation: &Operation) {
        self.target_operations.push(operation.inner());
    }

    /// Retuns the type of the tensor given an index.
    /// Returns `None` if the index is out of range or the output is not yet available.
    pub fn output_data_type(&self, output_idx: usize) -> Option<DataType> {
        if output_idx >= self.output_tensors.len() {
            return None;
        }
        if self.output_tensors[output_idx].is_null() {
            return None;
        }
        unsafe { Some(DataType::from_c(tf::TF_TensorType(self.output_tensors[output_idx]))) }
    }

    fn drop_output_tensors(&mut self) {
        for mut tensor in &mut self.output_tensors {
            // TODO: Is TF_DeleteTensor NULL safe?
            if !tensor.is_null() {
                unsafe {
                    tf::TF_DeleteTensor(*tensor);
                }
            }
            *tensor = ptr::null_mut();
        }
    }
}

impl<'l> Drop for StepWithGraph<'l> {
    fn drop(&mut self) {
        self.drop_output_tensors();
    }
}

////////////////////////

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::DataType;
    use super::super::Graph;
    use super::super::Operation;
    use super::super::Output;
    use super::super::SessionOptions;
    use super::super::Shape;
    use super::super::Tensor;

    fn create_session() -> (Session, Operation, Operation) {
        let mut g = Graph::new();
        let two = {
            let mut nd = g.new_operation("Const", "two").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            let mut value = Tensor::new(&[1]);
            value[0] = 2.0f32;
            nd.set_attr_tensor("value", value).unwrap();
            nd.finish().unwrap()
        };
        let x = {
            let mut nd = g.new_operation("Placeholder", "x").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            nd.set_attr_shape("shape", &Shape(Some(vec![]))).unwrap();
            nd.finish().unwrap()
        };
        let y = {
            let mut nd = g.new_operation("Mul", "y").unwrap();
            nd.add_input(Output {
                             operation: two,
                             index: 0,
                         });
            nd.add_input(Output {
                             operation: x.clone(),
                             index: 0,
                         });
            nd.finish().unwrap()
        };
        let options = SessionOptions::new();
        match Session::new(&options, &g) {
            Ok(session) => (session, x, y),
            Err(status) => panic!("Creating session failed with status: {}", status),
        }
    }

    #[test]
    fn smoke() {
        create_session();
    }

    #[test]
    fn test_close() {
        let (mut session, _, _) = create_session();
        let status = session.close();
        assert!(status.is_ok());
    }

    #[test]
    fn test_run() {
        let (mut session, x_operation, y_operation) = create_session();
        let mut x = <Tensor<f32>>::new(&[2]);
        x[0] = 2.0;
        x[1] = 3.0;
        let mut step = StepWithGraph::new();
        step.add_input(&x_operation, 0, &x);
        let output_token = step.request_output(&y_operation, 0);
        session.run(&mut step).unwrap();
        let output_tensor = step.take_output::<f32>(output_token).unwrap();
        assert_eq!(output_tensor.len(), 2);
        assert_eq!(output_tensor[0], 4.0);
        assert_eq!(output_tensor[1], 6.0);
    }

    #[test]
    fn test_savedmodelbundle() {
        let mut graph = Graph::new();
        let bundle = SavedModelBundle::load(
            &SessionOptions::new(),
            &["train", "serve"],
            &mut graph,
            "test_resources/regression-model",
        ).unwrap();

        let x_op = graph.operation_by_name_required("x").unwrap();
        let y_op = graph.operation_by_name_required("y").unwrap();
        let y_hat_op = graph.operation_by_name_required("y_hat").unwrap();
        let _train_op = graph.operation_by_name_required("train").unwrap();

        let SavedModelBundle {
            mut session,
            meta_graph_def,
        } = bundle;

        assert!(!meta_graph_def.is_empty());

        let mut x = <Tensor<f32>>::new(&[1]);
        x[0] = 2.0;
        let mut y = <Tensor<f32>>::new(&[1]);
        y[0] = 4.0;
        let mut step = StepWithGraph::new();
        step.add_input(&x_op, 0, &x);
        step.add_input(&y_op, 0, &y);
        let output_token = step.request_output(&y_hat_op, 0);
        session.run(&mut step).unwrap();
        let output_tensor = step.take_output::<f32>(output_token).unwrap();
        assert_eq!(output_tensor.len(), 1);
    }
}
<commit_msg>Add Session::device_list()<commit_after>use tf;
use libc::{c_char, c_int};
use std::ffi::CStr;
use std::ffi::CString;
use std::marker;
use std::path::Path;
use std::ptr;
use super::AnyTensor;
use super::{Buffer, BufferTrait};
use super::Code;
use super::DataType;
use super::Graph;
use super::GraphTrait;
use super::Operation;
use super::OperationTrait;
use super::Result;
use super::SessionOptions;
use super::Status;
use super::Tensor;
use super::TensorType;

/// Aggregation type for a saved model bundle.
#[derive(Debug)]
pub struct SavedModelBundle {
    /// The loaded session.
    pub session: Session,
    /// A meta graph definition as raw protocol buffer.
    pub meta_graph_def: Vec<u8>,
}

impl SavedModelBundle {

    /// Loads a session from an exported model, creating a bundle
    pub fn load<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
        (options: &SessionOptions,
         tags: Tags,
         graph: &mut Graph,
         export_dir: P)
         -> Result<SavedModelBundle> {
        let mut status = Status::new();

        let export_dir_cstr =
            export_dir.as_ref()
                     .to_str()
                     .and_then(|s| CString::new(s.as_bytes()).ok())
                     .ok_or_else(|| invalid_arg!("Invalid export directory path"))?;

        let tags_cstr: Vec<_> = tags.into_iter()
                                    .map(|t| CString::new(t.as_ref()))
                                    .collect::<::std::result::Result<_, _>>()
                                    .map_err(|_| invalid_arg!("Invalid tag name"))?;
        let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();

        // The empty TF_Buffer will be filled by LoadSessionFromSavedModel
        let mut meta = unsafe { Buffer::<u8>::from_ptr(ptr::null_mut(), 0) };

        let inner = unsafe {
            tf::TF_LoadSessionFromSavedModel(options.inner,
                                             ptr::null(),
                                             export_dir_cstr.as_ptr(),
                                             tags_ptr.as_ptr(),
                                             tags_ptr.len() as c_int,
                                             graph.inner(),
                                             meta.inner_mut(),
                                             status.inner())
        };
        if inner.is_null() {
            Err(status)
        } else {
            let session = Session { inner: inner };
            Ok(SavedModelBundle {
                session: session,
                meta_graph_def: Vec::from(meta.as_ref())
            })
        }
    }

}

/// Manages a single graph and execution.
#[derive(Debug)]
pub struct Session {
    inner: *mut tf::TF_Session,
}

impl Session {
    /// Creates a session.
    pub fn new(options: &SessionOptions, graph: &Graph) -> Result<Self> {
        let mut status = Status::new();
        let inner = unsafe { tf::TF_NewSession(graph.inner(), options.inner, status.inner()) };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Loads a session from an exported model.
    pub fn from_saved_model<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
        (options: &SessionOptions,
         tags: Tags,
         graph: &mut Graph,
         export_dir: P)
         -> Result<Self> {
        let mut status = Status::new();

        let export_dir_cstr = export_dir.as_ref()
                     .to_str()
                     .and_then(|s| CString::new(s.as_bytes()).ok())
                     .ok_or_else(|| invalid_arg!("Invalid export directory path"))?;

        let tags_cstr: Vec<_> = tags.into_iter()
                                    .map(|t| CString::new(t.as_ref()))
                                    .collect::<::std::result::Result<_, _>>()
                                    .map_err(|_| invalid_arg!("Invalid tag name"))?;
        // keeping tags_cstr to retain strings in memory
        let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();

        let inner = unsafe {
            tf::TF_LoadSessionFromSavedModel(options.inner,
                                             ptr::null(),
                                             export_dir_cstr.as_ptr(),
                                             tags_ptr.as_ptr(),
                                             tags_ptr.len() as c_int,
                                             graph.inner(),
                                             ptr::null_mut(),
                                             status.inner())
        };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Closes the session.
    pub fn close(&mut self) -> Result<()> {
        let mut status = Status::new();
        unsafe {
            tf::TF_CloseSession(self.inner, status.inner());
        }
        status.into_result()
    }

    /// Runs the graph, feeding the inputs and then fetching the outputs requested in the step.
    pub fn run(&mut self, step: &mut StepWithGraph) -> Result<()> {
        // In case we're running it a second time and not all outputs were taken out.
        step.drop_output_tensors();

        let mut status = Status::new();
        let maybe_tensors: Result<_> = step.input_tensors.iter().map(|t| t.inner()).collect();
        let input_tensors: Vec<_> = maybe_tensors?;
        unsafe {
            tf::TF_SessionRun(self.inner,
                              ptr::null(),
                              step.input_ports.as_ptr(),
                              input_tensors.as_ptr() as *const *const tf::TF_Tensor,
                              input_tensors.len() as c_int,
                              step.output_ports.as_ptr(),
                              step.output_tensors.as_mut_ptr(),
                              step.output_tensors.len() as c_int,
                              step.target_operations.as_mut_ptr(),
                              step.target_operations.len() as c_int,
                              ptr::null_mut(),
                              status.inner());
        };
        status.into_result()
    }

    /// Lists all devices in a session.
    pub fn device_list(&self) -> Result<Vec<Device>> {
        let status = Status::new();
        unsafe {
            let list = tf::TF_SessionListDevices(self.inner, status.inner);
            if !status.is_ok() {
                return Err(status);
            }
            let result = (|| {
                let n = tf::TF_DeviceListCount(list);
                let mut devices = Vec::with_capacity(n as usize);
                for i in 0..n {
                    let c_name = tf::TF_DeviceListName(list, i, status.inner);
                    if !status.is_ok() {
                        return Err(status);
                    }
                    let c_type = tf::TF_DeviceListType(list, i, status.inner);
                    if !status.is_ok() {
                        return Err(status);
                    }
                    let bytes = tf::TF_DeviceListMemoryBytes(list, i, status.inner);
                    if !status.is_ok() {
                        return Err(status);
                    }
                    devices.push(Device {
                        name: CStr::from_ptr(c_name).to_str()?.to_string(),
                        device_type: CStr::from_ptr(c_type).to_str()?.to_string(),
                        memory_bytes: bytes,
                    });
                }
                Ok(devices)
            })();
            tf::TF_DeleteDeviceList(list);
            result
        }
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        let mut status = Status::new();
        unsafe {
            tf::TF_DeleteSession(self.inner, status.inner());
        }
        // TODO: What do we do with the status?
    }
}

////////////////////////

/// An opaque token for retrieving an output from a computation.
#[derive(Copy,Clone,Debug)]
pub struct OutputToken {
    index: usize,
}

/// Manages the inputs and outputs for a single execution of a graph.
///
/// Typical usage involves creating an instance of this struct,
/// adding some inputs to it, requesting some outputs, passing it to `Session::run`
/// and then taking the outputs out of it.
///
/// This will be renamed to Step once the old API goes away.
#[derive(Debug)]
pub struct StepWithGraph<'l> {
    input_ports: Vec<tf::TF_Output>,
    input_tensors: Vec<&'l AnyTensor>,

    output_ports: Vec<tf::TF_Output>,
    output_tensors: Vec<*mut tf::TF_Tensor>,

    target_operations: Vec<*const tf::TF_Operation>,

    phantom: marker::PhantomData<&'l ()>,
}

impl<'l> StepWithGraph<'l> {
    /// Creates a StepWithGraph.
    pub fn new() -> Self {
        StepWithGraph {
            input_ports: vec![],
            input_tensors: vec![],

            output_ports: vec![],
            output_tensors: vec![],

            target_operations: vec![],

            phantom: marker::PhantomData,
        }
    }

    /// Adds an input to be fed to the graph.
    pub fn add_input<T: TensorType>(&mut self,
                                    operation: &Operation,
                                    index: c_int,
                                    tensor: &'l Tensor<T>) {
        self.input_ports.push(tf::TF_Output {
                                  oper: operation.inner(),
                                  index: index,
                              });
        self.input_tensors.push(tensor);
    }

    /// Requests that an output is fetched from the graph after running this step.
    /// Returns an index that you can then use to fetch this output from the step after running it.
    pub fn request_output(&mut self, operation: &Operation, index: c_int) -> OutputToken {
        self.output_ports.push(tf::TF_Output {
                                   oper: operation.inner(),
                                   index: index,
                               });
        self.output_tensors.push(ptr::null_mut());
        OutputToken { index: self.output_tensors.len() - 1 }
    }

    /// Extracts a tensor output given an index. A given index can only be extracted
    /// once per `Session::run`.
    /// Returns an error if output_idx is out of range, output is unavailable or the
    /// requested type does not match the type of the actual tensor.
    pub fn take_output<T: TensorType>(&mut self, token: OutputToken) -> Result<Tensor<T>> {
        let output_idx = token.index;
        if output_idx >= self.output_tensors.len() {
            return Err(Status::new_set(Code::OutOfRange,
                                       &format!("Requested output index is out of range: {} vs \
                                                 {}",
                                                output_idx,
                                                self.output_tensors.len()))
                               .unwrap());
        }
        if self.output_tensors[output_idx].is_null() {
            return Err(Status::new_set(Code::Unavailable,
                                       "Output not available. Either it was already taken, or \
                                        this step has not been sucessfully run yet.")
                               .unwrap());
        }
        let actual_data_type = self.output_data_type(output_idx).unwrap();
        if actual_data_type != T::data_type() {
            return Err(invalid_arg!("Requested tensor type does not match actual tensor type: \
                                     {} vs {}",
                                    actual_data_type,
                                    T::data_type()));
        }
        let tensor = unsafe { Tensor::from_tf_tensor(self.output_tensors[output_idx]).unwrap() };
        self.output_tensors[output_idx] = ptr::null_mut();
        Ok(tensor)
    }

    /// Adds a target operation to be executed when running the graph.
    pub fn add_target(&mut self, operation: &Operation) {
        self.target_operations.push(operation.inner());
    }

    /// Retuns the type of the tensor given an index.
    /// Returns `None` if the index is out of range or the output is not yet available.
    pub fn output_data_type(&self, output_idx: usize) -> Option<DataType> {
        if output_idx >= self.output_tensors.len() {
            return None;
        }
        if self.output_tensors[output_idx].is_null() {
            return None;
        }
        unsafe { Some(DataType::from_c(tf::TF_TensorType(self.output_tensors[output_idx]))) }
    }

    fn drop_output_tensors(&mut self) {
        for mut tensor in &mut self.output_tensors {
            // TODO: Is TF_DeleteTensor NULL safe?
            if !tensor.is_null() {
                unsafe {
                    tf::TF_DeleteTensor(*tensor);
                }
            }
            *tensor = ptr::null_mut();
        }
    }
}

impl<'l> Drop for StepWithGraph<'l> {
    fn drop(&mut self) {
        self.drop_output_tensors();
    }
}

////////////////////////

/// Metadata about a device.
#[derive(Debug,Eq,PartialEq,Clone,Hash)]
pub struct Device {
    /// Full name of the device (e.g. /job:worker/replica:0/...)
    pub name: String,

    /// Type of device.
    pub device_type: String,

    /// Amount of memory on the device.
    pub memory_bytes: i64,
}

////////////////////////

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::DataType;
    use super::super::Graph;
    use super::super::Operation;
    use super::super::Output;
    use super::super::SessionOptions;
    use super::super::Shape;
    use super::super::Tensor;

    fn create_session() -> (Session, Operation, Operation) {
        let mut g = Graph::new();
        let two = {
            let mut nd = g.new_operation("Const", "two").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            let mut value = Tensor::new(&[1]);
            value[0] = 2.0f32;
            nd.set_attr_tensor("value", value).unwrap();
            nd.finish().unwrap()
        };
        let x = {
            let mut nd = g.new_operation("Placeholder", "x").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            nd.set_attr_shape("shape", &Shape(Some(vec![]))).unwrap();
            nd.finish().unwrap()
        };
        let y = {
            let mut nd = g.new_operation("Mul", "y").unwrap();
            nd.add_input(Output {
                             operation: two,
                             index: 0,
                         });
            nd.add_input(Output {
                             operation: x.clone(),
                             index: 0,
                         });
            nd.finish().unwrap()
        };
        let options = SessionOptions::new();
        match Session::new(&options, &g) {
            Ok(session) => (session, x, y),
            Err(status) => panic!("Creating session failed with status: {}", status),
        }
    }

    #[test]
    fn smoke() {
        create_session();
    }

    #[test]
    fn test_close() {
        let (mut session, _, _) = create_session();
        let status = session.close();
        assert!(status.is_ok());
    }

    #[test]
    fn test_run() {
        let (mut session, x_operation, y_operation) = create_session();
        let mut x = <Tensor<f32>>::new(&[2]);
        x[0] = 2.0;
        x[1] = 3.0;
        let mut step = StepWithGraph::new();
        step.add_input(&x_operation, 0, &x);
        let output_token = step.request_output(&y_operation, 0);
        session.run(&mut step).unwrap();
        let output_tensor = step.take_output::<f32>(output_token).unwrap();
        assert_eq!(output_tensor.len(), 2);
        assert_eq!(output_tensor[0], 4.0);
        assert_eq!(output_tensor[1], 6.0);
    }

    #[test]
    fn test_savedmodelbundle() {
        let mut graph = Graph::new();
        let bundle = SavedModelBundle::load(
            &SessionOptions::new(),
            &["train", "serve"],
            &mut graph,
            "test_resources/regression-model",
        ).unwrap();

        let x_op = graph.operation_by_name_required("x").unwrap();
        let y_op = graph.operation_by_name_required("y").unwrap();
        let y_hat_op = graph.operation_by_name_required("y_hat").unwrap();
        let _train_op = graph.operation_by_name_required("train").unwrap();

        let SavedModelBundle {
            mut session,
            meta_graph_def,
        } = bundle;

        assert!(!meta_graph_def.is_empty());

        let mut x = <Tensor<f32>>::new(&[1]);
        x[0] = 2.0;
        let mut y = <Tensor<f32>>::new(&[1]);
        y[0] = 4.0;
        let mut step = StepWithGraph::new();
        step.add_input(&x_op, 0, &x);
        step.add_input(&y_op, 0, &y);
        let output_token = step.request_output(&y_hat_op, 0);
        session.run(&mut step).unwrap();
        let output_tensor = step.take_output::<f32>(output_token).unwrap();
        assert_eq!(output_tensor.len(), 1);
    }

    #[test]
    fn test_device_list() {
        let (session, _, _) = create_session();
        let devices = session.device_list().unwrap();
        assert!(devices.iter().any(|d| d.device_type == "CPU"), "devices: {:?}", devices);
    }
}
<|endoftext|>