<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Dictionaries of key-value pairs.

use base::{Boolean, CFAllocatorRef, CFIndex, CFIndexConvertible, CFRelease, CFRetain};
use base::{CFType, CFTypeID, CFTypeRef, TCFType, kCFAllocatorDefault};

use libc::c_void;
use std::mem;
use std::ptr;
use std::vec;

pub type CFDictionaryApplierFunction = *const u8;
pub type CFDictionaryCopyDescriptionCallBack = *const u8;
pub type CFDictionaryEqualCallBack = *const u8;
pub type CFDictionaryHashCallBack = *const u8;
pub type CFDictionaryReleaseCallBack = *const u8;
pub type CFDictionaryRetainCallBack = *const u8;

#[allow(dead_code)]
#[repr(C)]
#[deriving(Copy)]
pub struct CFDictionaryKeyCallBacks {
    version: CFIndex,
    retain: CFDictionaryRetainCallBack,
    release: CFDictionaryReleaseCallBack,
    copyDescription: CFDictionaryCopyDescriptionCallBack,
    equal: CFDictionaryEqualCallBack,
    hash: CFDictionaryHashCallBack
}

#[allow(dead_code)]
#[repr(C)]
#[deriving(Copy)]
pub struct CFDictionaryValueCallBacks {
    version: CFIndex,
    retain: CFDictionaryRetainCallBack,
    release: CFDictionaryReleaseCallBack,
    copyDescription: CFDictionaryCopyDescriptionCallBack,
    equal: CFDictionaryEqualCallBack
}

#[repr(C)]
struct __CFDictionary;

pub type CFDictionaryRef = *const __CFDictionary;

/// An immutable dictionary of key-value pairs.
///
/// FIXME(pcwalton): Should be a newtype struct, but that panics due to a Rust compiler bug.
pub struct CFDictionary {
    obj: CFDictionaryRef,
}

impl Drop for CFDictionary {
    fn drop(&mut self) {
        unsafe {
            CFRelease(self.as_CFTypeRef())
        }
    }
}

impl TCFType<CFDictionaryRef> for CFDictionary {
    #[inline]
    fn as_concrete_TypeRef(&self) -> CFDictionaryRef {
        self.obj
    }

    #[inline]
    unsafe fn wrap_under_get_rule(reference: CFDictionaryRef) -> CFDictionary {
        let reference: CFDictionaryRef = mem::transmute(CFRetain(mem::transmute(reference)));
        TCFType::wrap_under_create_rule(reference)
    }

    #[inline]
    fn as_CFTypeRef(&self) -> CFTypeRef {
        unsafe {
            mem::transmute(self.as_concrete_TypeRef())
        }
    }

    unsafe fn wrap_under_create_rule(obj: CFDictionaryRef) -> CFDictionary {
        CFDictionary {
            obj: obj,
        }
    }

    #[inline]
    fn type_id(_: Option<CFDictionary>) -> CFTypeID {
        unsafe {
            CFDictionaryGetTypeID()
        }
    }
}

impl CFDictionary {
    pub fn from_CFType_pairs(pairs: &[(CFType, CFType)]) -> CFDictionary {
        let (keys, values): (Vec<CFTypeRef>,Vec<CFTypeRef>) =
            pairs.iter()
            .map(|&(ref key, ref value)| (key.as_CFTypeRef(), value.as_CFTypeRef()))
            .unzip();

        unsafe {
            let dictionary_ref = CFDictionaryCreate(kCFAllocatorDefault,
                                                    mem::transmute(keys.as_ptr()),
                                                    mem::transmute(values.as_ptr()),
                                                    keys.len().to_CFIndex(),
                                                    &kCFTypeDictionaryKeyCallBacks,
                                                    &kCFTypeDictionaryValueCallBacks);
            TCFType::wrap_under_create_rule(dictionary_ref)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        unsafe {
            CFDictionaryGetCount(self.obj) as usize
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[inline]
    pub fn contains_key(&self, key: *const c_void) -> bool {
        unsafe {
            CFDictionaryContainsKey(self.obj, key) != 0
        }
    }

    #[inline]
    pub fn find(&self, key: *const c_void) -> Option<*const c_void> {
        unsafe {
            let mut value: *const c_void = ptr::null();
            if CFDictionaryGetValueIfPresent(self.obj, key, &mut value) != 0 {
                Some(value)
            } else {
                None
            }
        }
    }

    #[inline]
    pub fn get(&self, key: *const c_void) -> *const c_void {
        let value = self.find(key);
        if value.is_none() {
            panic!("No entry found for key [FIXME]"); // FIXME Not sure how to convert key to be printed
        }
        value.unwrap()
    }

    /// A convenience function to retrieve `CFType` instances.
    #[inline]
    pub unsafe fn get_CFType(&self, key: *const c_void) -> CFType {
        let value: CFTypeRef = mem::transmute(self.get(key));
        TCFType::wrap_under_get_rule(value)
    }
}

#[link(name = "CoreFoundation", kind = "framework")]
extern {
    /*
     * CFDictionary.h
     */

    static kCFTypeDictionaryKeyCallBacks: CFDictionaryKeyCallBacks;
    static kCFTypeDictionaryValueCallBacks: CFDictionaryValueCallBacks;

    fn CFDictionaryContainsKey(theDict: CFDictionaryRef, key: *const c_void) -> Boolean;
    fn CFDictionaryCreate(allocator: CFAllocatorRef, keys: *const *const c_void, values: *const *const c_void,
                          numValues: CFIndex, keyCallBacks: *const CFDictionaryKeyCallBacks,
                          valueCallBacks: *const CFDictionaryValueCallBacks)
                       -> CFDictionaryRef;
    fn CFDictionaryGetCount(theDict: CFDictionaryRef) -> CFIndex;
    fn CFDictionaryGetTypeID() -> CFTypeID;
    fn CFDictionaryGetValueIfPresent(theDict: CFDictionaryRef, key: *const c_void, value: *mut *const c_void)
                                     -> Boolean;
}
<commit_msg>Show key during panic<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Dictionaries of key-value pairs.

use base::{Boolean, CFAllocatorRef, CFIndex, CFIndexConvertible, CFRelease, CFRetain};
use base::{CFType, CFTypeID, CFTypeRef, TCFType, kCFAllocatorDefault};

use libc::c_void;
use std::mem;
use std::ptr;
use std::vec;

pub type CFDictionaryApplierFunction = *const u8;
pub type CFDictionaryCopyDescriptionCallBack = *const u8;
pub type CFDictionaryEqualCallBack = *const u8;
pub type CFDictionaryHashCallBack = *const u8;
pub type CFDictionaryReleaseCallBack = *const u8;
pub type CFDictionaryRetainCallBack = *const u8;

#[allow(dead_code)]
#[repr(C)]
#[deriving(Copy)]
pub struct CFDictionaryKeyCallBacks {
    version: CFIndex,
    retain: CFDictionaryRetainCallBack,
    release: CFDictionaryReleaseCallBack,
    copyDescription: CFDictionaryCopyDescriptionCallBack,
    equal: CFDictionaryEqualCallBack,
    hash: CFDictionaryHashCallBack
}

#[allow(dead_code)]
#[repr(C)]
#[deriving(Copy)]
pub struct CFDictionaryValueCallBacks {
    version: CFIndex,
    retain: CFDictionaryRetainCallBack,
    release: CFDictionaryReleaseCallBack,
    copyDescription: CFDictionaryCopyDescriptionCallBack,
    equal: CFDictionaryEqualCallBack
}

#[repr(C)]
struct __CFDictionary;

pub type CFDictionaryRef = *const __CFDictionary;

/// An immutable dictionary of key-value pairs.
///
/// FIXME(pcwalton): Should be a newtype struct, but that panics due to a Rust compiler bug.
pub struct CFDictionary {
    obj: CFDictionaryRef,
}

impl Drop for CFDictionary {
    fn drop(&mut self) {
        unsafe {
            CFRelease(self.as_CFTypeRef())
        }
    }
}

impl TCFType<CFDictionaryRef> for CFDictionary {
    #[inline]
    fn as_concrete_TypeRef(&self) -> CFDictionaryRef {
        self.obj
    }

    #[inline]
    unsafe fn wrap_under_get_rule(reference: CFDictionaryRef) -> CFDictionary {
        let reference: CFDictionaryRef = mem::transmute(CFRetain(mem::transmute(reference)));
        TCFType::wrap_under_create_rule(reference)
    }

    #[inline]
    fn as_CFTypeRef(&self) -> CFTypeRef {
        unsafe {
            mem::transmute(self.as_concrete_TypeRef())
        }
    }

    unsafe fn wrap_under_create_rule(obj: CFDictionaryRef) -> CFDictionary {
        CFDictionary {
            obj: obj,
        }
    }

    #[inline]
    fn type_id(_: Option<CFDictionary>) -> CFTypeID {
        unsafe {
            CFDictionaryGetTypeID()
        }
    }
}

impl CFDictionary {
    pub fn from_CFType_pairs(pairs: &[(CFType, CFType)]) -> CFDictionary {
        let (keys, values): (Vec<CFTypeRef>,Vec<CFTypeRef>) =
            pairs.iter()
            .map(|&(ref key, ref value)| (key.as_CFTypeRef(), value.as_CFTypeRef()))
            .unzip();

        unsafe {
            let dictionary_ref = CFDictionaryCreate(kCFAllocatorDefault,
                                                    mem::transmute(keys.as_ptr()),
                                                    mem::transmute(values.as_ptr()),
                                                    keys.len().to_CFIndex(),
                                                    &kCFTypeDictionaryKeyCallBacks,
                                                    &kCFTypeDictionaryValueCallBacks);
            TCFType::wrap_under_create_rule(dictionary_ref)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        unsafe {
            CFDictionaryGetCount(self.obj) as usize
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[inline]
    pub fn contains_key(&self, key: *const c_void) -> bool {
        unsafe {
            CFDictionaryContainsKey(self.obj, key) != 0
        }
    }

    #[inline]
    pub fn find(&self, key: *const c_void) -> Option<*const c_void> {
        unsafe {
            let mut value: *const c_void = ptr::null();
            if CFDictionaryGetValueIfPresent(self.obj, key, &mut value) != 0 {
                Some(value)
            } else {
                None
            }
        }
    }

    #[inline]
    pub fn get(&self, key: *const c_void) -> *const c_void {
        let value = self.find(key);
        if value.is_none() {
            panic!("No entry found for key {:p}", key);
        }
        value.unwrap()
    }

    /// A convenience function to retrieve `CFType` instances.
    #[inline]
    pub unsafe fn get_CFType(&self, key: *const c_void) -> CFType {
        let value: CFTypeRef = mem::transmute(self.get(key));
        TCFType::wrap_under_get_rule(value)
    }
}

#[link(name = "CoreFoundation", kind = "framework")]
extern {
    /*
     * CFDictionary.h
     */

    static kCFTypeDictionaryKeyCallBacks: CFDictionaryKeyCallBacks;
    static kCFTypeDictionaryValueCallBacks: CFDictionaryValueCallBacks;

    fn CFDictionaryContainsKey(theDict: CFDictionaryRef, key: *const c_void) -> Boolean;
    fn CFDictionaryCreate(allocator: CFAllocatorRef, keys: *const *const c_void, values: *const *const c_void,
                          numValues: CFIndex, keyCallBacks: *const CFDictionaryKeyCallBacks,
                          valueCallBacks: *const CFDictionaryValueCallBacks)
                       -> CFDictionaryRef;
    fn CFDictionaryGetCount(theDict: CFDictionaryRef) -> CFIndex;
    fn CFDictionaryGetTypeID() -> CFTypeID;
    fn CFDictionaryGetValueIfPresent(theDict: CFDictionaryRef, key: *const c_void, value: *mut *const c_void)
                                     -> Boolean;
}
<|endoftext|>


<commit_before>use futures::future::{Future,FutureResult};
use futures::future;
use futures::{Stream,Sink};
use std::collections::vec_deque::{VecDeque};
use std::default::Default;
use std::net::{SocketAddr};
use std::path::Path;
use std::time::Duration;
use tokio_core::net::{TcpStream};
use tokio_core::reactor;
use tokio_io::codec::{Framed};
use tokio_io::{AsyncRead};

use datastore::DataStore;
use error::{Error,Result};
use fillable::*;
use manifest::{ManifestWithFile};
use metainfo::{MetaInfo,InfoHash};
use peer_protocol::{PeerID,Message,BitTorrentPeerCodec};
use peer_protocol;
use tracker::{TrackerClient};
use util::{tcp_connect2,BxFuture,FutureEnhanced,VecDequeStream};

type PeerStream = Framed<TcpStream, BitTorrentPeerCodec>;

struct DownloaderState {
    info: MetaInfo,
    datastore: DataStore,
    manifest: ManifestWithFile,
    peer_state: PeerState,
    blah_state: BlahState,
}

type LoopState = (PeerStream, DownloaderState);

pub fn start<P: AsRef<Path>>(
    info: MetaInfo,
    peer_id: PeerID,
    store_path: P,
    manifest_path: P)
    -> Result<()>
{
    let mut core = reactor::Core::new()?;
    let handle = core.handle();

    let datastore = DataStore::create_or_open(&info, store_path)?;
    let manifest = ManifestWithFile::load_or_new(info.clone(), manifest_path)?;
    let mut tc = TrackerClient::new(info.clone(), peer_id.clone())?;

    let tracker_res = tc.easy_start()?;
    println!("tracker res: {:#?}", tracker_res);
    if let Some(reason) = tracker_res.failure_reason {
        return Err(Error::new_str(&format!("tracker failed: {}", reason)));
    }

    if tracker_res.peers.is_empty() {
        return Err(Error::new_str("tracker returned no peers"));
    }

    let peer = tracker_res.peers[0].clone();

    let num_pieces = info.num_pieces();
    let info_hash = info.info_hash.clone();

    let dstate = DownloaderState {
        info: info,
        datastore: datastore,
        manifest: manifest,
        peer_state: PeerState::new(num_pieces),
        blah_state: BlahState::default(),
    };

    let future_root = connect_peer(peer.address, info_hash, peer_id, &handle)
        .and_then(|(stream, remote_peer_id)| {
            let dstate = dstate;

            main_loop(stream, dstate)

                // .map_err(|err, _s| err)
                // .map(|(_, _s)| Ok(()))

            // Below lies ye ol' blocking implementation
            // for a single stream.

            // let mut state = PeerState::new(info.num_pieces());
            // let mut s = BlahState::default();

            // loop {
            //     let m = peer_protocol::read_message(&mut stream)?;
            //     println!("message: {}", m.summarize());
            //     match m {
            //         Message::Choke =>         state.peer_choking = true,
            //         Message::Unchoke =>       state.peer_choking = false,
            //         Message::Interested =>    state.peer_interested = true,
            //         Message::NotInterested => state.peer_interested = false,
            //         Message::Bitfield { bits } => {
            //             if bits.len() < info.num_pieces() {
            //                 println!("{}", Error::new_str(&format!("bitfield has less bits {} than pieces {}",
            //                                                 bits.len(), info.num_pieces())));
            //             }
            //             let mut i_start = 0;
            //             let mut in_interval = false;
            //             for b in 0..info.num_pieces() {
            //                 if bits[b] && !in_interval {
            //                     i_start = b;
            //                     in_interval = true;
            //                 } else if !bits[b] && in_interval {
            //                     state.has.add(i_start as u32, b as u32);
            //                     in_interval = false;
            //                 }
            //             }
            //             if in_interval {
            //                 state.has.add(i_start as u32, info.num_pieces() as u32);
            //             }
            //         },
            //         Message::Have { piece } => {
            //             state.has.add(piece, piece+1);
            //         },
            //         Message::Piece { piece, offset, block } => {
            //             datastore.write_block(piece as usize, offset, &block)?;
            //             manifest.manifest.add_block(piece as usize, offset, block.len() as u32)?;
            //             manifest.store()?;
            //             s.waiting = false;
            //         },
            //         _ => {},
            //     }
            //     s.nreceived += 1;
            //     if s.nreceived >= 1 && state.am_choking {
            //         let out = Message::Unchoke{};
            //         println!("sending message: {:?}", out);
            //         peer_protocol::send_message(&mut stream, &out)?;
            //         state.am_choking = false;
            //     }
            //     if s.nreceived >= 1 && !state.am_interested {
            //         let out = Message::Interested{};
            //         println!("sending message: {:?}", out);
            //         peer_protocol::send_message(&mut stream, &out)?;
            //         state.am_interested = true;
            //     }
            //     if !s.waiting && !state.peer_choking && state.am_interested {
            //         match manifest.manifest.next_desired_block() {
            //             None => {
            //                 return Err(Error::todo())
            //             },
            //             Some(desire) => {
            //                 let out = Message::Request{
            //                     piece: desire.piece,
            //                     offset: desire.offset,
            //                     length: desire.length,
            //                 };
            //                 println!("sending message: {:?}", out);
            //                 peer_protocol::send_message(&mut stream, &out)?;
            //                 s.waiting = true;
            //             }
            //         }
            //     }
            //     println!("state: {:?}", state);
            // }
        });
    core.run(future_root)?;
    Ok(())
}

/// Connect to a remote peer
fn connect_peer(
    addr: SocketAddr,
    info_hash: InfoHash,
    peer_id: PeerID,
    handle: &reactor::Handle)
    -> BxFuture<(PeerStream, PeerID), Error>
{
    let info_hash2 = info_hash.clone();

    println!("connecting to {} ...", addr);
    tcp_connect2(&addr, Duration::from_millis(3000), &handle)
        .map_err(|e| Error::annotate(e, "peer connection failed"))
        .and_then(|stream| -> FutureResult<TcpStream,Error> {
            match stream {
                Some(stream) => future::ok(stream),
                None => future::err(Error::new_str(&"peer connection timed out")),
            }
        })
        .and_then(move |stream| {
            println!("connected");
            peer_protocol::handshake_send_async(stream, info_hash.clone(), peer_id.clone())
        })
        .and_then(|stream| {
            peer_protocol::handshake_read_1_async(stream)
        })
        .and_then(move |(stream, remote_info_hash)| {
            println!("remote info hash: {:?}", remote_info_hash);
            if remote_info_hash != info_hash2 {
                return Err(Error::new_str(&format!("peer info hash mismatch peer:{:?} me:{:?}",
                    remote_info_hash, info_hash2)));
            }
            Ok(stream)
        })
        .and_then(|stream| {
            peer_protocol::handshake_read_2_async(stream)
                .map(move |(stream, remote_peer_id)| {
                    (stream, remote_peer_id)
                })
        })
        .and_then(|(stream, remote_peer_id)| {
            println!("remote peer id: {:?}", remote_peer_id);

            // let stream: Framed<TcpStream,BitTorrentPeerCodec> = stream.framed(BitTorrentPeerCodec);
            let stream: PeerStream = stream.framed(BitTorrentPeerCodec);
            Ok((stream, remote_peer_id))
        })
        .bxed()
}

/// Compose the main loop.
fn main_loop(stream: PeerStream, dstate: DownloaderState) -> BxFuture<(), Error> {
    let lstate: LoopState = (stream, dstate);
    future::loop_fn(lstate, loop_step)
        .map(|loop_result| {
            println!("loop finished: {:?}", loop_result);
        })
        .bxed()
}

/// One step of the main loop.
fn loop_step(lstate: LoopState) -> BxFuture<future::Loop<String, LoopState>, Error> {
    use futures::future::Loop::{Break,Continue};
    let (stream, mut dstate) = lstate;
    // Note: There are many `bxed` calls in here because match arms must return the same type.
    // And the long long types produced by future chaining are not equivalent.
    // So the boxing turns them into trait objects implementing the same specified Future.
    stream
        .into_future()
        .map_err(|(err, _stream)| err)
        .and_then(move |(omsg, stream)| { match omsg {
            None => future::ok(Break("event stream ended".to_owned())).bxed(),
            Some(msg) => {
                match single_step(msg, &mut dstate) {
                    Err(err) => future::err(err).bxed(),
                    Ok(outs) => {
                        let n_outs = outs.len();
                        stream.send_all(VecDequeStream::<Message, Error>::new(outs)).map(move |(stream, _)| {
                            println!("sent {}", n_outs);
                            Continue((stream, dstate))
                        })
                    }.bxed()
                }
            }
        }}).bxed()
}

/// One synchronous step.
/// Returns a message to send.
fn single_step(msg: Message, dstate: &mut DownloaderState) -> Result<VecDeque<Message>> {
    let mut outs = VecDeque::new();
    println!("recv message {}: {}", dstate.blah_state.nreceived, msg.summarize());
    dstate.blah_state.nreceived += 1;
    match msg {
        Message::Choke =>         dstate.peer_state.peer_choking = true,
        Message::Unchoke =>       dstate.peer_state.peer_choking = false,
        Message::Interested =>    dstate.peer_state.peer_interested = true,
        Message::NotInterested => dstate.peer_state.peer_interested = false,
        Message::Bitfield { bits } => {
            if bits.len() < dstate.info.num_pieces() {
                println!("{}", Error::new_str(&format!("bitfield has less bits {} than pieces {}",
                                                bits.len(), dstate.info.num_pieces())));
            }
            let mut i_start = 0;
            let mut in_interval = false;
            for b in 0..dstate.info.num_pieces() {
                if bits[b] && !in_interval {
                    i_start = b;
                    in_interval = true;
                } else if !bits[b] && in_interval {
                    dstate.peer_state.has.add(i_start as u32, b as u32)?;
                    in_interval = false;
                }
            }
            if in_interval {
                dstate.peer_state.has.add(i_start as u32, dstate.info.num_pieces() as u32)?;
            }
        },
        Message::Have { piece } => {
            dstate.peer_state.has.add(piece, piece+1)?;
        },
        Message::Piece { piece, offset, block } => {
            dstate.datastore.write_block(piece as usize, offset, &block)?;
            dstate.manifest.manifest.add_block(piece as usize, offset, block.len() as u32)?;
            dstate.manifest.store()?;
            dstate.blah_state.waiting = false;
        },
        _ => {},
    }
    if dstate.blah_state.nreceived >= 1 && dstate.peer_state.am_choking {
        let out = Message::Unchoke{};
        println!("sending message: {:?}", out);
        dstate.peer_state.am_choking = false;
        outs.push_back(out);
    }
    if dstate.blah_state.nreceived >= 1 && !dstate.peer_state.am_interested {
        let out = Message::Interested{};
        println!("sending message: {:?}", out);
        dstate.peer_state.am_interested = true;
        outs.push_back(out);
    }
    if !dstate.blah_state.waiting && !dstate.peer_state.peer_choking && dstate.peer_state.am_interested {
        match dstate.manifest.manifest.next_desired_block() {
            None => {
                return Err(Error::todo())
            },
            Some(desire) => {
                let out = Message::Request{
                    piece: desire.piece,
                    offset: desire.offset,
                    length: desire.length,
                };
                println!("sending message: {:?}", out);
                dstate.blah_state.waiting = true;
                outs.push_back(out);
            }
        }
    }

    Ok(outs)
}

#[derive(Debug)]
pub struct PeerState {
    /// Peer is interested in this client
    peer_interested: bool,
    /// Peer is choking this client
    peer_choking: bool,

    am_interested: bool,
    am_choking: bool,

    has: Fillable,
}

#[derive(Debug, Default)]
pub struct BlahState {
    nreceived: u64,
    requested: bool,
    waiting: bool,
}

impl PeerState {
    fn new(num_pieces: usize) -> Self {
        PeerState {
            peer_interested: false,
            peer_choking: true,
            am_interested: false,
            am_choking: true,
            has: Fillable::new(num_pieces as u32),
        }
    }
}
<commit_msg>delete sync code<commit_after>use futures::future::{Future,FutureResult};
use futures::future;
use futures::{Stream,Sink};
use std::collections::vec_deque::{VecDeque};
use std::default::Default;
use std::net::{SocketAddr};
use std::path::Path;
use std::time::Duration;
use tokio_core::net::{TcpStream};
use tokio_core::reactor;
use tokio_io::codec::{Framed};
use tokio_io::{AsyncRead};

use datastore::DataStore;
use error::{Error,Result};
use fillable::*;
use manifest::{ManifestWithFile};
use metainfo::{MetaInfo,InfoHash};
use peer_protocol::{PeerID,Message,BitTorrentPeerCodec};
use peer_protocol;
use tracker::{TrackerClient};
use util::{tcp_connect2,BxFuture,FutureEnhanced,VecDequeStream};

type PeerStream = Framed<TcpStream, BitTorrentPeerCodec>;

struct DownloaderState {
    info: MetaInfo,
    datastore: DataStore,
    manifest: ManifestWithFile,
    peer_state: PeerState,
    blah_state: BlahState,
}

type LoopState = (PeerStream, DownloaderState);

pub fn start<P: AsRef<Path>>(
    info: MetaInfo,
    peer_id: PeerID,
    store_path: P,
    manifest_path: P)
    -> Result<()>
{
    let mut core = reactor::Core::new()?;
    let handle = core.handle();

    let datastore = DataStore::create_or_open(&info, store_path)?;
    let manifest = ManifestWithFile::load_or_new(info.clone(), manifest_path)?;
    let mut tc = TrackerClient::new(info.clone(), peer_id.clone())?;

    let tracker_res = tc.easy_start()?;
    println!("tracker res: {:#?}", tracker_res);
    if let Some(reason) = tracker_res.failure_reason {
        return Err(Error::new_str(&format!("tracker failed: {}", reason)));
    }

    if tracker_res.peers.is_empty() {
        return Err(Error::new_str("tracker returned no peers"));
    }

    let peer = tracker_res.peers[0].clone();

    let num_pieces = info.num_pieces();
    let info_hash = info.info_hash.clone();

    let dstate = DownloaderState {
        info: info,
        datastore: datastore,
        manifest: manifest,
        peer_state: PeerState::new(num_pieces),
        blah_state: BlahState::default(),
    };

    let future_root = connect_peer(peer.address, info_hash, peer_id, &handle)
        .and_then(|(stream, remote_peer_id)| {
            let dstate = dstate;

            main_loop(stream, dstate)
        });
    core.run(future_root)?;
    Ok(())
}

/// Connect to a remote peer
fn connect_peer(
    addr: SocketAddr,
    info_hash: InfoHash,
    peer_id: PeerID,
    handle: &reactor::Handle)
    -> BxFuture<(PeerStream, PeerID), Error>
{
    let info_hash2 = info_hash.clone();

    println!("connecting to {} ...", addr);
    tcp_connect2(&addr, Duration::from_millis(3000), &handle)
        .map_err(|e| Error::annotate(e, "peer connection failed"))
        .and_then(|stream| -> FutureResult<TcpStream,Error> {
            match stream {
                Some(stream) => future::ok(stream),
                None => future::err(Error::new_str(&"peer connection timed out")),
            }
        })
        .and_then(move |stream| {
            println!("connected");
            peer_protocol::handshake_send_async(stream, info_hash.clone(), peer_id.clone())
        })
        .and_then(|stream| {
            peer_protocol::handshake_read_1_async(stream)
        })
        .and_then(move |(stream, remote_info_hash)| {
            println!("remote info hash: {:?}", remote_info_hash);
            if remote_info_hash != info_hash2 {
                return Err(Error::new_str(&format!("peer info hash mismatch peer:{:?} me:{:?}",
                    remote_info_hash, info_hash2)));
            }
            Ok(stream)
        })
        .and_then(|stream| {
            peer_protocol::handshake_read_2_async(stream)
                .map(move |(stream, remote_peer_id)| {
                    (stream, remote_peer_id)
                })
        })
        .and_then(|(stream, remote_peer_id)| {
            println!("remote peer id: {:?}", remote_peer_id);

            // let stream: Framed<TcpStream,BitTorrentPeerCodec> = stream.framed(BitTorrentPeerCodec);
            let stream: PeerStream = stream.framed(BitTorrentPeerCodec);
            Ok((stream, remote_peer_id))
        })
        .bxed()
}

/// Compose the main loop.
fn main_loop(stream: PeerStream, dstate: DownloaderState) -> BxFuture<(), Error> {
    let lstate: LoopState = (stream, dstate);
    future::loop_fn(lstate, loop_step)
        .map(|loop_result| {
            println!("loop finished: {:?}", loop_result);
        })
        .bxed()
}

/// One step of the main loop.
fn loop_step(lstate: LoopState) -> BxFuture<future::Loop<String, LoopState>, Error> {
    use futures::future::Loop::{Break,Continue};
    let (stream, mut dstate) = lstate;
    // Note: There are many `bxed` calls in here because match arms must return the same type.
    // And the long long types produced by future chaining are not equivalent.
    // So the boxing turns them into trait objects implementing the same specified Future.
    stream
        .into_future()
        .map_err(|(err, _stream)| err)
        .and_then(move |(omsg, stream)| { match omsg {
            None => future::ok(Break("event stream ended".to_owned())).bxed(),
            Some(msg) => {
                match single_step(msg, &mut dstate) {
                    Err(err) => future::err(err).bxed(),
                    Ok(outs) => {
                        let n_outs = outs.len();
                        stream.send_all(VecDequeStream::<Message, Error>::new(outs)).map(move |(stream, _)| {
                            println!("sent {}", n_outs);
                            Continue((stream, dstate))
                        })
                    }.bxed()
                }
            }
        }}).bxed()
}

/// One synchronous step.
/// Returns a message to send.
fn single_step(msg: Message, dstate: &mut DownloaderState) -> Result<VecDeque<Message>> {
    let mut outs = VecDeque::new();
    println!("recv message {}: {}", dstate.blah_state.nreceived, msg.summarize());
    dstate.blah_state.nreceived += 1;
    match msg {
        Message::Choke =>         dstate.peer_state.peer_choking = true,
        Message::Unchoke =>       dstate.peer_state.peer_choking = false,
        Message::Interested =>    dstate.peer_state.peer_interested = true,
        Message::NotInterested => dstate.peer_state.peer_interested = false,
        Message::Bitfield { bits } => {
            if bits.len() < dstate.info.num_pieces() {
                println!("{}", Error::new_str(&format!("bitfield has less bits {} than pieces {}",
                                                bits.len(), dstate.info.num_pieces())));
            }
            let mut i_start = 0;
            let mut in_interval = false;
            for b in 0..dstate.info.num_pieces() {
                if bits[b] && !in_interval {
                    i_start = b;
                    in_interval = true;
                } else if !bits[b] && in_interval {
                    dstate.peer_state.has.add(i_start as u32, b as u32)?;
                    in_interval = false;
                }
            }
            if in_interval {
                dstate.peer_state.has.add(i_start as u32, dstate.info.num_pieces() as u32)?;
            }
        },
        Message::Have { piece } => {
            dstate.peer_state.has.add(piece, piece+1)?;
        },
        Message::Piece { piece, offset, block } => {
            dstate.datastore.write_block(piece as usize, offset, &block)?;
            dstate.manifest.manifest.add_block(piece as usize, offset, block.len() as u32)?;
            dstate.manifest.store()?;
            dstate.blah_state.waiting = false;
        },
        _ => {},
    }
    if dstate.blah_state.nreceived >= 1 && dstate.peer_state.am_choking {
        let out = Message::Unchoke{};
        println!("sending message: {:?}", out);
        dstate.peer_state.am_choking = false;
        outs.push_back(out);
    }
    if dstate.blah_state.nreceived >= 1 && !dstate.peer_state.am_interested {
        let out = Message::Interested{};
        println!("sending message: {:?}", out);
        dstate.peer_state.am_interested = true;
        outs.push_back(out);
    }
    if !dstate.blah_state.waiting && !dstate.peer_state.peer_choking && dstate.peer_state.am_interested {
        match dstate.manifest.manifest.next_desired_block() {
            None => {
                return Err(Error::todo())
            },
            Some(desire) => {
                let out = Message::Request{
                    piece: desire.piece,
                    offset: desire.offset,
                    length: desire.length,
                };
                println!("sending message: {:?}", out);
                dstate.blah_state.waiting = true;
                outs.push_back(out);
            }
        }
    }

    Ok(outs)
}

#[derive(Debug)]
pub struct PeerState {
    /// Peer is interested in this client
    peer_interested: bool,
    /// Peer is choking this client
    peer_choking: bool,

    am_interested: bool,
    am_choking: bool,

    has: Fillable,
}

#[derive(Debug, Default)]
pub struct BlahState {
    nreceived: u64,
    requested: bool,
    waiting: bool,
}

impl PeerState {
    fn new(num_pieces: usize) -> Self {
        PeerState {
            peer_interested: false,
            peer_choking: true,
            am_interested: false,
            am_choking: true,
            has: Fillable::new(num_pieces as u32),
        }
    }
}
<|endoftext|>


<commit_before>//! CUDA driver
//!
//! Reference: http://docs.nvidia.com/cuda/cuda-driver-api/

use std::ffi::CStr;
use std::marker::PhantomData;
use std::{mem, ptr, result};

#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[allow(non_snake_case)]
#[allow(non_upper_case_globals)]
mod ll;

/// A CUDA "block"
pub struct Block {
    x: u32,
    y: u32,
    z: u32,
}

impl Block {
    /// One dimensional block
    pub fn x(x: u32) -> Self {
        Block { x: x, y: 1, z: 1 }
    }

    /// Two dimensional block
    pub fn xy(x: u32, y: u32) -> Self {
        Block { x: x, y: y, z: 1 }
    }

    /// Three dimensional block
    pub fn xyz(x: u32, y: u32, z: u32) -> Self {
        Block { x: x, y: y, z: z }
    }
}

/// A CUDA "context"
#[derive(Debug)]
pub struct Context {
    defused: bool,
    handle: ll::CUcontext,
}

impl Context {
    // TODO is this actually useful? Note that we are using "RAII" (cf. `drop`)
    // and ownership to manage `Context`es
    #[allow(dead_code)]
    fn current() -> Result<Option<Self>> {
        let mut handle = ptr::null_mut();

        unsafe { lift(ll::cuCtxGetCurrent(&mut handle))? }

        if handle.is_null() {
            Ok(None)
        } else {
            Ok(Some(Context {
                defused: true,
                handle: handle,
            }))
        }
    }

    /// Loads a PTX module
    pub fn load_module<'ctx>(&'ctx self, image: &CStr) -> Result<Module<'ctx>> {
        let mut handle = ptr::null_mut();

        unsafe {
            lift(ll::cuModuleLoadData(&mut handle, image.as_ptr() as *const _))?
        }

        Ok(Module {
            handle: handle,
            _context: PhantomData,
        })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        if !self.defused {
            unsafe { lift(ll::cuCtxDestroy_v2(self.handle)).unwrap() }
        }
    }
}

/// A CUDA device (a GPU)
#[derive(Debug)]
pub struct Device {
    handle: ll::CUdevice,
}

/// Binds to the `nth` device
#[allow(non_snake_case)]
pub fn Device(nth: u16) -> Result<Device> {
    let mut handle = 0;

    unsafe { lift(ll::cuDeviceGet(&mut handle, i32::from(nth)))? }

    Ok(Device { handle: handle })
}

impl Device {
    /// Returns the number of available devices
    pub fn count() -> Result<u32> {
        let mut count: i32 = 0;

        unsafe { lift(ll::cuDeviceGetCount(&mut count))? }

        Ok(count as u32)
    }

    /// Creates a CUDA context for this device
    pub fn create_context(&self) -> Result<Context> {
        let mut handle = ptr::null_mut();
        // TODO expose
        let flags = 0;

        unsafe { lift(ll::cuCtxCreate_v2(&mut handle, flags, self.handle))? }

        Ok(Context {
            defused: false,
            handle: handle,
        })
    }

    /// Returns the maximum number of threads a block can have
    pub fn max_threads_per_block(&self) -> Result<i32> {
        use self::ll::CUdevice_attribute_enum::*;

        self.get(CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK)
    }

    /// Returns the total amount of (non necessarily free) memory, in bytes,
    /// that the device has
    pub fn total_memory(&self) -> Result<usize> {
        let mut bytes = 0;

        unsafe { lift(ll::cuDeviceTotalMem_v2(&mut bytes, self.handle))? };

        Ok(bytes)
    }

    fn get(&self, attr: ll::CUdevice_attribute) -> Result<i32> {
        let mut value = 0;

        unsafe {
            lift(ll::cuDeviceGetAttribute(&mut value, attr, self.handle))?
        }

        Ok(value)
    }
}

/// A function that the CUDA device can execute. AKA a "kernel"
pub struct Function<'ctx: 'm, 'm> {
    handle: ll::CUfunction,
    _module: PhantomData<&'m Module<'ctx>>,
}

impl<'ctx, 'm> Function<'ctx, 'm> {
    /// Execute a function on the GPU
    ///
    /// NOTE This function blocks until the GPU has finished executing the
    /// kernel
    pub fn launch(&self,
                  args: &[&Any],
                  grid: Grid,
                  block: Block)
                  -> Result<()> {
        let stream = Stream::new()?;
        // TODO expose
        let shared_mem_bytes = 0;
        // TODO expose
        let extra = ptr::null_mut();

        unsafe {
            lift(ll::cuLaunchKernel(self.handle,
                                    grid.x,
                                    grid.y,
                                    grid.z,
                                    block.x,
                                    block.y,
                                    block.z,
                                    shared_mem_bytes,
                                    stream.handle,
                                    args.as_ptr() as *mut _,
                                    extra))?
        }

        stream.sync()?;
        stream.destroy()
    }
}

/// A CUDA "grid"
pub struct Grid {
    x: u32,
    y: u32,
    z: u32,
}

impl Grid {
    /// One dimensional grid
    pub fn x(x: u32) -> Self {
        Grid { x: x, y: 1, z: 1 }
    }

    /// Two dimensional grid
    pub fn xy(x: u32, y: u32) -> Self {
        Grid { x: x, y: y, z: 1 }
    }

    /// Three dimensional grid
    pub fn xyz(x: u32, y: u32, z: u32) -> Self {
        Grid { x: x, y: y, z: z }
    }
}

/// A PTX module
pub struct Module<'ctx> {
    handle: ll::CUmodule,
    _context: PhantomData<&'ctx Context>,
}

impl<'ctx> Module<'ctx> {
    /// Retrieves a function from the PTX module
    pub fn function<'m>(&'m self, name: &CStr) -> Result<Function<'ctx, 'm>> {
        let mut handle = ptr::null_mut();

        unsafe {
            lift(ll::cuModuleGetFunction(&mut handle,
                                         self.handle,
                                         name.as_ptr()))?
        }

        Ok(Function {
            handle: handle,
            _module: PhantomData,
        })
    }
}

impl<'ctx> Drop for Module<'ctx> {
    fn drop(&mut self) {
        unsafe { lift(ll::cuModuleUnload(self.handle)).unwrap() }
    }
}

// TODO expose
struct Stream {
    handle: ll::CUstream,
}

impl Stream {
    fn new() -> Result<Self> {
        let mut handle = ptr::null_mut();
        // TODO expose
        let flags = 0;

        unsafe { lift(ll::cuStreamCreate(&mut handle, flags))? }

        Ok(Stream { handle: handle })
    }

    fn destroy(self) -> Result<()> {
        unsafe { lift(ll::cuStreamDestroy_v2(self.handle)) }
    }

    fn sync(&self) -> Result<()> {
        unsafe { lift(ll::cuStreamSynchronize(self.handle)) }
    }
}

/// Value who's type has been erased
pub enum Any {}

/// Erase the type of a value
#[allow(non_snake_case)]
pub fn Any<T>(ref_: &T) -> &Any {
    unsafe { &*(ref_ as *const T as *const Any) }
}

/// `memcpy` direction
pub enum Direction {
    /// `src` points to device memory. `dst` points to host memory
    DeviceToHost,
    /// `src` points to host memory. `dst` points to device memory
    HostToDevice,
}

#[allow(missing_docs)]
#[derive(Debug)]
pub enum Error {
    AlreadyAcquired,
    AlreadyMapped,
    ArrayIsMapped,
    Assert,
    ContextAlreadyCurrent,
    ContextAlreadyInUse,
    ContextIsDestroyed,
    Deinitialized,
    EccUncorrectable,
    FileNotFound,
    HardwareStackError,
    HostMemoryAlreadyRegistered,
    HostMemoryNotRegistered,
    IllegalAddress,
    IllegalInstruction,
    InvalidAddressSpace,
    InvalidContext,
    InvalidDevice,
    InvalidGraphicsContext,
    InvalidHandle,
    InvalidImage,
    InvalidPc,
    InvalidPtx,
    InvalidSource,
    InvalidValue,
    LaunchFailed,
    LaunchIncompatibleTexturing,
    LaunchOutOfResources,
    LaunchTimeout,
    MapFailed,
    MisalignedAddress,
    NoBinaryForGpu,
    NoDevice,
    NotFound,
    NotInitialized,
    NotMapped,
    NotMappedAsArray,
    NotMappedAsPointer,
    NotPermitted,
    NotReady,
    NotSupported,
    NvlinkUncorrectable,
    OperatingSystem,
    OutOfMemory,
    PeerAccessAlreadyEnabled,
    PeerAccessNotEnabled,
    PeerAccessUnsupported,
    PrimaryContextActive,
    ProfilerAlreadyStarted,
    ProfilerAlreadyStopped,
    ProfilerDisabled,
    ProfilerNotInitialized,
    SharedObjectInitFailed,
    SharedObjectSymbolNotFound,
    TooManyPeers,
    Unknown,
    UnmapFailed,
    UnsupportedLimit,
}

// TODO should this be a method of `Context`?
/// Allocate `n` bytes of memory on the device
pub unsafe fn allocate(n: usize) -> Result<*mut u8> {
    let mut d_ptr = 0;

    lift(ll::cuMemAlloc_v2(&mut d_ptr, n))?;

    Ok(d_ptr as *mut u8)
}

/// Copy `n` bytes of memory from `src` to `dst`
///
/// `direction` indicates where `src` and `dst` are located (device or host)
pub unsafe fn copy<T>(src: *const T,
                      dst: *mut T,
                      count: usize,
                      direction: Direction)
                      -> Result<()> {
    use self::Direction::*;

    let bytes = count * mem::size_of::<T>();

    lift(match direction {
        DeviceToHost => ll::cuMemcpyDtoH_v2(dst as *mut _, src as u64, bytes),
        HostToDevice => ll::cuMemcpyHtoD_v2(dst as u64, src as *const _, bytes),
    })?;

    Ok(())
}

// TODO same question as `allocate`
/// Free the memory pointed to by `ptr`
pub unsafe fn deallocate(ptr: *mut u8) -> Result<()> {
    lift(ll::cuMemFree_v2(ptr as u64))
}

/// Initialize the CUDA runtime
pub fn initialize() -> Result<()> {
    // TODO expose
    let flags = 0;

    unsafe { lift(ll::cuInit(flags)) }
}

/// Returns the version of the CUDA runtime
pub fn version() -> Result<i32> {
    let mut version = 0;

    unsafe { lift(ll::cuDriverGetVersion(&mut version))? }

    Ok(version)
}

#[allow(missing_docs)]
pub type Result<T> = result::Result<T, Error>;

fn lift(e: ll::CUresult) -> Result<()> {
    use self::Error::*;
    use self::ll::cudaError_enum::*;

    Err(match e {
        CUDA_SUCCESS => return Ok(()),
        CUDA_ERROR_ALREADY_ACQUIRED => AlreadyAcquired,
        CUDA_ERROR_ALREADY_MAPPED => AlreadyMapped,
        CUDA_ERROR_ARRAY_IS_MAPPED => ArrayIsMapped,
        CUDA_ERROR_ASSERT => Assert,
        CUDA_ERROR_CONTEXT_ALREADY_CURRENT => ContextAlreadyCurrent,
        CUDA_ERROR_CONTEXT_ALREADY_IN_USE => ContextAlreadyInUse,
        CUDA_ERROR_CONTEXT_IS_DESTROYED => ContextIsDestroyed,
        CUDA_ERROR_DEINITIALIZED => Deinitialized,
        CUDA_ERROR_ECC_UNCORRECTABLE => EccUncorrectable,
        CUDA_ERROR_FILE_NOT_FOUND => FileNotFound,
        CUDA_ERROR_HARDWARE_STACK_ERROR => HardwareStackError,
        CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED => {
            HostMemoryAlreadyRegistered
        }
        CUDA_ERROR_HOST_MEMORY_NOT_REGISTERED => HostMemoryNotRegistered,
        CUDA_ERROR_ILLEGAL_ADDRESS => IllegalAddress,
        CUDA_ERROR_ILLEGAL_INSTRUCTION => IllegalInstruction,
        CUDA_ERROR_INVALID_ADDRESS_SPACE => InvalidAddressSpace,
        CUDA_ERROR_INVALID_CONTEXT => InvalidContext,
        CUDA_ERROR_INVALID_DEVICE => InvalidDevice,
        CUDA_ERROR_INVALID_GRAPHICS_CONTEXT => InvalidGraphicsContext,
        CUDA_ERROR_INVALID_HANDLE => InvalidHandle,
        CUDA_ERROR_INVALID_IMAGE => InvalidImage,
        CUDA_ERROR_INVALID_PC => InvalidPc,
        CUDA_ERROR_INVALID_PTX => InvalidPtx,
        CUDA_ERROR_INVALID_SOURCE => InvalidSource,
        CUDA_ERROR_INVALID_VALUE => InvalidValue,
        CUDA_ERROR_LAUNCH_FAILED => LaunchFailed,
        CUDA_ERROR_LAUNCH_INCOMPATIBLE_TEXTURING => LaunchIncompatibleTexturing,
        CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES => LaunchOutOfResources,
        CUDA_ERROR_LAUNCH_TIMEOUT => LaunchTimeout,
        CUDA_ERROR_MAP_FAILED => MapFailed,
        CUDA_ERROR_MISALIGNED_ADDRESS => MisalignedAddress,
        CUDA_ERROR_NOT_FOUND => NotFound,
        CUDA_ERROR_NOT_INITIALIZED => NotInitialized,
        CUDA_ERROR_NOT_MAPPED => NotMapped,
        CUDA_ERROR_NOT_MAPPED_AS_ARRAY => NotMappedAsArray,
        CUDA_ERROR_NOT_MAPPED_AS_POINTER => NotMappedAsPointer,
        CUDA_ERROR_NOT_PERMITTED => NotPermitted,
        CUDA_ERROR_NOT_READY => NotReady,
        CUDA_ERROR_NOT_SUPPORTED => NotSupported,
        CUDA_ERROR_NO_BINARY_FOR_GPU => NoBinaryForGpu,
        CUDA_ERROR_NO_DEVICE => NoDevice,
        CUDA_ERROR_OPERATING_SYSTEM => OperatingSystem,
        CUDA_ERROR_OUT_OF_MEMORY => OutOfMemory,
        CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED => PeerAccessAlreadyEnabled,
        CUDA_ERROR_PEER_ACCESS_NOT_ENABLED => PeerAccessNotEnabled,
        CUDA_ERROR_PEER_ACCESS_UNSUPPORTED => PeerAccessUnsupported,
        CUDA_ERROR_PRIMARY_CONTEXT_ACTIVE => PrimaryContextActive,
        CUDA_ERROR_PROFILER_ALREADY_STARTED => ProfilerAlreadyStarted,
        CUDA_ERROR_PROFILER_ALREADY_STOPPED => ProfilerAlreadyStopped,
        CUDA_ERROR_PROFILER_DISABLED => ProfilerDisabled,
        CUDA_ERROR_PROFILER_NOT_INITIALIZED => ProfilerNotInitialized,
        CUDA_ERROR_SHARED_OBJECT_INIT_FAILED => SharedObjectInitFailed,
        CUDA_ERROR_SHARED_OBJECT_SYMBOL_NOT_FOUND => SharedObjectSymbolNotFound,
        CUDA_ERROR_TOO_MANY_PEERS => TooManyPeers,
        CUDA_ERROR_UNKNOWN => Unknown,
        CUDA_ERROR_UNMAP_FAILED => UnmapFailed,
        CUDA_ERROR_UNSUPPORTED_LIMIT => UnsupportedLimit,
        CUDA_ERROR_NVLINK_UNCORRECTABLE => NvlinkUncorrectable,
    })
}
<commit_msg>Auto merge of #2 - denzp:master, r=japaric<commit_after>//! CUDA driver
//!
//! Reference: http://docs.nvidia.com/cuda/cuda-driver-api/

use std::ffi::CStr;
use std::marker::PhantomData;
use std::{mem, ptr, result};

#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[allow(non_snake_case)]
#[allow(non_upper_case_globals)]
mod ll;

/// A CUDA "block"
pub struct Block {
    x: u32,
    y: u32,
    z: u32,
}

impl Block {
    /// One dimensional block
    pub fn x(x: u32) -> Self {
        Block { x: x, y: 1, z: 1 }
    }

    /// Two dimensional block
    pub fn xy(x: u32, y: u32) -> Self {
        Block { x: x, y: y, z: 1 }
    }

    /// Three dimensional block
    pub fn xyz(x: u32, y: u32, z: u32) -> Self {
        Block { x: x, y: y, z: z }
    }
}

/// A CUDA "context"
#[derive(Debug)]
pub struct Context {
    defused: bool,
    handle: ll::CUcontext,
}

impl Context {
    // TODO is this actually useful? Note that we are using "RAII" (cf. `drop`)
    // and ownership to manage `Context`es
    #[allow(dead_code)]
    fn current() -> Result<Option<Self>> {
        let mut handle = ptr::null_mut();

        unsafe { lift(ll::cuCtxGetCurrent(&mut handle))? }

        if handle.is_null() {
            Ok(None)
        } else {
            Ok(Some(Context {
                defused: true,
                handle: handle,
            }))
        }
    }

    /// Binds context to the calling thread
    pub fn set_current(&self) -> Result<()> {
        unsafe {
            lift(ll::cuCtxSetCurrent(self.handle))
        }
    }

    /// Loads a PTX module
    pub fn load_module<'ctx>(&'ctx self, image: &CStr) -> Result<Module<'ctx>> {
        let mut handle = ptr::null_mut();

        unsafe {
            lift(ll::cuModuleLoadData(&mut handle, image.as_ptr() as *const _))?
        }

        Ok(Module {
            handle: handle,
            _context: PhantomData,
        })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        if !self.defused {
            unsafe { lift(ll::cuCtxDestroy_v2(self.handle)).unwrap() }
        }
    }
}

/// A CUDA device (a GPU)
#[derive(Debug)]
pub struct Device {
    handle: ll::CUdevice,
}

/// Binds to the `nth` device
#[allow(non_snake_case)]
pub fn Device(nth: u16) -> Result<Device> {
    let mut handle = 0;

    unsafe { lift(ll::cuDeviceGet(&mut handle, i32::from(nth)))? }

    Ok(Device { handle: handle })
}

impl Device {
    /// Returns the number of available devices
    pub fn count() -> Result<u32> {
        let mut count: i32 = 0;

        unsafe { lift(ll::cuDeviceGetCount(&mut count))? }

        Ok(count as u32)
    }

    /// Creates a CUDA context for this device
    pub fn create_context(&self) -> Result<Context> {
        let mut handle = ptr::null_mut();
        // TODO expose
        let flags = 0;

        unsafe { lift(ll::cuCtxCreate_v2(&mut handle, flags, self.handle))? }

        Ok(Context {
            defused: false,
            handle: handle,
        })
    }

    /// Returns the maximum number of threads a block can have
    pub fn max_threads_per_block(&self) -> Result<i32> {
        use self::ll::CUdevice_attribute_enum::*;

        self.get(CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK)
    }

    /// Returns the total amount of (non necessarily free) memory, in bytes,
    /// that the device has
    pub fn total_memory(&self) -> Result<usize> {
        let mut bytes = 0;

        unsafe { lift(ll::cuDeviceTotalMem_v2(&mut bytes, self.handle))? };

        Ok(bytes)
    }

    fn get(&self, attr: ll::CUdevice_attribute) -> Result<i32> {
        let mut value = 0;

        unsafe {
            lift(ll::cuDeviceGetAttribute(&mut value, attr, self.handle))?
        }

        Ok(value)
    }
}

/// A function that the CUDA device can execute. AKA a "kernel"
pub struct Function<'ctx: 'm, 'm> {
    handle: ll::CUfunction,
    _module: PhantomData<&'m Module<'ctx>>,
}

impl<'ctx, 'm> Function<'ctx, 'm> {
    /// Execute a function on the GPU
    ///
    /// NOTE This function blocks until the GPU has finished executing the
    /// kernel
    pub fn launch(&self,
                  args: &[&Any],
                  grid: Grid,
                  block: Block)
                  -> Result<()> {
        let stream = Stream::new()?;
        // TODO expose
        let shared_mem_bytes = 0;
        // TODO expose
        let extra = ptr::null_mut();

        unsafe {
            lift(ll::cuLaunchKernel(self.handle,
                                    grid.x,
                                    grid.y,
                                    grid.z,
                                    block.x,
                                    block.y,
                                    block.z,
                                    shared_mem_bytes,
                                    stream.handle,
                                    args.as_ptr() as *mut _,
                                    extra))?
        }

        stream.sync()?;
        stream.destroy()
    }
}

/// A CUDA "grid"
pub struct Grid {
    x: u32,
    y: u32,
    z: u32,
}

impl Grid {
    /// One dimensional grid
    pub fn x(x: u32) -> Self {
        Grid { x: x, y: 1, z: 1 }
    }

    /// Two dimensional grid
    pub fn xy(x: u32, y: u32) -> Self {
        Grid { x: x, y: y, z: 1 }
    }

    /// Three dimensional grid
    pub fn xyz(x: u32, y: u32, z: u32) -> Self {
        Grid { x: x, y: y, z: z }
    }
}

/// A PTX module
pub struct Module<'ctx> {
    handle: ll::CUmodule,
    _context: PhantomData<&'ctx Context>,
}

impl<'ctx> Module<'ctx> {
    /// Retrieves a function from the PTX module
    pub fn function<'m>(&'m self, name: &CStr) -> Result<Function<'ctx, 'm>> {
        let mut handle = ptr::null_mut();

        unsafe {
            lift(ll::cuModuleGetFunction(&mut handle,
                                         self.handle,
                                         name.as_ptr()))?
        }

        Ok(Function {
            handle: handle,
            _module: PhantomData,
        })
    }
}

impl<'ctx> Drop for Module<'ctx> {
    fn drop(&mut self) {
        unsafe { lift(ll::cuModuleUnload(self.handle)).unwrap() }
    }
}

// TODO expose
struct Stream {
    handle: ll::CUstream,
}

impl Stream {
    fn new() -> Result<Self> {
        let mut handle = ptr::null_mut();
        // TODO expose
        let flags = 0;

        unsafe { lift(ll::cuStreamCreate(&mut handle, flags))? }

        Ok(Stream { handle: handle })
    }

    fn destroy(self) -> Result<()> {
        unsafe { lift(ll::cuStreamDestroy_v2(self.handle)) }
    }

    fn sync(&self) -> Result<()> {
        unsafe { lift(ll::cuStreamSynchronize(self.handle)) }
    }
}

/// Value who's type has been erased
pub enum Any {}

/// Erase the type of a value
#[allow(non_snake_case)]
pub fn Any<T>(ref_: &T) -> &Any {
    unsafe { &*(ref_ as *const T as *const Any) }
}

/// `memcpy` direction
pub enum Direction {
    /// `src` points to device memory. `dst` points to host memory
    DeviceToHost,
    /// `src` points to host memory. `dst` points to device memory
    HostToDevice,
}

#[allow(missing_docs)]
#[derive(Debug)]
pub enum Error {
    AlreadyAcquired,
    AlreadyMapped,
    ArrayIsMapped,
    Assert,
    ContextAlreadyCurrent,
    ContextAlreadyInUse,
    ContextIsDestroyed,
    Deinitialized,
    EccUncorrectable,
    FileNotFound,
    HardwareStackError,
    HostMemoryAlreadyRegistered,
    HostMemoryNotRegistered,
    IllegalAddress,
    IllegalInstruction,
    InvalidAddressSpace,
    InvalidContext,
    InvalidDevice,
    InvalidGraphicsContext,
    InvalidHandle,
    InvalidImage,
    InvalidPc,
    InvalidPtx,
    InvalidSource,
    InvalidValue,
    LaunchFailed,
    LaunchIncompatibleTexturing,
    LaunchOutOfResources,
    LaunchTimeout,
    MapFailed,
    MisalignedAddress,
    NoBinaryForGpu,
    NoDevice,
    NotFound,
    NotInitialized,
    NotMapped,
    NotMappedAsArray,
    NotMappedAsPointer,
    NotPermitted,
    NotReady,
    NotSupported,
    NvlinkUncorrectable,
    OperatingSystem,
    OutOfMemory,
    PeerAccessAlreadyEnabled,
    PeerAccessNotEnabled,
    PeerAccessUnsupported,
    PrimaryContextActive,
    ProfilerAlreadyStarted,
    ProfilerAlreadyStopped,
    ProfilerDisabled,
    ProfilerNotInitialized,
    SharedObjectInitFailed,
    SharedObjectSymbolNotFound,
    TooManyPeers,
    Unknown,
    UnmapFailed,
    UnsupportedLimit,
}

// TODO should this be a method of `Context`?
/// Allocate `n` bytes of memory on the device
pub unsafe fn allocate(n: usize) -> Result<*mut u8> {
    let mut d_ptr = 0;

    lift(ll::cuMemAlloc_v2(&mut d_ptr, n))?;

    Ok(d_ptr as *mut u8)
}

/// Copy `n` bytes of memory from `src` to `dst`
///
/// `direction` indicates where `src` and `dst` are located (device or host)
pub unsafe fn copy<T>(src: *const T,
                      dst: *mut T,
                      count: usize,
                      direction: Direction)
                      -> Result<()> {
    use self::Direction::*;

    let bytes = count * mem::size_of::<T>();

    lift(match direction {
        DeviceToHost => ll::cuMemcpyDtoH_v2(dst as *mut _, src as u64, bytes),
        HostToDevice => ll::cuMemcpyHtoD_v2(dst as u64, src as *const _, bytes),
    })?;

    Ok(())
}

// TODO same question as `allocate`
/// Free the memory pointed to by `ptr`
pub unsafe fn deallocate(ptr: *mut u8) -> Result<()> {
    lift(ll::cuMemFree_v2(ptr as u64))
}

/// Initialize the CUDA runtime
pub fn initialize() -> Result<()> {
    // TODO expose
    let flags = 0;

    unsafe { lift(ll::cuInit(flags)) }
}

/// Returns the version of the CUDA runtime
pub fn version() -> Result<i32> {
    let mut version = 0;

    unsafe { lift(ll::cuDriverGetVersion(&mut version))? }

    Ok(version)
}

#[allow(missing_docs)]
pub type Result<T> = result::Result<T, Error>;

fn lift(e: ll::CUresult) -> Result<()> {
    use self::Error::*;
    use self::ll::cudaError_enum::*;

    Err(match e {
        CUDA_SUCCESS => return Ok(()),
        CUDA_ERROR_ALREADY_ACQUIRED => AlreadyAcquired,
        CUDA_ERROR_ALREADY_MAPPED => AlreadyMapped,
        CUDA_ERROR_ARRAY_IS_MAPPED => ArrayIsMapped,
        CUDA_ERROR_ASSERT => Assert,
        CUDA_ERROR_CONTEXT_ALREADY_CURRENT => ContextAlreadyCurrent,
        CUDA_ERROR_CONTEXT_ALREADY_IN_USE => ContextAlreadyInUse,
        CUDA_ERROR_CONTEXT_IS_DESTROYED => ContextIsDestroyed,
        CUDA_ERROR_DEINITIALIZED => Deinitialized,
        CUDA_ERROR_ECC_UNCORRECTABLE => EccUncorrectable,
        CUDA_ERROR_FILE_NOT_FOUND => FileNotFound,
        CUDA_ERROR_HARDWARE_STACK_ERROR => HardwareStackError,
        CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED => {
            HostMemoryAlreadyRegistered
        }
        CUDA_ERROR_HOST_MEMORY_NOT_REGISTERED => HostMemoryNotRegistered,
        CUDA_ERROR_ILLEGAL_ADDRESS => IllegalAddress,
        CUDA_ERROR_ILLEGAL_INSTRUCTION => IllegalInstruction,
        CUDA_ERROR_INVALID_ADDRESS_SPACE => InvalidAddressSpace,
        CUDA_ERROR_INVALID_CONTEXT => InvalidContext,
        CUDA_ERROR_INVALID_DEVICE => InvalidDevice,
        CUDA_ERROR_INVALID_GRAPHICS_CONTEXT => InvalidGraphicsContext,
        CUDA_ERROR_INVALID_HANDLE => InvalidHandle,
        CUDA_ERROR_INVALID_IMAGE => InvalidImage,
        CUDA_ERROR_INVALID_PC => InvalidPc,
        CUDA_ERROR_INVALID_PTX => InvalidPtx,
        CUDA_ERROR_INVALID_SOURCE => InvalidSource,
        CUDA_ERROR_INVALID_VALUE => InvalidValue,
        CUDA_ERROR_LAUNCH_FAILED => LaunchFailed,
        CUDA_ERROR_LAUNCH_INCOMPATIBLE_TEXTURING => LaunchIncompatibleTexturing,
        CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES => LaunchOutOfResources,
        CUDA_ERROR_LAUNCH_TIMEOUT => LaunchTimeout,
        CUDA_ERROR_MAP_FAILED => MapFailed,
        CUDA_ERROR_MISALIGNED_ADDRESS => MisalignedAddress,
        CUDA_ERROR_NOT_FOUND => NotFound,
        CUDA_ERROR_NOT_INITIALIZED => NotInitialized,
        CUDA_ERROR_NOT_MAPPED => NotMapped,
        CUDA_ERROR_NOT_MAPPED_AS_ARRAY => NotMappedAsArray,
        CUDA_ERROR_NOT_MAPPED_AS_POINTER => NotMappedAsPointer,
        CUDA_ERROR_NOT_PERMITTED => NotPermitted,
        CUDA_ERROR_NOT_READY => NotReady,
        CUDA_ERROR_NOT_SUPPORTED => NotSupported,
        CUDA_ERROR_NO_BINARY_FOR_GPU => NoBinaryForGpu,
        CUDA_ERROR_NO_DEVICE => NoDevice,
        CUDA_ERROR_OPERATING_SYSTEM => OperatingSystem,
        CUDA_ERROR_OUT_OF_MEMORY => OutOfMemory,
        CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED => PeerAccessAlreadyEnabled,
        CUDA_ERROR_PEER_ACCESS_NOT_ENABLED => PeerAccessNotEnabled,
        CUDA_ERROR_PEER_ACCESS_UNSUPPORTED => PeerAccessUnsupported,
        CUDA_ERROR_PRIMARY_CONTEXT_ACTIVE => PrimaryContextActive,
        CUDA_ERROR_PROFILER_ALREADY_STARTED => ProfilerAlreadyStarted,
        CUDA_ERROR_PROFILER_ALREADY_STOPPED => ProfilerAlreadyStopped,
        CUDA_ERROR_PROFILER_DISABLED => ProfilerDisabled,
        CUDA_ERROR_PROFILER_NOT_INITIALIZED => ProfilerNotInitialized,
        CUDA_ERROR_SHARED_OBJECT_INIT_FAILED => SharedObjectInitFailed,
        CUDA_ERROR_SHARED_OBJECT_SYMBOL_NOT_FOUND => SharedObjectSymbolNotFound,
        CUDA_ERROR_TOO_MANY_PEERS => TooManyPeers,
        CUDA_ERROR_UNKNOWN => Unknown,
        CUDA_ERROR_UNMAP_FAILED => UnmapFailed,
        CUDA_ERROR_UNSUPPORTED_LIMIT => UnsupportedLimit,
        CUDA_ERROR_NVLINK_UNCORRECTABLE => NvlinkUncorrectable,
    })
}
<|endoftext|>


<commit_before>extern crate libc;

use array::Array;
use defines::AfError;
use seq::Seq;
use self::libc::{c_int, c_uint, c_longlong};

type MutAfIndex    = *mut self::libc::c_longlong;
type MutAfArray    = *mut self::libc::c_longlong;
type AfArray       = self::libc::c_longlong;
type DimT          = self::libc::c_longlong;
type IndexT        = self::libc::c_longlong;

#[allow(dead_code)]
extern {
    fn af_create_indexers(indexers: MutAfIndex) -> c_int;
    fn af_set_array_indexer(indexer: MutAfIndex, idx: AfArray, dim: DimT) -> c_int;
    fn af_set_seq_indexer(indexer: MutAfIndex, idx: *const Seq, dim: DimT, is_batch: c_int) -> c_int;
    fn af_release_indexers(indexers: MutAfIndex) -> c_int;

    fn af_index(out: MutAfArray, input: AfArray, ndims: c_uint, index: *const Seq) -> c_int;
    fn af_lookup(out: MutAfArray, arr: AfArray, indices: AfArray, dim: c_uint) -> c_int;
    fn af_assign_seq(out: MutAfArray, lhs: AfArray, ndims: c_uint, indices: *const Seq, rhs: AfArray) -> c_int;
    fn af_index_gen(out: MutAfArray, input: AfArray, ndims: DimT, indices: *const IndexT) -> c_int;
    fn af_assign_gen(out: MutAfArray, lhs: AfArray, ndims: DimT, indices: *const IndexT, rhs: AfArray) -> c_int;
}

pub struct Indexer {
    handle: i64,
    count: u32,
}

pub trait Indexable {
    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError>;
}

impl Indexable for Array {
    #[allow(unused_variables)]
    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
        unsafe {
            let err_val = af_set_array_indexer(idxr.clone().get() as MutAfIndex,
                                             self.get() as AfArray,
                                             dim as DimT);
            match err_val {
                0 => Ok(()),
                _ => Err(AfError::from(err_val)),
            }
        }
    }
}

impl Indexable for Seq {
    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
        unsafe {
            let err_val = af_set_seq_indexer(idxr.clone().get() as MutAfIndex, self as *const Seq,
                                           dim as DimT, is_batch.unwrap() as c_int);
            match err_val {
                0 => Ok(()),
                _ => Err(AfError::from(err_val)),
            }
        }
    }
}

impl Indexer {
    #[allow(unused_mut)]
    pub fn new() -> Result<Indexer, AfError> {
        unsafe {
            let mut temp: i64 = 0;
            let err_val = af_create_indexers(&mut temp as MutAfIndex);
            match err_val {
                0 => Ok(Indexer{handle: temp, count: 0}),
                _ => Err(AfError::from(err_val)),
            }
        }
    }

    pub fn set_index<T: Indexable>(&mut self, idx: &T, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
        self.count = self.count + 1;
        idx.set(self, dim, is_batch)
    }

    pub fn get(&self) -> i64 {
        self.handle
    }

    pub fn len(&self) -> u32 {
        self.count
    }
}

impl Drop for Indexer {
    fn drop(&mut self) {
        unsafe {
            let ret_val = af_release_indexers(self.handle as MutAfIndex);
            match ret_val {
                0 => (),
                _ => panic!("Failed to release indexers resource: {}", ret_val),
            }
        }
    }
}

pub fn index(input: &Array, seqs: &[Seq]) -> Result<Array, AfError> {
    unsafe {
        let mut temp: i64 = 0;
        let err_val = af_index(&mut temp as MutAfArray
                               , input.get() as AfArray, seqs.len() as u32
                               , seqs.as_ptr() as *const Seq);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn row(input: &Array, row_num: u64) -> Result<Array, AfError> {
    let dims_err = input.dims();
    let dims = match dims_err {
        Ok(dim) =>  dim.clone(),
        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
    };

    index(input, &[Seq::new(row_num as f64, row_num as f64, 1.0)
                    ,Seq::new(0.0, dims[1] as f64 - 1.0, 1.0)])
}

pub fn col(input: &Array, col_num: u64) -> Result<Array, AfError> {
    let dims_err = input.dims();
    let dims = match dims_err {
        Ok(dim) =>  dim.clone(),
        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
    };

    index(input, &[Seq::new(0.0, dims[0] as f64 - 1.0, 1.0)
                    ,Seq::new(col_num as f64, col_num as f64, 1.0)])
}

pub fn lookup(input: &Array, indices: &Array, seq_dim: i32) -> Result<Array, AfError> {
    unsafe {
        let mut temp: i64 = 0;
        let err_val = af_lookup(&mut temp as MutAfArray, input.get() as AfArray,
                                indices.get() as AfArray, seq_dim as c_uint);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn assign_seq(lhs: &Array, ndims: usize, seqs: &[Seq], rhs: &Array) -> Result<Array, AfError> {
    unsafe{
        let mut temp: i64 = 0;
        let err_val = af_assign_seq(&mut temp as MutAfArray, lhs.get() as AfArray,
                                    ndims as c_uint, seqs.as_ptr() as *const Seq,
                                    rhs.get() as AfArray);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn index_gen(input: &Array, indices: Indexer) -> Result<Array, AfError> {
    unsafe{
        let mut temp: i64 = 0;
        let err_val = af_index_gen(&mut temp as MutAfArray, input.get() as AfArray,
                                   indices.len() as DimT, indices.get() as *const IndexT);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn assign_gen(lhs: &Array, indices: &Indexer, rhs: &Array) -> Result<Array, AfError> {
    unsafe{
        let mut temp: i64 = 0;
        let err_val = af_assign_gen(&mut temp as MutAfArray, lhs.get() as AfArray,
                                    indices.len() as DimT, indices.get() as *const IndexT,
                                    rhs.get() as AfArray);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}
<commit_msg>add rows and cols (note the plural) to parallel af mainline<commit_after>extern crate libc;

use array::Array;
use defines::AfError;
use seq::Seq;
use self::libc::{c_int, c_uint, c_longlong};

type MutAfIndex    = *mut self::libc::c_longlong;
type MutAfArray    = *mut self::libc::c_longlong;
type AfArray       = self::libc::c_longlong;
type DimT          = self::libc::c_longlong;
type IndexT        = self::libc::c_longlong;

#[allow(dead_code)]
extern {
    fn af_create_indexers(indexers: MutAfIndex) -> c_int;
    fn af_set_array_indexer(indexer: MutAfIndex, idx: AfArray, dim: DimT) -> c_int;
    fn af_set_seq_indexer(indexer: MutAfIndex, idx: *const Seq, dim: DimT, is_batch: c_int) -> c_int;
    fn af_release_indexers(indexers: MutAfIndex) -> c_int;

    fn af_index(out: MutAfArray, input: AfArray, ndims: c_uint, index: *const Seq) -> c_int;
    fn af_lookup(out: MutAfArray, arr: AfArray, indices: AfArray, dim: c_uint) -> c_int;
    fn af_assign_seq(out: MutAfArray, lhs: AfArray, ndims: c_uint, indices: *const Seq, rhs: AfArray) -> c_int;
    fn af_index_gen(out: MutAfArray, input: AfArray, ndims: DimT, indices: *const IndexT) -> c_int;
    fn af_assign_gen(out: MutAfArray, lhs: AfArray, ndims: DimT, indices: *const IndexT, rhs: AfArray) -> c_int;
}

pub struct Indexer {
    handle: i64,
    count: u32,
}

pub trait Indexable {
    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError>;
}

impl Indexable for Array {
    #[allow(unused_variables)]
    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
        unsafe {
            let err_val = af_set_array_indexer(idxr.clone().get() as MutAfIndex,
                                             self.get() as AfArray,
                                             dim as DimT);
            match err_val {
                0 => Ok(()),
                _ => Err(AfError::from(err_val)),
            }
        }
    }
}

impl Indexable for Seq {
    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
        unsafe {
            let err_val = af_set_seq_indexer(idxr.clone().get() as MutAfIndex, self as *const Seq,
                                           dim as DimT, is_batch.unwrap() as c_int);
            match err_val {
                0 => Ok(()),
                _ => Err(AfError::from(err_val)),
            }
        }
    }
}

impl Indexer {
    #[allow(unused_mut)]
    pub fn new() -> Result<Indexer, AfError> {
        unsafe {
            let mut temp: i64 = 0;
            let err_val = af_create_indexers(&mut temp as MutAfIndex);
            match err_val {
                0 => Ok(Indexer{handle: temp, count: 0}),
                _ => Err(AfError::from(err_val)),
            }
        }
    }

    pub fn set_index<T: Indexable>(&mut self, idx: &T, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
        self.count = self.count + 1;
        idx.set(self, dim, is_batch)
    }

    pub fn get(&self) -> i64 {
        self.handle
    }

    pub fn len(&self) -> u32 {
        self.count
    }
}

impl Drop for Indexer {
    fn drop(&mut self) {
        unsafe {
            let ret_val = af_release_indexers(self.handle as MutAfIndex);
            match ret_val {
                0 => (),
                _ => panic!("Failed to release indexers resource: {}", ret_val),
            }
        }
    }
}

pub fn index(input: &Array, seqs: &[Seq]) -> Result<Array, AfError> {
    unsafe {
        let mut temp: i64 = 0;
        let err_val = af_index(&mut temp as MutAfArray
                               , input.get() as AfArray, seqs.len() as u32
                               , seqs.as_ptr() as *const Seq);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn row(input: &Array, row_num: u64) -> Result<Array, AfError> {
    let dims_err = input.dims();
    let dims = match dims_err {
        Ok(dim) =>  dim.clone(),
        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
    };

    index(input, &[Seq::new(row_num as f64, row_num as f64, 1.0)
                    ,Seq::new(0.0, dims[1] as f64 - 1.0, 1.0)])
}

pub fn rows(input: &Array, first: u64, last: u64) -> Result<Array, AfError> {
    let dims_err = input.dims();
    let dims = match dims_err {
        Ok(dim) =>  dim.clone(),
        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
    };

    index(input, &[Seq::new(first as f64, last as f64, 1.0)
                    ,Seq::new(0.0, dims[1] as f64 - 1.0, 1.0)])
}

pub fn col(input: &Array, col_num: u64) -> Result<Array, AfError> {
    let dims_err = input.dims();
    let dims = match dims_err {
        Ok(dim) =>  dim.clone(),
        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
    };

    index(input, &[Seq::new(0.0, dims[0] as f64 - 1.0, 1.0)
                    ,Seq::new(col_num as f64, col_num as f64, 1.0)])
}

pub fn cols(input: &Array, first: u64, last: u64) -> Result<Array, AfError> {
    let dims_err = input.dims();
    let dims = match dims_err {
        Ok(dim) =>  dim.clone(),
        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
    };

    index(input, &[Seq::new(0.0, dims[0] as f64 - 1.0, 1.0)
                    ,Seq::new(first as f64, last as f64, 1.0)])
}

pub fn lookup(input: &Array, indices: &Array, seq_dim: i32) -> Result<Array, AfError> {
    unsafe {
        let mut temp: i64 = 0;
        let err_val = af_lookup(&mut temp as MutAfArray, input.get() as AfArray,
                                indices.get() as AfArray, seq_dim as c_uint);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn assign_seq(lhs: &Array, ndims: usize, seqs: &[Seq], rhs: &Array) -> Result<Array, AfError> {
    unsafe{
        let mut temp: i64 = 0;
        let err_val = af_assign_seq(&mut temp as MutAfArray, lhs.get() as AfArray,
                                    ndims as c_uint, seqs.as_ptr() as *const Seq,
                                    rhs.get() as AfArray);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn index_gen(input: &Array, indices: Indexer) -> Result<Array, AfError> {
    unsafe{
        let mut temp: i64 = 0;
        let err_val = af_index_gen(&mut temp as MutAfArray, input.get() as AfArray,
                                   indices.len() as DimT, indices.get() as *const IndexT);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}

pub fn assign_gen(lhs: &Array, indices: &Indexer, rhs: &Array) -> Result<Array, AfError> {
    unsafe{
        let mut temp: i64 = 0;
        let err_val = af_assign_gen(&mut temp as MutAfArray, lhs.get() as AfArray,
                                    indices.len() as DimT, indices.get() as *const IndexT,
                                    rhs.get() as AfArray);
        match err_val {
            0 => Ok(Array::from(temp)),
            _ => Err(AfError::from(err_val)),
        }
    }
}
<|endoftext|>


<commit_before>use super::*;
use byteorder::BigEndian;
use byteorder::WriteBytesExt;
use client;
use futures::prelude::*;
use index::btree;
use index::btree::test::u64_to_slice;
use index::key_with_id;
use index::lsmtree::tree::KeyRange;
use index::lsmtree::tree::LSMTree;
use index::Cursor;
use index::Ordering;
use itertools::Itertools;
use ram::types::Id;
use rand::distributions::Uniform;
use rand::thread_rng;
use rand::Rng;
use rayon::prelude::*;
use server::NebServer;
use server::ServerOptions;
use smallvec::SmallVec;
use std::env;
use std::io::Cursor as StdCursor;
use std::sync::Arc;
use std::thread;
use std::time::Duration;

fn dump_trees(lsm_tree: &LSMTree, name: &str) {
    for i in 0..lsm_tree.trees.len() {
        lsm_tree.trees[i].dump(&format!("{}_lsm_{}_dump.json", name, i));
    }
}

fn default_key_range() -> KeyRange {
    (smallvec!(), smallvec!())
}

#[test]
pub fn insertions() {
    env_logger::init();
    let tree = Arc::new(LSMTree::new(default_key_range(), Id::unit_id()));
    let num = env::var("LSM_TREE_TEST_ITEMS")
        // this value cannot do anything useful to the test
        // must arrange a long-term test to cover every levels
        .unwrap_or("331776".to_string())
        .parse::<u64>()
        .unwrap();

    let tree_clone = tree.clone();
    debug!("Testing LSM-tree");
    thread::spawn(move || loop {
        thread::sleep(Duration::from_secs(10));
        let tree_len = tree_clone.len();
        debug!(
            "LSM-Tree now have {}/{} elements, total {:.2}%",
            tree_len,
            num,
            tree_len as f32 / num as f32 * 100.0
        );
    });

    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        key_with_id(&mut key, &id);
        tree.insert(key);
    });

    dump_trees(&*tree, "stage_1_before_insertion");
    for _ in 0..20 {
        tree.check_and_merge();
    }
    dump_trees(&*tree, "stage_1_after_insertion");

    debug!("Start point search validations");
    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        debug!("checking: {}", i);
        let cursor = tree.seek(&key, Ordering::Forward);
        key_with_id(&mut key, &id);
        assert_eq!(cursor.current(), Some(&key), "{}", i);
    });

    LSMTree::start_sentinel(&tree);

    let mut test_data = (num..num * 2).collect_vec();
    thread_rng().shuffle(test_data.as_mut_slice());
    test_data.par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        key_with_id(&mut key, &id);
        tree.insert(key);
    });

    dump_trees(&*tree, "stage_2_after_insertion");
    thread::sleep(Duration::new(30, 0));
    dump_trees(&*tree, "stage_2_waited_insertion");

    (0..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        debug!("checking: {}", i);
        let cursor = tree.seek(&key, Ordering::Forward);
        key_with_id(&mut key, &id);
        assert_eq!(cursor.current(), Some(&key), "{}", i);
    });
}

#[test]
pub fn hybrid() {
    env_logger::init();
    let tree = Arc::new(LSMTree::new(default_key_range(), Id::unit_id()));
    let num = env::var("LSM_TREE_TEST_ITEMS")
        // this value cannot do anything useful to the test
        // must arrange a long-term test to cover every levels
        .unwrap_or("663552".to_string())
        .parse::<u64>()
        .unwrap();

    LSMTree::start_sentinel(&tree);
    let tree_clone = tree.clone();
    debug!("Testing LSM-tree");
    thread::spawn(move || loop {
        thread::sleep(Duration::from_secs(10));
        let tree_len = tree_clone.len();
        debug!(
            "LSM-Tree now have {}/{} elements, total {:.2}%",
            tree_len,
            num,
            tree_len as f32 / num as f32 * 100.0
        );
    });

    let mut test_data = (0..num).collect_vec();
    thread_rng().shuffle(test_data.as_mut_slice());
    test_data.par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        key_with_id(&mut key, &id);
        tree.insert(key);
    });

    dump_trees(&*tree, "hybird_after_insertion");
    thread::sleep(Duration::new(30, 0));
    dump_trees(&*tree, "hybird_after_sleep");

    debug!("Start point search validations");
    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        debug!("checking: {}", i);
        let cursor = tree.seek(&key, Ordering::Forward);
        key_with_id(&mut key, &id);
        assert_eq!(cursor.current(), Some(&key), "{}", i);
    });
}
<commit_msg>check insertions<commit_after>use super::*;
use byteorder::BigEndian;
use byteorder::WriteBytesExt;
use client;
use futures::prelude::*;
use index::btree;
use index::btree::test::u64_to_slice;
use index::key_with_id;
use index::lsmtree::tree::KeyRange;
use index::lsmtree::tree::LSMTree;
use index::Cursor;
use index::Ordering;
use itertools::Itertools;
use ram::types::Id;
use rand::distributions::Uniform;
use rand::thread_rng;
use rand::Rng;
use rayon::prelude::*;
use server::NebServer;
use server::ServerOptions;
use smallvec::SmallVec;
use std::env;
use std::io::Cursor as StdCursor;
use std::sync::Arc;
use std::thread;
use std::time::Duration;

fn dump_trees(lsm_tree: &LSMTree, name: &str) {
    for i in 0..lsm_tree.trees.len() {
        lsm_tree.trees[i].dump(&format!("{}_lsm_{}_dump.json", name, i));
    }
}

fn default_key_range() -> KeyRange {
    (smallvec!(), smallvec!())
}

#[test]
pub fn insertions() {
    env_logger::init();
    let tree = Arc::new(LSMTree::new(default_key_range(), Id::unit_id()));
    let num = env::var("LSM_TREE_TEST_ITEMS")
        // this value cannot do anything useful to the test
        // must arrange a long-term test to cover every levels
        .unwrap_or("331776".to_string())
        .parse::<u64>()
        .unwrap();

    let tree_clone = tree.clone();
    debug!("Testing LSM-tree");
    thread::spawn(move || loop {
        thread::sleep(Duration::from_secs(10));
        let tree_len = tree_clone.len();
        debug!(
            "LSM-Tree now have {}/{} elements, total {:.2}%",
            tree_len,
            num,
            tree_len as f32 / num as f32 * 100.0
        );
    });

    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        key_with_id(&mut key, &id);
        tree.insert(key);
    });

    debug!("Start point search validations");
    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        debug!("checking: {}", i);
        let cursor = tree.seek(&key, Ordering::Forward);
        key_with_id(&mut key, &id);
        assert_eq!(cursor.current(), Some(&key), "{}", i);
    });

    dump_trees(&*tree, "stage_1_before_insertion");
    for _ in 0..50 {
        tree.check_and_merge();
    }
    dump_trees(&*tree, "stage_1_after_insertion");

    debug!("Start point search validations");
    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        debug!("checking: {}", i);
        let cursor = tree.seek(&key, Ordering::Forward);
        key_with_id(&mut key, &id);
        assert_eq!(cursor.current(), Some(&key), "{}", i);
    });
}

#[test]
pub fn hybrid() {
    env_logger::init();
    let tree = Arc::new(LSMTree::new(default_key_range(), Id::unit_id()));
    let num = env::var("LSM_TREE_TEST_ITEMS")
        // this value cannot do anything useful to the test
        // must arrange a long-term test to cover every levels
        .unwrap_or("663552".to_string())
        .parse::<u64>()
        .unwrap();

    LSMTree::start_sentinel(&tree);
    let tree_clone = tree.clone();
    debug!("Testing LSM-tree");
    thread::spawn(move || loop {
        thread::sleep(Duration::from_secs(10));
        let tree_len = tree_clone.len();
        debug!(
            "LSM-Tree now have {}/{} elements, total {:.2}%",
            tree_len,
            num,
            tree_len as f32 / num as f32 * 100.0
        );
    });

    let mut test_data = (0..num).collect_vec();
    thread_rng().shuffle(test_data.as_mut_slice());
    test_data.par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        key_with_id(&mut key, &id);
        tree.insert(key);
    });

    dump_trees(&*tree, "hybird_after_insertion");
    thread::sleep(Duration::new(30, 0));
    dump_trees(&*tree, "hybird_after_sleep");

    debug!("Start point search validations");
    (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
        let i = *i;
        let id = Id::new(0, i);
        let key_slice = u64_to_slice(i);
        let mut key = SmallVec::from_slice(&key_slice);
        debug!("checking: {}", i);
        let cursor = tree.seek(&key, Ordering::Forward);
        key_with_id(&mut key, &id);
        assert_eq!(cursor.current(), Some(&key), "{}", i);
    });
}
<|endoftext|>


<commit_before>extern crate fragment;

use self::fragment::matching;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use std::clone::Clone;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

#[derive(Debug, PartialEq)]
pub struct IndexedPath(PathBuf);

impl ToString for IndexedPath {
    fn to_string(&self) -> String {
        self.0.to_string_lossy().into_owned()
    }
}

impl Clone for IndexedPath {
    fn clone(&self) -> Self {
        IndexedPath(self.0.clone())
    }

    fn clone_from(&mut self, source: &Self) {
        self.0 = source.0.clone()
    }
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new(),
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        for entry in WalkDir::new(&self.path) {
            match relative_entry_path(entry, prefix_length) {
                Some(entry_path) => self.entries.push(IndexedPath(PathBuf::from(entry_path))),
                _ => (),
            }
        }
    }

    pub fn find(&self, term: &str, limit: usize, case_sensitive: bool) -> Vec<PathBuf> {
        matching::find(term, &self.entries, limit, case_sensitive)
            .into_iter()
            .map(|r| {
                let IndexedPath(other_thing) = r.clone();
                other_thing
            })
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    match entry {
        Ok(e) => {
            match e.path().metadata() {
                Ok(metadata) => {
                    if metadata.is_file() {
                        match e.path().to_str() {
                            Some(path) => Some(path[prefix_length..].to_string()),
                            _ => None,
                        }
                    } else {
                        None
                    }
                }
                _ => None,
            }
        }
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    extern crate fragment;

    use super::{Index, IndexedPath};
    use std::path::PathBuf;

    #[test]
    fn populate_adds_all_files_to_entries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate();
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit, false)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }
}
<commit_msg>Clean up option/result handling<commit_after>extern crate fragment;

use self::fragment::matching;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use std::clone::Clone;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

#[derive(Debug, PartialEq)]
pub struct IndexedPath(PathBuf);

impl ToString for IndexedPath {
    fn to_string(&self) -> String {
        self.0.to_string_lossy().into_owned()
    }
}

impl Clone for IndexedPath {
    fn clone(&self) -> Self {
        IndexedPath(self.0.clone())
    }

    fn clone_from(&mut self, source: &Self) {
        self.0 = source.0.clone()
    }
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new(),
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        for entry in WalkDir::new(&self.path) {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    IndexedPath(PathBuf::from(entry_path))
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize, case_sensitive: bool) -> Vec<PathBuf> {
        matching::find(term, &self.entries, limit, case_sensitive)
            .into_iter()
            .map(|r| {
                let IndexedPath(other_thing) = r.clone();
                other_thing
            })
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    extern crate fragment;

    use super::{Index, IndexedPath};
    use std::path::PathBuf;

    #[test]
    fn populate_adds_all_files_to_entries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate();
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit, false)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }
}
<|endoftext|>