<commit_before>
//! Immediate mode gui. See [`ui::Ui`](struct.Ui.html) for more info.

use std::mem;
use gl;
use gl::types::*;
use cable_math::Vec2;

use color::Color;
use font::{Font, CachedFont};
use input::{InputManager, Key, State};
use matrix_stack::MatrixStack;
use shader::{Shader, ShaderPrototype};
use buffer::{Vertex, VertexBuffer, PrimitiveMode, BufferUsage};

const FONT_SIZE: f32 = 14.0;

/// A struct for using a imediate mode gui.
pub struct Ui {
    pub style: Style,

    mat_stack: MatrixStack,
    font: CachedFont,
    shader: Shader,
    draw_data: Vec<Vert>,
    draw_vbo: VertexBuffer<Vert>,

    held: Option<Id>,

    // Input state
    mouse_pos: Vec2<f32>,
    mouse_state: State,
}

impl Ui {
    /// Creates a new imediate mode gui system with the given font. Note that the font will be
    /// copied internally, so you can pass a reference to a font you are using elsewhere in your
    /// program.
    pub fn new(font: &Font) -> Ui {
        Ui {
            style: Default::default(),

            mat_stack: MatrixStack::new(),
            font: CachedFont::from_font(font.clone()),
            shader: build_shader(),
            draw_data: Vec::with_capacity(500),
            draw_vbo: VertexBuffer::with_capacity(PrimitiveMode::Triangles, BufferUsage::DynamicDraw, 500),

            held: None,

            mouse_pos: Vec2::zero(),
            mouse_state: State::Up,
        }
    }

    /// Updates this imgui system. This should be called once per frame, before using any of the
    /// gui creation functions.
    pub fn update(&mut self, input: &InputManager, window_size: Vec2<u32>) {
        self.mat_stack.ortho(0.0, window_size.x as f32, 0.0, window_size.y as f32, -1.0, 1.0);

        self.mouse_pos = input.mouse_pos();
        self.mouse_state = input.mouse_key(0);

        if self.mouse_state.up() && !self.mouse_state.released() {
            self.held = None;
        }
    }

    /// Shows all components added since the last call to `draw`. This function update the matrix
    /// buffers and binds new shaders. No special opengl state is required to be set when calling
    /// this function. Note that this function does not necessarily reset the state it changes.
    pub fn draw(&mut self) {
        self.mat_stack.update_buffer();

        self.draw_vbo.clear();
        self.draw_vbo.put(0, &self.draw_data);
        self.draw_data.clear();

        self.shader.bind();
        self.draw_vbo.draw();
        self.font.draw();
    }

    /// Shows a new button with the given text at the given location. Returns true if the button
    /// was pressed. Note that this function needs to be called every frame you want to see the
    /// button.
    pub fn button(&mut self, text: &str, pos: Vec2<f32>) -> bool {
        let width = self.font.font().width(text, FONT_SIZE) + self.style.padding.x;
        let height = self.font.font().line_height(FONT_SIZE) + self.style.padding.y;
        let text_start = self.style.padding.y/2.0 - self.font.font().descent(FONT_SIZE);

        let hovered = self.mouse_pos.x > pos.x && self.mouse_pos.y > pos.y &&
                      self.mouse_pos.x < pos.x + width && self.mouse_pos.y < pos.y + height;
        let held = hovered && self.mouse_state.down();

        let color = if held {
            self.style.hold_color
        } else if hovered {
            self.style.hover_color
        } else {
            self.style.button_color
        };

        quad(&mut self.draw_data, pos, Vec2::new(width, height), color);
        self.font.cache(text, FONT_SIZE, pos + Vec2::new(self.style.padding.x/2.0, height - text_start));

        let id = Id::from_str(text, CompType::Button);

        if hovered && self.mouse_state.pressed() {
            self.held = Some(id);
        }

        self.held == Some(id) && hovered && self.mouse_state.released()
    }
}

#[derive(Clone, Debug)]
pub struct Style {
    pub button_color: Color,
    pub hover_color: Color,
    pub hold_color: Color,

    pub padding: Vec2<f32>,
}
impl Default for Style {
    fn default() -> Style {
        Style {
            button_color: Color::hex("4c4665"),
            hover_color: Color::hex("575074"),
            hold_color: Color::hex("413c56"),
            padding: Vec2::new(10.0, 6.0),
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Id(u64, CompType);
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CompType {
    Button,
}

impl Id {
    fn from_str(text: &str, ty: CompType) -> Id {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hasher, Hash};

        let mut hasher = DefaultHasher::new();
        text.hash(&mut hasher);
        let id = hasher.finish();

        Id(id, ty)
    }
}

#[derive(Debug)]
#[repr(C)]
struct Vert {
    pos: Vec2<f32>,
    color: Color,
}

fn quad(buf: &mut Vec<Vert>, pos: Vec2<f32>, size: Vec2<f32>, color: Color){
    let min = pos;
    let max = pos + size;

    buf.push(Vert { pos: Vec2::new(min.x, min.y), color: color });
    buf.push(Vert { pos: Vec2::new(max.x, min.y), color: color });
    buf.push(Vert { pos: Vec2::new(max.x, max.y), color: color });

    buf.push(Vert { pos: Vec2::new(min.x, min.y), color: color });
    buf.push(Vert { pos: Vec2::new(max.x, max.y), color: color });
    buf.push(Vert { pos: Vec2::new(min.x, max.y), color: color });
}

// We cannot use the custom derive from within this crate
impl Vertex for Vert {
    fn bytes_per_vertex() -> usize { mem::size_of::<Vert>() }
    fn setup_attrib_pointers() {
        let stride = <Vert as Vertex>::bytes_per_vertex();
        let mut offset = 0;
        unsafe {
            gl::EnableVertexAttribArray(0);
            gl::VertexAttribPointer(0, 2, gl::FLOAT,
                                    false as GLboolean,
                                    stride as GLsizei, offset as *const GLvoid);
            offset += mem::size_of::<Vec2<f32>>();

            gl::EnableVertexAttribArray(1);
            gl::VertexAttribPointer(1, 4, gl::FLOAT,
                                    false as GLboolean,
                                    stride as GLsizei, offset as *const GLvoid);
        }
    }
    // Not used, we manualy declare inputs in the shader
    fn gen_shader_input_decl() -> String { String::new() }
}

const VERT_SRC: &'static str = "
    #version 330 core

    layout(location = 0) in vec2 pos;
    layout(location = 1) in vec4 color;

    out vec4 vert_col;

    // Matrix block is inserted automatically

    void main() {
        gl_Position = mvp * vec4(pos, 0.0, 1.0);
        vert_col = color;
    }
";
const FRAG_SRC: &'static str = "
    #version 330 core

    in vec4 vert_col;
    out vec4 color;

    void main() {
        color = vert_col;
    }
";

fn build_shader() -> Shader {
    let mut proto = ShaderPrototype::new_prototype(VERT_SRC, "", FRAG_SRC);
    proto.bind_to_matrix_storage();
    match proto.build() {
        Ok(shader) => shader,
        Err(err) => {
            println!("{}", err); // Print the error properly
            panic!();
        }
    }
}

<commit_msg>Added a internal caret to the ui, for sequential component placement<commit_after>
//! Immediate mode gui. See [`ui::Ui`](struct.Ui.html) for more info.

use std::mem;
use gl;
use gl::types::*;
use cable_math::Vec2;

use color::Color;
use font::{Font, CachedFont};
use input::{InputManager, Key, State};
use matrix_stack::MatrixStack;
use shader::{Shader, ShaderPrototype};
use buffer::{Vertex, VertexBuffer, PrimitiveMode, BufferUsage};

const FONT_SIZE: f32 = 14.0;

/// A struct for using a imediate mode gui.
pub struct Ui {
    pub style: Style,

    mat_stack: MatrixStack,
    font: CachedFont,
    shader: Shader,
    draw_data: Vec<Vert>,
    draw_vbo: VertexBuffer<Vert>,

    caret: Vec2<f32>,
    caret_start: Vec2<f32>,
    line_size: f32,
    line_dir: LineDir,
    held: Option<Id>,

    // Input state
    mouse_pos: Vec2<f32>,
    mouse_state: State,
}

impl Ui {
    /// Creates a new imediate mode gui system with the given font. Note that the font will be
    /// copied internally, so you can pass a reference to a font you are using elsewhere in your
    /// program.
    pub fn new(font: &Font) -> Ui {
        Ui {
            style: Default::default(),

            mat_stack: MatrixStack::new(),
            font: CachedFont::from_font(font.clone()),
            shader: build_shader(),
            draw_data: Vec::with_capacity(500),
            draw_vbo: VertexBuffer::with_capacity(PrimitiveMode::Triangles, BufferUsage::DynamicDraw, 500),

            caret: Vec2::zero(),
            caret_start: Vec2::zero(),
            line_size: 0.0,
            line_dir: LineDir::Vertical,
            held: None,

            mouse_pos: Vec2::zero(),
            mouse_state: State::Up,
        }
    }

    /// Updates this imgui system. This should be called once per frame, before using any of the
    /// gui creation functions.
    pub fn update(&mut self, input: &InputManager, window_size: Vec2<u32>) {
        self.mat_stack.ortho(0.0, window_size.x as f32, 0.0, window_size.y as f32, -1.0, 1.0);

        self.mouse_pos = input.mouse_pos();
        self.mouse_state = input.mouse_key(0);

        if self.mouse_state.up() && !self.mouse_state.released() {
            self.held = None;
        }

        self.caret = Vec2::zero();
    }

    /// Shows all components added since the last call to `draw`. This function update the matrix
    /// buffers and binds new shaders. No special opengl state is required to be set when calling
    /// this function. Note that this function does not necessarily reset the state it changes.
    pub fn draw(&mut self) {
        self.mat_stack.update_buffer();

        self.draw_vbo.clear();
        self.draw_vbo.put(0, &self.draw_data);
        self.draw_data.clear();

        self.shader.bind();
        self.draw_vbo.draw();
        self.font.draw();
    }

    /// Moves the internal caret to the given position. Consecutive items will be inserted at
    /// the caret.
    pub fn set_caret(&mut self, pos: Vec2<f32>, line_dir: LineDir) {
        self.caret = pos;
        self.caret_start = pos;
        self.line_dir = line_dir;
        self.line_size = 0.0;
    }

    /// Advances the caret to the next line. The direction of a line depends on the line direction
    /// set by [`set_caret`].
    ///
    /// [`set_caret`]: struct.Ui.html#method.set_caret
    pub fn next_line(&mut self) {
        match self.line_dir {
            LineDir::Horizontal => {
                self.caret.y += self.line_size + self.style.line_spacing;
                self.caret.x = self.caret_start.x;
                self.line_size = 0.0;
            },
            LineDir::Vertical => {
                self.caret.x += self.line_size + self.style.line_spacing;
                self.caret.y = self.caret_start.y;
                self.line_size = 0.0;
            },
        }
    }

    fn advance_caret(&mut self, comp_width: f32, comp_height: f32) {
        match self.line_dir {
            LineDir::Horizontal => {
                self.caret.x += comp_width + self.style.line_spacing;
                self.line_size = f32::max(comp_height, self.line_size);
            },
            LineDir::Vertical => {
                self.caret.y += comp_height + self.style.line_spacing;
                self.line_size = f32::max(comp_width, self.line_size);
            },
        }
    }

    /// Shows a new button with the given text at the given location. Returns true if the button
    /// was pressed. Note that this function needs to be called every frame you want to see the
    /// button.
    pub fn button(&mut self, text: &str) -> bool {
        let id = Id::from_str(text, CompType::Button);

        let width = self.font.font().width(text, FONT_SIZE) + self.style.internal_padding.x;
        let height = self.font.font().line_height(FONT_SIZE) + self.style.internal_padding.y;
        let pos = self.caret;
        self.advance_caret(width, height);

        let text_start = self.style.internal_padding.y/2.0 - self.font.font().descent(FONT_SIZE);

        let hovered = self.mouse_pos.x > pos.x && self.mouse_pos.y > pos.y &&
                      self.mouse_pos.x < pos.x + width && self.mouse_pos.y < pos.y + height;
        let held = hovered && self.mouse_state.down();

        let color = if held {
            self.style.hold_color
        } else if hovered {
            self.style.hover_color
        } else {
            self.style.button_color
        };
        quad(&mut self.draw_data, pos, Vec2::new(width, height), color);
        self.font.cache(text, FONT_SIZE, pos + Vec2::new(self.style.internal_padding.x/2.0, height - text_start));

        if hovered && self.mouse_state.pressed() {
            self.held = Some(id);
        }

        self.held == Some(id) && hovered && self.mouse_state.released()
    }
}

#[derive(Clone, Debug)]
pub struct Style {
    pub button_color: Color,
    pub hover_color: Color,
    pub hold_color: Color,

    pub internal_padding: Vec2<f32>,
    pub line_spacing: f32,
}
impl Default for Style {
    fn default() -> Style {
        Style {
            button_color: Color::hex("4c4665"),
            hover_color: Color::hex("575074"),
            hold_color: Color::hex("413c56"),
            internal_padding: Vec2::new(10.0, 6.0),
            line_spacing: 5.0,
        }
    }
}

pub enum LineDir {
    /// Components are layed out below each other
    Vertical,
    /// Components are layed out side by side
    Horizontal,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Id(u64, CompType);
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CompType {
    Button,
}

impl Id {
    fn from_str(text: &str, ty: CompType) -> Id {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hasher, Hash};

        let mut hasher = DefaultHasher::new();
        text.hash(&mut hasher);
        let id = hasher.finish();

        Id(id, ty)
    }
}

#[derive(Debug)]
#[repr(C)]
struct Vert {
    pos: Vec2<f32>,
    color: Color,
}

fn quad(buf: &mut Vec<Vert>, pos: Vec2<f32>, size: Vec2<f32>, color: Color){
    let min = pos;
    let max = pos + size;

    buf.push(Vert { pos: Vec2::new(min.x, min.y), color: color });
    buf.push(Vert { pos: Vec2::new(max.x, min.y), color: color });
    buf.push(Vert { pos: Vec2::new(max.x, max.y), color: color });

    buf.push(Vert { pos: Vec2::new(min.x, min.y), color: color });
    buf.push(Vert { pos: Vec2::new(max.x, max.y), color: color });
    buf.push(Vert { pos: Vec2::new(min.x, max.y), color: color });
}

// We cannot use the custom derive from within this crate
impl Vertex for Vert {
    fn bytes_per_vertex() -> usize { mem::size_of::<Vert>() }
    fn setup_attrib_pointers() {
        let stride = <Vert as Vertex>::bytes_per_vertex();
        let mut offset = 0;
        unsafe {
            gl::EnableVertexAttribArray(0);
            gl::VertexAttribPointer(0, 2, gl::FLOAT,
                                    false as GLboolean,
                                    stride as GLsizei, offset as *const GLvoid);
            offset += mem::size_of::<Vec2<f32>>();

            gl::EnableVertexAttribArray(1);
            gl::VertexAttribPointer(1, 4, gl::FLOAT,
                                    false as GLboolean,
                                    stride as GLsizei, offset as *const GLvoid);
        }
    }
    // Not used, we manualy declare inputs in the shader
    fn gen_shader_input_decl() -> String { String::new() }
}

const VERT_SRC: &'static str = "
    #version 330 core

    layout(location = 0) in vec2 pos;
    layout(location = 1) in vec4 color;

    out vec4 vert_col;

    // Matrix block is inserted automatically

    void main() {
        gl_Position = mvp * vec4(pos, 0.0, 1.0);
        vert_col = color;
    }
";
const FRAG_SRC: &'static str = "
    #version 330 core

    in vec4 vert_col;
    out vec4 color;

    void main() {
        color = vert_col;
    }
";

fn build_shader() -> Shader {
    let mut proto = ShaderPrototype::new_prototype(VERT_SRC, "", FRAG_SRC);
    proto.bind_to_matrix_storage();
    match proto.build() {
        Ok(shader) => shader,
        Err(err) => {
            println!("{}", err); // Print the error properly
            panic!();
        }
    }
}

<|endoftext|>


<commit_before>// FIXME: Currently, there is one VM. Ideally, we'd have multiple VMs that
// can be used for different purposes. For example, if we don't need captures
// (i.e., 'is_match'), then the VM can be drastically simplified and allocation
// reduced.
//
// So the TODO item here is to come up with a better abstraction for writing
// a VM. A nice start would be a trait, but I suspect that multiple VMs will
// share a lot of code. So find a way to reuse it.

use std::cmp;
use std::iter;
use std::mem;
use super::compile::{Inst, Char, CharClass, Any,
                     EmptyBegin, EmptyEnd, EmptyWordBoundary,
                     Match, Save, Jump, Split};

pub type CaptureIndices = Vec<Option<(uint, uint)>>;

pub fn run(insts: &[Inst], input: &[char], caps: bool) -> CaptureIndices {
    unflatten_capture_locations(Vm {
        insts: insts,
        input: input,
        caps: caps,
    }.run())
}

fn unflatten_capture_locations(locs: Vec<Option<uint>>) -> CaptureIndices {
    let mut caps = Vec::with_capacity(locs.len() / 2);
    for win in locs.as_slice().chunks(2) {
        match (win[0], win[1]) {
            (Some(s), Some(e)) => caps.push(Some((s, e))),
            (None, None) => caps.push(None),
            wins => fail!("BUG: Invalid capture group: {}", wins),
        }
    }
    caps
}

struct Thread {
    pc: uint,
    groups: Vec<Option<uint>>,
}

impl Thread {
    fn new(pc: uint, groups: Vec<Option<uint>>) -> Thread {
        Thread { pc: pc, groups: groups }
    }
}

struct Threads {
    queue: Vec<Thread>,
    sparse: Vec<uint>,
    size: uint,
}

impl Threads {
    // This is using a wicked neat trick to provide constant time lookup
    // for threads in the queue using a sparse set. A queue of threads is
    // allocated once with maximal size when the VM initializes and is reused
    // throughout execution.
    //
    // See http://research.swtch.com/sparse for the deets.
    fn new(num_insts: uint, num_caps: uint) -> Threads {
        Threads {
            queue: Vec::from_fn(num_insts, |_| {
                Thread::new(0, Vec::from_elem(num_caps, None))
            }),
            sparse: Vec::from_elem(num_insts, 0u),
            size: 0,
        }
    }

    fn add(&mut self, pc: uint, groups: &[Option<uint>], empty: bool) {
        let t = self.queue.get_mut(self.size);
        t.pc = pc;
        if !empty {
            let mut i = 0;
            while i < groups.len() {
                *t.groups.get_mut(i) = groups[i];
                i += 1;
            }
        }
        *self.sparse.get_mut(pc) = self.size;
        self.size += 1;
    }

    #[inline(always)]
    fn contains(&self, pc: uint) -> bool {
        let s = *self.sparse.get(pc);
        s < self.size && self.queue.get(s).pc == pc
    }

    fn empty(&mut self) {
        self.size = 0;
    }

    fn pc(&self, i: uint) -> uint {
        self.queue.get(i).pc
    }

    fn groups<'r>(&'r mut self, i: uint) -> &'r mut [Option<uint>] {
        self.queue.get_mut(i).groups.as_mut_slice()
    }
}

struct Vm<'r> {
    insts: &'r [Inst],
    input: &'r [char],
    caps: bool,
}

impl<'r> Vm<'r> {
    fn run(&self) -> Vec<Option<uint>> {
        let num_caps = numcaps(self.insts);
        let mut clist = Threads::new(self.insts.len(), num_caps);
        let mut nlist = Threads::new(self.insts.len(), num_caps);

        let mut groups = Vec::from_elem(num_caps, None);
        self.add(&mut clist, 0, 0, groups.as_mut_slice());

        for ic in iter::range_inclusive(0, self.input.len()) {
            if clist.size == 0 && nlist.size == 0 {
                break
            }
            let mut i = 0;
            while i < clist.size {
                let pc = clist.pc(i);
                match self.insts[pc] {
                    Match => {
                        if !self.caps {
                            // FIXME: This is a terrible hack that is used to
                            // indicate a match when the caller doesn't want
                            // any capture groups.
                            // The right way to do this, I think, is to create
                            // a separate VM for non-capturing search.
                            groups = vec!(Some(0), Some(0))
                        } else {
                            groups = Vec::from_slice(clist.groups(i))
                        }
                        clist.empty();
                    }
                    Char(c, casei) => {
                        if self.char_eq(casei, ic, c) {
                            self.add(&mut nlist, pc+1, ic+1, clist.groups(i));
                        }
                    }
                    CharClass(ref ranges, negate, casei) => {
                        if ic < self.input.len() {
                            let c = self.get(ic);
                            let found = ranges.as_slice();
                            let found = found.bsearch(|&rc| class_cmp(casei, c, rc));
                            let found = found.is_some();
                            if (found && !negate) || (!found && negate) {
                                self.add(&mut nlist, pc+1, ic+1, clist.groups(i));
                            }
                        }
                    }
                    Any(true) =>
                        self.add(&mut nlist, pc+1, ic+1, clist.groups(i)),
                    Any(false) => {
                        if !self.char_eq(false, ic, '\n') {
                            self.add(&mut nlist, pc+1, ic+1, clist.groups(i))
                        }
                    }
                    EmptyBegin(_) | EmptyEnd(_) | EmptyWordBoundary(_)
                    | Save(_) | Jump(_) | Split(_, _) => {},
                }
                i += 1;
            }
            mem::swap(&mut clist, &mut nlist);
            nlist.empty();
        }
        groups
    }

    fn add(&self, nlist: &mut Threads, pc: uint, ic: uint,
           groups: &mut [Option<uint>]) {
        if nlist.contains(pc) {
            return
        }
        // This is absolutely critical to the *correctness* of the VM.
        // TL;DR - It prevents cycles.
        // If we didn't care about cycles, we'd *only* add threads that
        // correspond to non-jumping instructions (Char, Any, Match, etc.).
        // But, it's possible for valid regexs (like '(a*)*') to result in
        // a cycle in the instruction list. e.g., We'll keep chasing the Split
        // instructions forever.
        // So we add these instructions to our thread queue, but in the main
        // VM loop, we look for them but simply ignore them.
        // Adding them to the queue prevents them from being revisited so we
        // can avoid cycles (and the inevitable stack overflow).
        match self.insts[pc] {
            EmptyBegin(multi) => {
                nlist.add(pc, groups, true);
                if self.is_begin(ic) || (multi && self.char_is(ic-1, '\n')) {
                    self.add(nlist, pc + 1, ic, groups)
                }
            }
            EmptyEnd(multi) => {
                nlist.add(pc, groups, true);
                if self.is_end(ic) || (multi && self.char_is(ic, '\n')) {
                    self.add(nlist, pc + 1, ic, groups)
                }
            }
            EmptyWordBoundary(yes) => {
                nlist.add(pc, groups, true);
                let wb = self.is_word_boundary(ic);
                if yes == wb {
                    self.add(nlist, pc + 1, ic, groups)
                }
            }
            Save(slot) => {
                nlist.add(pc, groups, true);
                if !self.caps {
                    self.add(nlist, pc + 1, ic, groups);
                } else {
                    let old = groups[slot];
                    groups[slot] = Some(ic);
                    self.add(nlist, pc + 1, ic, groups);
                    groups[slot] = old;
                }
            }
            Jump(to) => {
                nlist.add(pc, groups, true);
                self.add(nlist, to, ic, groups)
            }
            Split(x, y) => {
                nlist.add(pc, groups, true);
                self.add(nlist, x, ic, groups);
                self.add(nlist, y, ic, groups);
            }
            // Handled in 'run'
            Match | Char(_, _) | CharClass(_, _, _) | Any(_) => {
                nlist.add(pc, groups, !self.caps);
            }
        }
    }

    fn is_begin(&self, ic: uint) -> bool { ic == 0 }
    fn is_end(&self, ic: uint) -> bool { ic == self.input.len() }

    fn is_word_boundary(&self, ic: uint) -> bool {
        if self.is_begin(ic) {
            return self.is_word(ic)
        }
        if self.is_end(ic) {
            return self.is_word(self.input.len()-1)
        }
        (self.is_word(ic) && !self.is_word(ic-1))
        || (self.is_word(ic-1) && !self.is_word(ic))
    }

    fn is_word(&self, ic: uint) -> bool {
        if ic >= self.input.len() {
            return false
        }
        let c = self.input[ic];
        c == '_'
        || (c >= '0' && c <= '9')
        || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
    }

    // FIXME: For case insensitive comparisons, it uses the uppercase
    // character and tests for equality. IIUC, this does not generalize to
    // all of Unicode. I believe we need to check the entire fold for each
    // character. This will be easy to add if and when it gets added to Rust's
    // standard library.
    fn char_eq(&self, casei: bool, ic: uint, regc: char) -> bool {
        if ic >= self.input.len() {
            return false
        }
        let textc = self.get(ic);
        regc == textc || (casei && regc.to_uppercase() == textc.to_uppercase())
    }

    fn char_is(&self, ic: uint, c: char) -> bool {
        ic < self.input.len() && self.input[ic] == c
    }

    fn get(&self, ic: uint) -> char {
        self.input[ic]
    }
}

fn class_cmp(casei: bool, mut textc: char,
             (mut start, mut end): (char, char)) -> Ordering {
    if casei {
        // FIXME: This is pretty ridiculous. All of this case conversion
        // can be moved outside this function:
        // 1) textc should be uppercased outside the bsearch.
        // 2) the character class itself should be uppercased either in the
        //    parser or the compiler.
        // FIXME: This is too simplistic for correct Unicode support.
        //        See also: char_eq
        textc = textc.to_uppercase();
        start = start.to_uppercase();
        end = end.to_uppercase();
    }
    if textc >= start && textc <= end {
        Equal
    } else if start > textc {
        Greater
    } else {
        Less
    }
}

fn numcaps(insts: &[Inst]) -> uint {
    let mut n = 0;
    for inst in insts.iter() {
        match *inst {
            Save(c) => n = cmp::max(n, c+1),
            _ => {}
        }
    }
    n
}
<commit_msg>Fix some comments and add some.<commit_after>// FIXME: Currently, there is one VM. Ideally, we'd have multiple VMs that
// can be used for different purposes. For example, if we don't need captures
// (i.e., 'is_match'), then the VM can be drastically simplified and allocation
// reduced.
//
// So the TODO item here is to come up with a better abstraction for writing
// a VM. A nice start would be a trait, but I suspect that multiple VMs will
// share a lot of code. So find a way to reuse it.

use std::cmp;
use std::iter;
use std::mem;
use super::compile::{Inst, Char, CharClass, Any,
                     EmptyBegin, EmptyEnd, EmptyWordBoundary,
                     Match, Save, Jump, Split};

pub type CaptureIndices = Vec<Option<(uint, uint)>>;

pub fn run(insts: &[Inst], input: &[char], caps: bool) -> CaptureIndices {
    unflatten_capture_locations(Vm {
        insts: insts,
        input: input,
        caps: caps,
    }.run())
}

fn unflatten_capture_locations(locs: Vec<Option<uint>>) -> CaptureIndices {
    let mut caps = Vec::with_capacity(locs.len() / 2);
    for win in locs.as_slice().chunks(2) {
        match (win[0], win[1]) {
            (Some(s), Some(e)) => caps.push(Some((s, e))),
            (None, None) => caps.push(None),
            wins => fail!("BUG: Invalid capture group: {}", wins),
        }
    }
    caps
}

struct Thread {
    pc: uint,
    groups: Vec<Option<uint>>,
}

impl Thread {
    fn new(pc: uint, groups: Vec<Option<uint>>) -> Thread {
        Thread { pc: pc, groups: groups }
    }
}

struct Threads {
    queue: Vec<Thread>,
    sparse: Vec<uint>,
    size: uint,
}

impl Threads {
    // This is using a wicked neat trick to provide constant time lookup
    // for threads in the queue using a sparse set. A queue of threads is
    // allocated once with maximal size when the VM initializes and is reused
    // throughout execution.
    //
    // See http://research.swtch.com/sparse for the deets.
    fn new(num_insts: uint, num_caps: uint) -> Threads {
        Threads {
            queue: Vec::from_fn(num_insts, |_| {
                Thread::new(0, Vec::from_elem(num_caps, None))
            }),
            sparse: Vec::from_elem(num_insts, 0u),
            size: 0,
        }
    }

    fn add(&mut self, pc: uint, groups: &[Option<uint>], empty: bool) {
        let t = self.queue.get_mut(self.size);
        t.pc = pc;
        if !empty {
            let mut i = 0;
            while i < groups.len() {
                *t.groups.get_mut(i) = groups[i];
                i += 1;
            }
        }
        *self.sparse.get_mut(pc) = self.size;
        self.size += 1;
    }

    #[inline(always)]
    fn contains(&self, pc: uint) -> bool {
        let s = *self.sparse.get(pc);
        s < self.size && self.queue.get(s).pc == pc
    }

    fn empty(&mut self) {
        self.size = 0;
    }

    fn pc(&self, i: uint) -> uint {
        self.queue.get(i).pc
    }

    fn groups<'r>(&'r mut self, i: uint) -> &'r mut [Option<uint>] {
        self.queue.get_mut(i).groups.as_mut_slice()
    }
}

struct Vm<'r> {
    insts: &'r [Inst],
    input: &'r [char],
    caps: bool,
}

impl<'r> Vm<'r> {
    fn run(&self) -> Vec<Option<uint>> {
        let num_caps = numcaps(self.insts);
        let mut clist = Threads::new(self.insts.len(), num_caps);
        let mut nlist = Threads::new(self.insts.len(), num_caps);

        let mut groups = Vec::from_elem(num_caps, None);
        self.add(&mut clist, 0, 0, groups.as_mut_slice());

        for ic in iter::range_inclusive(0, self.input.len()) {
            if clist.size == 0 && nlist.size == 0 {
                break
            }
            let mut i = 0;
            while i < clist.size {
                let pc = clist.pc(i);
                match self.insts[pc] {
                    Match => {
                        if !self.caps {
                            // FIXME: This is a terrible hack that is used to
                            // indicate a match when the caller doesn't want
                            // any capture groups.
                            // The right way to do this, I think, is to create
                            // a separate VM for non-capturing search.
                            groups = vec!(Some(0), Some(0))
                        } else {
                            groups = Vec::from_slice(clist.groups(i))
                        }
                        clist.empty();
                    }
                    Char(c, casei) => {
                        if self.char_eq(casei, ic, c) {
                            self.add(&mut nlist, pc+1, ic+1, clist.groups(i));
                        }
                    }
                    CharClass(ref ranges, negate, casei) => {
                        if ic < self.input.len() {
                            let c = self.get(ic);
                            let found = ranges.as_slice();
                            let found = found.bsearch(|&rc| class_cmp(casei, c, rc));
                            let found = found.is_some();
                            if (found && !negate) || (!found && negate) {
                                self.add(&mut nlist, pc+1, ic+1, clist.groups(i));
                            }
                        }
                    }
                    Any(true) =>
                        self.add(&mut nlist, pc+1, ic+1, clist.groups(i)),
                    Any(false) => {
                        if !self.char_eq(false, ic, '\n') {
                            self.add(&mut nlist, pc+1, ic+1, clist.groups(i))
                        }
                    }
                    EmptyBegin(_) | EmptyEnd(_) | EmptyWordBoundary(_)
                    | Save(_) | Jump(_) | Split(_, _) => {},
                }
                i += 1;
            }
            mem::swap(&mut clist, &mut nlist);
            nlist.empty();
        }
        groups
    }

    fn add(&self, nlist: &mut Threads, pc: uint, ic: uint,
           groups: &mut [Option<uint>]) {
        if nlist.contains(pc) {
            return
        }
        // We have to add states to the threads list even if their empty.
        // TL;DR - It prevents cycles.
        // If we didn't care about cycles, we'd *only* add threads that
        // correspond to non-jumping instructions (Char, Any, Match, etc.).
        // But, it's possible for valid regexs (like '(a*)*') to result in
        // a cycle in the instruction list. e.g., We'll keep chasing the Split
        // instructions forever.
        // So we add these instructions to our thread queue, but in the main
        // VM loop, we look for them but simply ignore them.
        // Adding them to the queue prevents them from being revisited so we
        // can avoid cycles (and the inevitable stack overflow).
        //
        // We make a minor optimization by indicating that the state is "empty"
        // so that its capture groups are not filled in.
        match self.insts[pc] {
            EmptyBegin(multi) => {
                nlist.add(pc, groups, true);
                if self.is_begin(ic) || (multi && self.char_is(ic-1, '\n')) {
                    self.add(nlist, pc + 1, ic, groups)
                }
            }
            EmptyEnd(multi) => {
                nlist.add(pc, groups, true);
                if self.is_end(ic) || (multi && self.char_is(ic, '\n')) {
                    self.add(nlist, pc + 1, ic, groups)
                }
            }
            EmptyWordBoundary(yes) => {
                nlist.add(pc, groups, true);
                let wb = self.is_word_boundary(ic);
                if yes == wb {
                    self.add(nlist, pc + 1, ic, groups)
                }
            }
            Save(slot) => {
                nlist.add(pc, groups, true);
                if !self.caps {
                    self.add(nlist, pc + 1, ic, groups);
                } else {
                    let old = groups[slot];
                    groups[slot] = Some(ic);
                    self.add(nlist, pc + 1, ic, groups);
                    groups[slot] = old;
                }
            }
            Jump(to) => {
                nlist.add(pc, groups, true);
                self.add(nlist, to, ic, groups)
            }
            Split(x, y) => {
                nlist.add(pc, groups, true);
                self.add(nlist, x, ic, groups);
                self.add(nlist, y, ic, groups);
            }
            // Handled in 'run'
            Match | Char(_, _) | CharClass(_, _, _) | Any(_) => {
                // If captures are enabled, then we need to indicate that
                // this isn't an empty state.
                // Otherwise, we say it's an empty state (even though it isn't)
                // so that capture groups aren't copied.
                nlist.add(pc, groups, !self.caps);
            }
        }
    }

    fn is_begin(&self, ic: uint) -> bool { ic == 0 }
    fn is_end(&self, ic: uint) -> bool { ic == self.input.len() }

    fn is_word_boundary(&self, ic: uint) -> bool {
        if self.is_begin(ic) {
            return self.is_word(ic)
        }
        if self.is_end(ic) {
            return self.is_word(self.input.len()-1)
        }
        (self.is_word(ic) && !self.is_word(ic-1))
        || (self.is_word(ic-1) && !self.is_word(ic))
    }

    fn is_word(&self, ic: uint) -> bool {
        if ic >= self.input.len() {
            return false
        }
        let c = self.input[ic];
        c == '_'
        || (c >= '0' && c <= '9')
        || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
    }

    // FIXME: For case insensitive comparisons, it uses the uppercase
    // character and tests for equality. IIUC, this does not generalize to
    // all of Unicode. I believe we need to check the entire fold for each
    // character. This will be easy to add if and when it gets added to Rust's
    // standard library.
    fn char_eq(&self, casei: bool, ic: uint, regc: char) -> bool {
        if ic >= self.input.len() {
            return false
        }
        let textc = self.get(ic);
        regc == textc || (casei && regc.to_uppercase() == textc.to_uppercase())
    }

    fn char_is(&self, ic: uint, c: char) -> bool {
        ic < self.input.len() && self.input[ic] == c
    }

    fn get(&self, ic: uint) -> char {
        self.input[ic]
    }
}

fn class_cmp(casei: bool, mut textc: char,
             (mut start, mut end): (char, char)) -> Ordering {
    if casei {
        // FIXME: This is pretty ridiculous. All of this case conversion
        // can be moved outside this function:
        // 1) textc should be uppercased outside the bsearch.
        // 2) the character class itself should be uppercased either in the
        //    parser or the compiler.
        // FIXME: This is too simplistic for correct Unicode support.
        //        See also: char_eq
        textc = textc.to_uppercase();
        start = start.to_uppercase();
        end = end.to_uppercase();
    }
    if textc >= start && textc <= end {
        Equal
    } else if start > textc {
        Greater
    } else {
        Less
    }
}

fn numcaps(insts: &[Inst]) -> uint {
    let mut n = 0;
    for inst in insts.iter() {
        match *inst {
            Save(c) => n = cmp::max(n, c+1),
            _ => {}
        }
    }
    n
}
<|endoftext|>


<commit_before>use grammar::*;
use ast::*;
use std::collections::BTreeMap;
use std::rc::Rc;
use num::rational::BigRational;

#[derive(Clone, Eq, Debug, PartialEq)]
pub enum Value<'a> {
    Number(BigRational),
    CharString(String),
    Boolean(bool),
    Map(BTreeMap<Rc<Value<'a>>, Rc<Value<'a>>>),
    Fn(u8, Closure<'a>),
}

#[derive(Clone, Eq, Debug, PartialEq)]
pub enum Instruction<'a> {
    Push(Value<'a>),
    Fetch(String),
    LocalAssign(String),
    Assign(String),
    Call(u8),
    // Set(String, Value),
}

pub type InstructionSequence<'a> = Vec<Instruction<'a>>;

#[derive(Clone, Eq, Debug, PartialEq)]
pub struct Closure<'a> {
    instructions: Box<InstructionSequence<'a>>,
    parent_activation: &'a Activation<'a>,
}

impl<'a> Closure<'a> {
    pub fn new(instructions: InstructionSequence<'a>, parent_activation: &'a Activation) -> Closure<'a> {
        Closure {
            instructions: Box::new(instructions),
            parent_activation: parent_activation,
        }
    }
}

#[derive(Clone, Eq, Debug, PartialEq)]
pub struct Activation<'a> {
    // <String> sucks because I'm copying all the memory all the time.
    // Figure this shit out
    local_bindings: BTreeMap<String, Rc<Value<'a>>>,
}

impl<'a> Activation<'a> {
    pub fn new() -> Activation<'a> {
        Activation {
            local_bindings: BTreeMap::new(),
        }
    }
}

#[derive(Clone, Eq, Debug, PartialEq)]
struct Frame<'a> {
    activation: Activation<'a>,
    exception_handlers: Vec<()>,
}

impl<'a> Frame<'a> {
    pub fn new() -> Frame<'a> {
        Frame {
            activation: Activation::new(),
            exception_handlers: Vec::new(),
        }
    }

    pub fn local_assign(&mut self, binding_name: &str, value: Value<'a>) {
        self.activation.local_bindings.insert(binding_name.to_owned(), Rc::new(value));
    }
}

fn compile_statement<'a, 'b>(statement: &'a Statement) -> InstructionSequence<'b> {
    match statement {
        &Statement::Assign(local, ref binding_name, ref expression) => {
            let mut instructions = compile_expression(expression);
            let instruction = match local {
                true => Instruction::LocalAssign(binding_name.to_owned()),
                false => Instruction::Assign(binding_name.to_owned()),
            };
            instructions.push(instruction);
            instructions
        },
        &Statement::Call(ref binding_name, ref expressions) => {
            let mut instructions = expressions
                .iter()
                .flat_map(|exp| compile_expression(exp))
                .collect::<InstructionSequence>();
            instructions.push(Instruction::Fetch(binding_name.to_owned()));
            instructions.push(Instruction::Call(expressions.len() as u8));
            instructions
        },
        // _ => panic!("not implemented"),
    }
}

fn compile_expression<'a, 'b>(expression: &'a Expression) -> InstructionSequence<'b> {
    match expression {
        &Expression::Literal(ref literal) => {
            vec![Instruction::Push(compile_literal(literal))]
        },
        &Expression::Identifier(ref binding_name) => {
            vec![Instruction::Fetch(binding_name.to_owned())]
        }
        _ => panic!("not implemented"),
        // &Expression::Identifier(_) => { vec![] },
        // &Expression::BinOp(_, _, _) => { vec![] },
    }
}

fn compile_literal<'a, 'b>(literal: &'a Literal) -> Value<'b> {
    match literal {
        &Literal::Number(ref num) => Value::Number(num.to_owned()),
        &Literal::CharString(ref str) => Value::CharString(str.to_string()),
        // &Literal::Fn(ref args, ref statements) => {
        //     let closure = Closure::new(compile(&statements));
        //     Value::Fn(args.len() as u8, closure)
        // },
        _ => panic!("not implemented"),
    }
}

fn compile<'a, 'b>(statements: &'a Vec<Statement>) -> InstructionSequence<'b> {
    let mut instructions = InstructionSequence::new();

    for statement in statements.iter() {
        instructions.extend(compile_statement(&statement).iter().cloned());
    }

    instructions
}

#[derive(Clone, Eq, Debug, PartialEq)]
pub struct Vm<'a> {
    instructions: Rc<InstructionSequence<'a>>,
    pc: usize,
    stack: Vec<Value<'a>>,
    frames: Vec<Frame<'a>>,
}

impl<'a> Vm<'a> {
    pub fn new(source: &str) -> Vm<'a> {
        let stmts = statements(source);
        // println!("{:?}", stmts);
        let instructions = compile(&stmts.unwrap());
        // println!("{:?}", instructions);
        let mut frame = Frame::new();

        let mut vm = Vm{
            instructions: Rc::new(instructions),
            pc: 0,
            stack: Vec::new(),
            frames: vec![frame],
        };
        vm
    }

    pub fn run(&mut self) {
        for instruction in self.instructions.iter() {
            match instruction {
                &Instruction::Push(ref value) => {
                    self.stack.push(value.to_owned())
                },
                &Instruction::LocalAssign(ref binding_name) => {
                    let value = self.stack.pop().unwrap();
                    self.frames.last_mut().unwrap().local_assign(binding_name, value)
                },
                &Instruction::Call(arg_size) => {
                    let (closure_arg_size, closure) = match self.stack.pop() {
                        Some(Value::Fn(arg_size, closure)) => (arg_size, closure),
                        _ => panic!("expected a closure"),
                    };
                    // if arg_size != closure_arg_size {
                    //     panic!("wrong number of arguments")
                    // };
                    // make sure closure
                    // verify args
                    // push frame
                    // make new activation
                    // assign values
                    // change iseq
                },
                _ => ()
            };
        };
    }
}

#[cfg(test)]
mod test_vm {
    use super::*;
    use std::rc::Rc;
    use grammar::test_helpers::*;

    #[test]
    fn test_new_populates_basic_instructions() {
        let vm = Vm::new("let a = 1\nlet b = \"\"\n");
        assert_eq!(
            Rc::new(vec![
                Instruction::Push(Value::Number(build_ratio(1, 1))),
                Instruction::LocalAssign("a".to_owned()),
                Instruction::Push(Value::CharString("".to_string())),
                Instruction::LocalAssign("b".to_owned()),
            ]),
            vm.instructions
        )
    }

    fn test_new_populates_function_instructions() {
        let vm = Vm::new("let b = def(x)\nend\nb(1)");
        let activation = Activation::new();
        let closure = Closure::new(vec![], &activation);
        assert_eq!(
            Rc::new(vec![
                Instruction::Push(Value::Fn(1, closure)),
                Instruction::LocalAssign("b".to_owned()),
                Instruction::Push(Value::Number(build_ratio(1, 1))),
                Instruction::Fetch("b".to_owned()),
                Instruction::Call(1),
            ]),
            vm.instructions
        )
    }

    #[test]
    fn test_run_simple() {
        let mut vm = Vm::new("let a = 1");
        vm.run();
        assert_eq!(
            Rc::new(Value::Number(build_ratio(1, 1))),
            vm.frames.last().unwrap().activation.local_bindings.get("a").unwrap().to_owned()
        )
    }

    #[test]
    #[ignore]
    fn test_function_call() {
        let source =
            "let a = \"\"\
            let x = def() do
              a = 1
            end\
            x()";

        let mut vm = Vm::new(source);
        vm.run();
        assert_eq!(
            Rc::new(Value::Number(build_ratio(1, 1))),
            vm.frames.last().unwrap().activation.local_bindings.get("a").unwrap().to_owned()
        )
    }

    #[test]
    #[should_panic]
    fn test_calling_non_function() {
        let source =
            "let x = \"\"\
            x()";

        let mut vm = Vm::new(source);
        vm.run();
    }

    #[test]
    #[should_panic]
    fn test_function_with_wrong_arg_count() {
        let source =
            "let x = def(a, b) do
              a = 1
            end\
            x(1, 1)";

        let mut vm = Vm::new(source);
        vm.run();
    }
}
<commit_msg>Go back to Rc<RefCell> instead of static lifetimes<commit_after>use grammar::*;
use ast::*;
use std::collections::BTreeMap;
use std::rc::Rc;
use std::cell::RefCell;
use num::rational::BigRational;

#[derive(Clone, Eq, Debug, PartialEq)]
pub enum Value {
    Number(BigRational),
    CharString(String),
    Boolean(bool),
    Map(BTreeMap<Rc<Value>, Rc<Value>>),
    Closure(u8, Closure),
    Fn(u8, Closure),
}

#[derive(Clone, Eq, Debug, PartialEq)]
pub enum Instruction {
    Push(Literal),
    Fetch(String),
    LocalAssign(String),
    Assign(String),
    Call(u8),
    // Set(String, Value),
}

pub type InstructionSequence = Vec<Instruction>;
pub type BindingMap = Rc<RefCell<BTreeMap<String, Value>>>;

#[derive(Clone, Eq, Debug, PartialEq)]
pub struct Closure {
    instructions: Box<InstructionSequence>,
    parent_bindings: BindingMap,
}

impl Closure {
    pub fn new(instructions: InstructionSequence, parent_bindings: &BindingMap) -> Closure {
        Closure {
            instructions: Box::new(instructions),
            parent_bindings: parent_bindings.clone(),
        }
    }
}

#[derive(Clone, Eq, Debug, PartialEq)]
struct Frame {
    bindings: BindingMap,
    exception_handlers: Vec<()>,
}

impl Frame {
    pub fn new() -> Frame {
        Frame {
            bindings: Rc::new(RefCell::new(BTreeMap::new())),
            exception_handlers: Vec::new(),
        }
    }

    pub fn local_assign(&mut self, binding_name: &str, value: Value) {
        self.bindings.borrow_mut().insert(binding_name.to_owned(), value);
    }
}

fn compile_statement<'a>(statement: &'a Statement) -> InstructionSequence {
    match statement {
        &Statement::Assign(local, ref binding_name, ref expression) => {
            let mut instructions = compile_expression(expression);
            let instruction = match local {
                true => Instruction::LocalAssign(binding_name.to_owned()),
                false => Instruction::Assign(binding_name.to_owned()),
            };
            instructions.push(instruction);
            instructions
        },
        &Statement::Call(ref binding_name, ref expressions) => {
            let mut instructions = expressions
                .iter()
                .flat_map(|exp| compile_expression(exp))
                .collect::<InstructionSequence>();
            instructions.push(Instruction::Fetch(binding_name.to_owned()));
            instructions.push(Instruction::Call(expressions.len() as u8));
            instructions
        },
        // _ => panic!("not implemented"),
    }
}

fn compile_expression<'a>(expression: &'a Expression) -> InstructionSequence {
    match expression {
        &Expression::Literal(ref literal) => {
            vec![Instruction::Push(literal.to_owned())]
        },
        &Expression::Identifier(ref binding_name) => {
            vec![Instruction::Fetch(binding_name.to_owned())]
        }
        _ => panic!("not implemented"),
        // &Expression::Identifier(_) => { vec![] },
        // &Expression::BinOp(_, _, _) => { vec![] },
    }
}

fn compile<'a>(statements: &'a Vec<Statement>) -> InstructionSequence {
    let mut instructions = InstructionSequence::new();

    for statement in statements.iter() {
        instructions.extend(compile_statement(&statement).iter().cloned());
    }

    instructions
}

#[derive(Clone, Eq, Debug, PartialEq)]
pub struct Vm {
    instructions: Rc<InstructionSequence>,
    pc: usize,
    stack: Vec<Value>,
    frames: Vec<Frame>,
}

impl Vm {
    pub fn new(source: &str) -> Vm {
        let stmts = statements(source);
        // println!("{:?}", stmts);
        let instructions = compile(&stmts.unwrap());
        // println!("{:?}", instructions);
        let mut frame = Frame::new();

        let mut vm = Vm{
            instructions: Rc::new(instructions),
            pc: 0,
            stack: Vec::new(),
            frames: vec![frame],
        };
        vm
    }

    pub fn run<'b, 'c>(&'b mut self) {
        loop {
            let insn_result = Vm::next_instruction(self);
            let instruction;

            match insn_result {
                Some(i) => instruction = i,
                None => break,
            }

            match instruction {
                Instruction::Push(ref value) => {
                    let top_bindings = &mut self.frames.last_mut().unwrap().bindings;
                    self.stack.push(Vm::literal_to_value(value, top_bindings))
                    // Vm::push_value(&mut self.stack, top_activation, value)
                },
                Instruction::LocalAssign(ref binding_name) => {
                    let value = self.stack.pop().unwrap();
                    self.frames.last_mut().unwrap().local_assign(binding_name, value)
                },
                Instruction::Call(arg_size) => {
                    let (closure_arg_size, closure) = match self.stack.pop() {
                        Some(Value::Closure(arg_size, closure)) => (arg_size, closure),
                        _ => panic!("expected a closure"),
                    };
                    // if arg_size != closure_arg_size {
                    //     panic!("wrong number of arguments")
                    // };
                    // make sure closure
                    // verify args
                    // push frame
                    // make new activation
                    // assign values
                    // change iseq
                },
                _ => ()//panic!("unknown instruction {:?}", instruction)
            };
        };
    }

    fn next_instruction(vm: &mut Vm) -> Option<Instruction> {
        let instruction = match vm.instructions.get(vm.pc) {
            Some(i) => Some(i.clone()),
            None => None,
        };

        vm.pc += 1;

        instruction
    }

    fn literal_to_value<'b>(literal: &'b Literal, top_bindings: &BindingMap) -> Value {
        match literal {
            &Literal::Number(ref num) => Value::Number(num.to_owned()),
            &Literal::CharString(ref str) => Value::CharString(str.to_string()),
            &Literal::Fn(ref args, ref statements) => {
                // Statements should be compiled ahead of time
                let closure = Closure::new(compile(&statements), top_bindings);
                Value::Closure(args.len() as u8, closure)
            },
            _ => panic!("not implemented"),
        }
    }
}

#[cfg(test)]
mod test_vm {
    use super::*;
    use ast::*;
    use std::rc::Rc;
    use grammar::test_helpers::*;

    #[test]
    fn test_new_populates_basic_instructions() {
        let vm = Vm::new("let a = 1\nlet b = \"\"\n");
        assert_eq!(
            Rc::new(vec![
                Instruction::Push(Literal::Number(build_ratio(1, 1))),
                Instruction::LocalAssign("a".to_owned()),
                Instruction::Push(Literal::CharString("".to_string())),
                Instruction::LocalAssign("b".to_owned()),
            ]),
            vm.instructions
        )
    }

    fn test_new_populates_function_instructions() {
        let vm = Vm::new("let b = def(x)\nend\nb(1)");
        assert_eq!(
            Rc::new(vec![
                Instruction::Push(Literal::Fn(Box::new(vec!["x".to_owned()]), Box::new(vec![]))),
                Instruction::LocalAssign("b".to_owned()),
                Instruction::Push(Literal::Number(build_ratio(1, 1))),
                Instruction::Fetch("b".to_owned()),
                Instruction::Call(1),
            ]),
            vm.instructions
        )
    }

    #[test]
    fn test_run_simple() {
        let mut vm = Vm::new("let a = 1");
        vm.run();
        assert_eq!(
            Value::Number(build_ratio(1, 1)),
            vm.frames.last().unwrap().bindings.borrow().get("a").unwrap().to_owned()
        );
        ()
    }

    #[test]
    #[ignore]
    fn test_function_call() {
        let source =
            "let a = \"\"\
            let x = def() do
              a = 1
            end\
            x()";

        let mut vm = Vm::new(source);
        vm.run();
        assert_eq!(
            Value::Number(build_ratio(1, 1)),
            vm.frames.last().unwrap().bindings.borrow().get("a").unwrap().to_owned()
        );
        ()
    }

    #[test]
    #[should_panic]
    fn test_calling_non_function() {
        let source =
            "let x = \"\"\
            x()";

        let mut vm = Vm::new(source);
        vm.run();
    }

    #[test]
    #[should_panic]
    fn test_function_with_wrong_arg_count() {
        let source =
            "let x = def(a, b) do
              a = 1
            end\
            x(1, 1)";

        let mut vm = Vm::new(source);
        vm.run();
    }
}
<|endoftext|>


<commit_before>extern crate rand;

extern crate timely;
extern crate differential_dataflow;
extern crate dd_server;

use std::rc::Rc;
use std::cell::RefCell;

use rand::{Rng, SeedableRng, StdRng};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::nested::product::Product;
use timely::dataflow::operators::Probe;
use timely::dataflow::operators::generic::operator::source;

use differential_dataflow::AsCollection;
use differential_dataflow::operators::arrange::ArrangeByKey;
use differential_dataflow::trace::TraceReader;

use dd_server::{Environment, TraceHandle};

// load ./dataflows/random_graph/target/debug/librandom_graph.dylib build <graph_name> 1000 2000 1000000
// load ./dataflows/random_graph/target/release/librandom_graph.dylib build <graph_name> 10000000 100000000 1000000
// drop <graph_name>-capability

#[no_mangle]
pub fn build((dataflow, handles, probe, timer, args): Environment) -> Result<(), String> {

    // This call either starts or terminates the production of random graph edges.
    //
    // The arguments should be
    //
    //    <graph_name> <nodes> <edges> <rate>
    //
    // where <rate> is the target number of edge changes per second. The source
    // will play out changes to keep up with this, and timestamp them as if they
    // were emitted at the correct time. This currently uses a local timer, created
    // when this method is called, and should probably take / use a "system time"
    // from the worker so that timestamps align. Until then, it is a fine test that
    // things work even without synchronized timestamps.
    //
    // The method also registers a capability with name `<graph_name>-capability`,
    // and will continue to execute until this capability is dropped from `handles`.
    //
    // The operator also holds an `Weak<RefCell<Option<TraceHandle>>>` which it will
    // attempt to borrow and call `advance_by` in order to advance the capability
    // as it runs, to allow compaction and the maintenance of bounded state.

    if args.len() != 4 { return Err(format!("expected four arguments, instead: {:?}", args)); }

    let name = &args[0];
    let nodes: usize = args[1].parse().map_err(|_| format!("parse error, nodes: {:?}", args[1]))?;
    let edges: usize = args[2].parse().map_err(|_| format!("parse error, edges: {:?}", args[2]))?;
    let rate: usize = args[3].parse().map_err(|_| format!("parse error, rate: {:?}", args[3]))?;

    let requests_per_sec = rate;
    let ns_per_request = 1000000000 / requests_per_sec;

    // shared capability keeps graph generation going.
    let capability = Rc::new(RefCell::new(None));

    // shared (optional) trace handle, so that the operator can advance capabilities.
    let trace_handle: Rc<RefCell<Option<TraceHandle>>> = Rc::new(RefCell::new(None));
    let trace_handle_weak = Rc::downgrade(&trace_handle);

    let timer = timer.clone();

    // create a trace from a source of random graph edges.
    let mut trace =
        source(dataflow, "RandomGraph", |cap| {

            let index = dataflow.index();
            let peers = dataflow.peers();

            // RNGs for edge addition and deletion.
            let seed: &[_] = &[1, 2, 3, index];
            let mut rng1: StdRng = SeedableRng::from_seed(seed);
            let mut rng2: StdRng = SeedableRng::from_seed(seed);

            // numbers of times we've stepped each RNG.
            let mut additions = 0;
            let mut deletions = 0;

            // record delay between system start-up and operator start-up.
            let delay = timer.elapsed();
            let delay_ns = (delay.as_secs() as usize) * 1_000_000_000 + (delay.subsec_nanos() as usize);

            // stash capability in a rc::Weak.
            *capability.borrow_mut() = Some(cap);
            let capability = ::std::rc::Rc::downgrade(&capability);

            move |output| {

                // Open-loop latency-throughput test, parameterized by offered rate `ns_per_request`.
                let elapsed = timer.elapsed();
                let elapsed_ns = (elapsed.as_secs() as usize) * 1_000_000_000 + (elapsed.subsec_nanos() as usize);

                // attempt to advance the frontier of the trace handle.
                if let Some(trace_handle) = trace_handle_weak.upgrade() {
                    let mut borrow = trace_handle.borrow_mut();
                    if let Some(ref mut trace_handle) = borrow.as_mut() {
                        trace_handle.advance_by(&[Product::new(RootTimestamp, elapsed_ns)]);
                    }
                }

                // if our capability has not been cancelled ...
                if let Some(capability) = capability.upgrade() {

                    let mut borrow = capability.borrow_mut();
                    let capability = borrow.as_mut().unwrap();
                    let mut time = capability.time().clone();

                    {   // scope to allow session to drop, un-borrow.
                        let mut session = output.session(&capability);

                        // load initial graph.
                        while additions < edges + deletions {
                            time.inner = 0;
                            if additions % peers == index {
                                time.inner = 0;
                                let src = rng1.gen_range(0, nodes);
                                let dst = rng1.gen_range(0, nodes);
                                session.give(((src, dst), time, 1));
                            }
                            additions += 1;
                        }

                        // ship any scheduled edge additions.
                        while ns_per_request * (additions - edges) < (elapsed_ns - delay_ns) {
                            if additions % peers == index {
                                time.inner = ns_per_request * (additions - edges);
                                let src = rng1.gen_range(0, nodes);
                                let dst = rng1.gen_range(0, nodes);
                                session.give(((src, dst), time, 1));
                            }
                            additions += 1;
                        }

                        // ship any scheduled edge deletions.
                        while ns_per_request * deletions < (elapsed_ns - delay_ns) {
                            if deletions % peers == index {
                                time.inner = ns_per_request * deletions;
                                let src = rng2.gen_range(0, nodes);
                                let dst = rng2.gen_range(0, nodes);
                                session.give(((src, dst), time, -1));
                            }
                            deletions += 1;
                        }
                    }

                    time.inner = elapsed_ns;
                    capability.downgrade(&time);
                }
            }
        })
        .probe_with(probe)
        .as_collection()
        .arrange_by_key()
        .trace;

    // release all blocks on merging.
    trace.distinguish_since(&[]);
    *trace_handle.borrow_mut() = Some(trace);

    handles.set::<Rc<RefCell<Option<TraceHandle>>>>(name.to_owned(), trace_handle);
    handles.set(format!("{}-capability", name), capability);

    println!("handles set");

    Ok(())
}<commit_msg>comments improved<commit_after>extern crate rand;

extern crate timely;
extern crate differential_dataflow;
extern crate dd_server;

use std::rc::Rc;
use std::cell::RefCell;

use rand::{Rng, SeedableRng, StdRng};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::nested::product::Product;
use timely::dataflow::operators::Probe;
use timely::dataflow::operators::generic::operator::source;

use differential_dataflow::AsCollection;
use differential_dataflow::operators::arrange::ArrangeByKey;
use differential_dataflow::trace::TraceReader;

use dd_server::{Environment, TraceHandle};

// load ./dataflows/random_graph/target/debug/librandom_graph.dylib build <graph_name> 1000 2000 1000000
// load ./dataflows/random_graph/target/release/librandom_graph.dylib build <graph_name> 10000000 100000000 1000000
// drop <graph_name>-capability

#[no_mangle]
pub fn build((dataflow, handles, probe, timer, args): Environment) -> Result<(), String> {

    // This call either starts the production of random graph edges.
    //
    // The arguments should be
    //
    //    <graph_name> <nodes> <edges> <rate>
    //
    // where <rate> is the target number of edge changes per second. The source
    // will play out changes to keep up with this, and timestamp them as if they
    // were emitted at the correct time. The timestamps use the system `timer`,
    // but only start whenever the method is called. This means that the data are
    // not deterministic, but if you subtract the elapsed time between system start
    // up and method call, they should be deterministic.
    //
    // The method also registers a capability with name `<graph_name>-capability`,
    // and will continue to execute until this capability is dropped from `handles`.
    // To terminate the operator it is sufficient to drop the capability, as the
    // operator holds only a weak reference to it.
    //
    // The operator also holds an `Weak<RefCell<Option<TraceHandle>>>` which it will
    // attempt to borrow and call `advance_by` in order to advance the capability
    // as it runs, to allow compaction and the maintenance of bounded state.

    if args.len() != 4 { return Err(format!("expected four arguments, instead: {:?}", args)); }

    let name = &args[0];
    let nodes: usize = args[1].parse().map_err(|_| format!("parse error, nodes: {:?}", args[1]))?;
    let edges: usize = args[2].parse().map_err(|_| format!("parse error, edges: {:?}", args[2]))?;
    let rate: usize = args[3].parse().map_err(|_| format!("parse error, rate: {:?}", args[3]))?;

    let requests_per_sec = rate;
    let ns_per_request = 1000000000 / requests_per_sec;

    // shared capability keeps graph generation going.
    let capability = Rc::new(RefCell::new(None));

    // shared (optional) trace handle, so that the operator can advance capabilities.
    let trace_handle: Rc<RefCell<Option<TraceHandle>>> = Rc::new(RefCell::new(None));
    let trace_handle_weak = Rc::downgrade(&trace_handle);

    let timer = timer.clone();

    // create a trace from a source of random graph edges.
    let mut trace =
        source(dataflow, "RandomGraph", |cap| {

            let index = dataflow.index();
            let peers = dataflow.peers();

            // RNGs for edge addition and deletion.
            let seed: &[_] = &[1, 2, 3, index];
            let mut rng1: StdRng = SeedableRng::from_seed(seed);
            let mut rng2: StdRng = SeedableRng::from_seed(seed);

            // numbers of times we've stepped each RNG.
            let mut additions = 0;
            let mut deletions = 0;

            // record delay between system start-up and operator start-up.
            let delay = timer.elapsed();
            let delay_ns = (delay.as_secs() as usize) * 1_000_000_000 + (delay.subsec_nanos() as usize);

            // stash capability in a rc::Weak.
            *capability.borrow_mut() = Some(cap);
            let capability = ::std::rc::Rc::downgrade(&capability);

            move |output| {

                // Open-loop latency-throughput test, parameterized by offered rate `ns_per_request`.
                let elapsed = timer.elapsed();
                let elapsed_ns = (elapsed.as_secs() as usize) * 1_000_000_000 + (elapsed.subsec_nanos() as usize);

                // attempt to advance the frontier of the trace handle.
                if let Some(trace_handle) = trace_handle_weak.upgrade() {
                    let mut borrow = trace_handle.borrow_mut();
                    if let Some(ref mut trace_handle) = borrow.as_mut() {
                        trace_handle.advance_by(&[Product::new(RootTimestamp, elapsed_ns)]);
                    }
                }

                // if our capability has not been cancelled ...
                if let Some(capability) = capability.upgrade() {

                    let mut borrow = capability.borrow_mut();
                    let capability = borrow.as_mut().unwrap();
                    let mut time = capability.time().clone();

                    {   // scope to allow session to drop, un-borrow.
                        let mut session = output.session(&capability);

                        // load initial graph.
                        while additions < edges + deletions {
                            time.inner = 0;
                            if additions % peers == index {
                                time.inner = 0;
                                let src = rng1.gen_range(0, nodes);
                                let dst = rng1.gen_range(0, nodes);
                                session.give(((src, dst), time, 1));
                            }
                            additions += 1;
                        }

                        // ship any scheduled edge additions.
                        while ns_per_request * (additions - edges) < (elapsed_ns - delay_ns) {
                            if additions % peers == index {
                                time.inner = ns_per_request * (additions - edges);
                                let src = rng1.gen_range(0, nodes);
                                let dst = rng1.gen_range(0, nodes);
                                session.give(((src, dst), time, 1));
                            }
                            additions += 1;
                        }

                        // ship any scheduled edge deletions.
                        while ns_per_request * deletions < (elapsed_ns - delay_ns) {
                            if deletions % peers == index {
                                time.inner = ns_per_request * deletions;
                                let src = rng2.gen_range(0, nodes);
                                let dst = rng2.gen_range(0, nodes);
                                session.give(((src, dst), time, -1));
                            }
                            deletions += 1;
                        }
                    }

                    time.inner = elapsed_ns;
                    capability.downgrade(&time);
                }
            }
        })
        .probe_with(probe)
        .as_collection()
        .arrange_by_key()
        .trace;

    // release all blocks on merging.
    trace.distinguish_since(&[]);
    *trace_handle.borrow_mut() = Some(trace);

    handles.set::<Rc<RefCell<Option<TraceHandle>>>>(name.to_owned(), trace_handle);
    handles.set(format!("{}-capability", name), capability);

    println!("handles set");

    Ok(())
}<|endoftext|>


<commit_before>#![warn(rust_2018_idioms, single_use_lifetimes)]

use std::{env, process::Command};

fn main() {
    println!("cargo:rerun-if-changed=build.rs");

    if is_nightly() {
        println!("cargo:rustc-cfg=nightly");
    }
}

fn is_nightly() -> bool {
    env::var_os("RUSTC")
        .and_then(|rustc| Command::new(rustc).arg("--version").output().ok())
        .and_then(|output| String::from_utf8(output.stdout).ok())
        .map_or(false, |version| version.contains("nightly"))
}
<commit_msg>Update build script<commit_after>#![warn(rust_2018_idioms, single_use_lifetimes)]

use std::{env, process::Command};

fn main() {
    if is_nightly() {
        println!("cargo:rustc-cfg=nightly");
    }
}

fn is_nightly() -> bool {
    env::var_os("RUSTC")
        .and_then(|rustc| Command::new(rustc).arg("--version").output().ok())
        .and_then(|output| String::from_utf8(output.stdout).ok())
        .map_or(false, |version| version.contains("nightly"))
}
<|endoftext|>


<commit_before>
use std::collections::{HashMap, HashSet};
use std::sync::mpsc::{Sender, Receiver};

use futures::sync::mpsc;

use mqtt::packet::{Packet};
use mqtt::{TopicFilter, TopicName, QualityOfService};

use common::{Topic};
use store::{GlobalRetainMsg};
use super::{ClientSessionMsg, LocalRouterMsg};

pub fn run(
    local_router_rx: Receiver<super::LocalRouterMsg>,
    local_router_tx: Sender<super::LocalRouterMsg>,
    client_session_tx: Sender<super::ClientSessionMsg>,
    _router_follower_tx: mpsc::Sender<super::RouterFollowerMsg>,
    _router_leader_tx: mpsc::Sender<super::RouterLeaderMsg>,
    global_retain_tx: Sender<GlobalRetainMsg>,
) {
    let mut local_routes = LocalRoutes::new();
    loop {
        let msg = local_router_rx.recv().unwrap();
        match msg {
            LocalRouterMsg::ForwardPublish(user_id, packet) => {
                // Forward message to current receiver.
                local_router_tx.send(LocalRouterMsg::Publish(user_id, packet)).unwrap();
            }
            LocalRouterMsg::Publish(user_id, packet) => {
                let topic_name = TopicName::new(packet.topic_name().to_string()).unwrap();
                let clients: HashMap<_, _> = local_routes.search(user_id, &topic_name);
                for (addr, qos) in clients {
                    client_session_tx.send(ClientSessionMsg::Publish(user_id, addr, qos, packet.clone())).unwrap();
                }
            }
            LocalRouterMsg::Subscribe(user_id, client_identifier, addr, packet) => {
                for &(ref topic_filter, qos) in packet.payload().subscribes() {
                    let topic = Topic::from_filter(topic_filter);
                    let first_subscribe = local_routes.insert_topic(user_id, &topic, &client_identifier, qos);
                    if first_subscribe {
                        let msg = GlobalRetainMsg::MatchAll(user_id, addr, topic, qos);
                        global_retain_tx.send(msg).unwrap();
                    }
                }
            }
            LocalRouterMsg::Unsubscribe(user_id, client_identifier, packet) => {
                for topic_filter in packet.payload().subscribes() {
                    let topic = Topic::from_filter(topic_filter);
                    local_routes.remove_topic(user_id, &topic, &client_identifier);
                }
            }
            LocalRouterMsg::ClientDisconnect(user_id, client_identifier) => {
                local_routes.remove_all_topics(user_id, &client_identifier);
            }
        }
    }
}

/// Just for topic filters
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RouteKey {
    Normal(String),
    OneLevel,
    RestLevels,
}

impl<'a> From<&'a str> for RouteKey {
    fn from(token: &'a str) -> RouteKey {
        match token {
            "+" => RouteKey::OneLevel,
            "#" => RouteKey::RestLevels,
            s @ _ => RouteKey::Normal(s.to_string())
        }
    }
}

// type MqttClient = (u32, SocketAddr);

// TODO:
// ====
//  * Slow now, use `nikomatsakis/rayon` to speed up!

#[derive(Debug)]
pub struct LocalRouteNode {
    children: Option<HashMap<RouteKey, LocalRouteNode>>,
    clients: Option<HashMap<String, QualityOfService>>
}

impl LocalRouteNode {
    pub fn new () -> LocalRouteNode {
        LocalRouteNode {
            children: None,
            clients: None
        }
    }

    pub fn insert(&mut self, tf: &TopicFilter, client_identifier: &str, qos: QualityOfService) -> bool {
        let mut last_node = (*tf).split("/").fold(self, |current, token| {
            let token = RouteKey::from(token);
            if current.children.is_none() {
                current.children = Some(HashMap::<RouteKey, LocalRouteNode>::new());
            }
            match current.children {
                Some(ref mut map) => {
                    if !map.contains_key(&token) {
                        map.insert(token.clone(), LocalRouteNode::new());
                    }
                    map.get_mut(&token).unwrap()
                }
                _ => unreachable!()
            }
        });

        if last_node.clients.is_none() {
            last_node.clients = Some(HashMap::<String, QualityOfService>::new())
        }
        match last_node.clients {
            Some(ref mut client_map) => {
                client_map.insert(client_identifier.to_owned(), qos).is_none()
            }
            _ => unreachable!()
        }
    }

    pub fn is_empty(&self) -> bool {
        (match self.clients {
            Some(ref client_map) => client_map.is_empty(),
            None => true
        }) && (match self.children {
            Some(ref map) => map.is_empty(),
            None => true
        })
    }

    fn _remove<'a, I>(&mut self, mut tokens: I, client_identifier: &str) -> bool
        where I: Iterator<Item=&'a str>
    {
        match tokens.next() {
            Some(token) => {
                match self.children {
                    Some(ref mut map) => {
                        let mut child_empty = false;
                        let key = RouteKey::from(token);
                        let is_removed = match map.get_mut(&key) {
                            Some(node) => {
                                let is_removed = node._remove(tokens, client_identifier);
                                child_empty = node.is_empty();
                                is_removed
                            }
                            None => false
                        };
                        if child_empty {
                            map.remove(&key);
                        }
                        is_removed
                    }
                    None => false
                }
            }
            None => {
                match self.clients {
                    Some(ref mut client_map) => {
                        client_map.remove(client_identifier).is_some()
                    }
                    None => false
                }
            }
        }
    }

    pub fn remove(&mut self, tf: &TopicFilter, client_identifier: &str) -> bool {
        self._remove((*tf).split("/"), client_identifier)
    }

    fn _search(&self, tokens: &Vec<&str>, index: usize) -> HashMap<String, QualityOfService> {
        let mut clients = HashMap::<String, QualityOfService>::new();
        if let Some(token) = tokens.get(index) {
            match self.children {
                Some(ref map) => {
                    if let Some(node) = map.get(&RouteKey::Normal(token.to_string())) {
                        clients.extend(node._search(tokens, index+1));
                    }
                    if let Some(node) = map.get(&RouteKey::OneLevel) {
                        clients.extend(node._search(tokens, index+1));
                    }
                    if let Some(node) = map.get(&RouteKey::RestLevels) {
                        if let Some(ref new_clients) = node.clients {
                            clients.extend(new_clients.clone());
                        }
                    }
                }
                None => {}
            }
        } else {
            if let Some(ref new_clients) = self.clients {
                clients.extend(new_clients.clone());
            }
        }
        clients
    }

    pub fn search(&self, name: &TopicName) -> HashMap<String, QualityOfService> {
        self._search(&(*name).split("/").collect::<Vec<&str>>(), 0)
    }
}

pub struct LocalRoutes {
    // For topic filters
    filter_routes: HashMap<u32, LocalRouteNode>,
    // For topic names
    name_routes: HashMap<u32, HashMap<TopicName, HashMap<String, QualityOfService>>>,
    // For remove topic-names and topic-filters
    client_topics: HashMap<u32, HashMap<String, HashSet<Topic>>>,
}

/// Router thread
/// * subscribe a topic
/// * unsubscribe a topic
/// * unsubscribe all topics from a client
/// * match topics
impl LocalRoutes {
    pub fn new() -> LocalRoutes {
        LocalRoutes {
            filter_routes: HashMap::<u32, LocalRouteNode>::new(),
            name_routes: HashMap::<u32, HashMap<TopicName, HashMap<String, QualityOfService>>>::new(),
            client_topics: HashMap::<u32, HashMap<String, HashSet<Topic>>>::new()
        }
    }

    fn insert_topic(&mut self, user_id: u32, topic: &Topic,
                    client_identifier: &str, qos: QualityOfService) -> bool {
        if !self.filter_routes.contains_key(&user_id) {
            self.filter_routes.insert(user_id, LocalRouteNode::new());
        }
        if !self.name_routes.contains_key(&user_id) {
            self.name_routes.insert(user_id, HashMap::<TopicName, HashMap<String, QualityOfService>>::new());
        }
        if !self.client_topics.contains_key(&user_id) {
            self.client_topics.insert(user_id, HashMap::<String, HashSet<Topic>>::new());
        }
        let mut filter_routes = self.filter_routes.get_mut(&user_id).unwrap();
        let mut name_routes = self.name_routes.get_mut(&user_id).unwrap();
        let mut client_topics = self.client_topics.get_mut(&user_id).unwrap();

        if !client_topics.contains_key(client_identifier) {
            client_topics.insert(client_identifier.to_owned(), HashSet::<Topic>::new());
        }
        client_topics.get_mut(client_identifier).unwrap()
            .insert(topic.clone());

        match topic {
            &Topic::Filter(ref topic_filter) => {
                filter_routes.insert(topic_filter, client_identifier, qos)
            }
            &Topic::Name(ref topic_name) => {
                if !name_routes.contains_key(topic_name) {
                    name_routes.insert(topic_name.clone(), HashMap::<String, QualityOfService>::new());
                }
                name_routes.get_mut(topic_name).unwrap()
                    .insert(client_identifier.to_owned(), qos).is_none()
            }
        }
    }

    fn remove_topic(&mut self, user_id: u32, topic: &Topic, client_identifier: &str) {
        if let Some(ref mut client_topics) = self.client_topics.get_mut(&user_id) {
            match client_topics.get_mut(client_identifier) {
                Some(ref mut topic_name_set) => topic_name_set.remove(topic),
                None => false
            };
        }
        match topic {
            &Topic::Filter(ref topic_filter) => {
                if let Some(ref mut filter_routes) = self.filter_routes.get_mut(&user_id) {
                    filter_routes.remove(topic_filter, client_identifier);
                }
            }
            &Topic::Name(ref topic_name) => {
                if let Some(ref mut name_routes) = self.name_routes.get_mut(&user_id) {
                    match name_routes.get_mut(topic_name) {
                        Some(ref mut client_map) => client_map.remove(client_identifier).is_some(),
                        None => false
                    };
                }
            }
        }
    }

    fn remove_all_topics(&mut self, user_id: u32, client_identifier: &str) -> usize {
        if let Some(ref mut client_topics) = self.client_topics.get_mut(&user_id) {
            let removed_count = match client_topics.get_mut(client_identifier) {
                Some(ref mut topic_name_set) => {
                    let mut removed_count = 0;
                    for topic in topic_name_set.iter() {
                        removed_count += match topic {
                            &Topic::Filter(ref topic_filter) => {
                                if let Some(ref mut filter_routes) = self.filter_routes.get_mut(&user_id) {
                                    if filter_routes.remove(topic_filter, client_identifier) { 1 } else { 0 }
                                } else { 0 }
                            }
                            &Topic::Name(ref topic_name) => {
                                if let Some(ref mut name_routes) = self.name_routes.get_mut(&user_id) {
                                    match name_routes.get_mut(topic_name) {
                                        Some(ref mut client_map) => {
                                            if client_map.remove(client_identifier).is_some() { 1 } else { 0 }
                                        }
                                        None => 0
                                    }
                                } else { 0 }
                            }
                        }
                    }
                    removed_count
                }
                None => 0
            };
            client_topics.remove(client_identifier);
            removed_count
        } else { 0 }
    }

    fn search(&self, user_id: u32, topic_name: &TopicName) -> HashMap<String, QualityOfService> {
        let mut clients = HashMap::<String, QualityOfService>::new();
        if let Some(ref name_routes) = self.name_routes.get(&user_id) {
            if let Some(client_map) = name_routes.get(topic_name) {
                clients.extend(client_map.clone());
            }
        }
        if let Some(ref filter_routes) = self.filter_routes.get(&user_id) {
            clients.extend(filter_routes.search(topic_name));
        }
        clients
    }
}
<commit_msg>Fix [MQTT-3.8.4-3] (re-send retained messages)<commit_after>
use std::collections::{HashMap, HashSet};
use std::sync::mpsc::{Sender, Receiver};

use futures::sync::mpsc;

use mqtt::packet::{Packet};
use mqtt::{TopicFilter, TopicName, QualityOfService};

use common::{Topic};
use store::{GlobalRetainMsg};
use super::{ClientSessionMsg, LocalRouterMsg};

pub fn run(
    local_router_rx: Receiver<super::LocalRouterMsg>,
    local_router_tx: Sender<super::LocalRouterMsg>,
    client_session_tx: Sender<super::ClientSessionMsg>,
    _router_follower_tx: mpsc::Sender<super::RouterFollowerMsg>,
    _router_leader_tx: mpsc::Sender<super::RouterLeaderMsg>,
    global_retain_tx: Sender<GlobalRetainMsg>,
) {
    let mut local_routes = LocalRoutes::new();
    loop {
        let msg = local_router_rx.recv().unwrap();
        match msg {
            LocalRouterMsg::ForwardPublish(user_id, packet) => {
                // Forward message to current receiver.
                local_router_tx.send(LocalRouterMsg::Publish(user_id, packet)).unwrap();
            }
            LocalRouterMsg::Publish(user_id, packet) => {
                let topic_name = TopicName::new(packet.topic_name().to_string()).unwrap();
                let clients: HashMap<_, _> = local_routes.search(user_id, &topic_name);
                for (addr, qos) in clients {
                    client_session_tx.send(ClientSessionMsg::Publish(user_id, addr, qos, packet.clone())).unwrap();
                }
            }
            LocalRouterMsg::Subscribe(user_id, client_identifier, addr, packet) => {
                for &(ref topic_filter, qos) in packet.payload().subscribes() {
                    let topic = Topic::from_filter(topic_filter);
                    let _ = local_routes.insert_topic(user_id, &topic, &client_identifier, qos);
                    // <Spec>: [MQTT-3.8.4-3] "Any existing retained messages matching the Topic Filter MUST be re-sent"
                    let msg = GlobalRetainMsg::MatchAll(user_id, addr, topic, qos);
                    global_retain_tx.send(msg).unwrap();
                }
            }
            LocalRouterMsg::Unsubscribe(user_id, client_identifier, packet) => {
                for topic_filter in packet.payload().subscribes() {
                    let topic = Topic::from_filter(topic_filter);
                    local_routes.remove_topic(user_id, &topic, &client_identifier);
                }
            }
            LocalRouterMsg::ClientDisconnect(user_id, client_identifier) => {
                local_routes.remove_all_topics(user_id, &client_identifier);
            }
        }
    }
}

/// Just for topic filters
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RouteKey {
    Normal(String),
    OneLevel,
    RestLevels,
}

impl<'a> From<&'a str> for RouteKey {
    fn from(token: &'a str) -> RouteKey {
        match token {
            "+" => RouteKey::OneLevel,
            "#" => RouteKey::RestLevels,
            s @ _ => RouteKey::Normal(s.to_string())
        }
    }
}

// type MqttClient = (u32, SocketAddr);

// TODO:
// ====
//  * Slow now, use `nikomatsakis/rayon` to speed up!

#[derive(Debug)]
pub struct LocalRouteNode {
    children: Option<HashMap<RouteKey, LocalRouteNode>>,
    clients: Option<HashMap<String, QualityOfService>>
}

impl LocalRouteNode {
    pub fn new () -> LocalRouteNode {
        LocalRouteNode {
            children: None,
            clients: None
        }
    }

    pub fn insert(&mut self, tf: &TopicFilter, client_identifier: &str, qos: QualityOfService) -> bool {
        let mut last_node = (*tf).split("/").fold(self, |current, token| {
            let token = RouteKey::from(token);
            if current.children.is_none() {
                current.children = Some(HashMap::<RouteKey, LocalRouteNode>::new());
            }
            match current.children {
                Some(ref mut map) => {
                    if !map.contains_key(&token) {
                        map.insert(token.clone(), LocalRouteNode::new());
                    }
                    map.get_mut(&token).unwrap()
                }
                _ => unreachable!()
            }
        });

        if last_node.clients.is_none() {
            last_node.clients = Some(HashMap::<String, QualityOfService>::new())
        }
        match last_node.clients {
            Some(ref mut client_map) => {
                client_map.insert(client_identifier.to_owned(), qos).is_none()
            }
            _ => unreachable!()
        }
    }

    pub fn is_empty(&self) -> bool {
        (match self.clients {
            Some(ref client_map) => client_map.is_empty(),
            None => true
        }) && (match self.children {
            Some(ref map) => map.is_empty(),
            None => true
        })
    }

    fn _remove<'a, I>(&mut self, mut tokens: I, client_identifier: &str) -> bool
        where I: Iterator<Item=&'a str>
    {
        match tokens.next() {
            Some(token) => {
                match self.children {
                    Some(ref mut map) => {
                        let mut child_empty = false;
                        let key = RouteKey::from(token);
                        let is_removed = match map.get_mut(&key) {
                            Some(node) => {
                                let is_removed = node._remove(tokens, client_identifier);
                                child_empty = node.is_empty();
                                is_removed
                            }
                            None => false
                        };
                        if child_empty {
                            map.remove(&key);
                        }
                        is_removed
                    }
                    None => false
                }
            }
            None => {
                match self.clients {
                    Some(ref mut client_map) => {
                        client_map.remove(client_identifier).is_some()
                    }
                    None => false
                }
            }
        }
    }

    pub fn remove(&mut self, tf: &TopicFilter, client_identifier: &str) -> bool {
        self._remove((*tf).split("/"), client_identifier)
    }

    fn _search(&self, tokens: &Vec<&str>, index: usize) -> HashMap<String, QualityOfService> {
        let mut clients = HashMap::<String, QualityOfService>::new();
        if let Some(token) = tokens.get(index) {
            match self.children {
                Some(ref map) => {
                    if let Some(node) = map.get(&RouteKey::Normal(token.to_string())) {
                        clients.extend(node._search(tokens, index+1));
                    }
                    if let Some(node) = map.get(&RouteKey::OneLevel) {
                        clients.extend(node._search(tokens, index+1));
                    }
                    if let Some(node) = map.get(&RouteKey::RestLevels) {
                        if let Some(ref new_clients) = node.clients {
                            clients.extend(new_clients.clone());
                        }
                    }
                }
                None => {}
            }
        } else {
            if let Some(ref new_clients) = self.clients {
                clients.extend(new_clients.clone());
            }
        }
        clients
    }

    pub fn search(&self, name: &TopicName) -> HashMap<String, QualityOfService> {
        self._search(&(*name).split("/").collect::<Vec<&str>>(), 0)
    }
}

pub struct LocalRoutes {
    // For topic filters
    filter_routes: HashMap<u32, LocalRouteNode>,
    // For topic names
    name_routes: HashMap<u32, HashMap<TopicName, HashMap<String, QualityOfService>>>,
    // For remove topic-names and topic-filters
    client_topics: HashMap<u32, HashMap<String, HashSet<Topic>>>,
}

/// Router thread
/// * subscribe a topic
/// * unsubscribe a topic
/// * unsubscribe all topics from a client
/// * match topics
impl LocalRoutes {
    pub fn new() -> LocalRoutes {
        LocalRoutes {
            filter_routes: HashMap::<u32, LocalRouteNode>::new(),
            name_routes: HashMap::<u32, HashMap<TopicName, HashMap<String, QualityOfService>>>::new(),
            client_topics: HashMap::<u32, HashMap<String, HashSet<Topic>>>::new()
        }
    }

    fn insert_topic(&mut self, user_id: u32, topic: &Topic,
                    client_identifier: &str, qos: QualityOfService) -> bool {
        if !self.filter_routes.contains_key(&user_id) {
            self.filter_routes.insert(user_id, LocalRouteNode::new());
        }
        if !self.name_routes.contains_key(&user_id) {
            self.name_routes.insert(user_id, HashMap::<TopicName, HashMap<String, QualityOfService>>::new());
        }
        if !self.client_topics.contains_key(&user_id) {
            self.client_topics.insert(user_id, HashMap::<String, HashSet<Topic>>::new());
        }
        let mut filter_routes = self.filter_routes.get_mut(&user_id).unwrap();
        let mut name_routes = self.name_routes.get_mut(&user_id).unwrap();
        let mut client_topics = self.client_topics.get_mut(&user_id).unwrap();

        if !client_topics.contains_key(client_identifier) {
            client_topics.insert(client_identifier.to_owned(), HashSet::<Topic>::new());
        }
        client_topics.get_mut(client_identifier).unwrap()
            .insert(topic.clone());

        match topic {
            &Topic::Filter(ref topic_filter) => {
                filter_routes.insert(topic_filter, client_identifier, qos)
            }
            &Topic::Name(ref topic_name) => {
                if !name_routes.contains_key(topic_name) {
                    name_routes.insert(topic_name.clone(), HashMap::<String, QualityOfService>::new());
                }
                name_routes.get_mut(topic_name).unwrap()
                    .insert(client_identifier.to_owned(), qos).is_none()
            }
        }
    }

    fn remove_topic(&mut self, user_id: u32, topic: &Topic, client_identifier: &str) {
        if let Some(ref mut client_topics) = self.client_topics.get_mut(&user_id) {
            match client_topics.get_mut(client_identifier) {
                Some(ref mut topic_name_set) => topic_name_set.remove(topic),
                None => false
            };
        }
        match topic {
            &Topic::Filter(ref topic_filter) => {
                if let Some(ref mut filter_routes) = self.filter_routes.get_mut(&user_id) {
                    filter_routes.remove(topic_filter, client_identifier);
                }
            }
            &Topic::Name(ref topic_name) => {
                if let Some(ref mut name_routes) = self.name_routes.get_mut(&user_id) {
                    match name_routes.get_mut(topic_name) {
                        Some(ref mut client_map) => client_map.remove(client_identifier).is_some(),
                        None => false
                    };
                }
            }
        }
    }

    fn remove_all_topics(&mut self, user_id: u32, client_identifier: &str) -> usize {
        if let Some(ref mut client_topics) = self.client_topics.get_mut(&user_id) {
            let removed_count = match client_topics.get_mut(client_identifier) {
                Some(ref mut topic_name_set) => {
                    let mut removed_count = 0;
                    for topic in topic_name_set.iter() {
                        removed_count += match topic {
                            &Topic::Filter(ref topic_filter) => {
                                if let Some(ref mut filter_routes) = self.filter_routes.get_mut(&user_id) {
                                    if filter_routes.remove(topic_filter, client_identifier) { 1 } else { 0 }
                                } else { 0 }
                            }
                            &Topic::Name(ref topic_name) => {
                                if let Some(ref mut name_routes) = self.name_routes.get_mut(&user_id) {
                                    match name_routes.get_mut(topic_name) {
                                        Some(ref mut client_map) => {
                                            if client_map.remove(client_identifier).is_some() { 1 } else { 0 }
                                        }
                                        None => 0
                                    }
                                } else { 0 }
                            }
                        }
                    }
                    removed_count
                }
                None => 0
            };
            client_topics.remove(client_identifier);
            removed_count
        } else { 0 }
    }

    fn search(&self, user_id: u32, topic_name: &TopicName) -> HashMap<String, QualityOfService> {
        let mut clients = HashMap::<String, QualityOfService>::new();
        if let Some(ref name_routes) = self.name_routes.get(&user_id) {
            if let Some(client_map) = name_routes.get(topic_name) {
                clients.extend(client_map.clone());
            }
        }
        if let Some(ref filter_routes) = self.filter_routes.get(&user_id) {
            clients.extend(filter_routes.search(topic_name));
        }
        clients
    }
}
<|endoftext|>