<commit_before>#![feature(rustc_private)]
extern crate llvm_sys as llvm;
extern crate libc;
use std::ptr;
use std::ffi;

use self::llvm::prelude::{LLVMContextRef, LLVMModuleRef, LLVMBuilderRef, LLVMValueRef, LLVMTypeRef};
use self::llvm::core::*;

use std::collections::{HashMap};
use std::mem;

use syntax::ast::{Block, Expr, Decl, TType, OptionalTypeExprTupleList};
use syntax::ptr::{B};
//FIXME this import is for integration testing purposes
use syntax::parse::*;//{Parser};
use syntax::parse::parser::{Parser};

macro_rules! c_str_ptr {
    ($s:expr) => {
        ffi::CString::new($s).unwrap().as_ptr()
    };
}

pub struct Context{
    context : LLVMContextRef,
    module : LLVMModuleRef,
    builder : LLVMBuilderRef,
    named_values : HashMap<String, LLVMValueRef>
}

impl Context{
    fn new(module_name : &str) -> Self{
        unsafe{
            let llvm_context =  LLVMContextCreate();
            let llvm_module = LLVMModuleCreateWithNameInContext(c_str_ptr!(module_name),
                                                                llvm_context);
            let builder = LLVMCreateBuilderInContext(llvm_context);
            let named_values = HashMap::new();

            Context {
                context : llvm_context,
                module : llvm_module,
                builder : builder,
                named_values : named_values
            }
        }
    }

    pub fn dump(&self){
        unsafe{
            LLVMDumpModule(self.module);
        }
    }
}

impl Drop for Context{
    fn drop(&mut self){
        unsafe{
            LLVMDisposeBuilder(self.builder);
            LLVMDisposeModule(self.module);
            LLVMContextDispose(self.context);
        }
    }
}

//TODO move these in a seperate file
type IRBuildingResult = Result<LLVMValueRef, String>;

trait IRBuilder{
    fn codegen(&self, ctxt : &mut Context) -> IRBuildingResult;
}

fn std_functions_call_factory(fn_name : &str,
                              args : &OptionalTypeExprTupleList,
                              ctxt : &mut Context) -> Option<LLVMValueRef>{
    unsafe{
        match fn_name {
            "print" =>{
                debug_assert!(args.is_some(), "No args passed to print()");
                let lst = args.as_ref().unwrap();
                debug_assert!(lst.len() == 1, "One arg should be passed to print()");
                debug_assert!(lst[0].0 == TType::TString);
                let str_arg = match &*lst[0].1 {
                    &Expr::StringExpr(ref value) => c_str_ptr!(&*(value.clone())),
                    _ => panic!("Expected a string expr")
                };

                let print_ty = LLVMIntTypeInContext(ctxt.context, 32);
                let mut pf_type_args_vec = Vec::new();
                pf_type_args_vec.push(LLVMPointerType(LLVMIntTypeInContext(ctxt.context, 8),
                                                      0));
                let proto = LLVMFunctionType(print_ty, pf_type_args_vec.as_mut_ptr(), 1, 1);
                let print_function = LLVMAddFunction(ctxt.module,
                                                     c_str_ptr!("printf"),
                                                     proto);
                let gstr = LLVMBuildGlobalStringPtr(ctxt.builder,
                                                    str_arg,
                                                    c_str_ptr!(".str"));
                let mut pf_args = Vec::new();
                pf_args.push(gstr);

                Some(LLVMBuildCall(ctxt.builder,
                                   print_function,
                                   pf_args.as_mut_ptr(),
                                   1,
                                   c_str_ptr!("call")))
            },
            _ => {None}
        }
    }
}

fn get_llvm_type_for_ttype(ty : &TType, ctxt : &mut Context) -> LLVMTypeRef{
    unsafe{
        match ty {
            &TType::TVoid => LLVMVoidTypeInContext(ctxt.context),
            &TType::TInt32 => LLVMIntTypeInContext(ctxt.context, 32),
            _ => panic!("Other TTypes not mapped yet to the corresponding LLVM types")
        }
    }
}

impl IRBuilder for Expr{
    fn codegen(&self, ctxt : &mut Context) -> IRBuildingResult{
        unsafe{
            match self{
                &Expr::NumExpr(ref i) => {
                    let ty = LLVMIntTypeInContext(ctxt.context, 32);
                    Ok(LLVMConstInt(ty, *i as u64, 0))
                },
                &Expr::AddExpr(ref e1, ref e2) => {
                    let ev1 = try!(e1.codegen(ctxt));
                    let ev2 = try!(e2.codegen(ctxt));
                    Ok(LLVMBuildFAdd(ctxt.builder, ev1, ev2, "add_tmp".as_ptr() as *const i8))
                },
                &Expr::CallExpr(ref fn_name, ref optional_args) => {
                    match std_functions_call_factory(&*fn_name, optional_args, ctxt) {
                        Some(call) => Ok(call),
                        _ => {Err("Should handle non std functions".to_string())}
                    }
                },
                &Expr::LetExpr(ref decls, ref expr) => {
                    debug_assert!(!decls.is_empty(), "Declarations in a let block can't be empty");
                    debug_assert!(expr.is_some(), "Expr in a let block can't be empty");
                    for decl in &*decls {
                        match decl {
                            &Decl::FunDec(ref name, ref params, ref ty, ref body) => {
                                let ty = get_llvm_type_for_ttype(ty, ctxt);
                                let proto = LLVMFunctionType(ty, ptr::null_mut(), 0, 0);
                                let function = LLVMAddFunction(ctxt.module,
                                                               c_str_ptr!(&(*name.clone())),
                                                               proto);
                                let bb = LLVMAppendBasicBlockInContext(ctxt.context,
                                                                       function,
                                                                       c_str_ptr!("entry"));
                                LLVMPositionBuilderAtEnd(ctxt.builder, bb);
                                //trans_expr(body, &mut ctxt);
                                body.codegen(ctxt);
                            },
                            _ => panic!("More decl types should be covered")
                        }

                    }
                    //trans_expr(&*expr.unwrap(), &mut ctxt);
                    let e = &expr.as_ref().unwrap();
                    let v = try!(e.codegen(ctxt));
                    Ok(v)
                }
                _ => Err("error".to_string())
            }
        }
    }
}

pub fn translate(expr : &Expr) -> Option<Context>{
    let mut ctxt = Context::new("main_mod");
    unsafe{
        //build outer embedding main() fn
        let ty = LLVMIntTypeInContext(ctxt.context, 32);
        let proto = LLVMFunctionType(ty, ptr::null_mut(), 0, 0);
        let function = LLVMAddFunction(ctxt.module,
                                       c_str_ptr!("main"),
                                       proto);
        let bb = LLVMAppendBasicBlockInContext(ctxt.context,
                                               function,
                                               c_str_ptr!("entry"));
        LLVMPositionBuilderAtEnd(ctxt.builder, bb);
        trans_expr(expr, &mut ctxt);
        LLVMBuildRet(ctxt.builder,
                     LLVMConstInt(LLVMIntTypeInContext(ctxt.context, 32), 0 as u64, 0));

        //add translated code as part of the block
    }
    Some(ctxt)
}

fn trans_expr(expr: &Expr, ctxt : &mut Context){
    let result = expr.codegen(ctxt); //ctxt is already &mut
    match result {
        Ok(_) => {},
        Err(msg) => panic!(msg)
    }
}

//#[test]
fn test_translate_std_print_call() {
    let ctxt = translate(&Expr::CallExpr("print".to_string(),
                                  Some(vec![(TType::TString,
                                             B(Expr::StringExpr("abhi".to_string())))])));
    assert_eq!(ctxt.is_some(), true);
    ctxt.unwrap().dump();
}

//#[test]
fn test_prsr_bcknd_intgrtion_prnt_call() {
    let mut p = Parser::new("print(\"Grrrr!\n\")".to_string());
    p.start_lexer();
    let tup = p.expr();
    let (ty, b_expr) = tup.unwrap();
    let ctxt = translate(&*b_expr);
    // let ctxt = translate();&Expr::CallExpr("print".to_string(),
    //                               Some(vec![(TType::TString,
    //                                          B(Expr::StringExpr("abhi".to_string())))])));
    assert_eq!(ctxt.is_some(), true);
    ctxt.unwrap().dump();
}

#[test]
fn test_prsr_bcknd_intgrtion_let_blk() {
    let mut p = Parser::new("let function foo() = print(\"Grrrr!\n\") in foo() end".to_string());
    p.start_lexer();
    let tup = p.expr();
    let (ty, b_expr) = tup.unwrap();
    let ctxt = translate(&*b_expr);
    assert_eq!(ctxt.is_some(), true);
    ctxt.unwrap().dump();
}
<commit_msg>Add call expr translation for user defined functions, fix resolution of correct basic block for placing instructions<commit_after>#![feature(rustc_private)]
extern crate llvm_sys as llvm;
extern crate libc;
use std::ptr;
use std::ffi;

use self::llvm::prelude::{LLVMContextRef, LLVMModuleRef, LLVMBuilderRef, LLVMValueRef, LLVMTypeRef};
use self::llvm::core::*;

use std::collections::{HashMap};
use std::mem;

use syntax::ast::{Block, Expr, Decl, TType, OptionalTypeExprTupleList};
use syntax::ptr::{B};
//FIXME this import is for integration testing purposes
use syntax::parse::*;//{Parser};
use syntax::parse::parser::{Parser};

macro_rules! c_str_ptr {
    ($s:expr) => {
        ffi::CString::new($s).unwrap().as_ptr()
    };
}

pub struct Context{
    context : LLVMContextRef,
    module : LLVMModuleRef,
    builder : LLVMBuilderRef,
    sym_tab : HashMap<String, LLVMValueRef>,
    bb_stack : Vec<*mut llvm::LLVMBasicBlock>
}

impl Context{
    fn new(module_name : &str) -> Self{
        unsafe{
            let llvm_context =  LLVMContextCreate();
            let llvm_module = LLVMModuleCreateWithNameInContext(c_str_ptr!(module_name),
                                                                llvm_context);
            let builder = LLVMCreateBuilderInContext(llvm_context);
            let sym_tab = HashMap::new();
            let bb_stack = Vec::new();

            Context {
                context : llvm_context,
                module : llvm_module,
                builder : builder,
                sym_tab : sym_tab,
                bb_stack : bb_stack
            }
        }
    }

    pub fn dump(&self){
        unsafe{
            LLVMDumpModule(self.module);
        }
    }
}

impl Drop for Context{
    fn drop(&mut self){
        unsafe{
            LLVMDisposeBuilder(self.builder);
            LLVMDisposeModule(self.module);
            LLVMContextDispose(self.context);
        }
    }
}

//TODO move these in a seperate file
type IRBuildingResult = Result<LLVMValueRef, String>;

trait IRBuilder{
    fn codegen(&self, ctxt : &mut Context) -> IRBuildingResult;
}

fn std_functions_call_factory(fn_name : &str,
                              args : &OptionalTypeExprTupleList,
                              ctxt : &mut Context) -> Option<LLVMValueRef>{
    unsafe{
        match fn_name {
            "print" =>{
                debug_assert!(args.is_some(), "No args passed to print()");
                let lst = args.as_ref().unwrap();
                debug_assert!(lst.len() == 1, "One arg should be passed to print()");
                debug_assert!(lst[0].0 == TType::TString);
                let str_arg = match &*lst[0].1 {
                    &Expr::StringExpr(ref value) => c_str_ptr!(&*(value.clone())),
                    _ => panic!("Expected a string expr")
                };

                let print_ty = LLVMIntTypeInContext(ctxt.context, 32);
                let mut pf_type_args_vec = Vec::new();
                pf_type_args_vec.push(LLVMPointerType(LLVMIntTypeInContext(ctxt.context, 8),
                                                      0));
                let proto = LLVMFunctionType(print_ty, pf_type_args_vec.as_mut_ptr(), 1, 1);
                let print_function = LLVMAddFunction(ctxt.module,
                                                     c_str_ptr!("printf"),
                                                     proto);
                let gstr = LLVMBuildGlobalStringPtr(ctxt.builder,
                                                    str_arg,
                                                    c_str_ptr!(".str"));
                let mut pf_args = Vec::new();
                pf_args.push(gstr);

                Some(LLVMBuildCall(ctxt.builder,
                                   print_function,
                                   pf_args.as_mut_ptr(),
                                   1,
                                   c_str_ptr!("call")))
            },
            _ => {None}
        }
    }
}

fn get_llvm_type_for_ttype(ty : &TType, ctxt : &mut Context) -> LLVMTypeRef{
    unsafe{
        match ty {
            &TType::TVoid => LLVMVoidTypeInContext(ctxt.context),
            &TType::TInt32 => LLVMIntTypeInContext(ctxt.context, 32),
            _ => panic!("Other TTypes not mapped yet to the corresponding LLVM types")
        }
    }
}

impl IRBuilder for Expr{
    fn codegen(&self, ctxt : &mut Context) -> IRBuildingResult{
        unsafe{
            match self{
                &Expr::NumExpr(ref i) => {
                    let ty = LLVMIntTypeInContext(ctxt.context, 32);
                    Ok(LLVMConstInt(ty, *i as u64, 0))
                },
                &Expr::AddExpr(ref e1, ref e2) => {
                    let ev1 = try!(e1.codegen(ctxt));
                    let ev2 = try!(e2.codegen(ctxt));
                    Ok(LLVMBuildFAdd(ctxt.builder, ev1, ev2, "add_tmp".as_ptr() as *const i8))
                },
                &Expr::CallExpr(ref fn_name, ref optional_args) => {
                    //FIXME instead of directly passing to the factory
                    //fn_name can be checked in a map that records names of std functions
                    match std_functions_call_factory(&*fn_name, optional_args, ctxt) {
                        Some(call) => Ok(call),
                        _ => {
                            //user defined function call
                            let mut pf_args = Vec::new();
                            //FIXME pass args if present in the call
                            if optional_args.is_some() {

                            }
                            //pf_args.push(gstr);
                            let _fn = ctxt.sym_tab[&*fn_name];
                            Ok(LLVMBuildCall(ctxt.builder,
                                            _fn,
                                            pf_args.as_mut_ptr(),
                                            0,
                                            c_str_ptr!("")))
                        }
                    }
                },
                &Expr::LetExpr(ref decls, ref expr) => {
                    debug_assert!(!decls.is_empty(), "Declarations in a let block can't be empty");
                    debug_assert!(expr.is_some(), "Expr in a let block can't be empty");
                    for decl in &*decls {
                        match decl {
                            &Decl::FunDec(ref name, ref params, ref ty, ref body) => {
                                let ty = get_llvm_type_for_ttype(ty, ctxt);
                                let proto = LLVMFunctionType(ty, ptr::null_mut(), 0, 0);
                                let function = LLVMAddFunction(ctxt.module,
                                                               c_str_ptr!(&(*name.clone())),
                                                               proto);
                                let bb = LLVMAppendBasicBlockInContext(ctxt.context,
                                                                       function,
                                                                       c_str_ptr!("entry"));
                                ctxt.sym_tab.insert(name.clone(), function);
                                LLVMPositionBuilderAtEnd(ctxt.builder, bb);
                                //trans_expr(body, &mut ctxt);
                                body.codegen(ctxt);
                                LLVMBuildRetVoid(ctxt.builder);
                            },
                            _ => panic!("More decl types should be covered")
                        }

                    }
                    //trans_expr(&*expr.unwrap(), &mut ctxt);
                    //FIXME should the previous bb be popped here?
                    let bb = ctxt.bb_stack.pop().unwrap();
                    LLVMPositionBuilderAtEnd(ctxt.builder, bb);
                    let e = &expr.as_ref().unwrap();
                    let v = try!(e.codegen(ctxt));
                    Ok(v)
                }
                _ => Err("error".to_string())
            }
        }
    }
}

pub fn translate(expr : &Expr) -> Option<Context>{
    let mut ctxt = Context::new("main_mod");
    unsafe{
        //build outer embedding main() fn
        let ty = LLVMIntTypeInContext(ctxt.context, 32);
        let proto = LLVMFunctionType(ty, ptr::null_mut(), 0, 0);
        let function = LLVMAddFunction(ctxt.module,
                                       c_str_ptr!("main"),
                                       proto);
        let bb = LLVMAppendBasicBlockInContext(ctxt.context,
                                               function,
                                               c_str_ptr!("entry"));
        LLVMPositionBuilderAtEnd(ctxt.builder, bb);
        ctxt.bb_stack.push(bb);
        trans_expr(expr, &mut ctxt);
        LLVMBuildRet(ctxt.builder,
                     LLVMConstInt(LLVMIntTypeInContext(ctxt.context, 32), 0 as u64, 0));

        //add translated code as part of the block
    }
    Some(ctxt)
}

fn trans_expr(expr: &Expr, ctxt : &mut Context){
    let result = expr.codegen(ctxt); //ctxt is already &mut
    match result {
        Ok(_) => {},
        Err(msg) => panic!(msg)
    }
}

//#[test]
fn test_translate_std_print_call() {
    let ctxt = translate(&Expr::CallExpr("print".to_string(),
                                  Some(vec![(TType::TString,
                                             B(Expr::StringExpr("abhi".to_string())))])));
    assert_eq!(ctxt.is_some(), true);
    ctxt.unwrap().dump();
}

//#[test]
fn test_prsr_bcknd_intgrtion_prnt_call() {
    let mut p = Parser::new("print(\"Grrrr!\n\")".to_string());
    p.start_lexer();
    let tup = p.expr();
    let (ty, b_expr) = tup.unwrap();
    let ctxt = translate(&*b_expr);
    // let ctxt = translate();&Expr::CallExpr("print".to_string(),
    //                               Some(vec![(TType::TString,
    //                                          B(Expr::StringExpr("abhi".to_string())))])));
    assert_eq!(ctxt.is_some(), true);
    ctxt.unwrap().dump();
}

#[test]
fn test_prsr_bcknd_intgrtion_let_blk() {
    let mut p = Parser::new("let function foo() = print(\"Grrrr!\n\") in foo() end".to_string());
    p.start_lexer();
    let tup = p.expr();
    let (ty, b_expr) = tup.unwrap();
    let ctxt = translate(&*b_expr);
    assert_eq!(ctxt.is_some(), true);
    ctxt.unwrap().dump();
}
<|endoftext|>


<commit_before>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use std::fs;
use std::io::{Read, Result as IoResult, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use encryption::{DataKeyManager, DecrypterReader, EncrypterWriter};
use engine_traits::{
    CacheStats, RaftEngine, RaftEngineReadOnly, RaftLogBatch as RaftLogBatchTrait, RaftLogGCTask,
    Result,
};
use file_system::{IOOp, IORateLimiter, IOType};
use kvproto::raft_serverpb::RaftLocalState;
use raft::eraftpb::Entry;
use raft_engine::env::{DefaultFileSystem, FileSystem, Handle, WriteExt};
use raft_engine::{
    Command, Engine as RawRaftEngine, Error as RaftEngineError, LogBatch, MessageExt,
};
use tikv_util::Either;

pub use raft_engine::{Config as RaftEngineConfig, RecoveryMode};

#[derive(Clone)]
pub struct MessageExtTyped;

impl MessageExt for MessageExtTyped {
    type Entry = Entry;

    fn index(e: &Entry) -> u64 {
        e.index
    }
}

struct ManagedReader {
    inner: Either<
        <DefaultFileSystem as FileSystem>::Reader,
        DecrypterReader<<DefaultFileSystem as FileSystem>::Reader>,
    >,
    rate_limiter: Option<Arc<IORateLimiter>>,
}

impl Seek for ManagedReader {
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        match self.inner.as_mut() {
            Either::Left(reader) => reader.seek(pos),
            Either::Right(reader) => reader.seek(pos),
        }
    }
}

impl Read for ManagedReader {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let mut size = buf.len();
        if let Some(ref mut limiter) = self.rate_limiter {
            size = limiter.request(IOType::ForegroundRead, IOOp::Read, size);
        }
        match self.inner.as_mut() {
            Either::Left(reader) => reader.read(&mut buf[..size]),
            Either::Right(reader) => reader.read(&mut buf[..size]),
        }
    }
}

struct ManagedWriter {
    inner: Either<
        <DefaultFileSystem as FileSystem>::Writer,
        EncrypterWriter<<DefaultFileSystem as FileSystem>::Writer>,
    >,
    rate_limiter: Option<Arc<IORateLimiter>>,
}

impl Seek for ManagedWriter {
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        match self.inner.as_mut() {
            Either::Left(writer) => writer.seek(pos),
            Either::Right(writer) => writer.seek(pos),
        }
    }
}

impl Write for ManagedWriter {
    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        let mut size = buf.len();
        if let Some(ref mut limiter) = self.rate_limiter {
            size = limiter.request(IOType::ForegroundWrite, IOOp::Write, size);
        }
        match self.inner.as_mut() {
            Either::Left(writer) => writer.write(&buf[..size]),
            Either::Right(writer) => writer.write(&buf[..size]),
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        Ok(())
    }
}

impl WriteExt for ManagedWriter {
    fn truncate(&self, offset: usize) -> IoResult<()> {
        match self.inner.as_ref() {
            Either::Left(writer) => writer.truncate(offset),
            Either::Right(writer) => writer.inner().truncate(offset),
        }
    }

    fn sync(&self) -> IoResult<()> {
        match self.inner.as_ref() {
            Either::Left(writer) => writer.sync(),
            Either::Right(writer) => writer.inner().sync(),
        }
    }

    fn allocate(&self, offset: usize, size: usize) -> IoResult<()> {
        match self.inner.as_ref() {
            Either::Left(writer) => writer.allocate(offset, size),
            Either::Right(writer) => writer.inner().allocate(offset, size),
        }
    }
}

struct ManagedFileSystem {
    base_level_file_system: DefaultFileSystem,
    key_manager: Option<Arc<DataKeyManager>>,
    rate_limiter: Option<Arc<IORateLimiter>>,
}

impl ManagedFileSystem {
    fn new(
        key_manager: Option<Arc<DataKeyManager>>,
        rate_limiter: Option<Arc<IORateLimiter>>,
    ) -> Self {
        Self {
            base_level_file_system: DefaultFileSystem,
            key_manager,
            rate_limiter,
        }
    }
}

struct ManagedHandle {
    path: PathBuf,
    base: Arc<<DefaultFileSystem as FileSystem>::Handle>,
}

impl Handle for ManagedHandle {
    fn truncate(&self, offset: usize) -> IoResult<()> {
        self.base.truncate(offset)
    }

    fn file_size(&self) -> IoResult<usize> {
        self.base.file_size()
    }
}

impl FileSystem for ManagedFileSystem {
    type Handle = ManagedHandle;
    type Reader = ManagedReader;
    type Writer = ManagedWriter;

    fn create<P: AsRef<Path>>(&self, path: P) -> IoResult<Self::Handle> {
        Ok(ManagedHandle {
            path: path.as_ref().to_path_buf(),
            base: Arc::new(self.base_level_file_system.create(path.as_ref())?),
        })
    }

    fn open<P: AsRef<Path>>(&self, path: P) -> IoResult<Self::Handle> {
        Ok(ManagedHandle {
            path: path.as_ref().to_path_buf(),
            base: Arc::new(self.base_level_file_system.open(path.as_ref())?),
        })
    }

    fn new_reader(&self, handle: Arc<Self::Handle>) -> IoResult<Self::Reader> {
        let base_reader = self
            .base_level_file_system
            .new_reader(handle.base.clone())?;
        if let Some(ref key_manager) = self.key_manager {
            Ok(ManagedReader {
                inner: Either::Right(key_manager.open_file_with_reader(&handle.path, base_reader)?),
                rate_limiter: self.rate_limiter.clone(),
            })
        } else {
            Ok(ManagedReader {
                inner: Either::Left(base_reader),
                rate_limiter: self.rate_limiter.clone(),
            })
        }
    }

    fn new_writer(&self, handle: Arc<Self::Handle>) -> IoResult<Self::Writer> {
        let base_writer = self
            .base_level_file_system
            .new_writer(handle.base.clone())?;

        if let Some(ref key_manager) = self.key_manager {
            Ok(ManagedWriter {
                inner: Either::Right(key_manager.open_file_with_writer(
                    &handle.path,
                    base_writer,
                    false,
                )?),
                rate_limiter: self.rate_limiter.clone(),
            })
        } else {
            Ok(ManagedWriter {
                inner: Either::Left(base_writer),
                rate_limiter: self.rate_limiter.clone(),
            })
        }
    }
}

#[derive(Clone)]
pub struct RaftLogEngine(Arc<RawRaftEngine<ManagedFileSystem>>);

impl RaftLogEngine {
    pub fn new(
        config: RaftEngineConfig,
        key_manager: Option<Arc<DataKeyManager>>,
        rate_limiter: Option<Arc<IORateLimiter>>,
    ) -> Result<Self> {
        let file_system = Arc::new(ManagedFileSystem::new(key_manager, rate_limiter));
        Ok(RaftLogEngine(Arc::new(
            RawRaftEngine::open_with_file_system(config, file_system).map_err(transfer_error)?,
        )))
    }

    /// If path is not an empty directory, we say db exists.
    pub fn exists(path: &str) -> bool {
        let path = Path::new(path);
        if !path.exists() || !path.is_dir() {
            return false;
        }
        fs::read_dir(&path).unwrap().next().is_some()
    }

    pub fn raft_groups(&self) -> Vec<u64> {
        self.0.raft_groups()
    }

    pub fn first_index(&self, raft_id: u64) -> Option<u64> {
        self.0.first_index(raft_id)
    }

    pub fn last_index(&self, raft_id: u64) -> Option<u64> {
        self.0.last_index(raft_id)
    }
}

#[derive(Default)]
pub struct RaftLogBatch(LogBatch);

const RAFT_LOG_STATE_KEY: &[u8] = b"R";

impl RaftLogBatchTrait for RaftLogBatch {
    fn append(&mut self, raft_group_id: u64, entries: Vec<Entry>) -> Result<()> {
        self.0
            .add_entries::<MessageExtTyped>(raft_group_id, &entries)
            .map_err(transfer_error)
    }

    fn cut_logs(&mut self, _: u64, _: u64, _: u64) {
        // It's unnecessary because overlapped entries can be handled in `append`.
    }

    fn put_raft_state(&mut self, raft_group_id: u64, state: &RaftLocalState) -> Result<()> {
        self.0
            .put_message(raft_group_id, RAFT_LOG_STATE_KEY.to_vec(), state)
            .map_err(transfer_error)
    }

    fn persist_size(&self) -> usize {
        self.0.approximate_size()
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    fn merge(&mut self, mut src: Self) {
        self.0.merge(&mut src.0);
    }
}

impl RaftEngineReadOnly for RaftLogEngine {
    fn get_raft_state(&self, raft_group_id: u64) -> Result<Option<RaftLocalState>> {
        self.0
            .get_message(raft_group_id, RAFT_LOG_STATE_KEY)
            .map_err(transfer_error)
    }

    fn get_entry(&self, raft_group_id: u64, index: u64) -> Result<Option<Entry>> {
        self.0
            .get_entry::<MessageExtTyped>(raft_group_id, index)
            .map_err(transfer_error)
    }

    fn fetch_entries_to(
        &self,
        raft_group_id: u64,
        begin: u64,
        end: u64,
        max_size: Option<usize>,
        to: &mut Vec<Entry>,
    ) -> Result<usize> {
        self.0
            .fetch_entries_to::<MessageExtTyped>(raft_group_id, begin, end, max_size, to)
            .map_err(transfer_error)
    }

    fn get_all_entries_to(&self, raft_group_id: u64, buf: &mut Vec<Entry>) -> Result<()> {
        if let Some(first) = self.0.first_index(raft_group_id) {
            let last = self.0.last_index(raft_group_id).unwrap();
            buf.reserve((last - first + 1) as usize);
            self.fetch_entries_to(raft_group_id, first, last + 1, None, buf)?;
        }
        Ok(())
    }
}

impl RaftEngine for RaftLogEngine {
    type LogBatch = RaftLogBatch;

    fn log_batch(&self, capacity: usize) -> Self::LogBatch {
        RaftLogBatch(LogBatch::with_capacity(capacity))
    }

    fn sync(&self) -> Result<()> {
        self.0.sync().map_err(transfer_error)
    }

    fn consume(&self, batch: &mut Self::LogBatch, sync: bool) -> Result<usize> {
        self.0.write(&mut batch.0, sync).map_err(transfer_error)
    }

    fn consume_and_shrink(
        &self,
        batch: &mut Self::LogBatch,
        sync: bool,
        _: usize,
        _: usize,
    ) -> Result<usize> {
        self.0.write(&mut batch.0, sync).map_err(transfer_error)
    }

    fn clean(
        &self,
        raft_group_id: u64,
        _: u64,
        _: &RaftLocalState,
        batch: &mut RaftLogBatch,
    ) -> Result<()> {
        batch.0.add_command(raft_group_id, Command::Clean);
        Ok(())
    }

    fn append(&self, raft_group_id: u64, entries: Vec<Entry>) -> Result<usize> {
        let mut batch = Self::LogBatch::default();
        batch
            .0
            .add_entries::<MessageExtTyped>(raft_group_id, &entries)
            .map_err(transfer_error)?;
        self.0.write(&mut batch.0, false).map_err(transfer_error)
    }

    fn put_raft_state(&self, raft_group_id: u64, state: &RaftLocalState) -> Result<()> {
        let mut batch = Self::LogBatch::default();
        batch
            .0
            .put_message(raft_group_id, RAFT_LOG_STATE_KEY.to_vec(), state)
            .map_err(transfer_error)?;
        self.0.write(&mut batch.0, false).map_err(transfer_error)?;
        Ok(())
    }

    fn gc(&self, raft_group_id: u64, from: u64, to: u64) -> Result<usize> {
        self.batch_gc(vec![RaftLogGCTask {
            raft_group_id,
            from,
            to,
        }])
    }

    fn batch_gc(&self, tasks: Vec<RaftLogGCTask>) -> Result<usize> {
        let mut batch = self.log_batch(tasks.len());
        let mut old_first_index = Vec::with_capacity(tasks.len());
        for task in &tasks {
            batch
                .0
                .add_command(task.raft_group_id, Command::Compact { index: task.to });
            old_first_index.push(self.0.first_index(task.raft_group_id));
        }

        self.0.write(&mut batch.0, false).map_err(transfer_error)?;

        let mut total = 0;
        for (old_first_index, task) in old_first_index.iter().zip(tasks) {
            let new_first_index = self.0.first_index(task.raft_group_id);
            if let (Some(old), Some(new)) = (old_first_index, new_first_index) {
                total += new.saturating_sub(*old);
            }
        }
        Ok(total as usize)
    }

    fn purge_expired_files(&self) -> Result<Vec<u64>> {
        self.0.purge_expired_files().map_err(transfer_error)
    }

    fn has_builtin_entry_cache(&self) -> bool {
        false
    }

    fn gc_entry_cache(&self, _raft_group_id: u64, _to: u64) {}

    /// Flush current cache stats.
    fn flush_stats(&self) -> Option<CacheStats> {
        None
    }

    fn stop(&self) {}

    fn dump_stats(&self) -> Result<String> {
        // Raft engine won't dump anything.
        Ok("".to_owned())
    }

    fn get_engine_size(&self) -> Result<u64> {
        Ok(self.0.get_used_size() as u64)
    }
}

fn transfer_error(e: RaftEngineError) -> engine_traits::Error {
    match e {
        RaftEngineError::EntryCompacted => engine_traits::Error::EntriesCompacted,
        RaftEngineError::EntryNotFound => engine_traits::Error::EntriesUnavailable,
        e => {
            let e = box_err!(e);
            engine_traits::Error::Other(e)
        }
    }
}
<commit_msg>raft_log_engine: fix managedfs create encrypted file (#12163)<commit_after>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use std::fs;
use std::io::{Read, Result as IoResult, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use encryption::{DataKeyManager, DecrypterReader, EncrypterWriter};
use engine_traits::{
    CacheStats, EncryptionKeyManager, RaftEngine, RaftEngineReadOnly,
    RaftLogBatch as RaftLogBatchTrait, RaftLogGCTask, Result,
};
use file_system::{IOOp, IORateLimiter, IOType};
use kvproto::raft_serverpb::RaftLocalState;
use raft::eraftpb::Entry;
use raft_engine::env::{DefaultFileSystem, FileSystem, Handle, WriteExt};
use raft_engine::{
    Command, Engine as RawRaftEngine, Error as RaftEngineError, LogBatch, MessageExt,
};
use tikv_util::Either;

pub use raft_engine::{Config as RaftEngineConfig, RecoveryMode};

#[derive(Clone)]
pub struct MessageExtTyped;

impl MessageExt for MessageExtTyped {
    type Entry = Entry;

    fn index(e: &Entry) -> u64 {
        e.index
    }
}

struct ManagedReader {
    inner: Either<
        <DefaultFileSystem as FileSystem>::Reader,
        DecrypterReader<<DefaultFileSystem as FileSystem>::Reader>,
    >,
    rate_limiter: Option<Arc<IORateLimiter>>,
}

impl Seek for ManagedReader {
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        match self.inner.as_mut() {
            Either::Left(reader) => reader.seek(pos),
            Either::Right(reader) => reader.seek(pos),
        }
    }
}

impl Read for ManagedReader {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let mut size = buf.len();
        if let Some(ref mut limiter) = self.rate_limiter {
            size = limiter.request(IOType::ForegroundRead, IOOp::Read, size);
        }
        match self.inner.as_mut() {
            Either::Left(reader) => reader.read(&mut buf[..size]),
            Either::Right(reader) => reader.read(&mut buf[..size]),
        }
    }
}

struct ManagedWriter {
    inner: Either<
        <DefaultFileSystem as FileSystem>::Writer,
        EncrypterWriter<<DefaultFileSystem as FileSystem>::Writer>,
    >,
    rate_limiter: Option<Arc<IORateLimiter>>,
}

impl Seek for ManagedWriter {
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        match self.inner.as_mut() {
            Either::Left(writer) => writer.seek(pos),
            Either::Right(writer) => writer.seek(pos),
        }
    }
}

impl Write for ManagedWriter {
    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        let mut size = buf.len();
        if let Some(ref mut limiter) = self.rate_limiter {
            size = limiter.request(IOType::ForegroundWrite, IOOp::Write, size);
        }
        match self.inner.as_mut() {
            Either::Left(writer) => writer.write(&buf[..size]),
            Either::Right(writer) => writer.write(&buf[..size]),
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        Ok(())
    }
}

impl WriteExt for ManagedWriter {
    fn truncate(&self, offset: usize) -> IoResult<()> {
        match self.inner.as_ref() {
            Either::Left(writer) => writer.truncate(offset),
            Either::Right(writer) => writer.inner().truncate(offset),
        }
    }

    fn sync(&self) -> IoResult<()> {
        match self.inner.as_ref() {
            Either::Left(writer) => writer.sync(),
            Either::Right(writer) => writer.inner().sync(),
        }
    }

    fn allocate(&self, offset: usize, size: usize) -> IoResult<()> {
        match self.inner.as_ref() {
            Either::Left(writer) => writer.allocate(offset, size),
            Either::Right(writer) => writer.inner().allocate(offset, size),
        }
    }
}

struct ManagedFileSystem {
    base_level_file_system: DefaultFileSystem,
    key_manager: Option<Arc<DataKeyManager>>,
    rate_limiter: Option<Arc<IORateLimiter>>,
}

impl ManagedFileSystem {
    fn new(
        key_manager: Option<Arc<DataKeyManager>>,
        rate_limiter: Option<Arc<IORateLimiter>>,
    ) -> Self {
        Self {
            base_level_file_system: DefaultFileSystem,
            key_manager,
            rate_limiter,
        }
    }
}

struct ManagedHandle {
    path: PathBuf,
    base: Arc<<DefaultFileSystem as FileSystem>::Handle>,
}

impl Handle for ManagedHandle {
    fn truncate(&self, offset: usize) -> IoResult<()> {
        self.base.truncate(offset)
    }

    fn file_size(&self) -> IoResult<usize> {
        self.base.file_size()
    }
}

impl FileSystem for ManagedFileSystem {
    type Handle = ManagedHandle;
    type Reader = ManagedReader;
    type Writer = ManagedWriter;

    fn create<P: AsRef<Path>>(&self, path: P) -> IoResult<Self::Handle> {
        let base = Arc::new(self.base_level_file_system.create(path.as_ref())?);
        if let Some(ref manager) = self.key_manager {
            manager.new_file(path.as_ref().to_str().unwrap())?;
        }
        Ok(ManagedHandle {
            path: path.as_ref().to_path_buf(),
            base,
        })
    }

    fn open<P: AsRef<Path>>(&self, path: P) -> IoResult<Self::Handle> {
        Ok(ManagedHandle {
            path: path.as_ref().to_path_buf(),
            base: Arc::new(self.base_level_file_system.open(path.as_ref())?),
        })
    }

    fn new_reader(&self, handle: Arc<Self::Handle>) -> IoResult<Self::Reader> {
        let base_reader = self
            .base_level_file_system
            .new_reader(handle.base.clone())?;
        if let Some(ref key_manager) = self.key_manager {
            Ok(ManagedReader {
                inner: Either::Right(key_manager.open_file_with_reader(&handle.path, base_reader)?),
                rate_limiter: self.rate_limiter.clone(),
            })
        } else {
            Ok(ManagedReader {
                inner: Either::Left(base_reader),
                rate_limiter: self.rate_limiter.clone(),
            })
        }
    }

    fn new_writer(&self, handle: Arc<Self::Handle>) -> IoResult<Self::Writer> {
        let base_writer = self
            .base_level_file_system
            .new_writer(handle.base.clone())?;

        if let Some(ref key_manager) = self.key_manager {
            Ok(ManagedWriter {
                inner: Either::Right(key_manager.open_file_with_writer(
                    &handle.path,
                    base_writer,
                    false,
                )?),
                rate_limiter: self.rate_limiter.clone(),
            })
        } else {
            Ok(ManagedWriter {
                inner: Either::Left(base_writer),
                rate_limiter: self.rate_limiter.clone(),
            })
        }
    }
}

#[derive(Clone)]
pub struct RaftLogEngine(Arc<RawRaftEngine<ManagedFileSystem>>);

impl RaftLogEngine {
    pub fn new(
        config: RaftEngineConfig,
        key_manager: Option<Arc<DataKeyManager>>,
        rate_limiter: Option<Arc<IORateLimiter>>,
    ) -> Result<Self> {
        let file_system = Arc::new(ManagedFileSystem::new(key_manager, rate_limiter));
        Ok(RaftLogEngine(Arc::new(
            RawRaftEngine::open_with_file_system(config, file_system).map_err(transfer_error)?,
        )))
    }

    /// If path is not an empty directory, we say db exists.
    pub fn exists(path: &str) -> bool {
        let path = Path::new(path);
        if !path.exists() || !path.is_dir() {
            return false;
        }
        fs::read_dir(&path).unwrap().next().is_some()
    }

    pub fn raft_groups(&self) -> Vec<u64> {
        self.0.raft_groups()
    }

    pub fn first_index(&self, raft_id: u64) -> Option<u64> {
        self.0.first_index(raft_id)
    }

    pub fn last_index(&self, raft_id: u64) -> Option<u64> {
        self.0.last_index(raft_id)
    }
}

#[derive(Default)]
pub struct RaftLogBatch(LogBatch);

const RAFT_LOG_STATE_KEY: &[u8] = b"R";

impl RaftLogBatchTrait for RaftLogBatch {
    fn append(&mut self, raft_group_id: u64, entries: Vec<Entry>) -> Result<()> {
        self.0
            .add_entries::<MessageExtTyped>(raft_group_id, &entries)
            .map_err(transfer_error)
    }

    fn cut_logs(&mut self, _: u64, _: u64, _: u64) {
        // It's unnecessary because overlapped entries can be handled in `append`.
    }

    fn put_raft_state(&mut self, raft_group_id: u64, state: &RaftLocalState) -> Result<()> {
        self.0
            .put_message(raft_group_id, RAFT_LOG_STATE_KEY.to_vec(), state)
            .map_err(transfer_error)
    }

    fn persist_size(&self) -> usize {
        self.0.approximate_size()
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    fn merge(&mut self, mut src: Self) {
        self.0.merge(&mut src.0);
    }
}

impl RaftEngineReadOnly for RaftLogEngine {
    fn get_raft_state(&self, raft_group_id: u64) -> Result<Option<RaftLocalState>> {
        self.0
            .get_message(raft_group_id, RAFT_LOG_STATE_KEY)
            .map_err(transfer_error)
    }

    fn get_entry(&self, raft_group_id: u64, index: u64) -> Result<Option<Entry>> {
        self.0
            .get_entry::<MessageExtTyped>(raft_group_id, index)
            .map_err(transfer_error)
    }

    fn fetch_entries_to(
        &self,
        raft_group_id: u64,
        begin: u64,
        end: u64,
        max_size: Option<usize>,
        to: &mut Vec<Entry>,
    ) -> Result<usize> {
        self.0
            .fetch_entries_to::<MessageExtTyped>(raft_group_id, begin, end, max_size, to)
            .map_err(transfer_error)
    }

    fn get_all_entries_to(&self, raft_group_id: u64, buf: &mut Vec<Entry>) -> Result<()> {
        if let Some(first) = self.0.first_index(raft_group_id) {
            let last = self.0.last_index(raft_group_id).unwrap();
            buf.reserve((last - first + 1) as usize);
            self.fetch_entries_to(raft_group_id, first, last + 1, None, buf)?;
        }
        Ok(())
    }
}

impl RaftEngine for RaftLogEngine {
    type LogBatch = RaftLogBatch;

    fn log_batch(&self, capacity: usize) -> Self::LogBatch {
        RaftLogBatch(LogBatch::with_capacity(capacity))
    }

    fn sync(&self) -> Result<()> {
        self.0.sync().map_err(transfer_error)
    }

    fn consume(&self, batch: &mut Self::LogBatch, sync: bool) -> Result<usize> {
        self.0.write(&mut batch.0, sync).map_err(transfer_error)
    }

    fn consume_and_shrink(
        &self,
        batch: &mut Self::LogBatch,
        sync: bool,
        _: usize,
        _: usize,
    ) -> Result<usize> {
        self.0.write(&mut batch.0, sync).map_err(transfer_error)
    }

    fn clean(
        &self,
        raft_group_id: u64,
        _: u64,
        _: &RaftLocalState,
        batch: &mut RaftLogBatch,
    ) -> Result<()> {
        batch.0.add_command(raft_group_id, Command::Clean);
        Ok(())
    }

    fn append(&self, raft_group_id: u64, entries: Vec<Entry>) -> Result<usize> {
        let mut batch = Self::LogBatch::default();
        batch
            .0
            .add_entries::<MessageExtTyped>(raft_group_id, &entries)
            .map_err(transfer_error)?;
        self.0.write(&mut batch.0, false).map_err(transfer_error)
    }

    fn put_raft_state(&self, raft_group_id: u64, state: &RaftLocalState) -> Result<()> {
        let mut batch = Self::LogBatch::default();
        batch
            .0
            .put_message(raft_group_id, RAFT_LOG_STATE_KEY.to_vec(), state)
            .map_err(transfer_error)?;
        self.0.write(&mut batch.0, false).map_err(transfer_error)?;
        Ok(())
    }

    fn gc(&self, raft_group_id: u64, from: u64, to: u64) -> Result<usize> {
        self.batch_gc(vec![RaftLogGCTask {
            raft_group_id,
            from,
            to,
        }])
    }

    fn batch_gc(&self, tasks: Vec<RaftLogGCTask>) -> Result<usize> {
        let mut batch = self.log_batch(tasks.len());
        let mut old_first_index = Vec::with_capacity(tasks.len());
        for task in &tasks {
            batch
                .0
                .add_command(task.raft_group_id, Command::Compact { index: task.to });
            old_first_index.push(self.0.first_index(task.raft_group_id));
        }

        self.0.write(&mut batch.0, false).map_err(transfer_error)?;

        let mut total = 0;
        for (old_first_index, task) in old_first_index.iter().zip(tasks) {
            let new_first_index = self.0.first_index(task.raft_group_id);
            if let (Some(old), Some(new)) = (old_first_index, new_first_index) {
                total += new.saturating_sub(*old);
            }
        }
        Ok(total as usize)
    }

    fn purge_expired_files(&self) -> Result<Vec<u64>> {
        self.0.purge_expired_files().map_err(transfer_error)
    }

    fn has_builtin_entry_cache(&self) -> bool {
        false
    }

    fn gc_entry_cache(&self, _raft_group_id: u64, _to: u64) {}

    /// Flush current cache stats.
    fn flush_stats(&self) -> Option<CacheStats> {
        None
    }

    fn stop(&self) {}

    fn dump_stats(&self) -> Result<String> {
        // Raft engine won't dump anything.
        Ok("".to_owned())
    }

    fn get_engine_size(&self) -> Result<u64> {
        Ok(self.0.get_used_size() as u64)
    }
}

fn transfer_error(e: RaftEngineError) -> engine_traits::Error {
    match e {
        RaftEngineError::EntryCompacted => engine_traits::Error::EntriesCompacted,
        RaftEngineError::EntryNotFound => engine_traits::Error::EntriesUnavailable,
        e => {
            let e = box_err!(e);
            engine_traits::Error::Other(e)
        }
    }
}
<|endoftext|>


<commit_before>use super::ValidatorRegistration;

#[derive(Debug, Clone, PartialEq)]
pub struct ChainConfig {
    pub cycle_length: u8,
    pub shard_count: u16,
    pub min_committee_size: u64,
    pub genesis_time: u64,
    pub slot_duration_millis: u64,
    pub initial_validators: Vec<ValidatorRegistration>,
}

/*
 * Presently this is just some arbitrary time in Sept 2018.
 */
const GENESIS_TIME: u64 = 1_537_488_655;

impl ChainConfig {
    pub fn standard() -> Self {
        Self {
            cycle_length: 64,
            shard_count: 1024,
            min_committee_size: 128,
            genesis_time: GENESIS_TIME,
            slot_duration_millis: 16 * 1000,
            initial_validators: vec![],
        }
    }

    pub fn validate(&self) -> bool {
	    // criteria that ensure the config is valid

	    // shard_count / cycle_length > 0 otherwise validator delegation
	    // will fail.
	    if self.shard_count / self.cycle_length as u16 == 0  {
		    return false;
	    }

	    true
    }

    #[cfg(test)]
    pub fn super_fast_tests() -> Self {
        Self {
            cycle_length: 2,
            shard_count: 2,
            min_committee_size: 2,
            genesis_time: GENESIS_TIME,   // arbitrary
            slot_duration_millis: 16 * 1000,
            initial_validators: vec![],
        }
    }
}
<commit_msg>Extend chain config<commit_after>use super::ValidatorRegistration;

#[derive(Debug, Clone, PartialEq)]
pub struct ChainConfig {
    pub cycle_length: u8,
    pub deposit_size_gwei: u64,
    pub shard_count: u16,
    pub min_committee_size: u64,
    pub max_validator_churn_quotient: u64,
    pub genesis_time: u64,
    pub slot_duration_millis: u64,
    pub initial_validators: Vec<ValidatorRegistration>,
}

/*
 * Presently this is just some arbitrary time in Sept 2018.
 */
const TEST_GENESIS_TIME: u64 = 1_537_488_655;

impl ChainConfig {
    pub fn standard() -> Self {
        Self {
            cycle_length: 64,
            deposit_size_gwei: 32 * 10^9,
            shard_count: 1024,
            min_committee_size: 128,
            max_validator_churn_quotient: 32,
            genesis_time: TEST_GENESIS_TIME,
            slot_duration_millis: 16 * 1000,
            initial_validators: vec![],
        }
    }

    pub fn validate(&self) -> bool {
	    // criteria that ensure the config is valid

	    // shard_count / cycle_length > 0 otherwise validator delegation
	    // will fail.
	    if self.shard_count / self.cycle_length as u16 == 0  {
		    return false;
	    }

	    true
    }

    #[cfg(test)]
    pub fn super_fast_tests() -> Self {
        Self {
            cycle_length: 2,
            deposit_size_gwei: 32 * 10^9,
            shard_count: 2,
            min_committee_size: 2,
            max_validator_churn_quotient: 32,
            genesis_time: TEST_GENESIS_TIME,   // arbitrary
            slot_duration_millis: 16 * 1000,
            initial_validators: vec![],
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

mod from_str;
// mod relabel;

pub mod cli;

/// Variable or terminal symbol in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub enum VarT<T> {
    /// `Var(i, j)` represents the `j`th component of the `i`th successor.
    /// Indexing starts from `0`.
    Var(usize, usize),
    T(T),
}

/// Composition function in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub struct Composition<T> {
    pub composition: Vec<Vec<VarT<T>>>,
}

impl<T> From<Vec<Vec<VarT<T>>>> for Composition<T> {
    fn from(encapsulated_value: Vec<Vec<VarT<T>>>) -> Self {
        Composition { composition: encapsulated_value }
    }
}

/// Rule of a weighted MCFG.
#[derive(Debug, PartialOrd, Ord, Clone)]
pub struct PMCFGRule<N, T, W> {
    pub head: N,
    pub tail: Vec<N>,
    pub composition: Composition<T>,
    pub weight: W,
}

impl<N, T, W> PMCFGRule<N, T, W>
    where T: Clone,
          W: Clone,
{
    pub fn map_nonterminals<F, M>(&self, f: F) -> PMCFGRule<M, T, W>
        where F: Fn(&N) -> M,
    {
        PMCFGRule {
            head: f(&self.head),
            tail: self.tail.iter().map(f).collect(),
            composition: self.composition.clone(),
            weight: self.weight.clone()
        }
    }
}

/// A weighted MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct PMCFG<N, T, W> {
    pub _dummy: PhantomData<T>,
    pub initial: Vec<N>,
    pub rules: Vec<PMCFGRule<N, T, W>>,
}

impl<N: Hash, T: Hash, W> Hash for PMCFGRule<N, T, W> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.tail.hash(state);
        self.composition.hash(state);
    }
}

impl<N: PartialEq, T: PartialEq, W> PartialEq for PMCFGRule<N, T, W> {
    fn eq(&self, other: &Self) -> bool {
        self.head == other.head && self.tail == other.tail && self.composition == other.composition
    }
}

impl<N: Eq, T: Eq, W> Eq for PMCFGRule<N, T, W> {}

impl<T: fmt::Display> fmt::Display for Composition<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter0 = self.composition.iter().peekable();
        let mut iter1;

        buffer.push_str("[");
        while let Some(proj) = iter0.next() {
            iter1 = proj.into_iter().peekable();
            buffer.push_str("[");
            while let Some(vart) = iter1.next() {
                buffer.push_str(format!("{}", vart).as_str());
                if iter1.peek().is_some() {
                    buffer.push_str(", ");
                }
            }
            buffer.push_str("]");
            if iter0.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]");

        write!(f, "{}", buffer)
    }
}

impl<T: fmt::Display> fmt::Display for VarT<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VarT::Var(i, j) => {
                write!(f, "Var {} {}", i, j)
            },
            VarT::T(ref x) => {
                write!(f, "T \"{}\"", x)
            },
        }
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFGRule<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
           let mut buffer = "".to_string();

           let mut iter = self.tail.iter().peekable();

           buffer.push_str("(");
           while let Some(nt) = iter.next() {
               buffer.push_str(format!("\"{}\"", nt).as_str());
               if iter.peek().is_some() {
                   buffer.push_str(", ");
               }
           }
           buffer.push_str(")");

        write!(f, "\"{}\" → {} {}  # {}", self.head, self.composition, buffer, self.weight)
        // write!(f, "\"{}\" → {}", self.head, self.composition)
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFG<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter = self.initial.iter().peekable();

        buffer.push_str("initial: [");
        while let Some(nt) = iter.next() {
            buffer.push_str(format!("\"{}\"", nt).as_str());
            if iter.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]\n\n");

        for r in &self.rules {
            buffer.push_str(format!("{}\n", r).as_str());
        }

        write!(f, "{}", buffer)
    }
}

pub fn evaluate<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>) -> Composition<T> {
    evaluate_pos(term_map, vec![])
}

pub fn evaluate_pos<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>, address: Vec<usize>) -> Composition<T> {
    let unexpanded_composition = &term_map.get(&address).unwrap().composition;
    let mut expanded_nonterminals: BTreeMap<_, Vec<Vec<VarT<T>>>> = BTreeMap::new();
    let mut expanded_composition = Vec::new();

    for component in unexpanded_composition {
        let mut expanded_component = Vec::new();

        for variable in component {
            match variable {
                &VarT::Var(num_nonter, num_compon) => {
                    if !expanded_nonterminals.contains_key(&num_nonter) {
                        let mut child_address = address.clone();
                        child_address.push(num_nonter);
                        let nonter_compos = evaluate_pos(term_map, child_address).composition;
                        expanded_nonterminals.insert(num_nonter, nonter_compos);
                    }

                    let nonter_compos = expanded_nonterminals.get(&num_nonter).unwrap();

                    if let Some(compon) = nonter_compos.get(num_compon) {
                        for terminal in compon {
                            expanded_component.push(terminal.clone());
                        }
                    } else {
                        panic!("{:?}: use of {}-th component of nonterminal {} that has only {} components!",
                               unexpanded_composition, num_compon, num_nonter, nonter_compos.len());
                    }

                },
                &VarT::T(ref terminal) => {
                    expanded_component.push(VarT::T(terminal.clone()));
                },
            }
        }

        expanded_composition.push(expanded_component);
    }

    Composition::from(expanded_composition)
}

pub fn identify_terminals<A: Clone>(tree_map: &BTreeMap<Vec<usize>, Composition<A>>) -> (BTreeMap<Vec<usize>, Composition<(Vec<usize>, usize)>>, BTreeMap<(Vec<usize>, usize), A>) {
    let mut identified_tree_map = BTreeMap::new();
    let mut terminal_map = BTreeMap::new();

    for (address, composition) in tree_map {
        let vec_compos = &composition.composition;
        let mut identified_compos = Vec::new();
        let mut compos_var_pos = 0;

        for component in vec_compos {
            let mut identified_compon = Vec::new();

            for variable in component {
                match variable {
                    &VarT::Var(x, y) => {
                        identified_compon.push(VarT::Var(x, y));
                    },
                    &VarT::T(ref terminal) => {
                        let terminal_id = (address.clone(), compos_var_pos);
                        identified_compon.push(VarT::T(terminal_id.clone()));
                        terminal_map.insert(terminal_id, terminal.clone());
                    },
                };

                compos_var_pos += 1;
            }

            identified_compos.push(identified_compon);
        }

        identified_tree_map.insert(address.clone(), Composition::from(identified_compos));
    }

    (identified_tree_map, terminal_map)
}

#[cfg(test)]
mod tests {
    use super::*;
    use VarT::{Var, T};

    #[test]
    fn test_evaluate() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(1,0), Var(0,1), Var(1,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![0,0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![1], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![1,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![1,1], Composition::from(vec![
                vec![T('b')]
        ]));
        term_map.insert(vec![1,2], Composition::from(vec![
                vec![T('d')]
        ]));

        let expanded_compos = Composition::from(vec![
            vec![T('a'), T('a'), T('b'), T('c'), T('c'), T('d')]
        ]);

        assert_eq!(expanded_compos, evaluate(&term_map));
    }

    #[test]
    #[should_panic(expected =
        "[[Var(0, 0), Var(0, 1)]]: use of 1-th component of nonterminal 0 that has only 1 components!"
    )]
    fn test_evaluate_invalid_composition() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(0,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![T('a')]
        ]));

        evaluate(&term_map);
    }
}
<commit_msg>src/pmcfg/mod.rs: added test for identify_terminals()<commit_after>use std::collections::BTreeMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

mod from_str;
// mod relabel;

pub mod cli;

/// Variable or terminal symbol in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub enum VarT<T> {
    /// `Var(i, j)` represents the `j`th component of the `i`th successor.
    /// Indexing starts from `0`.
    Var(usize, usize),
    T(T),
}

/// Composition function in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub struct Composition<T> {
    pub composition: Vec<Vec<VarT<T>>>,
}

impl<T> From<Vec<Vec<VarT<T>>>> for Composition<T> {
    fn from(encapsulated_value: Vec<Vec<VarT<T>>>) -> Self {
        Composition { composition: encapsulated_value }
    }
}

/// Rule of a weighted MCFG.
#[derive(Debug, PartialOrd, Ord, Clone)]
pub struct PMCFGRule<N, T, W> {
    pub head: N,
    pub tail: Vec<N>,
    pub composition: Composition<T>,
    pub weight: W,
}

impl<N, T, W> PMCFGRule<N, T, W>
    where T: Clone,
          W: Clone,
{
    pub fn map_nonterminals<F, M>(&self, f: F) -> PMCFGRule<M, T, W>
        where F: Fn(&N) -> M,
    {
        PMCFGRule {
            head: f(&self.head),
            tail: self.tail.iter().map(f).collect(),
            composition: self.composition.clone(),
            weight: self.weight.clone()
        }
    }
}

/// A weighted MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct PMCFG<N, T, W> {
    pub _dummy: PhantomData<T>,
    pub initial: Vec<N>,
    pub rules: Vec<PMCFGRule<N, T, W>>,
}

impl<N: Hash, T: Hash, W> Hash for PMCFGRule<N, T, W> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.tail.hash(state);
        self.composition.hash(state);
    }
}

impl<N: PartialEq, T: PartialEq, W> PartialEq for PMCFGRule<N, T, W> {
    fn eq(&self, other: &Self) -> bool {
        self.head == other.head && self.tail == other.tail && self.composition == other.composition
    }
}

impl<N: Eq, T: Eq, W> Eq for PMCFGRule<N, T, W> {}

impl<T: fmt::Display> fmt::Display for Composition<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter0 = self.composition.iter().peekable();
        let mut iter1;

        buffer.push_str("[");
        while let Some(proj) = iter0.next() {
            iter1 = proj.into_iter().peekable();
            buffer.push_str("[");
            while let Some(vart) = iter1.next() {
                buffer.push_str(format!("{}", vart).as_str());
                if iter1.peek().is_some() {
                    buffer.push_str(", ");
                }
            }
            buffer.push_str("]");
            if iter0.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]");

        write!(f, "{}", buffer)
    }
}

impl<T: fmt::Display> fmt::Display for VarT<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VarT::Var(i, j) => {
                write!(f, "Var {} {}", i, j)
            },
            VarT::T(ref x) => {
                write!(f, "T \"{}\"", x)
            },
        }
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFGRule<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
           let mut buffer = "".to_string();

           let mut iter = self.tail.iter().peekable();

           buffer.push_str("(");
           while let Some(nt) = iter.next() {
               buffer.push_str(format!("\"{}\"", nt).as_str());
               if iter.peek().is_some() {
                   buffer.push_str(", ");
               }
           }
           buffer.push_str(")");

        write!(f, "\"{}\" → {} {}  # {}", self.head, self.composition, buffer, self.weight)
        // write!(f, "\"{}\" → {}", self.head, self.composition)
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFG<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter = self.initial.iter().peekable();

        buffer.push_str("initial: [");
        while let Some(nt) = iter.next() {
            buffer.push_str(format!("\"{}\"", nt).as_str());
            if iter.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]\n\n");

        for r in &self.rules {
            buffer.push_str(format!("{}\n", r).as_str());
        }

        write!(f, "{}", buffer)
    }
}

pub fn evaluate<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>) -> Composition<T> {
    evaluate_pos(term_map, vec![])
}

pub fn evaluate_pos<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>, address: Vec<usize>) -> Composition<T> {
    let unexpanded_composition = &term_map.get(&address).unwrap().composition;
    let mut expanded_nonterminals: BTreeMap<_, Vec<Vec<VarT<T>>>> = BTreeMap::new();
    let mut expanded_composition = Vec::new();

    for component in unexpanded_composition {
        let mut expanded_component = Vec::new();

        for variable in component {
            match variable {
                &VarT::Var(num_nonter, num_compon) => {
                    if !expanded_nonterminals.contains_key(&num_nonter) {
                        let mut child_address = address.clone();
                        child_address.push(num_nonter);
                        let nonter_compos = evaluate_pos(term_map, child_address).composition;
                        expanded_nonterminals.insert(num_nonter, nonter_compos);
                    }

                    let nonter_compos = expanded_nonterminals.get(&num_nonter).unwrap();

                    if let Some(compon) = nonter_compos.get(num_compon) {
                        for terminal in compon {
                            expanded_component.push(terminal.clone());
                        }
                    } else {
                        panic!("{:?}: use of {}-th component of nonterminal {} that has only {} components!",
                               unexpanded_composition, num_compon, num_nonter, nonter_compos.len());
                    }

                },
                &VarT::T(ref terminal) => {
                    expanded_component.push(VarT::T(terminal.clone()));
                },
            }
        }

        expanded_composition.push(expanded_component);
    }

    Composition::from(expanded_composition)
}

pub fn identify_terminals<A: Clone>(tree_map: &BTreeMap<Vec<usize>, Composition<A>>) -> (BTreeMap<Vec<usize>, Composition<(Vec<usize>, usize)>>, BTreeMap<(Vec<usize>, usize), A>) {
    let mut identified_tree_map = BTreeMap::new();
    let mut terminal_map = BTreeMap::new();

    for (address, composition) in tree_map {
        let vec_compos = &composition.composition;
        let mut identified_compos = Vec::new();
        let mut compos_var_pos = 0;

        for component in vec_compos {
            let mut identified_compon = Vec::new();

            for variable in component {
                match variable {
                    &VarT::Var(x, y) => {
                        identified_compon.push(VarT::Var(x, y));
                    },
                    &VarT::T(ref terminal) => {
                        let terminal_id = (address.clone(), compos_var_pos);
                        identified_compon.push(VarT::T(terminal_id.clone()));
                        terminal_map.insert(terminal_id, terminal.clone());
                    },
                };

                compos_var_pos += 1;
            }

            identified_compos.push(identified_compon);
        }

        identified_tree_map.insert(address.clone(), Composition::from(identified_compos));
    }

    (identified_tree_map, terminal_map)
}

#[cfg(test)]
mod tests {
    use super::*;
    use VarT::{Var, T};

    #[test]
    fn test_evaluate() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(1,0), Var(0,1), Var(1,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![0,0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![1], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![1,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![1,1], Composition::from(vec![
                vec![T('b')]
        ]));
        term_map.insert(vec![1,2], Composition::from(vec![
                vec![T('d')]
        ]));

        let expanded_compos = Composition::from(vec![
            vec![T('a'), T('a'), T('b'), T('c'), T('c'), T('d')]
        ]);

        assert_eq!(expanded_compos, evaluate(&term_map));
    }

    #[test]
    #[should_panic(expected =
        "[[Var(0, 0), Var(0, 1)]]: use of 1-th component of nonterminal 0 that has only 1 components!"
    )]
    fn test_evaluate_invalid_composition() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(0,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![T('a')]
        ]));

        evaluate(&term_map);
    }

    #[test]
    fn test_identify_terminals() {
        let mut tree_map = BTreeMap::new();
        tree_map.insert(vec![], Composition::from(vec![
            vec![Var(0,0), T("a"), Var(0,1), T("b")]
        ]));
        tree_map.insert(vec![0], Composition::from(vec![
            vec![Var(1,0)],
            vec![T("c")]
        ]));
        tree_map.insert(vec![1,0], Composition::from(vec![
            vec![T("d")]
        ]));

        let mut identified_tree_map = BTreeMap::new();
        identified_tree_map.insert(vec![], Composition::from(vec![
            vec![Var(0,0), T((vec![], 1)), Var(0,1), T((vec![], 3))]
        ]));
        identified_tree_map.insert(vec![0], Composition::from(vec![
            vec![Var(1,0)],
            vec![T((vec![0], 1))]
        ]));
        identified_tree_map.insert(vec![1,0], Composition::from(vec![
            vec![T((vec![1,0], 0))]
        ]));

        let mut terminal_map = BTreeMap::new();
        terminal_map.insert((vec![], 1), "a");
        terminal_map.insert((vec![], 3), "b");
        terminal_map.insert((vec![0], 1), "c");
        terminal_map.insert((vec![1,0], 0), "d");

        assert_eq!((identified_tree_map, terminal_map), identify_terminals(&tree_map));
    }
}
<|endoftext|>


<commit_before>use MESSAGES;

use channel::Archive;

use iron::prelude::*;
use iron::status;

use handlebars_iron::handlebars::*;

use url::Url;

use serde_json;

use std::path::PathBuf;
use std::fs::{self, File};

pub fn _refresh() {
  let archives = fs::read_dir("archives").unwrap();
  for server in archives {
    let server = server.unwrap().path();
    if server.is_dir() {
      let server_id: u64 = server.file_stem().unwrap().to_string_lossy().parse().unwrap();
      for channel in fs::read_dir(server).unwrap() {
        let channel = channel.unwrap().path();
        let channel_id: u64 = channel.file_stem().unwrap().to_string_lossy().parse().unwrap();
        add_messages(channel, server_id, channel_id);
      }
    }
  }
}

pub fn refresh(_: &mut Request) -> IronResult<Response> {
  _refresh();
  Ok(Response::with(("We gucci", status::Ok)))
}

fn add_messages(channel: PathBuf, server_id: u64, channel_id: u64) {
  let f = File::open(channel).unwrap();
  let mut archive: Archive = serde_json::from_reader(f).unwrap();
  for message in &mut archive.messages {
    if let Some(member) = archive.server.members.iter().find(|mem| mem.user.id == message.author.id) {
      if let Some(ref nick) = member.nick {
        message.author.name = nick.clone();
      }
    }
    // FIXME: continuing here doesn't escape and opens a vulnerability
    // can probably refactor these into individual methods that return, then always escape the end
    // result or base it on a return value
    let mut parts: Vec<String> = message.content.split(' ').map(ToOwned::to_owned).collect();
    for part in &mut parts {
      if part.starts_with("<@!") {
        let end = part.find('>').unwrap_or_else(|| part.len() - 1);
        let id: u64 = match part[3..end].parse() {
          Ok(u) => u,
          Err(_) => continue
        };
        if let Some(member) = archive.server.members.iter().find(|m| m.user.id.0 == id) {
          let name = member.nick.as_ref().unwrap_or(&member.user.name);
          *part = format!("@{}", html_escape(name));
        }
      } else if part.starts_with("<@&") {
        let end = part.find('>').unwrap_or_else(|| part.len() - 1);
        let id: u64 = match part[3..end].parse() {
          Ok(u) => u,
          Err(_) => continue
        };
        if let Some(role) = archive.server.roles.iter().find(|r| r.id.0 == id) {
          let name = if role.name == "@everyone" { role.name.clone() } else { format!("@{}", role.name) };
          *part = format!("{}", html_escape(&name));
        }
      } else if part.starts_with("<@") {
        let end = part.find('>').unwrap_or_else(|| part.len() - 1);
        let id: u64 = match part[2..end].parse() {
          Ok(u) => u,
          Err(_) => continue
        };
        if let Some(member) = archive.server.members.iter().find(|m| m.user.id.0 == id) {
          let name = member.nick.as_ref().unwrap_or(&member.user.name);
          *part = format!("@{}", html_escape(name));
        }
      } else if part.starts_with("<#") {
        let end = part.find('>').unwrap_or_else(|| part.len() - 1);
        let id: u64 = match part[2..end].parse() {
          Ok(u) => u,
          Err(_) => continue
        };
        if let Some(channel) = archive.server.channels.iter().find(|c| c.id.0 == id) {
          *part = format!("#{}", html_escape(&channel.name));
        }
      } else if part.starts_with("<:") {
        if let Some(index) = part[2..].find(':') {
          let end = part[2 + index..].find('>').map(|x| x + 2 + index).unwrap_or_else(|| part.len() - 1);
          let id: u64 = match part[3 + index..end].parse() {
            Ok(u) => u,
            Err(_) => continue
          };
          *part = format!("", &part[2..index], id);
        }
      } else {
        *part = html_escape(part);
      }

      if let Ok(url) = Url::parse(part) {
        if url.has_host() {
          *part = format!("{url}", url=part)
        }
      }
    }
    message.content = parts.join(" ");
  }
  let mut msgs = MESSAGES.write().unwrap();
  let server = msgs.entry(server_id).or_insert_with(Default::default);
  server.insert(channel_id, archive);
}
<commit_msg>refactor: parse in separate functions<commit_after>use MESSAGES;

use channel::{Archive, ArchiveServer};

use iron::prelude::*;
use iron::status;

use handlebars_iron::handlebars::*;

use url::Url;

use serde_json;

use std::path::PathBuf;
use std::fs::{self, File};

pub fn _refresh() {
  let archives = fs::read_dir("archives").unwrap();
  for server in archives {
    let server = server.unwrap().path();
    if server.is_dir() {
      let server_id: u64 = server.file_stem().unwrap().to_string_lossy().parse().unwrap();
      for channel in fs::read_dir(server).unwrap() {
        let channel = channel.unwrap().path();
        let channel_id: u64 = channel.file_stem().unwrap().to_string_lossy().parse().unwrap();
        add_messages(channel, server_id, channel_id);
      }
    }
  }
}

pub fn refresh(_: &mut Request) -> IronResult<Response> {
  _refresh();
  Ok(Response::with(("We gucci", status::Ok)))
}

fn parse_user_mention(server: &ArchiveServer, part: &mut String) -> bool {
  let end = part.find('>').unwrap_or_else(|| part.len() - 1);
  let id: u64 = match part[2..end].parse() {
    Ok(u) => u,
    Err(_) => return false
  };
  match server.members.iter().find(|m| m.user.id.0 == id) {
    Some(member) => {
      let name = member.nick.as_ref().unwrap_or(&member.user.name);
      *part = format!("@{}", html_escape(name));
      true
    },
    None => false
  }
}

fn parse_user_nick_mention(server: &ArchiveServer, part: &mut String) -> bool {
  let end = part.find('>').unwrap_or_else(|| part.len() - 1);
  let id: u64 = match part[3..end].parse() {
    Ok(u) => u,
    Err(_) => return false
  };
  match server.members.iter().find(|m| m.user.id.0 == id) {
    Some(member) => {
      let name = member.nick.as_ref().unwrap_or(&member.user.name);
      *part = format!("@{}", html_escape(name));
      true
    },
    None => false
  }
}

fn parse_channel_mention(server: &ArchiveServer, part: &mut String) -> bool {
  let end = part.find('>').unwrap_or_else(|| part.len() - 1);
  let id: u64 = match part[2..end].parse() {
    Ok(u) => u,
    Err(_) => return false
  };
  match server.channels.iter().find(|c| c.id.0 == id) {
    Some(channel) => {
      *part = format!("#{}", html_escape(&channel.name));
      true
    },
    None => false
  }
}

fn parse_role_mention(server: &ArchiveServer, part: &mut String) -> bool {
  let end = part.find('>').unwrap_or_else(|| part.len() - 1);
  let id: u64 = match part[3..end].parse() {
    Ok(u) => u,
    Err(_) => return false
  };
  match server.roles.iter().find(|r| r.id.0 == id) {
    Some(role) => {
      let name = if role.name == "@everyone" { role.name.clone() } else { format!("@{}", role.name) };
      *part = format!("{}", html_escape(&name));
      true
    },
    None => false
  }
}

fn parse_custom_emoji(part: &mut String) -> bool {
  match part[2..].find(':') {
    Some(index) => {
      let end = part[2 + index..].find('>').map(|x| x + 2 + index).unwrap_or_else(|| part.len() - 1);
      let id: u64 = match part[3 + index..end].parse() {
        Ok(u) => u,
        Err(_) => return false
      };
      *part = format!("", &part[2..index], id);
      true
    },
    None => false
  }
}

fn add_messages(channel: PathBuf, server_id: u64, channel_id: u64) {
  let f = File::open(channel).unwrap();
  let mut archive: Archive = serde_json::from_reader(f).unwrap();
  for message in &mut archive.messages {
    if let Some(member) = archive.server.members.iter().find(|mem| mem.user.id == message.author.id) {
      if let Some(ref nick) = member.nick {
        message.author.name = nick.clone();
      }
    }
    let mut parts: Vec<String> = message.content.split(' ').map(ToOwned::to_owned).collect();
    for part in &mut parts {
      let escaped = if part.starts_with("<@!") {
        parse_user_nick_mention(&archive.server, part)
      } else if part.starts_with("<@&") {
        parse_role_mention(&archive.server, part)
      } else if part.starts_with("<@") {
        parse_user_mention(&archive.server, part)
      } else if part.starts_with("<#") {
        parse_channel_mention(&archive.server, part)
      } else if part.starts_with("<:") {
        parse_custom_emoji(part)
      } else {
        false
      };

      if !escaped {
        *part = html_escape(part);
      }

      if let Ok(url) = Url::parse(part) {
        if url.has_host() {
          *part = format!("{url}", url=part)
        }
      }
    }
    message.content = parts.join(" ");
  }
  let mut msgs = MESSAGES.write().unwrap();
  let server = msgs.entry(server_id).or_insert_with(Default::default);
  server.insert(channel_id, archive);
}
<|endoftext|>


<commit_before>//! Serializes and deserializes Ardite Schema Definitions from different data
//! formats such as JSON and YAML.

mod types;

use std::io::BufReader;
use std::fs::File;
use std::path::PathBuf;

use regex::Regex;
use serde_json;
use serde_yaml;

use error::{Error, ErrorCode};
use schema::{Definition, Type, Schema};
use schema::serde::types::*;

/// Gets an Ardite Schema Definition from a file. Aims to support mainly the
/// JSON and YAML formats.
pub fn from_file(path: PathBuf) -> Result<Definition, Error> {
  let extension = path.extension().map_or("", |s| s.to_str().unwrap());
  let file = try!(File::open(&path));
  let reader = BufReader::new(file);
  match extension {
    "json" => {
      let definition: SerdeDefinition = try!(serde_json::from_reader(reader));
      Ok(try!(serde_definition_into_definition(definition)))
    },
    "yml" => {
      let definition: SerdeDefinition = try!(serde_yaml::from_reader(reader));
      Ok(try!(serde_definition_into_definition(definition)))
    },
    _ => Err(Error::new(
      ErrorCode::NotAcceptable,
      format!("File extension '{}' cannot be deserialized in '{}'.", extension, path.display()),
      Some("Use a recognizable file extension like '.json' or '.yml'.".to_owned())
    ))
  }
}

/// Transforms the intermediary type into the useful type.
fn serde_definition_into_definition(serde_definition: SerdeDefinition) -> Result<Definition, Error> {
  let mut definition = Definition::new();

  for (key, schema) in serde_definition.types.into_iter() {
    let mut type_ = Type::new(key);
    type_.set_boxed_schema(try!(serde_schema_into_schema(schema)));
    definition.add_type(type_);
  }

  Ok(definition)
}

/// Transforms the intermediary type into the useful type.
fn serde_schema_into_schema(serde_schema: SerdeSchema) -> Result<Box<Schema + 'static>, Error> {
  match serde_schema.type_ {
    Some(type_) => match type_.as_str() {
      "null" => Ok(Box::new(Schema::null())),
      "boolean" => Ok(Box::new(Schema::boolean())),
      type_ @ "number" | type_ @ "integer" => {
        let mut schema = Schema::number();
        if type_ == "integer" { schema.set_multiple_of(1.0); }
        else if let Some(multiple_of) = serde_schema.multiple_of { schema.set_multiple_of(multiple_of); }
        if let Some(minimum) = serde_schema.minimum { schema.set_minimum(minimum); }
        if let Some(maximum) = serde_schema.maximum { schema.set_maximum(maximum); }
        if serde_schema.exclusive_minimum.unwrap_or(false) { schema.enable_exclusive_minimum(); }
        if serde_schema.exclusive_maximum.unwrap_or(false) { schema.enable_exclusive_maximum(); }
        Ok(Box::new(schema))
      },
      "string" => {
        let mut schema = Schema::string();
        if let Some(min_length) = serde_schema.min_length { schema.set_min_length(min_length); }
        if let Some(max_length) = serde_schema.max_length { schema.set_max_length(max_length); }
        if let Some(pattern) = serde_schema.pattern.and_then(|p| Regex::new(&p).ok()) { schema.set_pattern(pattern); }
        Ok(Box::new(schema))
      },
      "array" => {
        let mut schema = Schema::array();
        if let Some(items) = serde_schema.items { schema.set_boxed_items(try!(serde_schema_into_schema(*items))); }
        Ok(Box::new(schema))
      },
      "object" => {
        let mut schema = Schema::object();
        schema.set_required(serde_schema.required.unwrap_or_default());
        if serde_schema.additional_properties.unwrap_or(false) { schema.enable_additional_properties(); }
        for (key, value) in serde_schema.properties.unwrap_or_default() {
          schema.add_boxed_property(key, try!(serde_schema_into_schema(value)));
        }
        Ok(Box::new(schema))
      },
      _ => Err(Error::invalid(
        format!("Invalid type '{}'.", type_),
        format!("Use a permitted type like 'string' and not '{}'.", type_)
      ))
    },
    None => {
      if let Some(enum_) = serde_schema.enum_ {
        Ok(Box::new(Schema::enum_(enum_)))
      } else {
        Err(Error::invalid("No schema type specified.", "Set a `type` property or an `enum` property."))
      }
    }
  }
}

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

  use std::path::PathBuf;

  use schema::definition::create_basic;

  #[test]
  fn test_basic_json() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.json")).unwrap(), create_basic());
  }

  #[test]
  fn test_basic_yaml() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.yml")).unwrap(), create_basic());
  }
}
<commit_msg>refactor(schema): remove some duplicated code<commit_after>//! Serializes and deserializes Ardite Schema Definitions from different data
//! formats such as JSON and YAML.

mod types;

use std::io::BufReader;
use std::fs::File;
use std::path::PathBuf;

use regex::Regex;
use serde_json;
use serde_yaml;

use error::{Error, ErrorCode};
use schema::{Definition, Type, Schema};
use schema::serde::types::*;

/// Gets an Ardite Schema Definition from a file. Aims to support mainly the
/// JSON and YAML formats.
// TODO: validate file against JSON schema.
pub fn from_file(path: PathBuf) -> Result<Definition, Error> {
  let extension = path.extension().map_or("", |s| s.to_str().unwrap());
  let file = try!(File::open(&path));
  let reader = BufReader::new(file);
  let definition: SerdeDefinition = match extension {
    "json" => try!(serde_json::from_reader(reader)),
    "yml" => try!(serde_yaml::from_reader(reader)),
    _ => {
      return Err(Error::new(
        ErrorCode::NotAcceptable,
        format!("File extension '{}' cannot be deserialized in '{}'.", extension, path.display()),
        Some("Use a recognizable file extension like '.json' or '.yml'.".to_owned())
      ))
    }
  };
  serde_definition_into_definition(definition)
}

/// Transforms the intermediary type into the useful type.
fn serde_definition_into_definition(serde_definition: SerdeDefinition) -> Result<Definition, Error> {
  let mut definition = Definition::new();

  for (key, schema) in serde_definition.types.into_iter() {
    let mut type_ = Type::new(key);
    type_.set_boxed_schema(try!(serde_schema_into_schema(schema)));
    definition.add_type(type_);
  }

  Ok(definition)
}

/// Transforms the intermediary type into the useful type.
fn serde_schema_into_schema(serde_schema: SerdeSchema) -> Result<Box<Schema + 'static>, Error> {
  match serde_schema.type_ {
    Some(type_) => match type_.as_str() {
      "null" => Ok(Box::new(Schema::null())),
      "boolean" => Ok(Box::new(Schema::boolean())),
      type_ @ "number" | type_ @ "integer" => {
        let mut schema = Schema::number();
        if type_ == "integer" { schema.set_multiple_of(1.0); }
        else if let Some(multiple_of) = serde_schema.multiple_of { schema.set_multiple_of(multiple_of); }
        if let Some(minimum) = serde_schema.minimum { schema.set_minimum(minimum); }
        if let Some(maximum) = serde_schema.maximum { schema.set_maximum(maximum); }
        if serde_schema.exclusive_minimum.unwrap_or(false) { schema.enable_exclusive_minimum(); }
        if serde_schema.exclusive_maximum.unwrap_or(false) { schema.enable_exclusive_maximum(); }
        Ok(Box::new(schema))
      },
      "string" => {
        let mut schema = Schema::string();
        if let Some(min_length) = serde_schema.min_length { schema.set_min_length(min_length); }
        if let Some(max_length) = serde_schema.max_length { schema.set_max_length(max_length); }
        if let Some(pattern) = serde_schema.pattern.and_then(|p| Regex::new(&p).ok()) { schema.set_pattern(pattern); }
        Ok(Box::new(schema))
      },
      "array" => {
        let mut schema = Schema::array();
        if let Some(items) = serde_schema.items { schema.set_boxed_items(try!(serde_schema_into_schema(*items))); }
        Ok(Box::new(schema))
      },
      "object" => {
        let mut schema = Schema::object();
        schema.set_required(serde_schema.required.unwrap_or_default());
        if serde_schema.additional_properties.unwrap_or(false) { schema.enable_additional_properties(); }
        for (key, value) in serde_schema.properties.unwrap_or_default() {
          schema.add_boxed_property(key, try!(serde_schema_into_schema(value)));
        }
        Ok(Box::new(schema))
      },
      _ => Err(Error::invalid(
        format!("Invalid type '{}'.", type_),
        format!("Use a permitted type like 'string' and not '{}'.", type_)
      ))
    },
    None => {
      if let Some(enum_) = serde_schema.enum_ {
        Ok(Box::new(Schema::enum_(enum_)))
      } else {
        Err(Error::invalid("No schema type specified.", "Set a `type` property or an `enum` property."))
      }
    }
  }
}

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

  use std::path::PathBuf;

  use schema::definition::create_basic;

  #[test]
  fn test_basic_json() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.json")).unwrap(), create_basic());
  }

  #[test]
  fn test_basic_yaml() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.yml")).unwrap(), create_basic());
  }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::iter::Fuse;
use std::fmt;

use super::adaptors::{PutBack, put_back};
use either_or_both::EitherOrBoth;

/// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](trait.Itertools.html#method.merge_join_by) for more information.
pub fn merge_join_by<I, J, F>(left: I, right: J, cmp_fn: F)
    -> MergeJoinBy<I::IntoIter, J::IntoIter, F>
    where I: IntoIterator,
          J: IntoIterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    MergeJoinBy {
        left: put_back(left.into_iter().fuse()),
        right: put_back(right.into_iter().fuse()),
        cmp_fn: cmp_fn
    }
}

/// An iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](../trait.Itertools.html#method.merge_join_by) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
    left: PutBack<Fuse<I>>,
    right: PutBack<Fuse<J>>,
    cmp_fn: F
}

impl<I, J, F> fmt::Debug for MergeJoinBy<I, J, F>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
          J: Iterator + fmt::Debug,
          J::Item: fmt::Debug,
{
    debug_fmt_fields!(MergeJoinBy, left, right);
}

impl<I, J, F> Iterator for MergeJoinBy<I, J, F>
    where I: Iterator,
          J: Iterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    type Item = EitherOrBoth<I::Item, J::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match (self.left.next(), self.right.next()) {
            (None, None) => None,
            (Some(left), None) =>
                Some(EitherOrBoth::Left(left)),
            (None, Some(right)) =>
                Some(EitherOrBoth::Right(right)),
            (Some(left), Some(right)) => {
                match (self.cmp_fn)(&left, &right) {
                    Ordering::Equal =>
                        Some(EitherOrBoth::Both(left, right)),
                    Ordering::Less => {
                        self.right.put_back(right);
                        Some(EitherOrBoth::Left(left))
                    },
                    Ordering::Greater => {
                        self.left.put_back(left);
                        Some(EitherOrBoth::Right(right))
                    }
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_lower, a_upper) = self.left.size_hint();
        let (b_lower, b_upper) = self.right.size_hint();

        let lower = ::std::cmp::max(a_lower, b_lower);

        let upper = match (a_upper, b_upper) {
            (Some(x), Some(y)) => Some(x + y),
            _ => None,
        };

        (lower, upper)
    }
}
<commit_msg>Fix potential overflow in MergeJoinBy::size_hint<commit_after>use std::cmp::Ordering;
use std::iter::Fuse;
use std::fmt;

use super::adaptors::{PutBack, put_back};
use either_or_both::EitherOrBoth;

/// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](trait.Itertools.html#method.merge_join_by) for more information.
pub fn merge_join_by<I, J, F>(left: I, right: J, cmp_fn: F)
    -> MergeJoinBy<I::IntoIter, J::IntoIter, F>
    where I: IntoIterator,
          J: IntoIterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    MergeJoinBy {
        left: put_back(left.into_iter().fuse()),
        right: put_back(right.into_iter().fuse()),
        cmp_fn: cmp_fn
    }
}

/// An iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](../trait.Itertools.html#method.merge_join_by) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
    left: PutBack<Fuse<I>>,
    right: PutBack<Fuse<J>>,
    cmp_fn: F
}

impl<I, J, F> fmt::Debug for MergeJoinBy<I, J, F>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
          J: Iterator + fmt::Debug,
          J::Item: fmt::Debug,
{
    debug_fmt_fields!(MergeJoinBy, left, right);
}

impl<I, J, F> Iterator for MergeJoinBy<I, J, F>
    where I: Iterator,
          J: Iterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    type Item = EitherOrBoth<I::Item, J::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match (self.left.next(), self.right.next()) {
            (None, None) => None,
            (Some(left), None) =>
                Some(EitherOrBoth::Left(left)),
            (None, Some(right)) =>
                Some(EitherOrBoth::Right(right)),
            (Some(left), Some(right)) => {
                match (self.cmp_fn)(&left, &right) {
                    Ordering::Equal =>
                        Some(EitherOrBoth::Both(left, right)),
                    Ordering::Less => {
                        self.right.put_back(right);
                        Some(EitherOrBoth::Left(left))
                    },
                    Ordering::Greater => {
                        self.left.put_back(left);
                        Some(EitherOrBoth::Right(right))
                    }
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_lower, a_upper) = self.left.size_hint();
        let (b_lower, b_upper) = self.right.size_hint();

        let lower = ::std::cmp::max(a_lower, b_lower);

        let upper = match (a_upper, b_upper) {
            (Some(x), Some(y)) => x.checked_add(y),
            _ => None,
        };

        (lower, upper)
    }
}
<|endoftext|>