<commit_before>extern mod cairo;
extern mod core_graphics;
extern mod core_text;

use azure_hl::{BackendType, CairoBackend, CoreGraphicsAcceleratedBackend, CoreGraphicsBackend};
use azure_hl::{Direct2DBackend, NoBackend, RecordingBackend, SkiaBackend};
use bindgen::{AzCreateScaledFontForNativeFont, AzReleaseScaledFont};
use cairo::{cairo_font_face_t, cairo_matrix_t, cairo_scaled_font_t};
use cairo::bindgen::{cairo_font_face_destroy, cairo_font_options_create};
use cairo::bindgen::{cairo_font_options_destroy, cairo_matrix_init_identity, cairo_matrix_scale};
use cairo::bindgen::{cairo_scaled_font_create};
use cairo::cairo_quartz::bindgen::cairo_quartz_font_face_create_for_cgfont;
use core_graphics::font::CGFontRef;
use core_text::font::{CTFont, CTFontRef};

#[cfg(target_os="macos")]
type NativeFont = CTFont;

type SkTypeface = *c_void;

pub struct ScaledFont {
    azure_scaled_font: AzScaledFontRef,

    drop {
        AzReleaseScaledFont(self.azure_scaled_font);
    }
}

impl ScaledFont {
    static priv fn create_cairo_font(face: *cairo_font_face_t, size: AzFloat)
                                  -> *cairo_scaled_font_t {
        // FIXME: error handling

        let idmatrix: cairo_matrix_t = {
            xx: 0 as c_double,
            yx: 0 as c_double,
            xy: 0 as c_double,
            yy: 0 as c_double,
            x0: 0 as c_double,
            y0: 0 as c_double
        };
        cairo_matrix_init_identity(ptr::to_unsafe_ptr(&idmatrix));

        let fontmatrix = idmatrix;
        cairo_matrix_scale(ptr::to_unsafe_ptr(&fontmatrix), size as c_double, size as c_double);
        let options = cairo_font_options_create();
        let cfont = cairo_scaled_font_create(face,
                                             ptr::to_unsafe_ptr(&fontmatrix),
                                             ptr::to_unsafe_ptr(&idmatrix),
                                             options);
        cairo_font_options_destroy(options);

        return cfont;
    }

    /// Mac-specific function to create a font for the given backend.
    #[cfg(target_os="macos")]
    static pub fn new(backend: BackendType, native_font: &const CTFont, size: AzFloat)
                   -> ScaledFont {
        let mut azure_native_font = {
            mType: 0,
            mFont: ptr::null()
        };

        match backend {
            CoreGraphicsBackend | CoreGraphicsAcceleratedBackend => {
                azure_native_font.mType = AZ_NATIVE_FONT_MAC_FONT_FACE;

                unsafe {
                    azure_native_font.mFont = cast::transmute(native_font.copy_cg_font());
                }
            }
            CairoBackend => {
                azure_native_font.mType = AZ_NATIVE_FONT_CAIRO_FONT_FACE;

                let cg_font = native_font.copy_cg_font();
                let face = cairo_quartz_font_face_create_for_cgfont(cg_font);
                if face == ptr::null() { fail; }

                let cairo_font = ScaledFont::create_cairo_font(face, size);
                cairo_font_face_destroy(face);

                unsafe {
                    azure_native_font.mFont = cast::transmute(cairo_font);
                }
            }
            SkiaBackend => {
                azure_native_font.mType = AZ_NATIVE_FONT_SKIA_FONT_FACE;
                let skia_typeface = _Z26SkCreateTypefaceFromCTFontPK8__CTFont(
                    native_font.get_ref());

                unsafe {
                    azure_native_font.mFont = cast::transmute(skia_typeface);
                }
            }
            NoBackend | Direct2DBackend | RecordingBackend => {
                fail ~"don't know how to make a scaled font for this backend";
            }
        }

        let azure_native_font_ptr = ptr::to_unsafe_ptr(&azure_native_font);
        let azure_scaled_font = AzCreateScaledFontForNativeFont(azure_native_font_ptr, size);
        ScaledFont { azure_scaled_font: move azure_scaled_font }
    }
}

// FIXME: Move this stuff to a rust-skia?
// FIXME: Demangle the names!!!
#[link_args="-lskia"]
extern {
    #[cfg(target_os="macos")]
    pub fn _Z26SkCreateTypefaceFromCTFontPK8__CTFont(font: CTFontRef) -> *SkTypeface;
}

<commit_msg>Fix Skia backend font loading on Mac<commit_after>extern mod cairo;
extern mod core_graphics;
extern mod core_text;

use azure_hl::{BackendType, CairoBackend, CoreGraphicsAcceleratedBackend, CoreGraphicsBackend};
use azure_hl::{Direct2DBackend, NoBackend, RecordingBackend, SkiaBackend};
use bindgen::{AzCreateScaledFontForNativeFont, AzReleaseScaledFont};
use cairo::{cairo_font_face_t, cairo_matrix_t, cairo_scaled_font_t};
use cairo::bindgen::{cairo_font_face_destroy, cairo_font_options_create};
use cairo::bindgen::{cairo_font_options_destroy, cairo_matrix_init_identity, cairo_matrix_scale};
use cairo::bindgen::{cairo_scaled_font_create};
use cairo::cairo_quartz::bindgen::cairo_quartz_font_face_create_for_cgfont;
use core_graphics::font::CGFontRef;
use core_text::font::{CTFont, CTFontRef};

#[cfg(target_os="macos")]
type NativeFont = CTFont;

type SkTypeface = *c_void;

pub struct ScaledFont {
    azure_scaled_font: AzScaledFontRef,

    drop {
        AzReleaseScaledFont(self.azure_scaled_font);
    }
}

impl ScaledFont {
    static priv fn create_cairo_font(face: *cairo_font_face_t, size: AzFloat)
                                  -> *cairo_scaled_font_t {
        // FIXME: error handling

        let idmatrix: cairo_matrix_t = {
            xx: 0 as c_double,
            yx: 0 as c_double,
            xy: 0 as c_double,
            yy: 0 as c_double,
            x0: 0 as c_double,
            y0: 0 as c_double
        };
        cairo_matrix_init_identity(ptr::to_unsafe_ptr(&idmatrix));

        let fontmatrix = idmatrix;
        cairo_matrix_scale(ptr::to_unsafe_ptr(&fontmatrix), size as c_double, size as c_double);
        let options = cairo_font_options_create();
        let cfont = cairo_scaled_font_create(face,
                                             ptr::to_unsafe_ptr(&fontmatrix),
                                             ptr::to_unsafe_ptr(&idmatrix),
                                             options);
        cairo_font_options_destroy(options);

        return cfont;
    }

    /// Mac-specific function to create a font for the given backend.
    #[cfg(target_os="macos")]
    static pub fn new(backend: BackendType, native_font: &const CTFont, size: AzFloat)
                   -> ScaledFont {
        let mut azure_native_font = {
            mType: 0,
            mFont: ptr::null()
        };

        match backend {
            CoreGraphicsBackend | CoreGraphicsAcceleratedBackend | SkiaBackend => {
                azure_native_font.mType = AZ_NATIVE_FONT_MAC_FONT_FACE;

                unsafe {
                    azure_native_font.mFont = cast::transmute(native_font.copy_cg_font());
                }
            }
            CairoBackend => {
                azure_native_font.mType = AZ_NATIVE_FONT_CAIRO_FONT_FACE;

                let cg_font = native_font.copy_cg_font();
                let face = cairo_quartz_font_face_create_for_cgfont(cg_font);
                if face == ptr::null() { fail; }

                let cairo_font = ScaledFont::create_cairo_font(face, size);
                cairo_font_face_destroy(face);

                unsafe {
                    azure_native_font.mFont = cast::transmute(cairo_font);
                }
            }
            NoBackend | Direct2DBackend | RecordingBackend => {
                fail ~"don't know how to make a scaled font for this backend";
            }
        }

        let azure_native_font_ptr = ptr::to_unsafe_ptr(&azure_native_font);
        let azure_scaled_font = AzCreateScaledFontForNativeFont(azure_native_font_ptr, size);
        ScaledFont { azure_scaled_font: move azure_scaled_font }
    }
}

// FIXME: Move this stuff to a rust-skia?
// FIXME: Demangle the names!!!
#[link_args="-lskia"]
extern {
    #[cfg(target_os="macos")]
    pub fn _Z26SkCreateTypefaceFromCTFontPK8__CTFont(font: CTFontRef) -> *SkTypeface;
}

<|endoftext|>


<commit_before>use std::ptr;
use std::fs::File;
use std::io::{Read, BufReader};
use std::path::{Path, PathBuf};
use std::ffi::{CStr, CString};

use rustc_llvm::archive_ro::ArchiveRO;
use cty::c_char;
use llvm;
use error::*;
use session::{Session, Output, Configuration};

pub struct Linker {
    session: Session,
    context: llvm::ContextRef,
    module: llvm::ModuleRef,
    pass_manager: llvm::PassManagerRef,
}

impl Linker {
    pub fn new(session: Session) -> Self {
        let module_name = CString::new("nvptx-module").unwrap();
        let context = unsafe { llvm::LLVMContextCreate() };

        Linker {
            session,
            context,
            pass_manager: unsafe { llvm::LLVMCreatePassManager() },
            module: unsafe {
                llvm::LLVMModuleCreateWithNameInContext(module_name.as_ptr(), context)
            },
        }
    }

    pub fn link(self) -> Result<()> {
        self.link_bitcode();
        self.link_rlibs();
        self.run_passes();

        // TODO: LLVMVerifyModule(mod, LLVMAbortProcessAction, &error);
        // TODO: LLVMDisposeMessage(error);

        for output in &self.session.emit {
            match *output {
                Output::PTXAssembly => self.emit_asm().chain_err(|| "Unable to emit PTX assembly")?,
                Output::Bitcode => self.emit_bc().chain_err(|| "Unable to emit LLVM bitcode")?,
                Output::IntermediateRepresentation => {
                    self.emit_ir().chain_err(|| "Unable to emit LLVM IR code")?
                }
            }
        }

        Ok(())
    }

    fn link_bitcode(&self) {
        for module_path in &self.session.include_bitcode_modules {
            debug!("Linking bitcode: {:?}", module_path);

            let mut bitcode_file = BufReader::new(File::open(&module_path).unwrap());
            let mut bitcode_bytes = vec![];

            bitcode_file.read_to_end(&mut bitcode_bytes).unwrap();
            unsafe {
                // TODO: check result
                llvm::LLVMRustLinkInExternalBitcode(self.module,
                                                    bitcode_bytes.as_ptr() as *const c_char,
                                                    bitcode_bytes.len());
            }
        }
    }

    fn link_rlibs(&self) {
        for rlib_path in &self.session.include_rlibs {
            debug!("Linking rlib: {:?}", rlib_path);

            let archive = ArchiveRO::open(rlib_path).unwrap();
            for item in archive.iter() {
                let name = Path::new(item.as_ref().unwrap().name().unwrap());

                if self.is_rlib_item_linkable(&name) {
                    debug!("  - linking archive item: {:?}", name);
                    let data = item.as_ref().unwrap().data();

                    unsafe {
                        // TODO: check result
                        llvm::LLVMRustLinkInExternalBitcode(self.module,
                                                            data.as_ptr() as *const c_char,
                                                            data.len());
                    }
                }
            }
        }
    }

    fn is_rlib_item_linkable(&self, name: &Path) -> bool {
        name.extension().unwrap() == "o"
    }

    fn run_passes(&self) {
        unsafe {
            let builder = llvm::LLVMPassManagerBuilderCreate();

            match self.session.configuration {
                Configuration::Debug => {
                    info!("Linking without optimisations");
                    llvm::LLVMPassManagerBuilderSetOptLevel(builder, 0);
                }

                Configuration::Release => {
                    info!("Linking with Link Time Optimisations");
                    llvm::LLVMPassManagerBuilderSetOptLevel(builder, 3);
                    llvm::LLVMPassManagerBuilderPopulateLTOPassManager(builder,
                                                                       self.pass_manager,
                                                                       llvm::True,
                                                                       llvm::True);
                }
            }

            llvm::LLVMPassManagerBuilderPopulateFunctionPassManager(builder, self.pass_manager);
            llvm::LLVMPassManagerBuilderPopulateModulePassManager(builder, self.pass_manager);
            llvm::LLVMPassManagerBuilderDispose(builder);

            llvm::LLVMAddStripSymbolsPass(self.pass_manager);
            llvm::LLVMAddStripDeadPrototypesPass(self.pass_manager);

            llvm::LLVMRunPassManager(self.pass_manager, self.module);
        }
    }

    fn emit_ir(&self) -> Result<()> {
        let path = CString::new(self.get_output_path_with_ext("ll")?.to_str().unwrap()).unwrap();

        unsafe {
            // TODO: check result
            let mut message: *const c_char = ptr::null();
            llvm::LLVMPrintModuleToFile(self.module, path.as_ptr(), &mut message);

            if message != ptr::null() {
                // TODO: stderr?
                println!("{}", CStr::from_ptr(message).to_str().unwrap());
            }

            llvm::LLVMDisposeMessage(message);
        }

        info!("LLVM IR has been written to {:?}", path);
        Ok(())
    }

    fn emit_bc(&self) -> Result<()> {
        let path = CString::new(self.get_output_path_with_ext("bc")?.to_str().unwrap()).unwrap();

        unsafe {
            // TODO: check result
            llvm::LLVMWriteBitcodeToFile(self.module, path.as_ptr());
        }

        info!("LLVM bitcode has been written to {:?}", path);
        Ok(())
    }

    fn emit_asm(&self) -> Result<()> {
        let path = CString::new(self.get_output_path()?.to_str().unwrap()).unwrap();

        let cpu = CString::new("sm_20").unwrap();
        let feature = CString::new("").unwrap();

        unsafe {
            llvm::LLVMInitializeNVPTXTargetInfo();
            llvm::LLVMInitializeNVPTXTarget();
            llvm::LLVMInitializeNVPTXTargetMC();
            llvm::LLVMInitializeNVPTXAsmPrinter();

            let triple = llvm::LLVMGetTarget(self.module);
            let target = llvm::LLVMRustCreateTargetMachine(triple,
                                                           cpu.as_ptr(),
                                                           feature.as_ptr(),
                                                           llvm::CodeModel::Default,
                                                           llvm::RelocMode::Default,
                                                           llvm::CodeGenOptLevel::Default,
                                                           false,
                                                           false,
                                                           false,
                                                           false);

            // TODO: check `target` != nullptr

            // TODO: check result
            llvm::LLVMRustWriteOutputFile(target,
                                          self.pass_manager,
                                          self.module,
                                          path.as_ptr(),
                                          llvm::FileType::AssemblyFile);

            llvm::LLVMRustDisposeTargetMachine(target);
        }

        info!("PTX assembly has been written to {:?}", path);
        Ok(())
    }

    fn get_output_path(&self) -> Result<PathBuf> {
        match self.session.output.as_ref() {
            Some(path) => Ok(path.clone()),
            None => Err(ErrorKind::NoOutputPathError.into()),
        }
    }

    fn get_output_path_with_ext(&self, extension: &str) -> Result<PathBuf> {
        let mut output_path = self.get_output_path()?;
        output_path.set_extension(extension);

        Ok(output_path)
    }
}

impl Drop for Linker {
    fn drop(&mut self) {
        unsafe {
            llvm::LLVMDisposeModule(self.module);
            llvm::LLVMContextDispose(self.context);
            llvm::LLVMDisposePassManager(self.pass_manager);
        }
    }
}
<commit_msg>Don't run passes twice<commit_after>use std::ptr;
use std::fs::File;
use std::io::{Read, BufReader};
use std::path::{Path, PathBuf};
use std::ffi::{CStr, CString};

use rustc_llvm::archive_ro::ArchiveRO;
use cty::c_char;
use llvm;
use error::*;
use session::{Session, Output, Configuration};

pub struct Linker {
    session: Session,
    context: llvm::ContextRef,
    module: llvm::ModuleRef,
}

impl Linker {
    pub fn new(session: Session) -> Self {
        let module_name = CString::new("nvptx-module").unwrap();
        let context = unsafe { llvm::LLVMContextCreate() };

        Linker {
            session,
            context,
            module: unsafe {
                llvm::LLVMModuleCreateWithNameInContext(module_name.as_ptr(), context)
            },
        }
    }

    pub fn link(self) -> Result<()> {
        self.link_bitcode();
        self.link_rlibs();
        self.run_passes();

        // TODO: LLVMVerifyModule(mod, LLVMAbortProcessAction, &error);
        // TODO: LLVMDisposeMessage(error);

        for output in &self.session.emit {
            match *output {
                Output::PTXAssembly => self.emit_asm().chain_err(|| "Unable to emit PTX assembly")?,
                Output::Bitcode => self.emit_bc().chain_err(|| "Unable to emit LLVM bitcode")?,
                Output::IntermediateRepresentation => {
                    self.emit_ir().chain_err(|| "Unable to emit LLVM IR code")?
                }
            }
        }

        Ok(())
    }

    fn link_bitcode(&self) {
        for module_path in &self.session.include_bitcode_modules {
            debug!("Linking bitcode: {:?}", module_path);

            let mut bitcode_file = BufReader::new(File::open(&module_path).unwrap());
            let mut bitcode_bytes = vec![];

            bitcode_file.read_to_end(&mut bitcode_bytes).unwrap();
            unsafe {
                // TODO: check result
                llvm::LLVMRustLinkInExternalBitcode(self.module,
                                                    bitcode_bytes.as_ptr() as *const c_char,
                                                    bitcode_bytes.len());
            }
        }
    }

    fn link_rlibs(&self) {
        for rlib_path in &self.session.include_rlibs {
            debug!("Linking rlib: {:?}", rlib_path);

            let archive = ArchiveRO::open(rlib_path).unwrap();
            for item in archive.iter() {
                let name = Path::new(item.as_ref().unwrap().name().unwrap());

                if self.is_rlib_item_linkable(&name) {
                    debug!("  - linking archive item: {:?}", name);
                    let data = item.as_ref().unwrap().data();

                    unsafe {
                        // TODO: check result
                        llvm::LLVMRustLinkInExternalBitcode(self.module,
                                                            data.as_ptr() as *const c_char,
                                                            data.len());
                    }
                }
            }
        }
    }

    fn is_rlib_item_linkable(&self, name: &Path) -> bool {
        name.extension().unwrap() == "o"
    }

    fn run_passes(&self) {
        unsafe {
            let builder = llvm::LLVMPassManagerBuilderCreate();
            let pass_manager = llvm::LLVMCreatePassManager();

            match self.session.configuration {
                Configuration::Debug => {
                    info!("Linking without optimisations");
                    llvm::LLVMPassManagerBuilderSetOptLevel(builder, 0);
                }

                Configuration::Release => {
                    info!("Linking with Link Time Optimisations");
                    llvm::LLVMPassManagerBuilderSetOptLevel(builder, 3);
                    llvm::LLVMPassManagerBuilderPopulateLTOPassManager(builder,
                                                                       pass_manager,
                                                                       llvm::True,
                                                                       llvm::True);
                }
            }

            llvm::LLVMPassManagerBuilderPopulateFunctionPassManager(builder, pass_manager);
            llvm::LLVMPassManagerBuilderPopulateModulePassManager(builder, pass_manager);
            llvm::LLVMPassManagerBuilderDispose(builder);

            llvm::LLVMAddStripSymbolsPass(pass_manager);
            llvm::LLVMAddStripDeadPrototypesPass(pass_manager);

            llvm::LLVMRunPassManager(pass_manager, self.module);
            llvm::LLVMDisposePassManager(pass_manager);
        }
    }

    fn emit_ir(&self) -> Result<()> {
        let path = CString::new(self.get_output_path_with_ext("ll")?.to_str().unwrap()).unwrap();

        unsafe {
            // TODO: check result
            let mut message: *const c_char = ptr::null();
            llvm::LLVMPrintModuleToFile(self.module, path.as_ptr(), &mut message);

            if message != ptr::null() {
                // TODO: stderr?
                println!("{}", CStr::from_ptr(message).to_str().unwrap());
            }

            llvm::LLVMDisposeMessage(message);
        }

        info!("LLVM IR has been written to {:?}", path);
        Ok(())
    }

    fn emit_bc(&self) -> Result<()> {
        let path = CString::new(self.get_output_path_with_ext("bc")?.to_str().unwrap()).unwrap();

        unsafe {
            // TODO: check result
            llvm::LLVMWriteBitcodeToFile(self.module, path.as_ptr());
        }

        info!("LLVM bitcode has been written to {:?}", path);
        Ok(())
    }

    fn emit_asm(&self) -> Result<()> {
        let path = CString::new(self.get_output_path()?.to_str().unwrap()).unwrap();

        let cpu = CString::new("sm_20").unwrap();
        let feature = CString::new("").unwrap();

        unsafe {
            llvm::LLVMInitializeNVPTXTargetInfo();
            llvm::LLVMInitializeNVPTXTarget();
            llvm::LLVMInitializeNVPTXTargetMC();
            llvm::LLVMInitializeNVPTXAsmPrinter();

            let triple = llvm::LLVMGetTarget(self.module);
            let target = llvm::LLVMRustCreateTargetMachine(triple,
                                                           cpu.as_ptr(),
                                                           feature.as_ptr(),
                                                           llvm::CodeModel::Default,
                                                           llvm::RelocMode::Default,
                                                           llvm::CodeGenOptLevel::Default,
                                                           false,
                                                           false,
                                                           false,
                                                           false);

            // TODO: check `target` != nullptr

            // We need an dummy pass manager because module is already optimised
            let dummy_pass_manager = llvm::LLVMCreatePassManager();

            // TODO: check result
            llvm::LLVMRustWriteOutputFile(target,
                                          dummy_pass_manager,
                                          self.module,
                                          path.as_ptr(),
                                          llvm::FileType::AssemblyFile);

            llvm::LLVMRustDisposeTargetMachine(target);
            llvm::LLVMDisposePassManager(dummy_pass_manager);
        }

        info!("PTX assembly has been written to {:?}", path);
        Ok(())
    }

    fn get_output_path(&self) -> Result<PathBuf> {
        match self.session.output.as_ref() {
            Some(path) => Ok(path.clone()),
            None => Err(ErrorKind::NoOutputPathError.into()),
        }
    }

    fn get_output_path_with_ext(&self, extension: &str) -> Result<PathBuf> {
        let mut output_path = self.get_output_path()?;
        output_path.set_extension(extension);

        Ok(output_path)
    }
}

impl Drop for Linker {
    fn drop(&mut self) {
        unsafe {
            llvm::LLVMDisposeModule(self.module);
            llvm::LLVMContextDispose(self.context);
        }
    }
}
<|endoftext|>


<commit_before>/*
    Copyright 2014-2015 Zumero, LLC

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
*/

#![feature(convert)]
#![feature(iter_arith)]

use std::collections::BTreeMap;

extern crate lsm;
use lsm::ICursor;
use lsm::IForwardCursor;

extern crate rand;
use rand::Rng;
use rand::SeedableRng;

fn dump_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let page = try!(db.get_page(pgnum));
    println!("{:?}", page);
    Ok(())
}

fn show_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let cursor = try!(db.open_cursor_on_page(pgnum));
    let pt = cursor.page_type();
    println!("page type: {:?}", pt);
    // TODO
    Ok(())
}

fn merge(name: &str, from_level: String) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let from_level =
        match from_level.as_str() {
            "fresh" => {
                lsm::FromLevel::Fresh
            },
            "young" => {
                lsm::FromLevel::Young
            },
            _ => {
                let level = from_level.parse::<usize>().unwrap();
                lsm::FromLevel::Other(level)
            },
        };
    try!(db.merge(from_level));
    Ok(())
}

fn show_leaf_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor_on_leaf_page(pgnum));
    try!(cursor.First());
    while cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.ValueRef());
            //println!("v: {:?}", v);
            //let q = try!(v.into_boxed_slice());
        }
        try!(cursor.Next());
    }
    Ok(())
}

fn show_parent_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let page = try!(db.read_parent_page(pgnum));
    println!("depth: {}", page.depth());
    println!("count_items: {}", page.count_items());
    let blocks = try!(page.blocklist_clean());
    println!("blocks ({} blocks, {} pages): {:?}", blocks.count_blocks(), blocks.count_pages(), blocks);
    println!("key range: {:?}", try!(page.range()));
    let count = page.count_items();
    println!("items ({}):", count);
    for i in 0 .. count {
        let p = page.get_child_pagenum(i);
        println!("    {}", p);
        let child_range = try!(page.child_range(i));
        println!("    {:?}", child_range);
    }
    try!(page.verify_child_keys());
    Ok(())
}

fn list_segments(name: &str) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let (fresh, young, levels) = try!(db.list_segments());
    println!("fresh ({}): ", fresh.len());

    fn print_seg(s: &lsm::SegmentLocation) {
        println!("    {}, {} pages", s.root_page, 1 + s.blocks.count_pages());
    }

    for s in fresh.iter() {
        print_seg(s);
    }
    println!("young ({}): ", young.len());
    for s in young.iter() {
        print_seg(s);
    }
    println!("levels ({}): ", levels.len());
    for s in levels.iter() {
        print_seg(s);
    }
    Ok(())
}

fn list_free_blocks(name: &str) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let blocks = try!(db.list_free_blocks());
    println!("{:?}", blocks);
    println!("total pages: {}", blocks.count_pages());
    Ok(())
}

fn list_keys(name: &str) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor());
    try!(cursor.First());
    while cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.LiveValueRef());
            //println!("v: {:?}", v);
            //let q = try!(v.into_boxed_slice());
        }
        try!(cursor.Next());
    }
    Ok(())
}

fn seek_string(name: &str, key: String, sop: String) -> Result<(),lsm::Error> {
    let sop =
        match sop.as_str() {
            "eq" => lsm::SeekOp::SEEK_EQ,
            "le" => lsm::SeekOp::SEEK_LE,
            "ge" => lsm::SeekOp::SEEK_GE,
            _ => return Err(lsm::Error::Misc(String::from("invalid sop"))),
        };
    let k = key.into_bytes().into_boxed_slice();
    let k = lsm::KeyRef::from_boxed_slice(k);
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor());
    let sr = try!(cursor.SeekRef(&k, sop));
    println!("sr: {:?}", sr);
    if cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.LiveValueRef());
            println!("v: {:?}", v);
        }
    }
    Ok(())
}

fn seek_bytes(name: &str, k: Box<[u8]>, sop: String) -> Result<(),lsm::Error> {
    let sop =
        match sop.as_str() {
            "eq" => lsm::SeekOp::SEEK_EQ,
            "le" => lsm::SeekOp::SEEK_LE,
            "ge" => lsm::SeekOp::SEEK_GE,
            _ => return Err(lsm::Error::Misc(String::from("invalid sop"))),
        };
    let k = lsm::KeyRef::from_boxed_slice(k);
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor());
    let sr = try!(cursor.SeekRef(&k, sop));
    println!("RESULT sr: {:?}", sr);
    if cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.LiveValueRef());
            println!("v: {:?}", v);
        }
        for x in 0 .. 20 {
            try!(cursor.Next());
            if cursor.IsValid() {
                let k = try!(cursor.KeyRef());
                println!("    k: {:?}", k);
                let v = try!(cursor.LiveValueRef());
                println!("    v: {:?}", v);
            } else {
                break;
            }
        }
    }
    Ok(())
}

fn add_numbers(name: &str, count: u64, start: u64, step: u64) -> Result<(),lsm::Error> {
    let mut pending = BTreeMap::new();
    for i in 0 .. count {
        let val = start + i * step;
        let k = format!("{:08}", val).into_bytes().into_boxed_slice();
        let v = format!("{}", val).into_bytes().into_boxed_slice();
        pending.insert(k, lsm::Blob::Boxed(v));
    }
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let seg = try!(db.write_segment(pending).map_err(lsm::wrap_err));
    if let Some(seg) = seg {
        let lck = try!(db.get_write_lock());
        try!(lck.commit_segment(seg).map_err(lsm::wrap_err));
    }
    Ok(())
}

fn add_random(name: &str, count: u64, seed: usize, klen: usize, vlen: usize) -> Result<(),lsm::Error> {
    fn make(rng: &mut rand::StdRng, max_len: usize) -> Box<[u8]> {
        let len = (rng.next_u64() as usize) % max_len + 1;
        let mut k = vec![0u8; len].into_boxed_slice();
        rng.fill_bytes(&mut k);
        k
    }

    let mut rng = rand::StdRng::from_seed(&[seed]);
    let mut pending = BTreeMap::new();
    for i in 0 .. count {
        let k = make(&mut rng, klen);
        let v = make(&mut rng, vlen);
        pending.insert(k, lsm::Blob::Boxed(v));
    }
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let seg = try!(db.write_segment(pending).map_err(lsm::wrap_err));
    if let Some(seg) = seg {
        let lck = try!(db.get_write_lock());
        try!(lck.commit_segment(seg).map_err(lsm::wrap_err));
    }
    Ok(())
}

fn result_main() -> Result<(),lsm::Error> {
    let args: Vec<_> = std::env::args().collect();
    println!("args: {:?}", args);
    if args.len() < 2 {
        return Err(lsm::Error::Misc(String::from("no filename given")));
    }
    if args.len() < 3 {
        return Err(lsm::Error::Misc(String::from("no command given")));
    }
    let name = args[1].as_str();
    let cmd = args[2].as_str();
    match cmd {
        "add_random" => {
            println!("usage: add_random count seed klen vlen");
            if args.len() < 7 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let count = args[3].parse::<u64>().unwrap();
            let seed = args[4].parse::<usize>().unwrap();
            let klen = args[5].parse::<usize>().unwrap();
            let vlen = args[6].parse::<usize>().unwrap();
            add_random(name, count, seed, klen, vlen)
        },
        "add_numbers" => {
            println!("usage: add_numbers count start step");
            if args.len() < 6 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let count = args[3].parse::<u64>().unwrap();
            let start = args[4].parse::<u64>().unwrap();
            let step = args[5].parse::<u64>().unwrap();
            if step == 0 {
                return Err(lsm::Error::Misc(String::from("step cannot be 0")));
            }
            add_numbers(name, count, start, step)
        },
        "show_page" => {
            println!("usage: show_page pagenum");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            show_page(name, pgnum)
        },
        "show_leaf_page" => {
            println!("usage: show_leaf_page pagenum");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            show_leaf_page(name, pgnum)
        },
        "show_parent_page" => {
            println!("usage: show_parent_page pagenum");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            show_parent_page(name, pgnum)
        },
        "seek_string" => {
            println!("usage: seek_string key sop");
            if args.len() < 5 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let key = args[3].clone();
            let sop = args[4].clone();
            seek_string(name, key, sop)
        },
        "merge" => {
            println!("usage: merge from_level");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let from_level = args[3].clone();
            merge(name, from_level)
        },
        "seek_bytes" => {
            println!("usage: seek_bytes sop numbytes b1 b2 b3 ... ");
            if args.len() < 5 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let sop = args[3].clone();
            let count = args[4].parse::<usize>().unwrap();
            if count == 0 {
                return Err(lsm::Error::Misc(String::from("count cannot be 0")));
            }
            if args.len() < 5 + count {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let mut k = Vec::with_capacity(count);
            for i in 0 .. count {
                let b = args[5 + i].parse::<u8>().unwrap();
                k.push(b);
            }
            seek_bytes(name, k.into_boxed_slice(), sop)
        },
        "dump_page" => {
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            dump_page(name, pgnum)
        },
        "list_keys" => {
            list_keys(name)
        },
        "list_segments" => {
            list_segments(name)
        },
        "list_free_blocks" => {
            list_free_blocks(name)
        },
        _ => {
            Err(lsm::Error::Misc(String::from("unknown command")))
        },
    }
}

pub fn main() {
    let r = result_main();
    if r.is_err() {
        println!("{:?}", r);
    }
}

<commit_msg>lsm_diag add_random allow multiple groups<commit_after>/*
    Copyright 2014-2015 Zumero, LLC

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
*/

#![feature(convert)]
#![feature(iter_arith)]

use std::collections::BTreeMap;

extern crate lsm;
use lsm::ICursor;
use lsm::IForwardCursor;

extern crate rand;
use rand::Rng;
use rand::SeedableRng;

fn dump_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let page = try!(db.get_page(pgnum));
    println!("{:?}", page);
    Ok(())
}

fn show_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let cursor = try!(db.open_cursor_on_page(pgnum));
    let pt = cursor.page_type();
    println!("page type: {:?}", pt);
    // TODO
    Ok(())
}

fn merge(name: &str, from_level: String) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let from_level =
        match from_level.as_str() {
            "fresh" => {
                lsm::FromLevel::Fresh
            },
            "young" => {
                lsm::FromLevel::Young
            },
            _ => {
                let level = from_level.parse::<usize>().unwrap();
                lsm::FromLevel::Other(level)
            },
        };
    try!(db.merge(from_level));
    Ok(())
}

fn show_leaf_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor_on_leaf_page(pgnum));
    try!(cursor.First());
    while cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.ValueRef());
            //println!("v: {:?}", v);
            //let q = try!(v.into_boxed_slice());
        }
        try!(cursor.Next());
    }
    Ok(())
}

fn show_parent_page(name: &str, pgnum: u32) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let page = try!(db.read_parent_page(pgnum));
    println!("depth: {}", page.depth());
    println!("count_items: {}", page.count_items());
    let blocks = try!(page.blocklist_clean());
    println!("blocks ({} blocks, {} pages): {:?}", blocks.count_blocks(), blocks.count_pages(), blocks);
    println!("key range: {:?}", try!(page.range()));
    let count = page.count_items();
    println!("items ({}):", count);
    for i in 0 .. count {
        let p = page.get_child_pagenum(i);
        println!("    {}", p);
        let child_range = try!(page.child_range(i));
        println!("    {:?}", child_range);
    }
    try!(page.verify_child_keys());
    Ok(())
}

fn list_segments(name: &str) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let (fresh, young, levels) = try!(db.list_segments());
    println!("fresh ({}): ", fresh.len());

    fn print_seg(s: &lsm::SegmentLocation) {
        println!("    {}, {} pages", s.root_page, 1 + s.blocks.count_pages());
    }

    for s in fresh.iter() {
        print_seg(s);
    }
    println!("young ({}): ", young.len());
    for s in young.iter() {
        print_seg(s);
    }
    println!("levels ({}): ", levels.len());
    for s in levels.iter() {
        print_seg(s);
    }
    Ok(())
}

fn list_free_blocks(name: &str) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let blocks = try!(db.list_free_blocks());
    println!("{:?}", blocks);
    println!("total pages: {}", blocks.count_pages());
    Ok(())
}

fn list_keys(name: &str) -> Result<(),lsm::Error> {
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor());
    try!(cursor.First());
    while cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.LiveValueRef());
            //println!("v: {:?}", v);
            //let q = try!(v.into_boxed_slice());
        }
        try!(cursor.Next());
    }
    Ok(())
}

fn seek_string(name: &str, key: String, sop: String) -> Result<(),lsm::Error> {
    let sop =
        match sop.as_str() {
            "eq" => lsm::SeekOp::SEEK_EQ,
            "le" => lsm::SeekOp::SEEK_LE,
            "ge" => lsm::SeekOp::SEEK_GE,
            _ => return Err(lsm::Error::Misc(String::from("invalid sop"))),
        };
    let k = key.into_bytes().into_boxed_slice();
    let k = lsm::KeyRef::from_boxed_slice(k);
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor());
    let sr = try!(cursor.SeekRef(&k, sop));
    println!("sr: {:?}", sr);
    if cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.LiveValueRef());
            println!("v: {:?}", v);
        }
    }
    Ok(())
}

fn seek_bytes(name: &str, k: Box<[u8]>, sop: String) -> Result<(),lsm::Error> {
    let sop =
        match sop.as_str() {
            "eq" => lsm::SeekOp::SEEK_EQ,
            "le" => lsm::SeekOp::SEEK_LE,
            "ge" => lsm::SeekOp::SEEK_GE,
            _ => return Err(lsm::Error::Misc(String::from("invalid sop"))),
        };
    let k = lsm::KeyRef::from_boxed_slice(k);
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let mut cursor = try!(db.open_cursor());
    let sr = try!(cursor.SeekRef(&k, sop));
    println!("RESULT sr: {:?}", sr);
    if cursor.IsValid() {
        {
            let k = try!(cursor.KeyRef());
            println!("k: {:?}", k);
            let v = try!(cursor.LiveValueRef());
            println!("v: {:?}", v);
        }
        for x in 0 .. 20 {
            try!(cursor.Next());
            if cursor.IsValid() {
                let k = try!(cursor.KeyRef());
                println!("    k: {:?}", k);
                let v = try!(cursor.LiveValueRef());
                println!("    v: {:?}", v);
            } else {
                break;
            }
        }
    }
    Ok(())
}

fn add_numbers(name: &str, count: u64, start: u64, step: u64) -> Result<(),lsm::Error> {
    let mut pending = BTreeMap::new();
    for i in 0 .. count {
        let val = start + i * step;
        let k = format!("{:08}", val).into_bytes().into_boxed_slice();
        let v = format!("{}", val).into_bytes().into_boxed_slice();
        pending.insert(k, lsm::Blob::Boxed(v));
    }
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::SETTINGS_NO_AUTOMERGE));
    let seg = try!(db.write_segment(pending).map_err(lsm::wrap_err));
    if let Some(seg) = seg {
        let lck = try!(db.get_write_lock());
        try!(lck.commit_segment(seg).map_err(lsm::wrap_err));
    }
    Ok(())
}

fn add_random(name: &str, seed: usize, count_groups: usize, count_per_group: usize, klen: usize, vlen: usize) -> Result<(),lsm::Error> {
    fn make(rng: &mut rand::StdRng, max_len: usize) -> Box<[u8]> {
        let len = (rng.next_u64() as usize) % max_len + 1;
        let mut k = vec![0u8; len].into_boxed_slice();
        rng.fill_bytes(&mut k);
        k
    }

    let mut rng = rand::StdRng::from_seed(&[seed]);
    let db = try!(lsm::DatabaseFile::new(String::from(name), lsm::DEFAULT_SETTINGS));

    for _ in 0 .. count_groups {
        let mut pending = BTreeMap::new();
        for _ in 0 .. count_per_group {
            let k = make(&mut rng, klen);
            let v = make(&mut rng, vlen);
            pending.insert(k, lsm::Blob::Boxed(v));
        }
        let seg = try!(db.write_segment(pending).map_err(lsm::wrap_err));
        if let Some(seg) = seg {
            let lck = try!(db.get_write_lock());
            try!(lck.commit_segment(seg).map_err(lsm::wrap_err));
        }
    }

    Ok(())
}

fn result_main() -> Result<(),lsm::Error> {
    let args: Vec<_> = std::env::args().collect();
    println!("args: {:?}", args);
    if args.len() < 2 {
        return Err(lsm::Error::Misc(String::from("no filename given")));
    }
    if args.len() < 3 {
        return Err(lsm::Error::Misc(String::from("no command given")));
    }
    let name = args[1].as_str();
    let cmd = args[2].as_str();
    match cmd {
        "add_random" => {
            println!("usage: add_random seed count_groups count_per_group klen vlen");
            if args.len() < 8 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let seed = args[3].parse::<usize>().unwrap();
            let count_groups = args[4].parse::<usize>().unwrap();
            let count_per_group = args[5].parse::<usize>().unwrap();
            let klen = args[6].parse::<usize>().unwrap();
            let vlen = args[7].parse::<usize>().unwrap();
            add_random(name, seed, count_groups, count_per_group, klen, vlen)
        },
        "add_numbers" => {
            println!("usage: add_numbers count start step");
            if args.len() < 6 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let count = args[3].parse::<u64>().unwrap();
            let start = args[4].parse::<u64>().unwrap();
            let step = args[5].parse::<u64>().unwrap();
            if step == 0 {
                return Err(lsm::Error::Misc(String::from("step cannot be 0")));
            }
            add_numbers(name, count, start, step)
        },
        "show_page" => {
            println!("usage: show_page pagenum");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            show_page(name, pgnum)
        },
        "show_leaf_page" => {
            println!("usage: show_leaf_page pagenum");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            show_leaf_page(name, pgnum)
        },
        "show_parent_page" => {
            println!("usage: show_parent_page pagenum");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            show_parent_page(name, pgnum)
        },
        "seek_string" => {
            println!("usage: seek_string key sop");
            if args.len() < 5 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let key = args[3].clone();
            let sop = args[4].clone();
            seek_string(name, key, sop)
        },
        "merge" => {
            println!("usage: merge from_level");
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let from_level = args[3].clone();
            merge(name, from_level)
        },
        "seek_bytes" => {
            println!("usage: seek_bytes sop numbytes b1 b2 b3 ... ");
            if args.len() < 5 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let sop = args[3].clone();
            let count = args[4].parse::<usize>().unwrap();
            if count == 0 {
                return Err(lsm::Error::Misc(String::from("count cannot be 0")));
            }
            if args.len() < 5 + count {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let mut k = Vec::with_capacity(count);
            for i in 0 .. count {
                let b = args[5 + i].parse::<u8>().unwrap();
                k.push(b);
            }
            seek_bytes(name, k.into_boxed_slice(), sop)
        },
        "dump_page" => {
            if args.len() < 4 {
                return Err(lsm::Error::Misc(String::from("too few args")));
            }
            let pgnum = args[3].parse::<u32>().unwrap();
            dump_page(name, pgnum)
        },
        "list_keys" => {
            list_keys(name)
        },
        "list_segments" => {
            list_segments(name)
        },
        "list_free_blocks" => {
            list_free_blocks(name)
        },
        _ => {
            Err(lsm::Error::Misc(String::from("unknown command")))
        },
    }
}

pub fn main() {
    let r = result_main();
    if r.is_err() {
        println!("{:?}", r);
    }
}

<|endoftext|>


<commit_before>use std;
use std::io::Read;
use std::time::Duration;
use std::collections::HashSet;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::net::UdpSocket;
use std::str;

use rustc_serialize::json;
use hyper::Client;

/// Returns a HashSet of hue bridge SocketAddr's
pub fn discover() -> HashSet<IpAddr> {

    let string_list = vec![
        "M-SEARCH * HTTP/1.1",
        "HOST:239.255.255.250:1900",
        "MAN:\"ssdp:discover\"",
        "ST:ssdp:all",
        "MX:1"
            ];
    let joined = string_list.join("\r\n");

    let socket =
        UdpSocket::bind("0.0.0.0:0").unwrap();

    let five_second_timeout = Duration::new(5, 0);
    let _ = socket.set_read_timeout(Some(five_second_timeout));
    socket.send_to(joined.as_bytes(), "239.255.255.250:1900").unwrap();

    let mut bridges = HashSet::new();
    loop {
        let mut buf = [0;255];
        let sockread = match socket.recv_from(&mut buf) {
            Ok(val) => val,
            Err(e) => {
                match e.kind() {
                    // a timeout on unix is considered a WouldBlock
                    std::io::ErrorKind::WouldBlock => break,
                    _ => panic!(e),
                }
            }
        };
        let _ = str::from_utf8(&buf).and_then(|s| {
            // Hue docs say to use "IpBridge" over "hue-bridgeid"
            if s.contains("IpBridge") {
                let bridge = sockread.1.ip();
                bridges.insert(bridge);
            }
            Ok(s)
        });
    }

    bridges
}

/// Hue Bridge
#[derive(Debug)]
pub struct Bridge {
    ip: Ipv4Addr,
}

impl Bridge {
    /// Returns a hue bridge with the given ip
    pub fn new(addr: Ipv4Addr) -> Bridge {
        Bridge {
            ip: addr,
        }
    }

    /// Attempt to register with the hue bridge
    pub fn register(&self, name: &str) {
        #[derive(RustcEncodable, RustcDecodable)]
        struct Devicetype {
            devicetype: String,
        }

        let client = Client::new();
        let url = format!("http://{}/api", self.ip);
        let payload = Devicetype { devicetype: name.to_owned() };
        let body = json::encode(&payload).unwrap();

        // TODO handle errors and return username
        let mut response = client.post(&url).body(&body).send().unwrap();
        let mut buf = String::new();
        response.read_to_string(&mut buf).unwrap();
        println!("{}", buf);
    }
}
<commit_msg>Hue only supports IPv4 bridges so lets do the check up front<commit_after>use std;
use std::io::Read;
use std::time::Duration;
use std::collections::HashSet;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::net::UdpSocket;
use std::str;

use rustc_serialize::json;
use hyper::Client;

/// Returns a HashSet of hue bridge SocketAddr's
pub fn discover() -> HashSet<Ipv4Addr> {

    let string_list = vec![
        "M-SEARCH * HTTP/1.1",
        "HOST:239.255.255.250:1900",
        "MAN:\"ssdp:discover\"",
        "ST:ssdp:all",
        "MX:1"
            ];
    let joined = string_list.join("\r\n");

    let socket =
        UdpSocket::bind("0.0.0.0:0").unwrap();

    let five_second_timeout = Duration::new(5, 0);
    let _ = socket.set_read_timeout(Some(five_second_timeout));
    socket.send_to(joined.as_bytes(), "239.255.255.250:1900").unwrap();

    let mut bridges = HashSet::new();
    loop {
        let mut buf = [0;255];
        let sockread = match socket.recv_from(&mut buf) {
            Ok(val) => val,
            Err(e) => {
                match e.kind() {
                    // a timeout on unix is considered a WouldBlock
                    std::io::ErrorKind::WouldBlock => break,
                    _ => panic!(e),
                }
            }
        };
        let _ = str::from_utf8(&buf).and_then(|s| {
            // Hue docs say to use "IpBridge" over "hue-bridgeid"
            if s.contains("IpBridge") {
                if let IpAddr::V4(addr) = sockread.1.ip() {
                    bridges.insert(addr);
                }
            }
            Ok(s)
        });
    }

    bridges
}

/// Hue Bridge
#[derive(Debug)]
pub struct Bridge {
    ip: Ipv4Addr,
}

impl Bridge {
    /// Returns a hue bridge with the given ip
    pub fn new(addr: Ipv4Addr) -> Bridge {
        Bridge {
            ip: addr,
        }
    }

    /// Attempt to register with the hue bridge
    pub fn register(&self, name: &str) {
        #[derive(RustcEncodable, RustcDecodable)]
        struct Devicetype {
            devicetype: String,
        }

        let client = Client::new();
        let url = format!("http://{}/api", self.ip);
        let payload = Devicetype { devicetype: name.to_owned() };
        let body = json::encode(&payload).unwrap();

        // TODO handle errors and return username
        let mut response = client.post(&url).body(&body).send().unwrap();
        let mut buf = String::new();
        response.read_to_string(&mut buf).unwrap();
        println!("{}", buf);
    }
}
<|endoftext|>


<commit_before>use std::rc::{Rc};
use std::cell::{RefCell};

pub struct Broker<T: Subscriber> {
    subscriptions: Subscriptions<T>,
}

impl<T: Subscriber> Broker<T> {
    pub fn new() -> Self {
        Broker{ subscriptions: Subscriptions::<T>::new() }
    }

    pub fn publish(&mut self, topic: &str, payload: &[u8]) {
        self.subscriptions.publish(topic, payload);
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        for sub in &self.subscriptions.subscribers {
            let sub1 = sub.clone();
            let sub2 = subscriber.clone();
            if sub1.borrow().id() == sub2.borrow().id() {
                sub2.borrow_mut().append_topics(topics);
                return;
            }
        }
        self.subscriptions.subscribers.push(subscriber);
        let last_subscriber = self.subscriptions.subscribers[self.subscriptions.subscribers.len() - 1].clone();
        last_subscriber.borrow_mut().append_topics(topics);
    }
}

pub struct Subscriptions<T: Subscriber> {
    subscribers: Vec<Rc<RefCell<T>>>,
}

impl<T: Subscriber> Subscriptions<T> {
    fn new() -> Self {
        return Subscriptions { subscribers: vec![], }
    }

    fn publish(&mut self, topic: &str, payload: &[u8]) {
        for s in &mut self.subscribers {
            let sub_clone = s.clone();
            let mut s = sub_clone.borrow_mut();

            for t in s.topics() {
                if t == topic {
                    s.new_message(payload);
                }
            }
        }
    }
}

pub trait Subscriber {
    fn new_message(&mut self, bytes: &[u8]);
    fn id(&self) -> i64;
    fn append_topics(&mut self, topics: &[&str]);
    fn topics(&self) -> Vec<String>;
}

#[cfg(test)]
struct TestSubscriber {
    msgs: Vec<Vec<u8>>,
    topics: Vec<String>,
}
#[cfg(test)]
impl TestSubscriber {
    fn new() -> Self {
        return TestSubscriber{msgs: vec![], topics: vec![]}
    }
}

#[cfg(test)]
impl Subscriber for TestSubscriber {

    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }

    fn id(&self) -> i64 {
        self as *const TestSubscriber as i64
    }

    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[test]
fn test_subscribe() {
    let mut broker = Broker::<TestSubscriber>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriber::new()));
    let subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}
<commit_msg>Got rid of id function in Subscriber trait<commit_after>use std::rc::{Rc};
use std::cell::{RefCell};

pub struct Broker<T: Subscriber> {
    subscriptions: Subscriptions<T>,
}

impl<T: Subscriber> Broker<T> {
    pub fn new() -> Self {
        Broker{ subscriptions: Subscriptions::<T>::new() }
    }

    pub fn publish(&mut self, topic: &str, payload: &[u8]) {
        self.subscriptions.publish(topic, payload);
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        for sub in &self.subscriptions.subscribers {
            let sub1 = sub.clone();
            let sub2 = subscriber.clone();
            //this is horrible
            //borrow doesn't return T, it returns Ref<T>
            //since Ref<T> itself doesn't satisfy the template contraint, we need to deref
            //it to get the to the T inside. And since is_same takes borrows, we reapply
            //the ampersand
            if is_same(&*sub1.borrow(), &*sub2.borrow()) {
                sub2.borrow_mut().append_topics(topics);
                return;
            }
        }
        self.subscriptions.subscribers.push(subscriber);
        let last_subscriber = self.subscriptions.subscribers[self.subscriptions.subscribers.len() - 1].clone();
        last_subscriber.borrow_mut().append_topics(topics);
    }
}

//is same identity
fn is_same<T>(lhs: &T, rhs: &T) -> bool {
    return lhs as *const T as i64 == rhs as *const T as i64
}

pub struct Subscriptions<T: Subscriber> {
    subscribers: Vec<Rc<RefCell<T>>>,
}

impl<T: Subscriber> Subscriptions<T> {
    fn new() -> Self {
        return Subscriptions { subscribers: vec![], }
    }

    fn publish(&mut self, topic: &str, payload: &[u8]) {
        for s in &mut self.subscribers {
            let sub_clone = s.clone();
            let mut s = sub_clone.borrow_mut();

            for t in s.topics() {
                if t == topic {
                    s.new_message(payload);
                }
            }
        }
    }
}

pub trait Subscriber {
    fn new_message(&mut self, bytes: &[u8]);
    fn append_topics(&mut self, topics: &[&str]);
    fn topics(&self) -> Vec<String>;
}

#[cfg(test)]
struct TestSubscriber {
    msgs: Vec<Vec<u8>>,
    topics: Vec<String>,
}
#[cfg(test)]
impl TestSubscriber {
    fn new() -> Self {
        return TestSubscriber{msgs: vec![], topics: vec![]}
    }
}

#[cfg(test)]
impl Subscriber for TestSubscriber {

    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }

    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[test]
fn test_subscribe() {
    let mut broker = Broker::<TestSubscriber>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriber::new()));
    let subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::io;
use rustyline::Editor;
use rustyline::error::ReadlineError as RLError;
use combine::{StreamOnce};
use error::*;

#[derive(Debug)]
pub struct Buffer {
    buffer: String,
    prompt: String,
    rl: Editor<()>,
}

impl Buffer {
    fn new<S>(prompt: S) -> Self
    where
        S: Into<String>
    {
        Buffer {
            buffer: String::with_capacity(128),
            prompt: prompt.into(),
            rl: Editor::new(),
        }
    }
}

impl Read for Buffer {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.push_str(&line);
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Interrupted) => return Err(io::ErrorKind::Interrupted.into()),
                Err(RLError::Eof) => return Ok(0),
                Err(RLError::Errno(num)) => return Err(io::Error::from_raw_os_error(num)),
                Err(err) => return Err(io::ErrorKind::Other.into()),
            }
        }

        // if buf > buffer, drain entire buffer
        if buf.len() >= self.buffer.len() {
            let result = buf.write(self.buffer.as_bytes());
            self.buffer.clear();
            result
        } else {
            let bytes = self.buffer.drain(0..buf.len()).collect::<String>();
            buf.write(bytes.into_bytes().as_slice())
        }
    }
}

impl BufRead for Buffer {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {

    }

    fn consume(&mut self, amt: usize) {

    }
}<commit_msg>Separate buffer into readline, buffer (generic over Read)<commit_after>use std::io;
use std::io::prelude::*;
use rustyline;
use rustyline::error::ReadlineError as RLError;
use combine;
use error::*;

pub struct Readline {
    rl: rustyline::Editor<()>,
    prompt: String,
    buffer: Vec<u8>,
}

impl Readline {
    fn new<S>(prompt: &str) -> Self {
        Readline {
            rl: rustyline::Editor::new(),
            prompt: prompt.to_owned(),
            buffer: Vec::with_capacity(128),
        }
    }
}

impl Read for Readline {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.write(line.as_bytes());
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Interrupted) => return Err(io::ErrorKind::Interrupted.into()),
                Err(RLError::Eof) => return Ok(0),
                Err(RLError::Errno(num)) => return Err(io::Error::from_raw_os_error(num)),
                Err(err) => return Err(io::ErrorKind::Other.into()),
            }
        }

        // if buf > buffer, drain entire buffer
        if buf.len() >= self.buffer.len() {
            let result = buf.write(&self.buffer);
            self.buffer.clear();
            result
        } else {
            let drained = self.buffer.drain(0..buf.len()).collect::<Vec<_>>();
            buf.write(&drained)
        }
    }
}

impl io::BufRead for Readline {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        unimplemented!()
    }

    fn consume(&mut self, amt: usize) {
        unimplemented!()
    }
}

#[derive(Clone, Debug)]
pub struct StringStream<R: io::Read> {
    src: R,
    buffer: String,
}

impl<R: io::Read> StringStream<R> {
    fn new(src: R) -> Self {
        StringStream {
            src: src,
            buffer: String::with_capacity(128),
        }
    }
}

impl<R: io::Read> combine::StreamOnce for StringStream<R> {
    type Item = char;
    type Range = str;
    type Position usize;
    fn uncons(&mut self) -> Result<Self::Item, Error<Self::Item, Self::Range>>;
    fn position(&self) -> Self::Position;
}
<|endoftext|>


<commit_before>use std::cmp;
use std::io::{self, Read, BufRead};

pub struct BufReader<R> {
    inner: R,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
}

const INIT_BUFFER_SIZE: usize = 4096;
const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100;

impl<R: Read> BufReader<R> {
    #[inline]
    pub fn new(rdr: R) -> BufReader<R> {
        BufReader::with_capacity(rdr, INIT_BUFFER_SIZE)
    }

    #[inline]
    pub fn with_capacity(rdr: R, cap: usize) -> BufReader<R> {
        let mut buf = Vec::with_capacity(cap);
        unsafe {
            grow_zerofill(&mut buf, cap);
        }
        BufReader {
            inner: rdr,
            buf: buf,
            pos: 0,
            cap: 0,
        }
    }

    #[inline]
    pub fn get_ref(&self) -> &R { &self.inner }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R { &mut self.inner }

    #[inline]
    pub fn get_buf(&self) -> &[u8] {
        if self.pos < self.cap {
            trace!("slicing {:?}", (self.pos, self.cap, self.buf.len()));
            &self.buf[self.pos..self.cap]
        } else {
            &[]
        }
    }

    #[inline]
    pub fn into_inner(self) -> R { self.inner }

    #[inline]
    pub fn read_into_buf(&mut self) -> io::Result<usize> {
        self.maybe_reserve();
        let v = &mut self.buf;
        trace!("read_into_buf pos={}, cap={}", self.cap, v.capacity());
        if self.cap < v.capacity() {
            let nread = try!(self.inner.read(&mut v[self.cap..]));
            self.cap += nread;
            Ok(nread)
        } else {
            trace!("read_into_buf at full capacity");
            Ok(0)
        }
    }

    #[inline]
    fn maybe_reserve(&mut self) {
        let cap = self.buf.capacity();
        if self.cap == cap && cap < MAX_BUFFER_SIZE {
            self.buf.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap);
            let new = self.buf.capacity() - self.buf.len();
            trace!("reserved {}", new);
            unsafe { grow_zerofill(&mut self.buf, new) }
        }
    }
}

#[inline]
unsafe fn grow_zerofill(buf: &mut Vec<u8>, additional: usize) {
    use std::ptr;
    let len = buf.len();
    buf.set_len(len + additional);
    ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len());
}

impl<R: Read> Read for BufReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.cap == self.pos && buf.len() >= self.buf.len() {
            return self.inner.read(buf);
        }
        let nread = {
           let mut rem = try!(self.fill_buf());
           try!(rem.read(buf))
        };
        self.consume(nread);
        Ok(nread)
    }
}

impl<R: Read> BufRead for BufReader<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.pos == self.cap {
            self.cap = try!(self.inner.read(&mut self.buf));
            self.pos = 0;
        }
        Ok(&self.buf[self.pos..self.cap])
    }

    #[inline]
    fn consume(&mut self, amt: usize) {
        self.pos = cmp::min(self.pos + amt, self.cap);
        if self.pos == self.cap {
            self.pos = 0;
            self.cap = 0;
        }
    }
}

#[cfg(test)]
mod tests {

    use std::io::{self, Read, BufRead};
    use super::BufReader;

    struct SlowRead(u8);

    impl Read for SlowRead {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            let state = self.0;
            self.0 += 1;
            (&match state % 3 {
                0 => b"foo",
                1 => b"bar",
                _ => b"baz",
            }[..]).read(buf)
        }
    }

    #[test]
    fn test_consume_and_get_buf() {
        let mut rdr = BufReader::new(SlowRead(0));
        rdr.read_into_buf().unwrap();
        rdr.consume(1);
        assert_eq!(rdr.get_buf(), b"oo");
        rdr.read_into_buf().unwrap();
        rdr.read_into_buf().unwrap();
        assert_eq!(rdr.get_buf(), b"oobarbaz");
        rdr.consume(5);
        assert_eq!(rdr.get_buf(), b"baz");
        rdr.consume(3);
        assert_eq!(rdr.get_buf(), b"");
        assert_eq!(rdr.pos, 0);
        assert_eq!(rdr.cap, 0);
    }
}
<commit_msg>refactor(buffer): use vec macro in constructor, change trace logs<commit_after>use std::cmp;
use std::io::{self, Read, BufRead};

pub struct BufReader<R> {
    inner: R,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
}

const INIT_BUFFER_SIZE: usize = 4096;
const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100;

impl<R: Read> BufReader<R> {
    #[inline]
    pub fn new(rdr: R) -> BufReader<R> {
        BufReader::with_capacity(rdr, INIT_BUFFER_SIZE)
    }

    #[inline]
    pub fn with_capacity(rdr: R, cap: usize) -> BufReader<R> {
        BufReader {
            inner: rdr,
            buf: vec![0; cap],
            pos: 0,
            cap: 0,
        }
    }

    #[inline]
    pub fn get_ref(&self) -> &R { &self.inner }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R { &mut self.inner }

    #[inline]
    pub fn get_buf(&self) -> &[u8] {
        if self.pos < self.cap {
            trace!("get_buf [u8; {}][{}..{}]", self.buf.len(), self.pos, self.cap);
            &self.buf[self.pos..self.cap]
        } else {
            trace!("get_buf []");
            &[]
        }
    }

    #[inline]
    pub fn into_inner(self) -> R { self.inner }

    #[inline]
    pub fn read_into_buf(&mut self) -> io::Result<usize> {
        self.maybe_reserve();
        let v = &mut self.buf;
        trace!("read_into_buf buf[{}..{}]", self.cap, v.len());
        if self.cap < v.capacity() {
            let nread = try!(self.inner.read(&mut v[self.cap..]));
            self.cap += nread;
            Ok(nread)
        } else {
            trace!("read_into_buf at full capacity");
            Ok(0)
        }
    }

    #[inline]
    fn maybe_reserve(&mut self) {
        let cap = self.buf.capacity();
        if self.cap == cap && cap < MAX_BUFFER_SIZE {
            self.buf.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap);
            let new = self.buf.capacity() - self.buf.len();
            trace!("reserved {}", new);
            unsafe { grow_zerofill(&mut self.buf, new) }
        }
    }
}

#[inline]
unsafe fn grow_zerofill(buf: &mut Vec<u8>, additional: usize) {
    use std::ptr;
    let len = buf.len();
    buf.set_len(len + additional);
    ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len());
}

impl<R: Read> Read for BufReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.cap == self.pos && buf.len() >= self.buf.len() {
            return self.inner.read(buf);
        }
        let nread = {
           let mut rem = try!(self.fill_buf());
           try!(rem.read(buf))
        };
        self.consume(nread);
        Ok(nread)
    }
}

impl<R: Read> BufRead for BufReader<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.pos == self.cap {
            self.cap = try!(self.inner.read(&mut self.buf));
            self.pos = 0;
        }
        Ok(&self.buf[self.pos..self.cap])
    }

    #[inline]
    fn consume(&mut self, amt: usize) {
        self.pos = cmp::min(self.pos + amt, self.cap);
        if self.pos == self.cap {
            self.pos = 0;
            self.cap = 0;
        }
    }
}

#[cfg(test)]
mod tests {

    use std::io::{self, Read, BufRead};
    use super::BufReader;

    struct SlowRead(u8);

    impl Read for SlowRead {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            let state = self.0;
            self.0 += 1;
            (&match state % 3 {
                0 => b"foo",
                1 => b"bar",
                _ => b"baz",
            }[..]).read(buf)
        }
    }

    #[test]
    fn test_consume_and_get_buf() {
        let mut rdr = BufReader::new(SlowRead(0));
        rdr.read_into_buf().unwrap();
        rdr.consume(1);
        assert_eq!(rdr.get_buf(), b"oo");
        rdr.read_into_buf().unwrap();
        rdr.read_into_buf().unwrap();
        assert_eq!(rdr.get_buf(), b"oobarbaz");
        rdr.consume(5);
        assert_eq!(rdr.get_buf(), b"baz");
        rdr.consume(3);
        assert_eq!(rdr.get_buf(), b"");
        assert_eq!(rdr.pos, 0);
        assert_eq!(rdr.cap, 0);
    }
}
<|endoftext|>


<commit_before>use std::str;
use std::cmp;
use std::io::{Read, Result};

#[allow(dead_code)]
pub struct IBuffer {
    buf: Vec<u8>,
    offset: usize,
}

#[allow(dead_code)]
impl IBuffer {
    pub fn from_str(s: &str) -> IBuffer {
        IBuffer {
            buf: s.as_bytes().to_vec(),
            offset: 0,
        }
    }
}

impl Read for IBuffer {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let len = cmp::min(buf.len(), self.buf.len() - self.offset);
        buf[0..len].copy_from_slice(&self.buf[self.offset..(self.offset + len)]);
        self.offset += len;
        Ok(len)
    }
}

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

    #[test]
    fn read_from_string() {
        let mut in_buf = IBuffer::from_str("Hello world");

        let mut buf = [0u8; 6];

        let res1 = in_buf.read(&mut buf);
        assert!(match res1 {
                    Ok(6) => true,
                    _ => false,
                });
        assert_eq!([b'H', b'e', b'l', b'l', b'o', b' '],
                   buf);
    }

    #[test]
    fn read_from_string_overflow() {
        let mut in_buf = IBuffer::from_str("Hello world");

        let mut buf = [0u8; 6];

        let _ = in_buf.read(&mut buf);
        let res = in_buf.read(&mut buf);
        assert!(match res {
                    Ok(5) => true,
                    _ => false,
                });
        assert_eq!([b'w', b'o', b'r', b'l', b'd'],
                   buf[0..5]);
    }

    #[test]
    fn read_from_string_eof() {
        let mut in_buf = IBuffer::from_str("Hello world");

        let mut buf = [0u8; 6];

        let _ = in_buf.read(&mut buf);
        let _ = in_buf.read(&mut buf);
        let res = in_buf.read(&mut buf);
        assert!(match res {
                    Ok(0) => true,
                    _ => false,
                });
    }

}
<commit_msg>rust format<commit_after>use std::str;
use std::cmp;
use std::io::{Read, Result};

#[allow(dead_code)]
pub struct IBuffer {
    buf: Vec<u8>,
    offset: usize,
}

#[allow(dead_code)]
impl IBuffer {
    pub fn from_str(s: &str) -> IBuffer {
        IBuffer {
            buf: s.as_bytes().to_vec(),
            offset: 0,
        }
    }
}

impl Read for IBuffer {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let len = cmp::min(buf.len(), self.buf.len() - self.offset);
        buf[0..len].copy_from_slice(&self.buf[self.offset..(self.offset + len)]);
        self.offset += len;
        Ok(len)
    }
}

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

    #[test]
    fn read_from_string() {
        let mut in_buf = IBuffer::from_str("Hello world");

        let mut buf = [0u8; 6];

        let res1 = in_buf.read(&mut buf);
        assert!(match res1 {
                    Ok(6) => true,
                    _ => false,
                });
        assert_eq!([b'H', b'e', b'l', b'l', b'o', b' '], buf);
    }

    #[test]
    fn read_from_string_overflow() {
        let mut in_buf = IBuffer::from_str("Hello world");

        let mut buf = [0u8; 6];

        let _ = in_buf.read(&mut buf);
        let res = in_buf.read(&mut buf);
        assert!(match res {
                    Ok(5) => true,
                    _ => false,
                });
        assert_eq!([b'w', b'o', b'r', b'l', b'd'], buf[0..5]);
    }

    #[test]
    fn read_from_string_eof() {
        let mut in_buf = IBuffer::from_str("Hello world");

        let mut buf = [0u8; 6];

        let _ = in_buf.read(&mut buf);
        let _ = in_buf.read(&mut buf);
        let res = in_buf.read(&mut buf);
        assert!(match res {
                    Ok(0) => true,
                    _ => false,
                });
    }

}
<|endoftext|>


<commit_before>use std::ops::{Add, Sub, Div};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::collections::VecDeque;
use std::iter::repeat;
use std::fmt::{Debug, Formatter, Error};
use std::cmp::Ordering::Equal;
use std::f32;

#[derive(Clone,)]
pub struct WordVec {
    pub word: String,
    pub count: u64,
    vec: Vec<f32>,
}

impl PartialEq for WordVec {
    fn eq(&self, other: &WordVec) -> bool {
        self.vec == other.vec
    }
}

impl Debug for WordVec {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        try!(self.word.fmt(fmt));
        try!(fmt.write_str(": "));
        try!(self.count.fmt(fmt));
        try!(fmt.write_str(", "));
        self.vec.iter().take(5).collect::<Vec<_>>().fmt(fmt)
    }
}

impl WordVec {

    fn dot_prod(&self, other: &WordVec) -> f32 {
        self.vec.iter()
                .zip(other.vec.iter())
                .map(|(x, y)| x * y)
                .fold(0.0, |x, y| x + y)
    }

    fn magnitude(&self) -> f32 {
        self.vec.iter()
                .map(|x| x * x)
                .fold(0.0, |x, y| x + y)
                .sqrt()
    }

    pub fn distance(&self, other: &WordVec) -> f32 {
        self.vec.iter().zip(other.vec.iter())
                .map(|(x, y)| (x * x - y * y).abs().sqrt())
                .fold(0.0, |x, y| x + y)
    }

    fn new(word: String, num_words: usize) -> WordVec {
        WordVec {
            word: word,
            count: 0,
            vec: repeat(0.0).take(num_words).collect(),
        }
    }

    fn normalize(&mut self) {
        let count = self.count as f32;
        for f in &mut self.vec {
            *f /= count;
        }
    }

    fn inc(&mut self, i: usize) {
        self.vec[i] += 1.0;
    }

}

impl Div<i32> for WordVec {
    type Output = WordVec;

    fn div(mut self, i: i32) -> WordVec {
        for f in &mut self.vec {
            *f /= i as f32;
        }

        self
    }
}

impl<'a> Add<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn add(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} + {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..newvec.vec.len() {
            newvec.vec[i] += other.vec[i];
        }

        newvec
    }
}

impl<'a> Sub<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn sub(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} - {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..(newvec.vec.len()) {
            newvec.vec[i] -= other.vec[i];
        }

        newvec
    }
}

pub struct LanguageModel {
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct LanguageModelBuilder {
    window_width: usize,
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct WordAcceptor<'a> {
    window: VecDeque<(usize, String)>,
    builder: &'a mut LanguageModelBuilder,
}

impl LanguageModelBuilder {
    pub fn new(window_radius: usize, words: Vec<String>) -> LanguageModelBuilder {

        let word_vecs = words.iter()
                             .map(|s| WordVec::new(s.clone(), words.len()))
                             .collect();

        let words = HashMap::from_iter(
                        words.into_iter()
                             .enumerate()
                             .map(|(a, b)| (b, a)));

        LanguageModelBuilder {
            window_width: 2 * window_radius + 1,
            words: words,
            word_vecs: word_vecs,
        }
    }

    pub fn build(mut self) -> LanguageModel {
        for vec in self.word_vecs.iter_mut() {
            vec.normalize();
        }

        LanguageModel {
            words: self.words,
            word_vecs: self.word_vecs,
        }
    }

    pub fn new_file<'a>(&'a mut self) -> WordAcceptor<'a> {
        WordAcceptor {
            window: VecDeque::new(),
            builder: self,
        }
    }

    fn get_vec(&mut self, word: &str) -> &mut WordVec {
        let i = self.words[word];
        &mut self.word_vecs[i]
    }

    fn add_word(&mut self, from: usize, to: usize) {
        self.word_vecs[from].inc(to);
    }

    fn word_seen(&mut self, word: &str) {
        self.get_vec(word).count += 1;
    }
}

impl<'a> WordAcceptor<'a> {
    pub fn add_word(&mut self, word: String) {
        let ww = self.builder.window_width;
        let idx_opt = self.builder.words.get(&word).map(|w| *w);

        if let Some(next_idx) = idx_opt {
            self.window.push_back((next_idx, word));
            if self.window.len() > ww {
                self.window.pop_front();
            }

            if self.window.len() == ww {
                let (center_idx, ref center_word) = self.window[ww / 2];
                self.builder.word_seen(center_word);
                for i in (0..ww).filter(|a| *a != ww / 2) {
                    self.builder.add_word(center_idx, self.window[i].0);
                }
            }
        }
    }
}

impl LanguageModel {
    pub fn get(&self, word: &str) -> Option<&WordVec> {
        self.words.get(word).map(|i| &self.word_vecs[*i])
    }

    pub fn nearest_words(&self, word: &WordVec) -> Vec<&WordVec> {
        let mut vec_refs = self.word_vecs.iter().filter_map(|w| {
            let dist = w.distance(word);
            if w.word != word.word {
                Some((dist, w))
            } else {
                None
            }
        }).collect::<Vec<_>>();
        vec_refs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or_else(|| {
            println!("bad vector: {:?}, {:?}", a, b);
            Equal
        }));
        vec_refs.into_iter().map(|x| x.1).collect()
    }
}<commit_msg>Make model use every work fed to an acceptor rather than ignoring the ends.<commit_after>use std::ops::{Add, Sub, Div};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::iter::repeat;
use std::fmt::{Debug, Formatter, Error};
use std::cmp::Ordering::Equal;
use std::ops::Drop;
use std::cmp;

#[derive(Clone,)]
pub struct WordVec {
    pub word: String,
    pub count: u64,
    vec: Vec<f32>,
}

impl PartialEq for WordVec {
    fn eq(&self, other: &WordVec) -> bool {
        self.vec == other.vec
    }
}

impl Debug for WordVec {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        try!(self.word.fmt(fmt));
        try!(fmt.write_str(": "));
        try!(self.count.fmt(fmt));
        try!(fmt.write_str(", "));
        self.vec.iter().take(5).collect::<Vec<_>>().fmt(fmt)
    }
}

impl WordVec {

    fn dot_prod(&self, other: &WordVec) -> f32 {
        self.vec.iter()
                .zip(other.vec.iter())
                .map(|(x, y)| x * y)
                .fold(0.0, |x, y| x + y)
    }

    fn magnitude(&self) -> f32 {
        self.vec.iter()
                .map(|x| x * x)
                .fold(0.0, |x, y| x + y)
                .sqrt()
    }

    pub fn distance(&self, other: &WordVec) -> f32 {
        self.vec.iter().zip(other.vec.iter())
                .map(|(x, y)| (x * x - y * y).abs().sqrt())
                .fold(0.0, |x, y| x + y)
    }

    fn new(word: String, num_words: usize) -> WordVec {
        WordVec {
            word: word,
            count: 0,
            vec: repeat(0.0).take(num_words).collect(),
        }
    }

    fn normalize(&mut self) {
        let count = self.count as f32;
        for f in &mut self.vec {
            *f /= count;
        }
    }

    fn inc(&mut self, i: usize) {
        self.vec[i] += 1.0;
    }

}

impl Div<i32> for WordVec {
    type Output = WordVec;

    fn div(mut self, i: i32) -> WordVec {
        for f in &mut self.vec {
            *f /= i as f32;
        }

        self
    }
}

impl<'a> Add<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn add(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} + {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..newvec.vec.len() {
            newvec.vec[i] += other.vec[i];
        }

        newvec
    }
}

impl<'a> Sub<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn sub(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} - {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..(newvec.vec.len()) {
            newvec.vec[i] -= other.vec[i];
        }

        newvec
    }
}

pub struct LanguageModel {
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct LanguageModelBuilder {
    window_radius: usize,
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct WordAcceptor<'a> {
    sentence: Vec<Option<usize>>,
    builder: &'a mut LanguageModelBuilder,
}

impl LanguageModelBuilder {
    pub fn new(window_radius: usize, words: Vec<String>) -> LanguageModelBuilder {

        let word_vecs = words.iter()
                             .map(|s| WordVec::new(s.clone(), words.len()))
                             .collect();

        let words = HashMap::from_iter(
                        words.into_iter()
                             .enumerate()
                             .map(|(a, b)| (b, a)));

        LanguageModelBuilder {
            window_radius: window_radius,
            words: words,
            word_vecs: word_vecs,
        }
    }

    pub fn build(mut self) -> LanguageModel {
        for vec in self.word_vecs.iter_mut() {
            vec.normalize();
        }

        LanguageModel {
            words: self.words,
            word_vecs: self.word_vecs,
        }
    }

    pub fn new_sentence<'a>(&'a mut self) -> WordAcceptor<'a> {
        WordAcceptor {
            sentence: Vec::new(),
            builder: self,
        }
    }

    fn get_vec(&mut self, word: &str) -> &mut WordVec {
        let i = self.words[word];
        &mut self.word_vecs[i]
    }

    fn add_word(&mut self, from: usize, to: usize) {
        self.word_vecs[from].inc(to);
    }

    fn word_seen(&mut self, word: &str) {
        self.get_vec(word).count += 1;
    }
}

impl<'a> WordAcceptor<'a> {
    pub fn add_word(&mut self, word: &str) {
        let idx_opt = self.builder.words.get(word).map(|w| *w);

        if let Some(next_idx) = idx_opt {
            self.builder.word_vecs[next_idx].count += 1;
            self.sentence.push(Some(next_idx));
        }
        else {
            self.sentence.push(None);
        }
    }
}

impl<'a> Drop for WordAcceptor<'a> {
    fn drop(&mut self) {
        let WordAcceptor {
            ref sentence,
            ref mut builder
        } = *self;

        println!("{:?}", sentence);

        for (i, &word_idx) in self.sentence.iter().enumerate() {
            let start = cmp::max(0, i as isize- builder.window_radius as isize) as usize;
            let end = cmp::min(sentence.len(), i + builder.window_radius + 1);
            for &j in &sentence[start..end] {
                match (word_idx, j) {
                    (Some(w), Some(j)) if w != j =>
                        builder.word_vecs[w].inc(j),
                    _ => {},
                }
            }
        }

        drop(sentence); // not actually sure if this is necessary
    }
}

impl LanguageModel {
    pub fn get(&self, word: &str) -> Option<&WordVec> {
        self.words.get(word).map(|i| &self.word_vecs[*i])
    }

    pub fn nearest_words(&self, word: &WordVec) -> Vec<&WordVec> {
        let mut vec_refs = self.word_vecs.iter().filter_map(|w| {
            let dist = w.distance(word);
            if w.word != word.word {
                Some((dist, w))
            } else {
                None
            }
        }).collect::<Vec<_>>();
        vec_refs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or_else(|| {
            println!("bad vector: {:?}, {:?}", a, b);
            Equal
        }));
        vec_refs.into_iter().map(|x| x.1).collect()
    }
}

#[cfg(test)]
mod test {
    use super::{WordVec, LanguageModel, LanguageModelBuilder};

    #[test]
    fn test_accept_sentences() {
        let words = "foo bar baz blort".words().map(|w| w.to_string()).collect::<Vec<_>>();
        let mut builder = LanguageModelBuilder::new(1, words);

        let input = "x foo bar baz x x x x x x x blort".words();

        {
            let mut acc = builder.new_sentence();
            for word in input {
                acc.add_word(word);
            }
        }

        let model = builder.build();
        let foo = model.get("foo").unwrap();
        let baz = model.get("baz").unwrap();
        let bar = model.get("bar").unwrap();
        let blort = model.get("blort").unwrap();

        assert!(foo.distance(baz) < foo.distance(blort));
        assert!(foo.distance(bar) == bar.distance(baz));

    }
}<|endoftext|>