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

#[derive(Debug,PartialEq)]
pub enum Token {
    Number(String),
    String(String),
    Name(String),
    ParenOpening,
    ParenClosing,
}

pub fn tokenizer(input: &str) -> Result<Vec<Token>, String> {
    let mut tokens: Vec<Token> = vec![];

    let mut char_iter = input.chars().peekable();
    while let Some(c) = char_iter.next() {
        match c {
            c if c.is_whitespace() => (),
            '(' => tokens.push(Token::ParenOpening),
            ')' => tokens.push(Token::ParenClosing),
            '0'...'9' => {
                let mut value = String::new();
                value.push(c);

                while match char_iter.peek() {
                    Some(&'0'...'9') => true,
                    _ => false,
                } {
                    if let Some(c) = char_iter.next() {
                        value.push(c);
                    }
                }
                tokens.push(Token::Number(value));
            }
            'a'...'z' => {
                let mut value = String::new();
                value.push(c);

                while match char_iter.peek() {
                    Some(&'a'...'z') => true,
                    _ => false,
                } {
                    if let Some(c) = char_iter.next() {
                        value.push(c);
                    }
                }
                tokens.push(Token::Name(value));
            }
            '"' => {
                let mut value = String::new();
                // skip opening "

                while match char_iter.peek() {
                    Some(&'"') | None => false,
                    _ => true,
                } {
                    if let Some(c) = char_iter.next() {
                        value.push(c);
                    }
                }
                tokens.push(Token::String(value));

                // skip closing ""
                if char_iter.peek() == Some(&'"') {
                    char_iter.next();
                }
            }
            _ => return Err(format!("I dont know what this character is: {}", c)),
        }
    }

    Ok(tokens)
}

#[derive(Debug,PartialEq)]
pub enum Ast {
    Programm { body: Vec<Ast> },
    CallExpression { name: String, params: Vec<Ast> },
    StringLiteral(String),
    NumberLiteral(String),
    SomethingElse,
}

pub fn parser(tokens: Vec<Token>) -> Result<Ast, String> {
    fn walk(token: &Token, token_iter: &mut Peekable<Iter<Token>>) -> Result<Ast, String> {
        match token {
            &Token::Number(ref value @ _) => Ok(Ast::NumberLiteral(value.to_string())),
            &Token::String(ref value @ _) => Ok(Ast::StringLiteral(value.to_string())),
            &Token::ParenOpening => {
                if let Some(token) = token_iter.next() {
                    match token {
                        &Token::Name(ref value @ _) => {
                            let name = value.to_string();
                            let mut params: Vec<Ast> = vec![];

                            while match token_iter.peek() {
                                Some(&&Token::ParenClosing) => {
                                    // skip closing )
                                    token_iter.next();
                                    false
                                }
                                None => false,
                                _ => true,
                            } {
                                if let Some(token) = token_iter.next() {
                                    match walk(token, token_iter) {
                                        Ok(nodes) => params.push(nodes),
                                        Err(value) => return Err(value),
                                    }
                                }
                            }

                            Ok(Ast::CallExpression {
                                name: name,
                                params: params,
                            })
                        }
                        _ => {
                            return Err(format!("{:?} isn't followed by a {:?}.",
                                               Token::ParenOpening,
                                               Token::Name("example".to_string())))
                        }
                    }
                } else {
                    return Err(format!("{:?} isn't followed by a node.", Token::ParenOpening));
                }
            }
            _ => return Err(format!("I dont know what this token is: {:?}", token)),
        }
    }

    let mut body: Vec<Ast> = vec![];

    let mut token_iter = tokens.iter().peekable();
    while let Some(token) = token_iter.next() {
        match walk(token, &mut token_iter) {
            Ok(nodes) => body.push(nodes),
            Err(value) => return Err(value),
        }
    }

    Ok(Ast::Programm { body: body })
}
<commit_msg>make intent clearer<commit_after>use std::iter::Peekable;
use std::slice::Iter;

#[derive(Debug,PartialEq)]
pub enum Token {
    Number(String),
    String(String),
    Name(String),
    ParenOpening,
    ParenClosing,
}

pub fn tokenizer(input: &str) -> Result<Vec<Token>, String> {
    let mut tokens: Vec<Token> = vec![];

    let mut char_iter = input.chars().peekable();
    while let Some(c) = char_iter.next() {
        match c {
            c if c.is_whitespace() => continue,
            '(' => tokens.push(Token::ParenOpening),
            ')' => tokens.push(Token::ParenClosing),
            '0'...'9' => {
                let mut value = String::new();
                value.push(c);

                while match char_iter.peek() {
                    Some(&'0'...'9') => true,
                    _ => false,
                } {
                    if let Some(c) = char_iter.next() {
                        value.push(c);
                    }
                }
                tokens.push(Token::Number(value));
            }
            'a'...'z' => {
                let mut value = String::new();
                value.push(c);

                while match char_iter.peek() {
                    Some(&'a'...'z') => true,
                    _ => false,
                } {
                    if let Some(c) = char_iter.next() {
                        value.push(c);
                    }
                }
                tokens.push(Token::Name(value));
            }
            '"' => {
                let mut value = String::new();
                // skip opening "

                while match char_iter.peek() {
                    Some(&'"') | None => false,
                    _ => true,
                } {
                    if let Some(c) = char_iter.next() {
                        value.push(c);
                    }
                }
                tokens.push(Token::String(value));

                // skip closing ""
                if char_iter.peek() == Some(&'"') {
                    char_iter.next();
                }
            }
            _ => return Err(format!("I dont know what this character is: {}", c)),
        }
    }

    Ok(tokens)
}

#[derive(Debug,PartialEq)]
pub enum Ast {
    Programm { body: Vec<Ast> },
    CallExpression { name: String, params: Vec<Ast> },
    StringLiteral(String),
    NumberLiteral(String),
    SomethingElse,
}

pub fn parser(tokens: Vec<Token>) -> Result<Ast, String> {
    fn walk(token: &Token, token_iter: &mut Peekable<Iter<Token>>) -> Result<Ast, String> {
        match token {
            &Token::Number(ref value @ _) => Ok(Ast::NumberLiteral(value.to_string())),
            &Token::String(ref value @ _) => Ok(Ast::StringLiteral(value.to_string())),
            &Token::ParenOpening => {
                if let Some(token) = token_iter.next() {
                    match token {
                        &Token::Name(ref value @ _) => {
                            let name = value.to_string();
                            let mut params: Vec<Ast> = vec![];

                            while match token_iter.peek() {
                                Some(&&Token::ParenClosing) => {
                                    // skip closing )
                                    token_iter.next();
                                    false
                                }
                                None => false,
                                _ => true,
                            } {
                                if let Some(token) = token_iter.next() {
                                    match walk(token, token_iter) {
                                        Ok(nodes) => params.push(nodes),
                                        Err(value) => return Err(value),
                                    }
                                }
                            }

                            Ok(Ast::CallExpression {
                                name: name,
                                params: params,
                            })
                        }
                        _ => {
                            return Err(format!("{:?} isn't followed by a {:?}.",
                                               Token::ParenOpening,
                                               Token::Name("example".to_string())))
                        }
                    }
                } else {
                    return Err(format!("{:?} isn't followed by a node.", Token::ParenOpening));
                }
            }
            _ => return Err(format!("I dont know what this token is: {:?}", token)),
        }
    }

    let mut body: Vec<Ast> = vec![];

    let mut token_iter = tokens.iter().peekable();
    while let Some(token) = token_iter.next() {
        match walk(token, &mut token_iter) {
            Ok(nodes) => body.push(nodes),
            Err(value) => return Err(value),
        }
    }

    Ok(Ast::Programm { body: body })
}
<|endoftext|>


<commit_before>#![crate_name = "rustsocks"]
#![crate_type = "rlib"]

pub mod socks4a;
pub mod socks4;
pub mod socks5;

mod util {
  use std::io::{IoResult, IoError, IoErrorKind};

  pub fn io_err<T>(kind: IoErrorKind, desc: &'static str) -> IoResult<T> {
      Err(IoError { kind: kind, desc: desc, detail: None })
  }
}
<commit_msg>Reexport for convenience<commit_after>#![crate_name = "rustsocks"]
#![crate_type = "rlib"]

pub use socks4a::Socks4a;
pub use socks4::Socks4;
pub use socks5::Socks5;

pub mod socks4a;
pub mod socks4;
pub mod socks5;

mod util {
  use std::io::{IoResult, IoError, IoErrorKind};

  pub fn io_err<T>(kind: IoErrorKind, desc: &'static str) -> IoResult<T> {
      Err(IoError { kind: kind, desc: desc, detail: None })
  }
}
<|endoftext|>


<commit_before>extern crate x264_sys as ffi;

use std::mem;
use ffi::x264::*;
use std::ptr::null;
use std::os::raw::c_int;

// TODO: Provide a builder API instead?
pub struct Param {
    par: x264_param_t,
}

impl Param {
    pub fn new() -> Param {
        let mut par = unsafe { mem::uninitialized() };

        unsafe {
            x264_param_default(&mut par as *mut x264_param_t);
        }

        Param { par: par }
    }
    pub fn default_preset(tune: Option<&str>, preset: Option<&str>) -> Result<Param, &'static str> {
        let mut par = unsafe { mem::uninitialized() };

        match unsafe {
                  x264_param_default_preset(&mut par as *mut x264_param_t,
                                            tune.map_or_else(|| null(),
                                                             |v| v.as_ptr() as *const i8),
                                            preset.map_or_else(|| null(), |v| v.as_ptr()) as
                                            *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(Param { par: par }),
            _ => Err("Unexpected"),
        }
    }
    pub fn apply_profile(mut self, profile: &str) -> Result<Param, &'static str> {
        match unsafe { x264_param_apply_profile(&mut self.par, profile.as_ptr() as *const i8) } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn param_parse<'a>(mut self, name: &str, value: &str) -> Result<Param, &'static str> {
        match unsafe {
                  x264_param_parse(&mut self.par,
                                   name.as_ptr() as *const i8,
                                   value.as_ptr() as *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn set_dimension(mut self, height: usize, width: usize) -> Param {
        self.par.i_height = height as c_int;
        self.par.i_width = width as c_int;

        self
    }
}

// TODO: Expose a NAL abstraction
pub struct NalData {
    vec: Vec<u8>,
}

impl NalData {
    fn with_capacity(capacity: usize) -> NalData {
        NalData { vec: Vec::with_capacity(capacity) }
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.vec.as_slice()
    }
}

pub struct Encoder {
    enc: *mut x264_t,
}

impl Encoder {
    pub fn open(par: &mut Param) -> Result<Encoder, &'static str> {
        let enc = unsafe { x264_encoder_open(&mut par.par as *mut x264_param_t) };

        if enc.is_null() {
            Err("Out of Memory")
        } else {
            Ok(Encoder { enc: enc })
        }
    }

    pub fn get_headers(&mut self) -> Result<NalData, &'static str> {
        let mut nb_nal: c_int = 0;
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };

        let bytes = unsafe {
            x264_encoder_headers(self.enc,
                                 &mut c_nals as *mut *mut x264_nal_t,
                                 &mut nb_nal as *mut c_int)
        };

        if bytes < 0 {
            Err("Encoding Headers Failed")
        } else {
            let mut data = NalData::with_capacity(bytes as usize);

            for i in 0..nb_nal {
                let nal = unsafe { Box::from_raw(c_nals.offset(i as isize)) };

                let payload =
                    unsafe { std::slice::from_raw_parts(nal.p_payload, nal.i_payload as usize) };

                data.vec.extend_from_slice(payload);

                mem::forget(payload);
                mem::forget(nal);
            }

            Ok(data)
        }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { x264_encoder_close(self.enc) };
    }
}

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

    #[test]
    fn test_open() {
        let mut par = Param::new().set_dimension(640, 480);

        let mut enc = Encoder::open(&mut par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());
    }
}
<commit_msg>picture: Basic safe wrapper around x264_picture_t<commit_after>extern crate x264_sys as ffi;

use std::mem;
use ffi::x264::*;
use std::ptr::null;
use std::os::raw::c_int;

pub struct Picture {
    pic: x264_picture_t,
    height: i32, // to compute the slice dimension
    native: bool,
}

impl Picture {
    /*
    pub fn new() -> Picture {
        let mut pic = unsafe { mem::uninitialized() };

        unsafe { x264_picture_init(&mut pic as *mut x264_picture_t) };

        Picture { pic: pic }
    }
*/
    pub fn from_param(param: &Param) -> Result<Picture, &'static str> {
        let mut pic = unsafe { mem::uninitialized() };

        let ret = unsafe {
            x264_picture_alloc(&mut pic as *mut x264_picture_t,
                               param.par.i_csp,
                               param.par.i_width,
                               param.par.i_height)
        };
        if ret < 0 {
            Err("Allocation Failure")
        } else {
            Ok(Picture {
                   pic: pic,
                   height: param.par.i_height,
                   native: true,
               })
        }
    }

    pub fn as_slice<'a>(&'a self, plane: usize) -> Result<&'a [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn as_slice_mut<'a>(&'a mut self, plane: usize) -> Result<&'a mut [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts_mut(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn set_timestamp(mut self, pts: i64) -> Picture {
        self.pic.i_pts = pts;
        self
    }
}

impl Drop for Picture {
    fn drop(&mut self) {
        if self.native {
            unsafe { x264_picture_clean(&mut self.pic as *mut x264_picture_t) };
        }
    }
}

// TODO: Provide a builder API instead?
pub struct Param {
    par: x264_param_t,
}

impl Param {
    pub fn new() -> Param {
        let mut par = unsafe { mem::uninitialized() };

        unsafe {
            x264_param_default(&mut par as *mut x264_param_t);
        }

        Param { par: par }
    }
    pub fn default_preset(tune: Option<&str>, preset: Option<&str>) -> Result<Param, &'static str> {
        let mut par = unsafe { mem::uninitialized() };

        match unsafe {
                  x264_param_default_preset(&mut par as *mut x264_param_t,
                                            tune.map_or_else(|| null(),
                                                             |v| v.as_ptr() as *const i8),
                                            preset.map_or_else(|| null(), |v| v.as_ptr()) as
                                            *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(Param { par: par }),
            _ => Err("Unexpected"),
        }
    }
    pub fn apply_profile(mut self, profile: &str) -> Result<Param, &'static str> {
        match unsafe { x264_param_apply_profile(&mut self.par, profile.as_ptr() as *const i8) } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn param_parse<'a>(mut self, name: &str, value: &str) -> Result<Param, &'static str> {
        match unsafe {
                  x264_param_parse(&mut self.par,
                                   name.as_ptr() as *const i8,
                                   value.as_ptr() as *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn set_dimension(mut self, height: usize, width: usize) -> Param {
        self.par.i_height = height as c_int;
        self.par.i_width = width as c_int;

        self
    }
}

// TODO: Expose a NAL abstraction
pub struct NalData {
    vec: Vec<u8>,
}

impl NalData {
    fn with_capacity(capacity: usize) -> NalData {
        NalData { vec: Vec::with_capacity(capacity) }
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.vec.as_slice()
    }
}

pub struct Encoder {
    enc: *mut x264_t,
}

impl Encoder {
    pub fn open(par: &mut Param) -> Result<Encoder, &'static str> {
        let enc = unsafe { x264_encoder_open(&mut par.par as *mut x264_param_t) };

        if enc.is_null() {
            Err("Out of Memory")
        } else {
            Ok(Encoder { enc: enc })
        }
    }

    pub fn get_headers(&mut self) -> Result<NalData, &'static str> {
        let mut nb_nal: c_int = 0;
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };

        let bytes = unsafe {
            x264_encoder_headers(self.enc,
                                 &mut c_nals as *mut *mut x264_nal_t,
                                 &mut nb_nal as *mut c_int)
        };

        if bytes < 0 {
            Err("Encoding Headers Failed")
        } else {
            let mut data = NalData::with_capacity(bytes as usize);

            for i in 0..nb_nal {
                let nal = unsafe { Box::from_raw(c_nals.offset(i as isize)) };

                let payload =
                    unsafe { std::slice::from_raw_parts(nal.p_payload, nal.i_payload as usize) };

                data.vec.extend_from_slice(payload);

                mem::forget(payload);
                mem::forget(nal);
            }

            Ok(data)
        }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { x264_encoder_close(self.enc) };
    }
}

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

    #[test]
    fn test_open() {
        let mut par = Param::new().set_dimension(640, 480);

        let mut enc = Encoder::open(&mut par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());
    }

    #[test]
    fn test_picture() {
        let par = Param::new().set_dimension(640, 480);
        {
            let mut pic = Picture::from_param(&par).unwrap();
            {
                let p = pic.as_slice_mut(0).unwrap();
                p[0] = 1;
            }
            let p = pic.as_slice(0).unwrap();

            assert_eq!(p[0], 1);
        }
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;

// What does this do? You'll never guess!
type HeapId<'a> = PhantomData<std::cell::Cell<&'a mut ()>>;

// The heap is (for now) just a big Vec of Pairs
pub struct Heap<'a> {
    id: HeapId<'a>,
    storage: Vec<PairStorage<'a>>,
    freelist: *mut PairStorage<'a>,
    pins: RefCell<HashMap<*mut PairStorage<'a>, usize>>
}

pub const HEAP_SIZE: usize = 10000;

pub fn with_heap<F, O>(f: F) -> O
    where F: for<'a> FnOnce(&mut Heap<'a>) -> O
{
    f(&mut Heap::new())
}

unsafe trait Mark {
    unsafe fn mark(ptr: *mut Self);
}

impl<'a> Heap<'a> {
    fn new() -> Heap<'a> {
        // Allocate with the full capacity so that allocated Pairs never move
        let mut h = Heap {
            id: PhantomData,
            storage: Vec::with_capacity(HEAP_SIZE),
            freelist: std::ptr::null_mut(),
            pins: RefCell::new(HashMap::new())
        };
        for i in 0 .. HEAP_SIZE {
            h.storage.push(PairStorage {
                marked: false,
                head: XValue::Null,
                tail: XValue::Null
            });
            let p = &mut h.storage[i] as *mut PairStorage<'a>;
            unsafe {
                h.add_to_free_list(p);
            }
        }
        h
    }

    fn pin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        let entry = pins.entry(p).or_insert(0);
        *entry += 1;
    }

    fn unpin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        if {
            let entry = pins.entry(p).or_insert(0);
            assert!(*entry != 0);
            *entry -= 1;
            *entry == 0
        } {
            pins.remove(&p);
        }
    }

    unsafe fn add_to_free_list(&mut self, p: *mut PairStorage<'a>) {
        let listp: *mut *mut PairStorage<'a> = std::mem::transmute(p);
        *listp = self.freelist;
        assert_eq!(*listp, self.freelist);
        self.freelist = p;
    }

    pub fn try_alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> Option<PairRoot<'a>> {
        if self.freelist.is_null() {
            unsafe {
                self.gc();
            }
            if self.freelist.is_null() {
                return None;
            }
        }

        let p = self.freelist;
        unsafe {
            let listp: *mut *mut PairStorage<'a> = std::mem::transmute(p);
            self.freelist = *listp;

            let (h, t) = pair;
            *p = PairStorage {
                marked: false,
                head: value_to_heap(h),
                tail: value_to_heap(t)
            };
            Some(PairRoot::new(self, p))
        }
    }

    pub fn alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> PairRoot<'a> {
        self.try_alloc(pair).expect("out of memory (gc did not collect anything)")
    }

    pub fn alloc_null(&mut self) -> PairRoot<'a> {
        self.alloc((Value::Null, Value::Null))
    }

    unsafe fn gc(&mut self) {
        // clear mark bits
        for p in &mut self.storage {
            p.marked = false;
        }

        // mark phase
        for (&p, _) in self.pins.borrow().iter() {
            Mark::mark(p);
        }

        // sweep phase
        self.freelist = std::ptr::null_mut();
        for i in 0 .. HEAP_SIZE {
            let p = &mut self.storage[i] as *mut PairStorage<'a>;
            if !(*p).marked {
                self.add_to_free_list(p);
            }
        }
    }

    #[cfg(test)]
    pub fn force_gc(&mut self) {
        unsafe { self.gc(); }
    }
}

// === Pair, the reference type

// PairStorage is the only type that is actually allocated inside the heap (private)
struct PairStorage<'a> {
    marked: bool,
    head: XValue<'a>,
    tail: XValue<'a>
}

unsafe impl<'a> Mark for PairStorage<'a> {
    unsafe fn mark(ptr: *mut PairStorage<'a>) {
        if !ptr.is_null() && !(*ptr).marked {
            (*ptr).marked = true;
            Mark::mark(&mut (*ptr).head as *mut XValue<'a>);
            Mark::mark(&mut (*ptr).tail as *mut XValue<'a>);
        }
    }
}

// Handle to a PairStorage that lives in the heap.
#[allow(raw_pointer_derive)]
#[derive(Debug, PartialEq)]
pub struct PairRoot<'a> {
    ptr: *mut PairStorage<'a>,
    heap: *const Heap<'a>,
    heap_id: HeapId<'a>
}

impl<'a> Drop for PairRoot<'a> {
    fn drop(&mut self) {
        unsafe {
            (*self.heap).unpin(self.ptr);
        }
    }
}

impl<'a> Clone for PairRoot<'a> {
    fn clone(&self) -> PairRoot<'a> {
        let &PairRoot { ptr, heap, heap_id } = self;
        unsafe {
            (*heap).pin(ptr);
        }
        PairRoot {
            ptr: ptr,
            heap: heap,
            heap_id: heap_id
        }
    }
}

impl<'a> PairRoot<'a> {
    unsafe fn new(heap: &Heap<'a>, p: *mut PairStorage<'a>) -> PairRoot<'a> {
        heap.pin(p);
        PairRoot {
            ptr: p,
            heap: heap as *const Heap<'a>,
            heap_id: heap.id
        }
    }

    pub fn head(&self) -> Value<'a> {
        let ptr = self.ptr;
        unsafe { value_from_heap(&*self.heap, &(*ptr).head) }
    }

    pub fn tail(&self) -> Value<'a> {
        let ptr = self.ptr;
        unsafe { value_from_heap(&*self.heap, &(*ptr).tail) }
    }

    pub fn set_head(&self, v: Value<'a>) {
        let ptr = self.ptr;
        unsafe { (*ptr).head = value_to_heap(v); }
    }

    pub fn set_tail(&self, v: Value<'a>) {
        let ptr = self.ptr;
        unsafe { (*ptr).tail = value_to_heap(v); }
    }

    #[cfg(test)]
    pub fn address(&self) -> usize {
        unsafe { std::mem::transmute(self.ptr) }
    }
}

// === Values (a heap-inline enum)

// Values inside the heap (GC heap-to-heap cross-references) (private)
enum XValue<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),
    Pair(*mut PairStorage<'a>, HeapId<'a>)
}

unsafe impl<'a> Mark for XValue<'a> {
    unsafe fn mark(ptr: *mut XValue<'a>) {
        match *ptr {
            XValue::Pair(p, _) => Mark::mark(p),
            _ => {}
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Value<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),  // <-- equality is by value
    Pair(PairRoot<'a>)  // <-- equality is by pointer
}

// === Getting data into and out of the heap

fn value_to_heap<'a>(v: Value<'a>) -> XValue<'a> {
    match v {
        Value::Null => XValue::Null,
        Value::Int(n) => XValue::Int(n),
        Value::Str(rcstr) => XValue::Str(rcstr),
        Value::Pair(PairRoot{ptr, heap_id, ..}) => XValue::Pair(ptr, heap_id)
    }
}

unsafe fn value_from_heap<'a>(heap: &Heap<'a>, v: &XValue<'a>) -> Value<'a> {
    match v {
        &XValue::Null => Value::Null,
        &XValue::Int(n) => Value::Int(n),
        &XValue::Str(ref rcstr) => Value::Str(rcstr.clone()),
        &XValue::Pair(ptr, _) => Value::Pair(PairRoot::new(heap, ptr))
    }
}

#[cfg(test)]
mod test;
<commit_msg>Rename `XValue` -> `ValueStorage`.<commit_after>use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;

// What does this do? You'll never guess!
type HeapId<'a> = PhantomData<std::cell::Cell<&'a mut ()>>;

// The heap is (for now) just a big Vec of Pairs
pub struct Heap<'a> {
    id: HeapId<'a>,
    storage: Vec<PairStorage<'a>>,
    freelist: *mut PairStorage<'a>,
    pins: RefCell<HashMap<*mut PairStorage<'a>, usize>>
}

pub const HEAP_SIZE: usize = 10000;

pub fn with_heap<F, O>(f: F) -> O
    where F: for<'a> FnOnce(&mut Heap<'a>) -> O
{
    f(&mut Heap::new())
}

unsafe trait Mark {
    unsafe fn mark(ptr: *mut Self);
}

impl<'a> Heap<'a> {
    fn new() -> Heap<'a> {
        // Allocate with the full capacity so that allocated Pairs never move
        let mut h = Heap {
            id: PhantomData,
            storage: Vec::with_capacity(HEAP_SIZE),
            freelist: std::ptr::null_mut(),
            pins: RefCell::new(HashMap::new())
        };
        for i in 0 .. HEAP_SIZE {
            h.storage.push(PairStorage {
                marked: false,
                head: ValueStorage::Null,
                tail: ValueStorage::Null
            });
            let p = &mut h.storage[i] as *mut PairStorage<'a>;
            unsafe {
                h.add_to_free_list(p);
            }
        }
        h
    }

    fn pin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        let entry = pins.entry(p).or_insert(0);
        *entry += 1;
    }

    fn unpin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        if {
            let entry = pins.entry(p).or_insert(0);
            assert!(*entry != 0);
            *entry -= 1;
            *entry == 0
        } {
            pins.remove(&p);
        }
    }

    unsafe fn add_to_free_list(&mut self, p: *mut PairStorage<'a>) {
        let listp: *mut *mut PairStorage<'a> = std::mem::transmute(p);
        *listp = self.freelist;
        assert_eq!(*listp, self.freelist);
        self.freelist = p;
    }

    pub fn try_alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> Option<PairRoot<'a>> {
        if self.freelist.is_null() {
            unsafe {
                self.gc();
            }
            if self.freelist.is_null() {
                return None;
            }
        }

        let p = self.freelist;
        unsafe {
            let listp: *mut *mut PairStorage<'a> = std::mem::transmute(p);
            self.freelist = *listp;

            let (h, t) = pair;
            *p = PairStorage {
                marked: false,
                head: value_to_heap(h),
                tail: value_to_heap(t)
            };
            Some(PairRoot::new(self, p))
        }
    }

    pub fn alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> PairRoot<'a> {
        self.try_alloc(pair).expect("out of memory (gc did not collect anything)")
    }

    pub fn alloc_null(&mut self) -> PairRoot<'a> {
        self.alloc((Value::Null, Value::Null))
    }

    unsafe fn gc(&mut self) {
        // clear mark bits
        for p in &mut self.storage {
            p.marked = false;
        }

        // mark phase
        for (&p, _) in self.pins.borrow().iter() {
            Mark::mark(p);
        }

        // sweep phase
        self.freelist = std::ptr::null_mut();
        for i in 0 .. HEAP_SIZE {
            let p = &mut self.storage[i] as *mut PairStorage<'a>;
            if !(*p).marked {
                self.add_to_free_list(p);
            }
        }
    }

    #[cfg(test)]
    pub fn force_gc(&mut self) {
        unsafe { self.gc(); }
    }
}

// === Pair, the reference type

// PairStorage is the only type that is actually allocated inside the heap (private)
struct PairStorage<'a> {
    marked: bool,
    head: ValueStorage<'a>,
    tail: ValueStorage<'a>
}

unsafe impl<'a> Mark for PairStorage<'a> {
    unsafe fn mark(ptr: *mut PairStorage<'a>) {
        if !ptr.is_null() && !(*ptr).marked {
            (*ptr).marked = true;
            Mark::mark(&mut (*ptr).head as *mut ValueStorage<'a>);
            Mark::mark(&mut (*ptr).tail as *mut ValueStorage<'a>);
        }
    }
}

// Handle to a PairStorage that lives in the heap.
#[allow(raw_pointer_derive)]
#[derive(Debug, PartialEq)]
pub struct PairRoot<'a> {
    ptr: *mut PairStorage<'a>,
    heap: *const Heap<'a>,
    heap_id: HeapId<'a>
}

impl<'a> Drop for PairRoot<'a> {
    fn drop(&mut self) {
        unsafe {
            (*self.heap).unpin(self.ptr);
        }
    }
}

impl<'a> Clone for PairRoot<'a> {
    fn clone(&self) -> PairRoot<'a> {
        let &PairRoot { ptr, heap, heap_id } = self;
        unsafe {
            (*heap).pin(ptr);
        }
        PairRoot {
            ptr: ptr,
            heap: heap,
            heap_id: heap_id
        }
    }
}

impl<'a> PairRoot<'a> {
    unsafe fn new(heap: &Heap<'a>, p: *mut PairStorage<'a>) -> PairRoot<'a> {
        heap.pin(p);
        PairRoot {
            ptr: p,
            heap: heap as *const Heap<'a>,
            heap_id: heap.id
        }
    }

    pub fn head(&self) -> Value<'a> {
        let ptr = self.ptr;
        unsafe { value_from_heap(&*self.heap, &(*ptr).head) }
    }

    pub fn tail(&self) -> Value<'a> {
        let ptr = self.ptr;
        unsafe { value_from_heap(&*self.heap, &(*ptr).tail) }
    }

    pub fn set_head(&self, v: Value<'a>) {
        let ptr = self.ptr;
        unsafe { (*ptr).head = value_to_heap(v); }
    }

    pub fn set_tail(&self, v: Value<'a>) {
        let ptr = self.ptr;
        unsafe { (*ptr).tail = value_to_heap(v); }
    }

    #[cfg(test)]
    pub fn address(&self) -> usize {
        unsafe { std::mem::transmute(self.ptr) }
    }
}

// === Values (a heap-inline enum)

// Values inside the heap (GC heap-to-heap cross-references) (private)
enum ValueStorage<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),
    Pair(*mut PairStorage<'a>, HeapId<'a>)
}

unsafe impl<'a> Mark for ValueStorage<'a> {
    unsafe fn mark(ptr: *mut ValueStorage<'a>) {
        match *ptr {
            ValueStorage::Pair(p, _) => Mark::mark(p),
            _ => {}
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Value<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),  // <-- equality is by value
    Pair(PairRoot<'a>)  // <-- equality is by pointer
}

// === Getting data into and out of the heap

fn value_to_heap<'a>(v: Value<'a>) -> ValueStorage<'a> {
    match v {
        Value::Null => ValueStorage::Null,
        Value::Int(n) => ValueStorage::Int(n),
        Value::Str(rcstr) => ValueStorage::Str(rcstr),
        Value::Pair(PairRoot{ptr, heap_id, ..}) => ValueStorage::Pair(ptr, heap_id)
    }
}

unsafe fn value_from_heap<'a>(heap: &Heap<'a>, v: &ValueStorage<'a>) -> Value<'a> {
    match v {
        &ValueStorage::Null => Value::Null,
        &ValueStorage::Int(n) => Value::Int(n),
        &ValueStorage::Str(ref rcstr) => Value::Str(rcstr.clone()),
        &ValueStorage::Pair(ptr, _) => Value::Pair(PairRoot::new(heap, ptr))
    }
}

#[cfg(test)]
mod test;
<|endoftext|>


<commit_before>#![feature(read_exact, associated_consts, associated_type_defaults)]

#[macro_use]
extern crate quick_error;
extern crate byteorder;
#[macro_use]
extern crate log;

mod gamedata_io;

pub use gamedata_io::{ReadError, StringReadError};

use std::io::{self, Read, Write};
use std::path;

/// The data of a Game Maker Studio game.
///
/// This is the collective information acquired from "data.win".
pub struct GameData {
    pub metadata: MetaData,
    optn: Optn,
    extn: Extn,
    sounds: Sounds,
    audio_groups: Option<AudioGroups>,
    sprites: Sprites,
    backgrounds: Backgrounds,
    paths: Paths,
    scripts: Scripts,
    shaders: Shaders,
    fonts: Fonts,
    timelines: Timelines,
    objects: Objects,
    rooms: Rooms,
    dafl: Dafl,
    tpag: Tpag,
    code: Code,
    pub variables: Variables,
    pub functions: Functions,
    pub strings: Strings,
    textures: Textures,
    audio: Audio,
}

impl GameData {
    pub fn from_reader<R: Read>(reader: &mut R) -> Result<GameData, ReadError> {
        gamedata_io::read(reader)
    }
    pub fn from_file<P: AsRef<path::Path>>(path: &P) -> Result<GameData, ReadError> {
        use std::fs::File;
        use std::io::BufReader;
        let file = try!(File::open(path));
        GameData::from_reader(&mut BufReader::new(file))
    }
    pub fn write_to_writer<W: Write>(&self, writer: &mut W) -> Result<(), io::Error> {
        gamedata_io::write(self, writer)
    }
    pub fn save_to_file<P: AsRef<path::Path>>(&self, path: &P) -> Result<(), io::Error> {
        use std::fs::File;
        use std::io::BufWriter;
        let file = try!(File::create(path));
        self.write_to_writer(&mut BufWriter::new(file))
    }
    pub fn string_at(&self, index: usize) -> &str {
        &self.strings.strings[index]
    }
    pub fn replace_string_at<S: Into<String>>(&mut self, index: usize, with: S) {
        self.strings.strings[index] = with.into();
    }
    pub fn window_width(&self) -> u32 {
        self.metadata.window_width
    }
    pub fn window_height(&self) -> u32 {
        self.metadata.window_height
    }
    pub fn set_window_dimensions(&mut self, width: u32, height: u32) {
        self.metadata.window_width = width;
        self.metadata.window_height = height;
    }
    pub fn window_title(&self) -> &str {
        &self.strings.strings[self.metadata.window_title_index]
    }
    pub fn set_window_title<S: Into<String>>(&mut self, new: S) {
        self.strings.strings[self.metadata.window_title_index] = new.into();
    }
}

/// Contains various metadata.
pub struct MetaData {
    unk1: u32, // Purpose unknown
    game_id_1_index: usize, // Some kind of game id
    default_index: usize, // Points to "Default"
    unk2: u32,
    unk3: u32,
    unk4: u32,
    unk5: u32,
    unk6: u32,
    unk7: u32,
    unk8: u32,
    game_id_2_index: usize, // Some kind of game id, identical to game_id_1
    unk9: u32,
    unk10: u32,
    unk11: u32,
    unk12: u32,
    window_width: u32,
    window_height: u32,
    unk13: u32,
    unk14: u32,
    unk15: u32,
    unk16: u32,
    unk17: u32,
    unk18: u32,
    unk19: u32,
    unk20: u32,
    window_title_index: usize,
    unknown: Vec<u32>,
}

/// Purpose unknown.
pub struct Optn {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Purpose unknown.
pub struct Extn {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of sounds.
pub struct Sounds {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Collection of audio groups. Not present in all games.
pub struct AudioGroups {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of sprites.
pub struct Sprites {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of backgrounds.
pub struct Backgrounds {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of paths.
pub struct Paths {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of scripts.
pub struct Scripts {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of shaders.
pub struct Shaders {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of fonts.
pub struct Fonts {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of timelines.
pub struct Timelines {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of objects.
pub struct Objects {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of rooms.
pub struct Rooms {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Purpose unknown.
pub struct Dafl {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Purpose unknown.
pub struct Tpag {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Something to do with code. Maybe code of scripts?
pub struct Code {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A game maker variable.
pub struct Variable {
    /// Index of the name of the variable in the strings section.
    pub name_index: usize,
    unknown1: u32, // Purpose unknown
    unknown2: u32, // Purpose unknown
}

/// A collection of variables.
pub struct Variables {
    pub variables: Vec<Variable>,
}

/// A game maker function.
pub struct Function {
    /// Index of the name of the function in the strings section.
    pub name_index: usize,
    unknown1: u32, // Purpose unknown
    unknown2: u32, // Purpose unknown
}

/// A collection of functions.
pub struct Functions {
    pub functions: Vec<Function>,
}

/// A collection of strings.
///
/// All strings are assumed to be valid UTF-8.
pub struct Strings {
    pub strings: Vec<String>,
}

/// A collection of textures.
pub struct Textures {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of audio data.
pub struct Audio {
    raw: Vec<u8>, // Data not analyzed yet
}
<commit_msg>Replace some leftover Result<(), io::Error>s<commit_after>#![feature(read_exact, associated_consts, associated_type_defaults)]

#[macro_use]
extern crate quick_error;
extern crate byteorder;
#[macro_use]
extern crate log;

mod gamedata_io;

pub use gamedata_io::{ReadError, StringReadError};

use std::io::{self, Read, Write};
use std::path;

/// The data of a Game Maker Studio game.
///
/// This is the collective information acquired from "data.win".
pub struct GameData {
    pub metadata: MetaData,
    optn: Optn,
    extn: Extn,
    sounds: Sounds,
    audio_groups: Option<AudioGroups>,
    sprites: Sprites,
    backgrounds: Backgrounds,
    paths: Paths,
    scripts: Scripts,
    shaders: Shaders,
    fonts: Fonts,
    timelines: Timelines,
    objects: Objects,
    rooms: Rooms,
    dafl: Dafl,
    tpag: Tpag,
    code: Code,
    pub variables: Variables,
    pub functions: Functions,
    pub strings: Strings,
    textures: Textures,
    audio: Audio,
}

impl GameData {
    pub fn from_reader<R: Read>(reader: &mut R) -> Result<GameData, ReadError> {
        gamedata_io::read(reader)
    }
    pub fn from_file<P: AsRef<path::Path>>(path: &P) -> Result<GameData, ReadError> {
        use std::fs::File;
        use std::io::BufReader;
        let file = try!(File::open(path));
        GameData::from_reader(&mut BufReader::new(file))
    }
    pub fn write_to_writer<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        gamedata_io::write(self, writer)
    }
    pub fn save_to_file<P: AsRef<path::Path>>(&self, path: &P) -> io::Result<()> {
        use std::fs::File;
        use std::io::BufWriter;
        let file = try!(File::create(path));
        self.write_to_writer(&mut BufWriter::new(file))
    }
    pub fn string_at(&self, index: usize) -> &str {
        &self.strings.strings[index]
    }
    pub fn replace_string_at<S: Into<String>>(&mut self, index: usize, with: S) {
        self.strings.strings[index] = with.into();
    }
    pub fn window_width(&self) -> u32 {
        self.metadata.window_width
    }
    pub fn window_height(&self) -> u32 {
        self.metadata.window_height
    }
    pub fn set_window_dimensions(&mut self, width: u32, height: u32) {
        self.metadata.window_width = width;
        self.metadata.window_height = height;
    }
    pub fn window_title(&self) -> &str {
        &self.strings.strings[self.metadata.window_title_index]
    }
    pub fn set_window_title<S: Into<String>>(&mut self, new: S) {
        self.strings.strings[self.metadata.window_title_index] = new.into();
    }
}

/// Contains various metadata.
pub struct MetaData {
    unk1: u32, // Purpose unknown
    game_id_1_index: usize, // Some kind of game id
    default_index: usize, // Points to "Default"
    unk2: u32,
    unk3: u32,
    unk4: u32,
    unk5: u32,
    unk6: u32,
    unk7: u32,
    unk8: u32,
    game_id_2_index: usize, // Some kind of game id, identical to game_id_1
    unk9: u32,
    unk10: u32,
    unk11: u32,
    unk12: u32,
    window_width: u32,
    window_height: u32,
    unk13: u32,
    unk14: u32,
    unk15: u32,
    unk16: u32,
    unk17: u32,
    unk18: u32,
    unk19: u32,
    unk20: u32,
    window_title_index: usize,
    unknown: Vec<u32>,
}

/// Purpose unknown.
pub struct Optn {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Purpose unknown.
pub struct Extn {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of sounds.
pub struct Sounds {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Collection of audio groups. Not present in all games.
pub struct AudioGroups {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of sprites.
pub struct Sprites {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of backgrounds.
pub struct Backgrounds {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of paths.
pub struct Paths {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of scripts.
pub struct Scripts {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of shaders.
pub struct Shaders {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of fonts.
pub struct Fonts {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of timelines.
pub struct Timelines {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of objects.
pub struct Objects {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of rooms.
pub struct Rooms {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Purpose unknown.
pub struct Dafl {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Purpose unknown.
pub struct Tpag {
    raw: Vec<u8>, // Data not analyzed yet
}

/// Something to do with code. Maybe code of scripts?
pub struct Code {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A game maker variable.
pub struct Variable {
    /// Index of the name of the variable in the strings section.
    pub name_index: usize,
    unknown1: u32, // Purpose unknown
    unknown2: u32, // Purpose unknown
}

/// A collection of variables.
pub struct Variables {
    pub variables: Vec<Variable>,
}

/// A game maker function.
pub struct Function {
    /// Index of the name of the function in the strings section.
    pub name_index: usize,
    unknown1: u32, // Purpose unknown
    unknown2: u32, // Purpose unknown
}

/// A collection of functions.
pub struct Functions {
    pub functions: Vec<Function>,
}

/// A collection of strings.
///
/// All strings are assumed to be valid UTF-8.
pub struct Strings {
    pub strings: Vec<String>,
}

/// A collection of textures.
pub struct Textures {
    raw: Vec<u8>, // Data not analyzed yet
}

/// A collection of audio data.
pub struct Audio {
    raw: Vec<u8>, // Data not analyzed yet
}
<|endoftext|>


<commit_before>extern crate emailaddress;
#[macro_use]
extern crate log;
extern crate hyper;
extern crate iron;
extern crate lettre;
extern crate redis;
extern crate serde;
extern crate serde_json;
extern crate time;
extern crate url;
extern crate urlencoded;

use emailaddress::EmailAddress;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::modifiers;
use iron::method::Method;
use iron::prelude::*;
use iron::status;
use serde_json::builder::{ArrayBuilder, ObjectBuilder};
use serde_json::value::Value;
use std::env;
use std::fs::File;
use std::io::{BufReader, Read};
use std::sync::Arc;
use std::error::Error;
use time::now_utc;
use urlencoded::{UrlEncodedBody, UrlEncodedQuery};

pub mod error;
pub mod config;
pub use config::AppConfig;
pub mod crypto;
pub mod email;
pub mod oidc;
pub mod store;
pub mod store_cache;

use error::{BrokerResult, BrokerError};

/// Macro that creates Handler implementations that log the request,
/// and keep a reference to the AppConfig.
macro_rules! broker_handler {
    ( $name:ident , | $app:ident, $req:ident | $body:block ) => {
        pub struct $name {
            pub app: Arc<AppConfig>,
        }
        impl $name {
            pub fn new(app: &Arc<AppConfig>) -> Self {
                $name { app: app.clone() }
            }
            fn handle_body($app: &AppConfig, $req: &mut Request)
                -> BrokerResult<Response> $body
        }
        impl Handler for $name {
            fn handle(&self, req: &mut Request) -> IronResult<Response> {
                info!("{} {}", req.method, req.url);
                Self::handle_body(&self.app, req)
                    .or_else(|e| handle_error(&self.app, req, e))
            }
        }
    }
}

/// Macro used to extract a parameter from a QueryMap.
///
/// Will return from the caller with a BrokerError if the parameter is missing.
///
/// ```
/// let foo = try_get_param!(params, "foo");
/// ```
macro_rules! try_get_param {
    ( $input:expr , $param:tt ) => {
        try!($input.get($param).map(|list| &list[0]).ok_or_else(|| {
            BrokerError::Input(concat!("missing request parameter ", $param).to_string())
        }))
    }
}

/// Handle an BrokerError and create an IronResult.
///
/// The `broker_handler!` macro calls this on error. We don't use a `From`
/// implementation, because this way we get app and request context, and we
/// don't necessarily have to pass the error on to Iron.
fn handle_error(app: &AppConfig, _: &mut Request, err: BrokerError) -> IronResult<Response> {
    match err {
        BrokerError::Input(_) => {
            let obj = ObjectBuilder::new()
                .insert("error", err.description())
                .build();
            let content = serde_json::to_string(&obj).unwrap();
            let content_type = modifiers::Header(ContentType::json());
            Ok(Response::with((status::BadRequest, content_type, content)))
        }
        BrokerError::Provider(_) => {
            // TODO: Redirect to RP with the error description
            Err(IronError::new(err, status::ServiceUnavailable))
        }
        _ => {
            Err(IronError::new(err, (status::InternalServerError,
                                     modifiers::Header(ContentType::html()),
                                     app.templates.error.render(&[]))))
        }
    }
}

/// Helper function for returning an Iron response with JSON data.
///
/// Serializes the argument value to JSON and returns a HTTP 200 response
/// code with the serialized JSON as the body.
fn json_response(obj: &Value) -> BrokerResult<Response> {
    let content = serde_json::to_string(&obj).unwrap();
    Ok(Response::with((status::Ok,
                       modifiers::Header(ContentType::json()),
                       content)))
}

/// Iron handler for the root path, returns human-friendly message.
///
/// This is not actually used in the protocol.
broker_handler!(WelcomeHandler, |_app, _req| {
    json_response(&ObjectBuilder::new()
        .insert("ladaemon", "Welcome")
        .insert("version", env!("CARGO_PKG_VERSION"))
        .build())
});

/// Iron handler for files in .well-known.
///
/// Mainly directed at Let's Encrypt verification. Returns text/plain always.
broker_handler!(WellKnownHandler, |_app, req| {
    let mut file_name = env::current_dir().unwrap();
    file_name.push(req.url.path().join("/"));
    let file_res = File::open(file_name);
    if file_res.is_err() {
        return Ok(Response::with((status::NotFound)));
    }
    let mut bytes = Vec::<u8>::new();
    let mut reader = BufReader::new(file_res.unwrap());
    let _ = reader.read_to_end(&mut bytes).unwrap();
    Ok(Response::with((status::Ok,
                       modifiers::Header(ContentType::plaintext()),
                       bytes)))
});

/// Iron handler to return the OpenID Discovery document.
///
/// Most of this is hard-coded for now, although the URLs are constructed by
/// using the base URL as configured in the `base_url` configuration value.
broker_handler!(OIDConfigHandler, |app, _req| {
    json_response(&ObjectBuilder::new()
        .insert("issuer", &app.base_url)
        .insert("authorization_endpoint",
                format!("{}/auth", app.base_url))
        .insert("jwks_uri", format!("{}/keys.json", app.base_url))
        .insert("scopes_supported", vec!["openid", "email"])
        .insert("claims_supported",
                vec!["aud", "email", "email_verified", "exp", "iat", "iss", "sub"])
        .insert("response_types_supported", vec!["id_token"])
        .insert("response_modes_supported", vec!["form_post"])
        .insert("grant_types_supported", vec!["implicit"])
        .insert("subject_types_supported", vec!["public"])
        .insert("id_token_signing_alg_values_supported", vec!["RS256"])
        .build())
});

/// Iron handler for the JSON Web Key Set document.
///
/// Respond with the JWK key set containing all of the configured keys.
///
/// Relying Parties will need to fetch this data to be able to verify identity
/// tokens issued by this daemon instance.
broker_handler!(KeysHandler, |app, _req| {
    let mut keys = ArrayBuilder::new();
    for key in &app.keys {
        keys = keys.push(key.public_jwk())
    }
    json_response(&ObjectBuilder::new()
                      .insert("keys", keys.build())
                      .build())
});

/// Iron handler for authentication requests from the RP.
///
/// Calls the `oidc::request()` function if the provided email address's
/// domain matches one of the configured famous providers. Otherwise, calls the
/// `email::request()` function to allow authentication through the email loop.
broker_handler!(AuthHandler, |app, req| {
    let params = try!(
        match req.method {
            Method::Get => {
                req.get_ref::<UrlEncodedQuery>()
                    .map_err(|_| BrokerError::Input("no query string in GET request".to_string()))
            },
            Method::Post => {
                req.get_ref::<UrlEncodedBody>()
                    .map_err(|_| BrokerError::Input("no query string in POST data".to_string()))
            },
            _ => {
                panic!("Unsupported method: {}", req.method)
            }
        }
    );
    let client_id = try_get_param!(params, "client_id");
    let nonce = try_get_param!(params, "nonce");
    let redirect_uri = try_get_param!(params, "redirect_uri");
    let email_addr = try!(
        EmailAddress::new(try_get_param!(params, "login_hint"))
            .map_err(|_| BrokerError::Input("login_hint is not a valid email address".to_string()))
    );
    if app.providers.contains_key(&email_addr.domain) {

        // OIDC authentication. Using 302 Found for redirection here. Note
        // that, per RFC 7231, a user agent MAY change the request method
        // from POST to GET for the subsequent request.
        let auth_url = try!(
            oidc::request(app, email_addr, client_id, nonce, redirect_uri)
        );
        Ok(Response::with((status::Found, modifiers::Redirect(auth_url))))

    } else {

        // Email loop authentication. Render a message and form.
        let session_id = try!(email::request(app, email_addr, client_id, nonce, redirect_uri));
        Ok(Response::with((status::Ok,
                           modifiers::Header(ContentType::html()),
                           app.templates.confirm_email.render(&[
                               ("client_id", &client_id),
                               ("session_id", &session_id),
                           ]))))

    }
});

/// Helper method to create a JWT for a given email address and origin.
///
/// Builds the JSON payload, then signs it using the last key provided in
/// the configuration object.
fn create_jwt(app: &AppConfig, email: &str, origin: &str, nonce: &str) -> String {
    let now = now_utc().to_timespec().sec;
    let payload = &ObjectBuilder::new()
        .insert("aud", origin)
        .insert("email", email)
        .insert("email_verified", email)
        .insert("exp", now + app.token_validity as i64)
        .insert("iat", now)
        .insert("iss", &app.base_url)
        .insert("sub", email)
        .insert("nonce", nonce)
        .build();
    app.keys.last().unwrap().sign_jws(&payload)
}

/// Helper function for returning result to the Relying Party.
///
/// Takes a `(jwt, redirect)` pair from one of the verification functions and
/// embeds it in a form in `tmpl/forward.html`, from where it's POSTed to
/// the RP's `redirect` as soon as the page has loaded.
fn return_to_relier(app: &AppConfig, result: (String, String)) -> BrokerResult<Response> {
    let (jwt, redirect) = result;
    Ok(Response::with((status::Ok,
                       modifiers::Header(ContentType::html()),
                       app.templates.forward.render(&[
                           ("return_url", &redirect),
                           ("jwt", &jwt),
                       ]))))
}

/// Iron handler for one-time pad email loop confirmation.
///
/// Retrieves the session based session ID and the expected one-time pad.
/// Verify the code and return the resulting token or error to the RP.
broker_handler!(ConfirmHandler, |app, req| {
    let params = try!(
        req.get_ref::<UrlEncodedQuery>()
            .map_err(|_| BrokerError::Input("no query string in GET request".to_string()))
    );
    let session_id = try_get_param!(params, "session");
    let code = try_get_param!(params, "code");
    return_to_relier(app, try!(email::verify(app, session_id, code)))
});

/// Iron handler for OAuth callbacks
///
/// After the user allows or denies the Authentication Request with the famous
/// identity provider, they will be redirected back to the callback handler.
/// Verify the callback data and return the resulting token or error.
broker_handler!(CallbackHandler, |app, req| {
    let params = try!(
        req.get_ref::<UrlEncodedQuery>()
            .map_err(|_| BrokerError::Input("no query string in GET request".to_string()))
    );
    let session = try_get_param!(params, "state");
    let code = try_get_param!(params, "code");
    return_to_relier(app, try!(oidc::verify(app, session, code)))
});
<commit_msg>Redirect with 302 instead of 303.<commit_after>extern crate emailaddress;
#[macro_use]
extern crate log;
extern crate hyper;
extern crate iron;
extern crate lettre;
extern crate redis;
extern crate serde;
extern crate serde_json;
extern crate time;
extern crate url;
extern crate urlencoded;

use emailaddress::EmailAddress;
use iron::headers::{ContentType, Location};
use iron::middleware::Handler;
use iron::modifiers;
use iron::method::Method;
use iron::prelude::*;
use iron::status;
use serde_json::builder::{ArrayBuilder, ObjectBuilder};
use serde_json::value::Value;
use std::env;
use std::fs::File;
use std::io::{BufReader, Read};
use std::sync::Arc;
use std::error::Error;
use time::now_utc;
use urlencoded::{UrlEncodedBody, UrlEncodedQuery};

pub mod error;
pub mod config;
pub use config::AppConfig;
pub mod crypto;
pub mod email;
pub mod oidc;
pub mod store;
pub mod store_cache;

use error::{BrokerResult, BrokerError};

/// Macro that creates Handler implementations that log the request,
/// and keep a reference to the AppConfig.
macro_rules! broker_handler {
    ( $name:ident , | $app:ident, $req:ident | $body:block ) => {
        pub struct $name {
            pub app: Arc<AppConfig>,
        }
        impl $name {
            pub fn new(app: &Arc<AppConfig>) -> Self {
                $name { app: app.clone() }
            }
            fn handle_body($app: &AppConfig, $req: &mut Request)
                -> BrokerResult<Response> $body
        }
        impl Handler for $name {
            fn handle(&self, req: &mut Request) -> IronResult<Response> {
                info!("{} {}", req.method, req.url);
                Self::handle_body(&self.app, req)
                    .or_else(|e| handle_error(&self.app, req, e))
            }
        }
    }
}

/// Macro used to extract a parameter from a QueryMap.
///
/// Will return from the caller with a BrokerError if the parameter is missing.
///
/// ```
/// let foo = try_get_param!(params, "foo");
/// ```
macro_rules! try_get_param {
    ( $input:expr , $param:tt ) => {
        try!($input.get($param).map(|list| &list[0]).ok_or_else(|| {
            BrokerError::Input(concat!("missing request parameter ", $param).to_string())
        }))
    }
}

/// Handle an BrokerError and create an IronResult.
///
/// The `broker_handler!` macro calls this on error. We don't use a `From`
/// implementation, because this way we get app and request context, and we
/// don't necessarily have to pass the error on to Iron.
fn handle_error(app: &AppConfig, _: &mut Request, err: BrokerError) -> IronResult<Response> {
    match err {
        BrokerError::Input(_) => {
            let obj = ObjectBuilder::new()
                .insert("error", err.description())
                .build();
            let content = serde_json::to_string(&obj).unwrap();
            let content_type = modifiers::Header(ContentType::json());
            Ok(Response::with((status::BadRequest, content_type, content)))
        }
        BrokerError::Provider(_) => {
            // TODO: Redirect to RP with the error description
            Err(IronError::new(err, status::ServiceUnavailable))
        }
        _ => {
            Err(IronError::new(err, (status::InternalServerError,
                                     modifiers::Header(ContentType::html()),
                                     app.templates.error.render(&[]))))
        }
    }
}

/// Helper function for returning an Iron response with JSON data.
///
/// Serializes the argument value to JSON and returns a HTTP 200 response
/// code with the serialized JSON as the body.
fn json_response(obj: &Value) -> BrokerResult<Response> {
    let content = serde_json::to_string(&obj).unwrap();
    Ok(Response::with((status::Ok,
                       modifiers::Header(ContentType::json()),
                       content)))
}

/// Iron handler for the root path, returns human-friendly message.
///
/// This is not actually used in the protocol.
broker_handler!(WelcomeHandler, |_app, _req| {
    json_response(&ObjectBuilder::new()
        .insert("ladaemon", "Welcome")
        .insert("version", env!("CARGO_PKG_VERSION"))
        .build())
});

/// Iron handler for files in .well-known.
///
/// Mainly directed at Let's Encrypt verification. Returns text/plain always.
broker_handler!(WellKnownHandler, |_app, req| {
    let mut file_name = env::current_dir().unwrap();
    file_name.push(req.url.path().join("/"));
    let file_res = File::open(file_name);
    if file_res.is_err() {
        return Ok(Response::with((status::NotFound)));
    }
    let mut bytes = Vec::<u8>::new();
    let mut reader = BufReader::new(file_res.unwrap());
    let _ = reader.read_to_end(&mut bytes).unwrap();
    Ok(Response::with((status::Ok,
                       modifiers::Header(ContentType::plaintext()),
                       bytes)))
});

/// Iron handler to return the OpenID Discovery document.
///
/// Most of this is hard-coded for now, although the URLs are constructed by
/// using the base URL as configured in the `base_url` configuration value.
broker_handler!(OIDConfigHandler, |app, _req| {
    json_response(&ObjectBuilder::new()
        .insert("issuer", &app.base_url)
        .insert("authorization_endpoint",
                format!("{}/auth", app.base_url))
        .insert("jwks_uri", format!("{}/keys.json", app.base_url))
        .insert("scopes_supported", vec!["openid", "email"])
        .insert("claims_supported",
                vec!["aud", "email", "email_verified", "exp", "iat", "iss", "sub"])
        .insert("response_types_supported", vec!["id_token"])
        .insert("response_modes_supported", vec!["form_post"])
        .insert("grant_types_supported", vec!["implicit"])
        .insert("subject_types_supported", vec!["public"])
        .insert("id_token_signing_alg_values_supported", vec!["RS256"])
        .build())
});

/// Iron handler for the JSON Web Key Set document.
///
/// Respond with the JWK key set containing all of the configured keys.
///
/// Relying Parties will need to fetch this data to be able to verify identity
/// tokens issued by this daemon instance.
broker_handler!(KeysHandler, |app, _req| {
    let mut keys = ArrayBuilder::new();
    for key in &app.keys {
        keys = keys.push(key.public_jwk())
    }
    json_response(&ObjectBuilder::new()
                      .insert("keys", keys.build())
                      .build())
});

/// Iron handler for authentication requests from the RP.
///
/// Calls the `oidc::request()` function if the provided email address's
/// domain matches one of the configured famous providers. Otherwise, calls the
/// `email::request()` function to allow authentication through the email loop.
broker_handler!(AuthHandler, |app, req| {
    let params = try!(
        match req.method {
            Method::Get => {
                req.get_ref::<UrlEncodedQuery>()
                    .map_err(|_| BrokerError::Input("no query string in GET request".to_string()))
            },
            Method::Post => {
                req.get_ref::<UrlEncodedBody>()
                    .map_err(|_| BrokerError::Input("no query string in POST data".to_string()))
            },
            _ => {
                panic!("Unsupported method: {}", req.method)
            }
        }
    );
    let client_id = try_get_param!(params, "client_id");
    let nonce = try_get_param!(params, "nonce");
    let redirect_uri = try_get_param!(params, "redirect_uri");
    let email_addr = try!(
        EmailAddress::new(try_get_param!(params, "login_hint"))
            .map_err(|_| BrokerError::Input("login_hint is not a valid email address".to_string()))
    );
    if app.providers.contains_key(&email_addr.domain) {

        // OIDC authentication. Redirect to the identity provider.
        let auth_url = try!(oidc::request(app, email_addr, client_id, nonce, redirect_uri));
        Ok(Response::with((status::SeeOther, modifiers::Header(Location(auth_url.to_string())))))

    } else {

        // Email loop authentication. Render a message and form.
        let session_id = try!(email::request(app, email_addr, client_id, nonce, redirect_uri));
        Ok(Response::with((status::Ok,
                           modifiers::Header(ContentType::html()),
                           app.templates.confirm_email.render(&[
                               ("client_id", &client_id),
                               ("session_id", &session_id),
                           ]))))

    }
});

/// Helper method to create a JWT for a given email address and origin.
///
/// Builds the JSON payload, then signs it using the last key provided in
/// the configuration object.
fn create_jwt(app: &AppConfig, email: &str, origin: &str, nonce: &str) -> String {
    let now = now_utc().to_timespec().sec;
    let payload = &ObjectBuilder::new()
        .insert("aud", origin)
        .insert("email", email)
        .insert("email_verified", email)
        .insert("exp", now + app.token_validity as i64)
        .insert("iat", now)
        .insert("iss", &app.base_url)
        .insert("sub", email)
        .insert("nonce", nonce)
        .build();
    app.keys.last().unwrap().sign_jws(&payload)
}

/// Helper function for returning result to the Relying Party.
///
/// Takes a `(jwt, redirect)` pair from one of the verification functions and
/// embeds it in a form in `tmpl/forward.html`, from where it's POSTed to
/// the RP's `redirect` as soon as the page has loaded.
fn return_to_relier(app: &AppConfig, result: (String, String)) -> BrokerResult<Response> {
    let (jwt, redirect) = result;
    Ok(Response::with((status::Ok,
                       modifiers::Header(ContentType::html()),
                       app.templates.forward.render(&[
                           ("return_url", &redirect),
                           ("jwt", &jwt),
                       ]))))
}

/// Iron handler for one-time pad email loop confirmation.
///
/// Retrieves the session based session ID and the expected one-time pad.
/// Verify the code and return the resulting token or error to the RP.
broker_handler!(ConfirmHandler, |app, req| {
    let params = try!(
        req.get_ref::<UrlEncodedQuery>()
            .map_err(|_| BrokerError::Input("no query string in GET request".to_string()))
    );
    let session_id = try_get_param!(params, "session");
    let code = try_get_param!(params, "code");
    return_to_relier(app, try!(email::verify(app, session_id, code)))
});

/// Iron handler for OAuth callbacks
///
/// After the user allows or denies the Authentication Request with the famous
/// identity provider, they will be redirected back to the callback handler.
/// Verify the callback data and return the resulting token or error.
broker_handler!(CallbackHandler, |app, req| {
    let params = try!(
        req.get_ref::<UrlEncodedQuery>()
            .map_err(|_| BrokerError::Input("no query string in GET request".to_string()))
    );
    let session = try_get_param!(params, "state");
    let code = try_get_param!(params, "code");
    return_to_relier(app, try!(oidc::verify(app, session, code)))
});
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::fmt;
use std::iter::FromIterator;

#[derive(Clone, PartialEq, Eq)]
pub enum Token {
    Output,
    Input,
    Loop(Vec<Token>),
    Move(i32),
    Add(i32, i32),
    Set(i32, i32),
    MulCopy(i32, i32, i32),
    Scan(i32),
    LoadOut(i32, i32),
    LoadOutSet(i32),
    If(i32, Vec<Token>),
}

impl fmt::Debug for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Output => write!(f, "Output"),
            Input => write!(f, "Input"),
            Move(offset) => write!(f, "Move(offset={})", offset),
            Add(offset, value) => write!(f, "Add(offset={}, value={})", offset, value),
            Set(offset, value) => write!(f, "Set(offset={}, value={})", offset, value),
            MulCopy(src, dest, mul) => {
                write!(f, "MulCopy(src={}, dest={}, mul={})", src, dest, mul)
            }
            Scan(offset) => write!(f, "Scan(offset={})", offset),
            LoadOut(offset, add) => write!(f, "LoadOut(offset={}, add={})", offset, add),
            LoadOutSet(value) => write!(f, "LoadOutSet(value={})", value),
            Loop(ref content) => {
                if f.alternate() {
                    write!(f, "Loop(content={:#?})", content)
                } else {
                    write!(f, "Loop(content={:?})", content)
                }
            }
            If(offset, ref content) => {
                if f.alternate() {
                    write!(f, "If(offset={}, content={:#?})", offset, content)
                } else {
                    write!(f, "If(offset={}, content={:?})", offset, content)
                }
            }
        }
    }
}

use Token::*;

#[derive(Default)]
struct OptimizeState {
    tokens: Vec<Token>,
    // With HashMap, the order sometimes switches
    // in recursion, and the optimizer never exits.
    adds: BTreeMap<i32, i32>,
    sets: BTreeMap<i32, i32>,
    shift: i32,
}

impl OptimizeState {
    fn apply_shift(&mut self) {
        if self.shift != 0 {
            self.tokens.push(Move(self.shift));
            self.shift = 0;
        }
    }

    fn apply_adds_sets(&mut self) {
        for (offset, value) in &self.sets {
            self.tokens.push(Set(*offset, *value));
        }
        for (offset, value) in &self.adds {
            self.tokens.push(Add(*offset, *value));
        }
        self.sets.clear();
        self.adds.clear();
    }

    fn add(&mut self, offset: i32, mut value: i32) {
        if self.sets.contains_key(&offset) {
            value = self.sets.get(&offset).unwrap() + value;
            self.sets.insert(offset, value);
        } else {
            value = self.adds.get(&offset).unwrap_or(&0) + value;
            if value != 0 {
                self.adds.insert(offset, value);
            } else {
                self.adds.remove(&offset);
            }
        }
    }

    fn set(&mut self, offset: i32, value: i32) {
        // Add before Set does nothing; remove it
        self.adds.remove(&offset);
        self.sets.insert(offset, value);
    }
}

pub fn parse(code: &str) -> Vec<Token> {
    _parse(&mut code.chars())
}

fn _parse(chars: &mut std::str::Chars) -> Vec<Token> {
    let mut tokens = Vec::new();
    while let Some(i) = chars.next() {
        match i {
            '+' => tokens.push(Add(0, 1)),
            '-' => tokens.push(Add(0, -1)),
            '>' => tokens.push(Move(1)),
            '<' => tokens.push(Move(-1)),
            '[' => tokens.push(Loop(_parse(chars))),
            ']' => {
                break;
            }
            ',' => tokens.push(Input),
            '.' => {
                tokens.push(LoadOut(0, 0));
                tokens.push(Output);
            }
            _ => (),
        };
    }

    tokens
}

fn _optimize(tokens: &Vec<Token>) -> OptimizeState {
    let mut do_output = false;
    let mut state = OptimizeState::default();

    for token in tokens {
        match *token {
            Set(..) | Add(..) | Move(_) | LoadOut(..) | LoadOutSet(_) | Output | MulCopy(..) => {}
            _ => {
                if do_output {
                    state.tokens.push(Output);
                    do_output = false;
                }

                state.apply_adds_sets();
            }
        }

        match *token {
            Loop(_) | Input | Scan(_) => {
                state.apply_shift();
            }
            _ => {}
        }

        match *token {
            Set(mut offset, value) => {
                offset += state.shift;
                state.set(offset, value);
            }
            Add(mut offset, value) => {
                offset += state.shift;
                state.add(offset, value);
            }
            MulCopy(mut src, mut dest, mul) => {
                src += state.shift;
                dest += state.shift;
                if let Some(value) = state.sets.remove(&src) {
                    state.set(dest, value * mul);
                } else {
                    state.apply_adds_sets();
                    state.tokens.push(MulCopy(src, dest, mul));
                }
            }
            If(offset, ref contents) => {
                let mut newcontents = Vec::new();
                for i in contents {
                    newcontents.push(match *i {
                        Set(offset, value) => Set(offset + state.shift, value),
                        MulCopy(src, dest, mul) => {
                            MulCopy(src + state.shift, dest + state.shift, mul)
                        }
                        _ => unreachable!(),
                    });
                }
                state.tokens.push(If(offset + state.shift, newcontents));
            }
            Move(offset) => state.shift += offset,
            Output => do_output = true,
            LoadOut(mut offset, add) => {
                offset += state.shift;
                if state.sets.contains_key(&offset) {
                    state.tokens.push(LoadOutSet(state.sets.get(&offset).unwrap() + add));
                } else {
                    state.tokens.push(LoadOut(offset, state.adds.get(&offset).unwrap_or(&0) + add));
                }
            }
            Loop(ref contents) => _optimize_loop(contents, &mut state),
            LoadOutSet(value) => state.tokens.push(LoadOutSet(value)),
            Input => state.tokens.push(Input),
            Scan(offset) => state.tokens.push(Scan(offset + state.shift)),
        }
    }

    if do_output {
        state.tokens.push(Output);
    }

    state
}

fn _optimize_loop(tokens: &Vec<Token>, outer: &mut OptimizeState) {
    let mut inner = _optimize(tokens);

    if inner.shift != 0 && inner.sets.is_empty() && inner.adds.is_empty() &&
       inner.tokens.is_empty() {
        outer.tokens.push(Scan(inner.shift));
    } else if inner.shift == 0 && inner.tokens.is_empty() && inner.adds.contains_key(&0) &&
              inner.adds.len() == 1 {
        if !inner.sets.is_empty() {
            let mut iftokens = Vec::new();
            for (offset, value) in &inner.sets {
                iftokens.push(Set(*offset, *value));
            }
            iftokens.push(Set(0, 0));
            outer.tokens.push(If(0, iftokens));
        } else {
            outer.set(0, 0);
        }
    } else if inner.shift == 0 && inner.tokens.is_empty() && inner.adds.get(&0) == Some(&-1) {
        let contents = inner.adds.iter().filter_map(|(offset, value)| {
            if *offset != 0 {
                let src = 0;
                let dest = *offset;
                let mul = *value;
                Some(MulCopy(src, dest, mul))
            } else {
                None
            }
        });

        if !inner.sets.is_empty() {
            let iftokens = Vec::from_iter(inner.sets
                .iter()
                .map(|(offset, value)| Set(*offset, *value))
                .chain(contents));
            outer.tokens.push(If(0, iftokens));
        } else {
            outer.tokens.extend(contents);
        }

        outer.set(0, 0);
    } else {
        inner.apply_adds_sets();
        inner.apply_shift();

        outer.tokens.push(Loop(inner.tokens));
    }
}

pub fn optimize(tokens: Vec<Token>) -> Vec<Token> {
    let mut oldtokens = tokens;
    // Ignore sets/adds/shifts at end of file
    let mut newtokens = _optimize(&oldtokens).tokens;
    while newtokens != oldtokens {
        oldtokens = newtokens;
        newtokens = _optimize(&oldtokens).tokens;
    }
    newtokens
}
<commit_msg>Make MulCopy optimization ignore set/add at irrelevant offsets<commit_after>use std::collections::BTreeMap;
use std::fmt;
use std::iter::FromIterator;

#[derive(Clone, PartialEq, Eq)]
pub enum Token {
    Output,
    Input,
    Loop(Vec<Token>),
    Move(i32),
    Add(i32, i32),
    Set(i32, i32),
    MulCopy(i32, i32, i32),
    Scan(i32),
    LoadOut(i32, i32),
    LoadOutSet(i32),
    If(i32, Vec<Token>),
}

impl fmt::Debug for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Output => write!(f, "Output"),
            Input => write!(f, "Input"),
            Move(offset) => write!(f, "Move(offset={})", offset),
            Add(offset, value) => write!(f, "Add(offset={}, value={})", offset, value),
            Set(offset, value) => write!(f, "Set(offset={}, value={})", offset, value),
            MulCopy(src, dest, mul) => {
                write!(f, "MulCopy(src={}, dest={}, mul={})", src, dest, mul)
            }
            Scan(offset) => write!(f, "Scan(offset={})", offset),
            LoadOut(offset, add) => write!(f, "LoadOut(offset={}, add={})", offset, add),
            LoadOutSet(value) => write!(f, "LoadOutSet(value={})", value),
            Loop(ref content) => {
                if f.alternate() {
                    write!(f, "Loop(content={:#?})", content)
                } else {
                    write!(f, "Loop(content={:?})", content)
                }
            }
            If(offset, ref content) => {
                if f.alternate() {
                    write!(f, "If(offset={}, content={:#?})", offset, content)
                } else {
                    write!(f, "If(offset={}, content={:?})", offset, content)
                }
            }
        }
    }
}

use Token::*;

#[derive(Default)]
struct OptimizeState {
    tokens: Vec<Token>,
    // With HashMap, the order sometimes switches
    // in recursion, and the optimizer never exits.
    adds: BTreeMap<i32, i32>,
    sets: BTreeMap<i32, i32>,
    shift: i32,
}

impl OptimizeState {
    fn apply_shift(&mut self) {
        if self.shift != 0 {
            self.tokens.push(Move(self.shift));
            self.shift = 0;
        }
    }

    fn apply_adds_sets(&mut self) {
        for (offset, value) in &self.sets {
            self.tokens.push(Set(*offset, *value));
        }
        for (offset, value) in &self.adds {
            self.tokens.push(Add(*offset, *value));
        }
        self.sets.clear();
        self.adds.clear();
    }

    fn add(&mut self, offset: i32, mut value: i32) {
        if self.sets.contains_key(&offset) {
            value = self.sets.get(&offset).unwrap() + value;
            self.sets.insert(offset, value);
        } else {
            value = self.adds.get(&offset).unwrap_or(&0) + value;
            if value != 0 {
                self.adds.insert(offset, value);
            } else {
                self.adds.remove(&offset);
            }
        }
    }

    fn set(&mut self, offset: i32, value: i32) {
        // Add before Set does nothing; remove it
        self.adds.remove(&offset);
        self.sets.insert(offset, value);
    }
}

pub fn parse(code: &str) -> Vec<Token> {
    _parse(&mut code.chars())
}

fn _parse(chars: &mut std::str::Chars) -> Vec<Token> {
    let mut tokens = Vec::new();
    while let Some(i) = chars.next() {
        match i {
            '+' => tokens.push(Add(0, 1)),
            '-' => tokens.push(Add(0, -1)),
            '>' => tokens.push(Move(1)),
            '<' => tokens.push(Move(-1)),
            '[' => tokens.push(Loop(_parse(chars))),
            ']' => {
                break;
            }
            ',' => tokens.push(Input),
            '.' => {
                tokens.push(LoadOut(0, 0));
                tokens.push(Output);
            }
            _ => (),
        };
    }

    tokens
}

fn _optimize(tokens: &Vec<Token>) -> OptimizeState {
    let mut do_output = false;
    let mut state = OptimizeState::default();

    for token in tokens {
        match *token {
            Set(..) | Add(..) | Move(_) | LoadOut(..) | LoadOutSet(_) | Output | MulCopy(..) => {}
            _ => {
                if do_output {
                    state.tokens.push(Output);
                    do_output = false;
                }

                state.apply_adds_sets();
            }
        }

        match *token {
            Loop(_) | Input | Scan(_) => {
                state.apply_shift();
            }
            _ => {}
        }

        match *token {
            Set(mut offset, value) => {
                offset += state.shift;
                state.set(offset, value);
            }
            Add(mut offset, value) => {
                offset += state.shift;
                state.add(offset, value);
            }
            MulCopy(mut src, mut dest, mul) => {
                src += state.shift;
                dest += state.shift;
                if let Some(value) = state.sets.remove(&src) {
                    state.set(dest, value * mul);
                } else {
                    if state.sets.contains_key(&dest) || state.adds.contains_key(&src) ||
                       state.adds.contains_key(&dest) {
                        state.apply_adds_sets();
                    }
                    state.tokens.push(MulCopy(src, dest, mul));
                }
            }
            If(offset, ref contents) => {
                let mut newcontents = Vec::new();
                for i in contents {
                    newcontents.push(match *i {
                        Set(offset, value) => Set(offset + state.shift, value),
                        MulCopy(src, dest, mul) => {
                            MulCopy(src + state.shift, dest + state.shift, mul)
                        }
                        _ => unreachable!(),
                    });
                }
                state.tokens.push(If(offset + state.shift, newcontents));
            }
            Move(offset) => state.shift += offset,
            Output => do_output = true,
            LoadOut(mut offset, add) => {
                offset += state.shift;
                if state.sets.contains_key(&offset) {
                    state.tokens.push(LoadOutSet(state.sets.get(&offset).unwrap() + add));
                } else {
                    state.tokens.push(LoadOut(offset, state.adds.get(&offset).unwrap_or(&0) + add));
                }
            }
            Loop(ref contents) => _optimize_loop(contents, &mut state),
            LoadOutSet(value) => state.tokens.push(LoadOutSet(value)),
            Input => state.tokens.push(Input),
            Scan(offset) => state.tokens.push(Scan(offset + state.shift)),
        }
    }

    if do_output {
        state.tokens.push(Output);
    }

    state
}

fn _optimize_loop(tokens: &Vec<Token>, outer: &mut OptimizeState) {
    let mut inner = _optimize(tokens);

    if inner.shift != 0 && inner.sets.is_empty() && inner.adds.is_empty() &&
       inner.tokens.is_empty() {
        outer.tokens.push(Scan(inner.shift));
    } else if inner.shift == 0 && inner.tokens.is_empty() && inner.adds.contains_key(&0) &&
              inner.adds.len() == 1 {
        if !inner.sets.is_empty() {
            let mut iftokens = Vec::new();
            for (offset, value) in &inner.sets {
                iftokens.push(Set(*offset, *value));
            }
            iftokens.push(Set(0, 0));
            outer.tokens.push(If(0, iftokens));
        } else {
            outer.set(0, 0);
        }
    } else if inner.shift == 0 && inner.tokens.is_empty() && inner.adds.get(&0) == Some(&-1) {
        let contents = inner.adds.iter().filter_map(|(offset, value)| {
            if *offset != 0 {
                let src = 0;
                let dest = *offset;
                let mul = *value;
                Some(MulCopy(src, dest, mul))
            } else {
                None
            }
        });

        if !inner.sets.is_empty() {
            let iftokens = Vec::from_iter(inner.sets
                .iter()
                .map(|(offset, value)| Set(*offset, *value))
                .chain(contents));
            outer.tokens.push(If(0, iftokens));
        } else {
            outer.tokens.extend(contents);
        }

        outer.set(0, 0);
    } else {
        inner.apply_adds_sets();
        inner.apply_shift();

        outer.tokens.push(Loop(inner.tokens));
    }
}

pub fn optimize(tokens: Vec<Token>) -> Vec<Token> {
    let mut oldtokens = tokens;
    // Ignore sets/adds/shifts at end of file
    let mut newtokens = _optimize(&oldtokens).tokens;
    while newtokens != oldtokens {
        oldtokens = newtokens;
        newtokens = _optimize(&oldtokens).tokens;
    }
    newtokens
}
<|endoftext|>