<commit_before>//! An append-only collection of update batches.
//!
//! The `Spine` is a general-purpose trace implementation based on collection and merging
//! immutable batches of updates. It is generic with respect to the batch type, and can be
//! instantiated for any implementor of `trace::Batch`.

use ::Diff;
use lattice::Lattice;
use trace::{Batch, BatchReader, Trace, TraceReader};
use trace::cursor::cursor_list::CursorList;
use trace::cursor::Cursor;
use trace::Merger;

enum MergeState<K, V, T, R, B: Batch<K, V, T, R>> {
    Merging(B, B, Option<Vec<T>>, <B as Batch<K,V,T,R>>::Merger),
    Complete(B),
}

impl<K, V, T: Eq, R, B: Batch<K, V, T, R>> MergeState<K, V, T, R, B> {
    fn is_complete(&self) -> bool {
        match *self {
            MergeState::Complete(_) => true,
            _ => false,
        }
    }
    fn begin_merge(batch1: B, batch2: B, frontier: Option<Vec<T>>) -> Self {
        assert!(batch1.upper() == batch2.lower());
        let begin_merge = <B as Batch<K, V, T, R>>::begin_merge(&batch1, &batch2);
        MergeState::Merging(batch1, batch2, frontier, begin_merge)
    }
    fn work(mut self, fuel: &mut usize) -> Self {
        if let MergeState::Merging(ref source1, ref source2, ref frontier, ref mut in_progress) = self {
            in_progress.work(source1, source2, frontier, fuel);
        }
        if *fuel > 0 {
            match self {
                // ALLOC: Here is where we may de-allocate batches.
                MergeState::Merging(_, _, _, finished) => MergeState::Complete(finished.done()),
                MergeState::Complete(x) => MergeState::Complete(x),
            }
        }
        else { self }
    }
    fn len(&self) -> usize {
        match *self {
            MergeState::Merging(ref batch1, ref batch2, _, _) => batch1.len() + batch2.len(),
            MergeState::Complete(ref batch) => batch.len(),
        }
    }
}

/// An append-only collection of update tuples.
///
/// A spine maintains a small number of immutable collections of update tuples, merging the collections when
/// two have similar sizes. In this way, it allows the addition of more tuples, which may then be merged with
/// other immutable collections.
pub struct Spine<K, V, T: Lattice+Ord, R: Diff, B: Batch<K, V, T, R>> {
    phantom: ::std::marker::PhantomData<(K, V, R)>,
    advance_frontier: Vec<T>,            // Times after which the trace must accumulate correctly.
    through_frontier: Vec<T>,            // Times after which the trace must be able to subset its inputs.
    merging: Vec<MergeState<K,V,T,R,B>>, // Several possibly shared collections of updates.
    pending: Vec<B>,                     // Batches at times in advance of `frontier`.
    // max_len: usize,
    // max_dur: ::std::time::Duration,
    reserve: usize,                      // fuel reserves.
}

impl<K, V, T, R, B> TraceReader<K, V, T, R> for Spine<K, V, T, R, B>
where
    K: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    V: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    T: Lattice+Ord+Clone,   // Clone is required by `advance_by` and `batch::advance_*`.
    R: Diff,
    B: Batch<K, V, T, R>+Clone+'static,
{
    type Batch = B;
    type Cursor = CursorList<K, V, T, R, <B as BatchReader<K, V, T, R>>::Cursor>;

    fn cursor_through(&mut self, upper: &[T]) -> Option<(Self::Cursor, <Self::Cursor as Cursor<K, V, T, R>>::Storage)> {

        // we shouldn't grab a cursor into a closed trace, right?
        assert!(self.advance_frontier.len() > 0);

        // Check that `upper` is greater or equal to `self.through_frontier`.
        // Otherwise, the cut could be in `self.merging` and it is user error anyhow.
        if upper.iter().all(|t1| self.through_frontier.iter().any(|t2| t2.less_equal(t1))) {

            let mut cursors = Vec::new();
            let mut storage = Vec::new();

            for merge_state in self.merging.iter() {
                match *merge_state {
                    MergeState::Merging(ref batch1, ref batch2, _, _) => {
                        cursors.push(batch1.cursor());
                        storage.push(batch1.clone());
                        cursors.push(batch2.cursor());
                        storage.push(batch2.clone());
                    },
                    MergeState::Complete(ref batch) => {
                        cursors.push(batch.cursor());
                        storage.push(batch.clone());
                    }
                }
            }

            for batch in &self.pending {
                let include_lower = upper.iter().all(|t1| batch.lower().iter().any(|t2| t2.less_equal(t1)));
                let include_upper = upper.iter().all(|t1| batch.upper().iter().any(|t2| t2.less_equal(t1)));

                if include_lower != include_upper && upper != batch.lower() {
                    panic!("`cursor_through`: `upper` straddles batch");
                    // return None;
                }

                // include pending batches
                if include_upper {
                    cursors.push(batch.cursor());
                    storage.push(batch.clone());
                }
            }
            Some((CursorList::new(cursors, &storage), storage))
        }
        else {
            None
        }
    }
    fn advance_by(&mut self, frontier: &[T]) {
        self.advance_frontier = frontier.to_vec();
        if self.advance_frontier.len() == 0 {
            self.pending.clear();
            self.merging.clear();
        }
    }
    fn advance_frontier(&mut self) -> &[T] { &self.advance_frontier[..] }
    fn distinguish_since(&mut self, frontier: &[T]) {
        self.through_frontier = frontier.to_vec();
        self.consider_merges();
    }
    fn distinguish_frontier(&mut self) -> &[T] { &self.through_frontier[..] }

    fn map_batches<F: FnMut(&Self::Batch)>(&mut self, mut f: F) {
        for batch in self.merging.iter() {
            match *batch {
                MergeState::Merging(ref batch1, ref batch2, _, _) => { f(batch1); f(batch2); },
                MergeState::Complete(ref batch) => { f(batch); },
            }
        }
        for batch in self.pending.iter() {
            f(batch);
        }
    }
}

// A trace implementation for any key type that can be borrowed from or converted into `Key`.
// TODO: Almost all this implementation seems to be generic with respect to the trace and batch types.
impl<K, V, T, R, B> Trace<K, V, T, R> for Spine<K, V, T, R, B>
where
    K: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    V: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    T: Lattice+Ord+Clone,   // Clone is required by `advance_by` and `batch::advance_*`.
    R: Diff,
    B: Batch<K, V, T, R>+Clone+'static,
{

    fn new() -> Self {
        Spine {
            phantom: ::std::marker::PhantomData,
            advance_frontier: vec![<T as Lattice>::minimum()],
            through_frontier: vec![<T as Lattice>::minimum()],
            merging: Vec::new(),
            pending: Vec::new(),
            // max_len: 0,
            // max_dur: Default::default(),
            reserve: 0,
        }
    }

    // Ideally, this method acts as insertion of `batch`, even if we are not yet able to begin
    // merging the batch. This means it is a good time to perform amortized work proportional
    // to the size of batch.
    fn insert(&mut self, batch: Self::Batch) {

        // we can ignore degenerate batches (TODO: learn where they come from; suppress them?)
        if batch.lower() != batch.upper() {
            self.pending.push(batch);
            self.consider_merges();
        }
        else {
            // degenerate batches had best be empty.
            assert!(batch.len() == 0);
        }
    }
}

impl<K, V, T, R, B> Spine<K, V, T, R, B>
where
    K: Ord+Clone,           // Clone is required by `advance_mut`.
    V: Ord+Clone,           // Clone is required by `advance_mut`.
    T: Lattice+Ord+Clone,   // Clone is required by `advance_mut`.
    R: Diff,
    B: Batch<K, V, T, R>,
{
    // Migrate data from `self.pending` into `self.merging`.
    #[inline(never)]
    fn consider_merges(&mut self) {

        // let start = ::std::time::Instant::now();

        // TODO: We could consider merging in batches here, rather than in sequence.
        //       Little is currently known about whether this is important ...
        while self.pending.len() > 0 &&
              self.through_frontier.iter().all(|t1| self.pending[0].upper().iter().any(|t2| t2.less_equal(t1)))
        {
            // this could be a VecDeque, if we ever notice this.
            let batch = self.pending.remove(0);

            // First, we want to complete any merges of batches smaller than `batch`.
            if self.merging.len() > 1 {

                let mut most_recent = self.merging.pop().unwrap();

                while self.merging.last_mut().map(|x| x.len() < batch.len()) == Some(true) {

                    // finish all smaller merges unconditionally.
                    let mut less_recent = self.merging.pop().unwrap();
                    most_recent = most_recent.work(&mut usize::max_value());
                    less_recent = less_recent.work(&mut usize::max_value());

                    match (less_recent, most_recent) {
                        (MergeState::Complete(less), MergeState::Complete(most)) => {
                            assert!(less.upper() == most.lower());
                            most_recent = MergeState::begin_merge(less, most, None);
                        },
                        _ => panic!("spine_fueled: unmerged small data discovered; logic bug!"),
                    }
                }

                self.merging.push(most_recent);
            }

            // add some fuel to reserve, and get ready to spend some.
            self.reserve += 4 * (self.merging.len()) * batch.len();
            let mut fuel = self.reserve / 10;
            self.reserve -= fuel;

            let mut most_recent = MergeState::Complete(batch);

            while fuel > 0 && self.merging.last_mut().map(|x| x.len() < 2 * most_recent.len()) == Some(true) {

                let mut less_recent = self.merging.pop().unwrap();
                most_recent = most_recent.work(&mut fuel);
                less_recent = less_recent.work(&mut fuel);

                match (less_recent, most_recent) {

                    (MergeState::Complete(less), MergeState::Complete(most)) => {
                        assert!(less.upper() == most.lower());
                        let frontier = if self.merging.len() == 0 { Some(self.advance_frontier.clone()) } else { None };
                        most_recent = MergeState::begin_merge(less, most, frontier);
                    }
                    (less, most) => {
                        // Can't merge; stash `less` and assume stash `most` next round.
                        assert!(fuel == 0);
                        most_recent = most;
                        self.merging.push(less);
                    }
                }
            }

            self.merging.push(most_recent);

            // as long as we have fuel to spend, and not all merges complete ...
            while fuel > 0 && self.merging.iter().any(|b| !b.is_complete()) {

                // Spend any remaining fuel.
                for index in (0 .. self.merging.len()).rev() {
                    if fuel > 0 {
                        let temp = self.merging.remove(index);
                        self.merging.insert(index, temp.work(&mut fuel));
                    }
                }

                // Scan forward, looking for possible merges to start.
                let mut index = 1;
                while index < self.merging.len() {
                    if self.merging[index-1].len() < 2 * self.merging[index].len() && self.merging[index-1].is_complete() && self.merging[index].is_complete() {
                        let less_recent = self.merging.remove(index-1);
                        let more_recent = self.merging.remove(index-1);
                        match (less_recent, more_recent) {
                            (MergeState::Complete(less), MergeState::Complete(more)) => {
                                let frontier = if index == 1 { Some(self.advance_frontier.clone()) } else { None };
                                self.merging.insert(index - 1, MergeState::begin_merge(less, more, frontier));
                            },
                            _ => panic!("unreachable"),
                        }
                    }
                    else {
                        index += 1;
                    }
                }
            }

            if fuel > 0 {
                self.reserve = 0;
            }
        }

        // let elapsed = start.elapsed();
        // if elapsed > self.max_dur {
        //     self.max_dur = elapsed;
        //     // println!("max_dur increased to {:?}", self.max_dur);
        // }

        // // if self.merging.len() + self.pending.len() > self.max_len {
        // //     self.max_len = self.merging.len() + self.pending.len();
        // //     println!("max_len increased to {:?}", self.max_len);
        // // }
    }
}

// impl<K, V, T: Lattice+Ord, R: Diff, B: Batch<K, V, T, R>> Drop for Spine<K, V, T, R, B> {
//     fn drop(&mut self) {
//         println!("max_duration: {:?}", self.max_dur);
//     }
// }
<commit_msg>spine_fueled update<commit_after>//! An append-only collection of update batches.
//!
//! The `Spine` is a general-purpose trace implementation based on collection and merging
//! immutable batches of updates. It is generic with respect to the batch type, and can be
//! instantiated for any implementor of `trace::Batch`.

use ::Diff;
use lattice::Lattice;
use trace::{Batch, BatchReader, Trace, TraceReader};
use trace::cursor::cursor_list::CursorList;
use trace::cursor::Cursor;
use trace::Merger;

enum MergeState<K, V, T, R, B: Batch<K, V, T, R>> {
    Merging(B, B, Option<Vec<T>>, <B as Batch<K,V,T,R>>::Merger),
    Complete(B),
}

impl<K, V, T: Eq, R, B: Batch<K, V, T, R>> MergeState<K, V, T, R, B> {
    fn complete(mut self) -> B {
        if let MergeState::Merging(ref source1, ref source2, ref frontier, ref mut in_progress) = self {
            let mut fuel = usize::max_value();
            in_progress.work(source1, source2, frontier, &mut fuel);
            assert!(fuel > 0);
        }
        match self {
            // ALLOC: Here is where we may de-allocate batches.
            MergeState::Merging(_, _, _, finished) => finished.done(),
            MergeState::Complete(x) => x,
        }
    }

    fn is_complete(&self) -> bool {
        match *self {
            MergeState::Complete(_) => true,
            _ => false,
        }
    }
    fn begin_merge(batch1: B, batch2: B, frontier: Option<Vec<T>>) -> Self {
        assert!(batch1.upper() == batch2.lower());
        let begin_merge = <B as Batch<K, V, T, R>>::begin_merge(&batch1, &batch2);
        MergeState::Merging(batch1, batch2, frontier, begin_merge)
    }
    fn work(mut self, fuel: &mut usize) -> Self {
        if let MergeState::Merging(ref source1, ref source2, ref frontier, ref mut in_progress) = self {
            in_progress.work(source1, source2, frontier, fuel);
        }
        if *fuel > 0 {
            match self {
                // ALLOC: Here is where we may de-allocate batches.
                MergeState::Merging(_, _, _, finished) => MergeState::Complete(finished.done()),
                MergeState::Complete(x) => MergeState::Complete(x),
            }
        }
        else { self }
    }
    fn len(&self) -> usize {
        match *self {
            MergeState::Merging(ref batch1, ref batch2, _, _) => batch1.len() + batch2.len(),
            MergeState::Complete(ref batch) => batch.len(),
        }
    }
}

/// An append-only collection of update tuples.
///
/// A spine maintains a small number of immutable collections of update tuples, merging the collections when
/// two have similar sizes. In this way, it allows the addition of more tuples, which may then be merged with
/// other immutable collections.
pub struct Spine<K, V, T: Lattice+Ord, R: Diff, B: Batch<K, V, T, R>> {
    phantom: ::std::marker::PhantomData<(K, V, R)>,
    advance_frontier: Vec<T>,            // Times after which the trace must accumulate correctly.
    through_frontier: Vec<T>,            // Times after which the trace must be able to subset its inputs.
    merging: Vec<Option<MergeState<K,V,T,R,B>>>, // Several possibly shared collections of updates.
    pending: Vec<B>,                     // Batches at times in advance of `frontier`.
    // max_len: usize,
    // max_dur: ::std::time::Duration,
    // reserve: usize,                      // fuel reserves.
}

impl<K, V, T, R, B> TraceReader<K, V, T, R> for Spine<K, V, T, R, B>
where
    K: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    V: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    T: Lattice+Ord+Clone,   // Clone is required by `advance_by` and `batch::advance_*`.
    R: Diff,
    B: Batch<K, V, T, R>+Clone+'static,
{
    type Batch = B;
    type Cursor = CursorList<K, V, T, R, <B as BatchReader<K, V, T, R>>::Cursor>;

    fn cursor_through(&mut self, upper: &[T]) -> Option<(Self::Cursor, <Self::Cursor as Cursor<K, V, T, R>>::Storage)> {

        // we shouldn't grab a cursor into a closed trace, right?
        assert!(self.advance_frontier.len() > 0);

        // Check that `upper` is greater or equal to `self.through_frontier`.
        // Otherwise, the cut could be in `self.merging` and it is user error anyhow.
        if upper.iter().all(|t1| self.through_frontier.iter().any(|t2| t2.less_equal(t1))) {

            let mut cursors = Vec::new();
            let mut storage = Vec::new();

            for merge_state in self.merging.iter().rev() {
                match *merge_state {
                    Some(MergeState::Merging(ref batch1, ref batch2, _, _)) => {
                        cursors.push(batch1.cursor());
                        storage.push(batch1.clone());
                        cursors.push(batch2.cursor());
                        storage.push(batch2.clone());
                    },
                    Some(MergeState::Complete(ref batch)) => {
                        cursors.push(batch.cursor());
                        storage.push(batch.clone());
                    },
                    None => { }
                }
            }

            for batch in &self.pending {
                let include_lower = upper.iter().all(|t1| batch.lower().iter().any(|t2| t2.less_equal(t1)));
                let include_upper = upper.iter().all(|t1| batch.upper().iter().any(|t2| t2.less_equal(t1)));

                if include_lower != include_upper && upper != batch.lower() {
                    panic!("`cursor_through`: `upper` straddles batch");
                    // return None;
                }

                // include pending batches
                if include_upper {
                    cursors.push(batch.cursor());
                    storage.push(batch.clone());
                }
            }
            Some((CursorList::new(cursors, &storage), storage))
        }
        else {
            None
        }
    }
    fn advance_by(&mut self, frontier: &[T]) {
        self.advance_frontier = frontier.to_vec();
        if self.advance_frontier.len() == 0 {
            self.pending.clear();
            self.merging.clear();
        }
    }
    fn advance_frontier(&mut self) -> &[T] { &self.advance_frontier[..] }
    fn distinguish_since(&mut self, frontier: &[T]) {
        self.through_frontier = frontier.to_vec();
        self.consider_merges();
    }
    fn distinguish_frontier(&mut self) -> &[T] { &self.through_frontier[..] }

    fn map_batches<F: FnMut(&Self::Batch)>(&mut self, mut f: F) {
        for batch in self.merging.iter().rev() {
            match *batch {
                Some(MergeState::Merging(ref batch1, ref batch2, _, _)) => { f(batch1); f(batch2); },
                Some(MergeState::Complete(ref batch)) => { f(batch); },
                None => { },
            }
        }
        for batch in self.pending.iter() {
            f(batch);
        }
    }
}

// A trace implementation for any key type that can be borrowed from or converted into `Key`.
// TODO: Almost all this implementation seems to be generic with respect to the trace and batch types.
impl<K, V, T, R, B> Trace<K, V, T, R> for Spine<K, V, T, R, B>
where
    K: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    V: Ord+Clone,           // Clone is required by `batch::advance_*` (in-place could remove).
    T: Lattice+Ord+Clone,   // Clone is required by `advance_by` and `batch::advance_*`.
    R: Diff,
    B: Batch<K, V, T, R>+Clone+'static,
{

    fn new() -> Self {
        Spine {
            phantom: ::std::marker::PhantomData,
            advance_frontier: vec![<T as Lattice>::minimum()],
            through_frontier: vec![<T as Lattice>::minimum()],
            merging: Vec::new(),
            pending: Vec::new(),
            // max_len: 0,
            // max_dur: Default::default(),
            // reserve: 0,
        }
    }

    // Ideally, this method acts as insertion of `batch`, even if we are not yet able to begin
    // merging the batch. This means it is a good time to perform amortized work proportional
    // to the size of batch.
    fn insert(&mut self, batch: Self::Batch) {

        // we can ignore degenerate batches (TODO: learn where they come from; suppress them?)
        if batch.lower() != batch.upper() {
            self.pending.push(batch);
            self.consider_merges();
        }
        else {
            // degenerate batches had best be empty.
            assert!(batch.len() == 0);
        }
    }
}

impl<K, V, T, R, B> Spine<K, V, T, R, B>
where
    K: Ord+Clone,           // Clone is required by `advance_mut`.
    V: Ord+Clone,           // Clone is required by `advance_mut`.
    T: Lattice+Ord+Clone,   // Clone is required by `advance_mut`.
    R: Diff,
    B: Batch<K, V, T, R>,
{
    // Migrate data from `self.pending` into `self.merging`.
    #[inline(never)]
    fn consider_merges(&mut self) {

        // We have a new design here, in an attempt to rationalize progressive merging of batches.
        //
        // Batches arrive with a number of records, and are assigned a power-of-two "size", which is this
        // number rounded up. Batches are placed in a slot based on their size, and each slot can either be
        //
        //  i. empty,
        //  ii. occupied by a batch,
        //  iii. occupied by a merge in progress.
        //
        // As we introduce a new batch, we need to do a few things.
        //
        //  0. We determine a size for the batch, and an implied target slot.
        //  1. For each slot smaller than the size of the batch, we should merge the batches so that the first
        //     occupied slot is no less than the target slot of the new batch.
        //  2. We perform size units of work for each merge in progress, starting from the large batches and
        //     working backwards (so that we have had the chance to complete work before starting a new merge).
        //  3. We install the new batch in the target slot, initiating a merge if necessary.
        //
        // At various points we may find that the number of updates in a batch is out of line with the size
        // associated with the slot the batch currently occupies. In this case, we can slide a batch (or a merge
        // in progress) down through empty slots, as long as the number of updates is no more than the associated
        // size.

        while self.pending.len() > 0 &&
              self.through_frontier.iter().all(|t1| self.pending[0].upper().iter().any(|t2| t2.less_equal(t1)))
        {
            // this could be a VecDeque, if we ever notice this.
            let batch = self.pending.remove(0);

            // println!("batch.len(): {:?}", batch.len());

            // Step 0: Determine batch size and target slot.
            let batch_size = batch.len().next_power_of_two();
            let batch_index = batch_size.trailing_zeros() as usize;
            while self.merging.len() <= batch_index { self.merging.push(None); }

            if self.merging.len() > 32 { println!("len: {:?}", self.merging.len()); }

            // Step 1: Forcibly merge batches in lower slots.
            for position in 0 .. batch_index {
                if let Some(mut batch) = self.merging[position].take() {
                    let batch = batch.complete();
                    if let Some(mut batch2) = self.merging[position+1].take() {
                        let batch2 = batch2.complete();
                        self.merging[position+1] = Some(MergeState::begin_merge(batch2, batch, None));
                    }
                    else {
                        self.merging[position+1] = Some(MergeState::Complete(batch));
                    };
                }
            }

            // Step 2: Perform `size` work on each in-progress merge, from large to small.
            let mut fuel;// = 0; //8 * batch_size * self.merging.len();
            for position in (batch_index .. self.merging.len()).rev() {
                fuel = 16 * batch_size;
                if let Some(mut batch) = self.merging[position].take() {
                    let batch = batch.work(&mut fuel);
                    if batch.is_complete() {
                        let batch = batch.complete();
                        let intended_position = batch.len().next_power_of_two().trailing_zeros() as usize;
                        if intended_position > position {
                            while self.merging.len() <= intended_position { self.merging.push(None); }
                            if let Some(batch2) = self.merging[position+1].take() {
                                if !batch2.is_complete() {
                                    println!("batch[{}] not complete (size: {:?})", position+1, batch2.len());
                                    println!("sizes: {:?}", self.merging.iter().map(|x| x.as_ref().map(|y|y.len()).unwrap_or(0)).collect::<Vec<_>>());
                                }
                                assert!(batch2.is_complete());
                                let batch2 = batch2.complete();
                                let frontier = if position + 1 == self.merging.len() { Some(self.advance_frontier.clone()) } else { None };
                                self.merging[position+1] = Some(MergeState::begin_merge(batch2, batch, frontier));
                            }
                            else {
                                self.merging[position+1] = Some(MergeState::Complete(batch));
                            };
                        }
                        else {
                            self.merging[position] = Some(MergeState::Complete(batch));
                        }
                    }
                    else {
                        self.merging[position] = Some(batch);
                    }
                }
            }

            // Step 3: Insert new batch at target position
            if let Some(batch2) = self.merging[batch_index].take() {
                assert!(batch2.is_complete());
                let batch2 = batch2.complete();
                self.merging[batch_index] = Some(MergeState::begin_merge(batch2, batch, None));
            }
            else {
                self.merging[batch_index] = Some(MergeState::Complete(batch));
            }

            // Step 4: Consider migrating complete batches to lower bins, if appropriate.
            for index in (self.merging.len() .. 1).rev() {
                if self.merging[index].as_ref().map(|x| x.is_complete() && x.len() < (1 << (index-1))).unwrap_or(false) {
                    if self.merging[index-1].is_none() {
                        self.merging[index-1] = self.merging[index].take();
                    }
                }
            }
        }
    }
}<|endoftext|>


<commit_before>use token::Token;
use token::Token::*;

#[derive(Default)]
struct CompileState {
    output: String,
    loopnum: i32,
    ifnum: i32,
    outbuffsize: i32,
}

/// Takes an offset from the current cell, and returns a string in assembly code
/// representing the register or memory region it is stored in
fn offset_to_operand(offset: i32) -> String {
    if offset == 0 {
        "%r12".to_string()
    } else {
        format!("{}(%rbx)", (offset * 8))
    }
}

fn compile_iter(state: &mut CompileState, tokens: Vec<Token>) {
    let mut outbuffpos = 0;
    for token in tokens {
        match token {
            Add(offset, value) => {
                let dest = offset_to_operand(offset);
                if value == 1 && dest == "%r12" {
                    state.output.push_str("    inc %r12\n");
                } else if value >= 1 {
                    state.output.push_str(&format!("    addq ${}, {}\n", value, dest));
                } else if value == -1 && dest == "%r12" {
                    state.output.push_str("    dec %r12\n");
                } else if value <= -1 {
                    state.output.push_str(&format!("    subq ${}, {}\n", -value, dest));
                }
            }
            MulCopy(src_idx, dest_idx, mul) => {
                let mut src = offset_to_operand(src_idx);
                let dest = offset_to_operand(dest_idx);

                if mul != -1 && mul != 1 {
                    state.output.push_str(&format!(concat!("    movq {}, %rax\n",
                                                           "    movq ${}, %rdx\n",
                                                           "    mulq %rdx\n"),
                                                   src,
                                                   mul.abs()));
                    src = "%rax".to_string();
                } else if src != "%r12" && dest != "%r12" {
                    // x86 cannot move memory to memory
                    state.output.push_str(&format!("    movq {}, %rax\n", src));
                    src = "%rax".to_string();
                }

                if mul > 0 {
                    state.output.push_str(&format!("    addq {}, {}\n", src, dest));
                } else {
                    state.output.push_str(&format!("    subq {}, {}\n", src, dest));
                }
            }
            Set(offset, value) => {
                if offset == 0 && value == 0 {
                    state.output.push_str("    xor %r12, %r12\n");
                } else if offset == 0 {
                    state.output.push_str(&format!("    movq ${}, %r12\n", value));
                } else {
                    state.output.push_str(&format!("    movq ${}, {}(%rbx)\n", value, offset * 8));
                }
            }
            Move(offset) => {
                if offset != 0 {
                    state.output.push_str("    movq %r12, (%rbx)\n");
                    if offset > 0 {
                        state.output.push_str(&format!("    addq ${}, %rbx\n", offset * 8));
                    } else {
                        state.output.push_str(&format!("    subq ${}, %rbx\n", -offset * 8));
                    }
                    state.output.push_str("    movq (%rbx), %r12\n");
                }
            }
            Loop(content) => {
                state.loopnum += 1;
                let curloop = state.loopnum;
                state.output.push_str(&format!(concat!("    jmp endloop{}\n", "    loop{}:\n"),
                                               curloop,
                                               curloop));

                compile_iter(state, content);

                state.output.push_str(&format!(concat!("    endloop{}:\n",
                                                       "    test %r12, %r12\n",
                                                       "    jnz loop{}\n"),
                                               curloop,
                                               curloop))
            }
            If(offset, content) => {
                state.ifnum += 1;
                let curif = state.ifnum;
                if offset == 0 {
                    state.output.push_str("    test %r12, %r12\n");
                } else {
                    state.output.push_str(&format!("    cmpq $0, {}(%rbx)\n", offset * 8));
                }
                state.output.push_str(&format!("    jz endif{}\n", curif));

                compile_iter(state, content);

                state.output.push_str(&format!("    endif{}:\n", curif))
            }
            Scan(offset) => {
                // Slighly more optimal than normal loop and move
                state.loopnum += 1;
                state.output.push_str(&format!(concat!("    movq %r12, (%rbx)\n",
                                                       "    jmp endloop{num}\n",
                                                       "    loop{num}:\n",
                                                       "    {add_sub} ${shift}, %rbx\n",
                                                       "    endloop{num}:\n",
                                                       "    cmp $0, (%rbx)\n",
                                                       "    jnz loop{num}\n",
                                                       "    movq (%rbx), %r12\n"),
                                               num = state.loopnum,
                                               add_sub = if offset > 0 { "addq" } else { "subq" },
                                               shift = offset.abs() * 8));
            }
            Input => {
                state.output.push_str(concat!("\n    xor %rax, %rax\n",
                                              "    xor %rdi, %rdi\n",
                                              "    movq %rbx, %rsi\n",
                                              "    movq $1, %rdx\n",
                                              "    syscall\n",
                                              "    movq (%rbx), %r12\n\n"))
            }
            LoadOut(offset, add) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                if offset == 0 {
                    state.output.push_str(&format!("    movq %r12, {}\n", outaddr));
                } else {
                    state.output.push_str(&format!("    movq {}(%rbx), %rax\n", offset * 8));
                    state.output.push_str(&format!("    movq %rax, {}\n", outaddr));
                }
                if add > 0 {
                    state.output.push_str(&format!("    addb ${}, {}\n", add, outaddr));
                } else if add < 0 {
                    state.output.push_str(&format!("    subb ${}, {}\n", -add, outaddr));
                }
                outbuffpos += 1;
            }
            LoadOutSet(value) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                state.output.push_str(&format!("    movq ${}, {}\n", value, outaddr));
                outbuffpos += 1;
            }
            Output => {
                state.output.push_str(&format!(concat!("    movq $1, %rax\n",
                                                       "    movq $1, %rdi\n",
                                                       "    movq $strbuff, %rsi\n",
                                                       "    movq ${}, %rdx\n",
                                                       "    syscall\n\n"),
                                               outbuffpos));

                if state.outbuffsize < outbuffpos + 8 {
                    state.outbuffsize = outbuffpos + 8;
                }
                outbuffpos = 0;
            }
        }
    }
}

pub fn compile(tokens: Vec<Token>, tape_size: i32) -> String {
    let mut state = CompileState::default();

    compile_iter(&mut state, tokens);

    format!(concat!(".section .bss\n",
                    "    .lcomm strbuff, {outbuffsize}\n",
                    "    .lcomm mem, {}\n",
                    "    .set startidx, mem + {}\n",
                    ".section .text\n",
                    ".global _start\n",
                    "_start:\n",
                    "    xor %r12, %r12\n",
                    "    movq $startidx, %rbx\n\n",
                    "{}\n",
                    // Exit syscall
                    "    movq $60, %rax\n",
                    "    movq $0, %rdi\n",
                    "    syscall\n"),
            tape_size,
            tape_size / 2,
            state.output,
            outbuffsize = state.outbuffsize)
}
<commit_msg>Use offset_to_operand() in Set<commit_after>use token::Token;
use token::Token::*;

#[derive(Default)]
struct CompileState {
    output: String,
    loopnum: i32,
    ifnum: i32,
    outbuffsize: i32,
}

/// Takes an offset from the current cell, and returns a string in assembly code
/// representing the register or memory region it is stored in
fn offset_to_operand(offset: i32) -> String {
    if offset == 0 {
        "%r12".to_string()
    } else {
        format!("{}(%rbx)", (offset * 8))
    }
}

fn compile_iter(state: &mut CompileState, tokens: Vec<Token>) {
    let mut outbuffpos = 0;
    for token in tokens {
        match token {
            Add(offset, value) => {
                let dest = offset_to_operand(offset);
                if value == 1 && dest == "%r12" {
                    state.output.push_str("    inc %r12\n");
                } else if value >= 1 {
                    state.output.push_str(&format!("    addq ${}, {}\n", value, dest));
                } else if value == -1 && dest == "%r12" {
                    state.output.push_str("    dec %r12\n");
                } else if value <= -1 {
                    state.output.push_str(&format!("    subq ${}, {}\n", -value, dest));
                }
            }
            MulCopy(src_idx, dest_idx, mul) => {
                let mut src = offset_to_operand(src_idx);
                let dest = offset_to_operand(dest_idx);

                if mul != -1 && mul != 1 {
                    state.output.push_str(&format!(concat!("    movq {}, %rax\n",
                                                           "    movq ${}, %rdx\n",
                                                           "    mulq %rdx\n"),
                                                   src,
                                                   mul.abs()));
                    src = "%rax".to_string();
                } else if src != "%r12" && dest != "%r12" {
                    // x86 cannot move memory to memory
                    state.output.push_str(&format!("    movq {}, %rax\n", src));
                    src = "%rax".to_string();
                }

                if mul > 0 {
                    state.output.push_str(&format!("    addq {}, {}\n", src, dest));
                } else {
                    state.output.push_str(&format!("    subq {}, {}\n", src, dest));
                }
            }
            Set(offset, value) => {
                if offset == 0 && value == 0 {
                    state.output.push_str("    xor %r12, %r12\n");
                } else {
                    state.output.push_str(&format!("    movq ${}, {}\n", value, offset_to_operand(offset)));
                }
            }
            Move(offset) => {
                if offset != 0 {
                    state.output.push_str("    movq %r12, (%rbx)\n");
                    if offset > 0 {
                        state.output.push_str(&format!("    addq ${}, %rbx\n", offset * 8));
                    } else {
                        state.output.push_str(&format!("    subq ${}, %rbx\n", -offset * 8));
                    }
                    state.output.push_str("    movq (%rbx), %r12\n");
                }
            }
            Loop(content) => {
                state.loopnum += 1;
                let curloop = state.loopnum;
                state.output.push_str(&format!(concat!("    jmp endloop{}\n", "    loop{}:\n"),
                                               curloop,
                                               curloop));

                compile_iter(state, content);

                state.output.push_str(&format!(concat!("    endloop{}:\n",
                                                       "    test %r12, %r12\n",
                                                       "    jnz loop{}\n"),
                                               curloop,
                                               curloop))
            }
            If(offset, content) => {
                state.ifnum += 1;
                let curif = state.ifnum;
                if offset == 0 {
                    state.output.push_str("    test %r12, %r12\n");
                } else {
                    state.output.push_str(&format!("    cmpq $0, {}(%rbx)\n", offset * 8));
                }
                state.output.push_str(&format!("    jz endif{}\n", curif));

                compile_iter(state, content);

                state.output.push_str(&format!("    endif{}:\n", curif))
            }
            Scan(offset) => {
                // Slighly more optimal than normal loop and move
                state.loopnum += 1;
                state.output.push_str(&format!(concat!("    movq %r12, (%rbx)\n",
                                                       "    jmp endloop{num}\n",
                                                       "    loop{num}:\n",
                                                       "    {add_sub} ${shift}, %rbx\n",
                                                       "    endloop{num}:\n",
                                                       "    cmp $0, (%rbx)\n",
                                                       "    jnz loop{num}\n",
                                                       "    movq (%rbx), %r12\n"),
                                               num = state.loopnum,
                                               add_sub = if offset > 0 { "addq" } else { "subq" },
                                               shift = offset.abs() * 8));
            }
            Input => {
                state.output.push_str(concat!("\n    xor %rax, %rax\n",
                                              "    xor %rdi, %rdi\n",
                                              "    movq %rbx, %rsi\n",
                                              "    movq $1, %rdx\n",
                                              "    syscall\n",
                                              "    movq (%rbx), %r12\n\n"))
            }
            LoadOut(offset, add) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                if offset == 0 {
                    state.output.push_str(&format!("    movq %r12, {}\n", outaddr));
                } else {
                    state.output.push_str(&format!("    movq {}(%rbx), %rax\n", offset * 8));
                    state.output.push_str(&format!("    movq %rax, {}\n", outaddr));
                }
                if add > 0 {
                    state.output.push_str(&format!("    addb ${}, {}\n", add, outaddr));
                } else if add < 0 {
                    state.output.push_str(&format!("    subb ${}, {}\n", -add, outaddr));
                }
                outbuffpos += 1;
            }
            LoadOutSet(value) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                state.output.push_str(&format!("    movq ${}, {}\n", value, outaddr));
                outbuffpos += 1;
            }
            Output => {
                state.output.push_str(&format!(concat!("    movq $1, %rax\n",
                                                       "    movq $1, %rdi\n",
                                                       "    movq $strbuff, %rsi\n",
                                                       "    movq ${}, %rdx\n",
                                                       "    syscall\n\n"),
                                               outbuffpos));

                if state.outbuffsize < outbuffpos + 8 {
                    state.outbuffsize = outbuffpos + 8;
                }
                outbuffpos = 0;
            }
        }
    }
}

pub fn compile(tokens: Vec<Token>, tape_size: i32) -> String {
    let mut state = CompileState::default();

    compile_iter(&mut state, tokens);

    format!(concat!(".section .bss\n",
                    "    .lcomm strbuff, {outbuffsize}\n",
                    "    .lcomm mem, {}\n",
                    "    .set startidx, mem + {}\n",
                    ".section .text\n",
                    ".global _start\n",
                    "_start:\n",
                    "    xor %r12, %r12\n",
                    "    movq $startidx, %rbx\n\n",
                    "{}\n",
                    // Exit syscall
                    "    movq $60, %rax\n",
                    "    movq $0, %rdi\n",
                    "    syscall\n"),
            tape_size,
            tape_size / 2,
            state.output,
            outbuffsize = state.outbuffsize)
}
<|endoftext|>


<commit_before>use llvm_sys::core::LLVMDisposeMessage;
use llvm_sys::execution_engine::{LLVMGetExecutionEngineTargetData, LLVMExecutionEngineRef, LLVMRunFunction, LLVMRunFunctionAsMain, LLVMDisposeExecutionEngine, LLVMGetFunctionAddress, LLVMAddModule, LLVMFindFunction, LLVMLinkInMCJIT, LLVMLinkInInterpreter, LLVMRemoveModule, LLVMGenericValueRef, LLVMFreeMachineCodeForFunction, LLVMAddGlobalMapping};

use module::Module;
use targets::TargetData;
use values::{AnyValue, AsValueRef, FunctionValue, GenericValue};

use std::rc::Rc;
use std::ops::Deref;
use std::ffi::{CStr, CString};
use std::mem::{forget, uninitialized, zeroed, transmute_copy, size_of};
use std::fmt::{self, Debug, Formatter};

#[derive(Debug, PartialEq, Eq)]
pub enum FunctionLookupError {
    JITNotEnabled,
    FunctionNotFound, // 404!
}

#[derive(PartialEq, Eq, Debug)]
pub struct ExecutionEngine {
    pub(crate) execution_engine: Rc<LLVMExecutionEngineRef>,
    target_data: Option<TargetData>,
    jit_mode: bool,
}

impl ExecutionEngine {
    pub(crate) fn new(execution_engine: Rc<LLVMExecutionEngineRef>, jit_mode: bool) -> ExecutionEngine {
        assert!(!execution_engine.is_null());

        // REVIEW: Will we have to do this for LLVMGetExecutionEngineTargetMachine too?
        let target_data = unsafe {
            LLVMGetExecutionEngineTargetData(*execution_engine)
        };

        ExecutionEngine {
            execution_engine: execution_engine,
            target_data: Some(TargetData::new(target_data)),
            jit_mode: jit_mode,
        }
    }

    pub fn link_in_mc_jit() {
        unsafe {
            LLVMLinkInMCJIT()
        }
    }

    pub fn link_in_interpreter() {
        unsafe {
            LLVMLinkInInterpreter();
        }
    }

    /// Maps the specified value to an address.
    ///
    /// # Example
    /// ```no_run
    /// use inkwell::targets::{InitializationConfig, Target};
    /// use inkwell::context::Context;
    /// use inkwell::OptimizationLevel;
    ///
    /// Target::initialize_native(&InitializationConfig::default()).unwrap();
    ///
    /// extern fn sumf(a: f64, b: f64) -> f64 {
    ///     a + b
    /// }
    ///
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let builder = context.create_builder();
    ///
    /// let ft = context.f64_type();
    /// let fnt = ft.fn_type(&[], false);
    ///
    /// let f = module.add_function("test_fn", &fnt, None);
    /// let b = context.append_basic_block(&f, "entry");
    /// builder.position_at_end(&b);
    ///
    /// let extf = module.add_function("sumf", &ft.fn_type(&[ &ft, &ft ], false), None);
    ///
    /// let argf = ft.const_float(64.);
    /// let retv = builder.build_call(&extf, &[ &argf, &argf ], "retv", false).left().unwrap().into_float_value();
    ///
    /// builder.build_return(Some(&retv));
    ///
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    /// ee.add_global_mapping(&extf, sumf as usize);
    ///
    /// let result = unsafe { ee.run_function(&f, &[]) }.as_float(&ft);
    ///
    /// assert_eq!(result, 128.);
    /// ```
    pub fn add_global_mapping(&self, value: &AnyValue, addr: usize) {
        unsafe {
            LLVMAddGlobalMapping(*self.execution_engine, value.as_value_ref(), addr as *mut _)
        }
    }

    /// Adds a module to an `ExecutionEngine`.
    ///
    /// The method will be `Ok(())` if the module does not belong to an `ExecutionEngine` already and `Err(())` otherwise.
    ///
    /// ```rust,no_run
    /// use inkwell::targets::{InitializationConfig, Target};
    /// use inkwell::context::Context;
    /// use inkwell::OptimizationLevel;
    ///
    /// Target::initialize_native(&InitializationConfig::default()).unwrap();
    ///
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    ///
    /// assert!(ee.add_module(&module).is_err());
    /// ```
    pub fn add_module(&self, module: &Module) -> Result<(), ()> {
        unsafe {
            LLVMAddModule(*self.execution_engine, module.module.get())
        }

        if module.owned_by_ee.borrow().is_some() {
            return Err(());
        }

        *module.owned_by_ee.borrow_mut() = Some(ExecutionEngine::new(self.execution_engine.clone(), self.jit_mode));

        Ok(())
    }

    pub fn remove_module(&self, module: &Module) -> Result<(), String> {
        match *module.owned_by_ee.borrow() {
            Some(ref ee) if *ee.execution_engine != *self.execution_engine => return Err("Module is not owned by this Execution Engine".into()),
            None => return Err("Module is not owned by an Execution Engine".into()),
            _ => ()
        }

        let mut new_module = unsafe { uninitialized() };
        let mut err_str = unsafe { zeroed() };

        let code = unsafe {
            LLVMRemoveModule(*self.execution_engine, module.module.get(), &mut new_module, &mut err_str)
        };

        if code == 1 {
            let rust_str = unsafe {
                let rust_str = CStr::from_ptr(err_str).to_string_lossy().into_owned();

                LLVMDisposeMessage(err_str);

                rust_str
            };

            return Err(rust_str);
        }

        module.module.set(new_module);
        *module.owned_by_ee.borrow_mut() = None;

        Ok(())
    }

    /// Try to load a function from the execution engine.
    ///
    /// If a target hasn't already been initialized, spurious "function not
    /// found" errors may be encountered.
    ///
    /// The [`UnsafeFunctionPointer`] trait is designed so only `unsafe extern
    /// "C"` functions can be retrieved via the `get_function()` method. If you
    /// get funny type errors then it's probably because you have specified the
    /// wrong calling convention or forgotten to specify the retrieved function
    /// is `unsafe`.
    ///
    /// # Examples
    ///
    ///
    /// ```rust
    /// # use inkwell::targets::{InitializationConfig, Target};
    /// # use inkwell::context::Context;
    /// # use inkwell::OptimizationLevel;
    /// # Target::initialize_native(&InitializationConfig::default()).unwrap();
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let builder = context.create_builder();
    ///
    /// // Set up the function signature
    /// let double = context.f64_type();
    /// let sig = double.fn_type(&[], false);
    ///
    /// // Add the function to our module
    /// let f = module.add_function("test_fn", &sig, None);
    /// let b = context.append_basic_block(&f, "entry");
    /// builder.position_at_end(&b);
    ///
    /// // Insert a return statement
    /// let ret = double.const_float(64.0);
    /// builder.build_return(Some(&ret));
    ///
    /// // create the JIT engine
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    ///
    /// // fetch our JIT'd function and execute it
    /// unsafe {
    ///     let test_fn = ee.get_function::<fn() -> f64>("test_fn").unwrap();
    ///     let return_value = test_fn();
    ///     assert_eq!(return_value, 64.0);
    /// }
    /// ```
    ///
    /// # Safety
    ///
    /// It is the caller's responsibility to ensure they call the function with
    /// the correct signature and calling convention.
    ///
    /// The `Symbol` wrapper ensures a function won't accidentally outlive the
    /// execution engine it came from, but adding functions after calling this
    /// method *may* invalidate the function pointer.
    ///
    /// [`UnsafeFunctionPointer`]: trait.UnsafeFunctionPointer.html
    pub unsafe fn get_function<F>(&self, fn_name: &str) -> Result<Symbol<F>, FunctionLookupError>
    where F: UnsafeFunctionPointer
    {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");

        let address = LLVMGetFunctionAddress(*self.execution_engine, c_string.as_ptr());

        // REVIEW: Can also return 0 if no targets are initialized.
        // One option might be to set a global to true if any at all of the targets have been
        // initialized (maybe we could figure out which config in particular is the trigger)
        // and if not return an "NoTargetsInitialized" error, instead of not found.
        if address == 0 {
            return Err(FunctionLookupError::FunctionNotFound);
        }

        assert_eq!(size_of::<F>(), size_of::<usize>(),
            "The type `F` must have the same size as a function pointer");

        Ok(Symbol {
            execution_engine: self.execution_engine.clone(),
            inner: transmute_copy(&address),
        })
    }

    // REVIEW: Not sure if an EE's target data can change.. if so we might want to update the value
    // when making this call
    pub fn get_target_data(&self) -> &TargetData {
        self.target_data.as_ref().expect("TargetData should always exist until Drop")
    }

    // REVIEW: Can also find nothing if no targeting is initialized. Maybe best to
    // do have a global flag for anything initialized. Catch is that it must be initialized
    // before EE is created
    pub fn get_function_value(&self, fn_name: &str) -> Result<FunctionValue, FunctionLookupError> {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");
        let mut function = unsafe { zeroed() };

        let code = unsafe {
            LLVMFindFunction(*self.execution_engine, c_string.as_ptr(), &mut function)
        };

        if code == 0 {
            return FunctionValue::new(function).ok_or(FunctionLookupError::FunctionNotFound)
        };

        Err(FunctionLookupError::FunctionNotFound)
    }

    // TODOC: Marked as unsafe because input function could very well do something unsafe. It's up to the caller
    // to ensure that doesn't happen by defining their function correctly.
    pub unsafe fn run_function(&self, function: &FunctionValue, args: &[&GenericValue]) -> GenericValue {
        let mut args: Vec<LLVMGenericValueRef> = args.iter()
                                                     .map(|val| val.generic_value)
                                                     .collect();

        let value = LLVMRunFunction(*self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_mut_ptr()); // REVIEW: usize to u32 ok??

        GenericValue::new(value)
    }

    pub fn run_function_as_main(&self, function: &FunctionValue) {
        let args = vec![]; // TODO: Support argc, argv
        let env_p = vec![]; // REVIEW: No clue what this is

        unsafe {
            LLVMRunFunctionAsMain(*self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_ptr(), env_p.as_ptr()); // REVIEW: usize to u32 cast ok??
        }
    }

    pub fn free_fn_machine_code(&self, function: &FunctionValue) {
        unsafe {
            LLVMFreeMachineCodeForFunction(*self.execution_engine, function.as_value_ref())
        }
    }
}

// Modules owned by the EE will be discarded by the EE so we don't
// want owned modules to drop.
impl Drop for ExecutionEngine {
    fn drop(&mut self) {
        forget(self.target_data.take().expect("TargetData should always exist until Drop"));

        if Rc::strong_count(&self.execution_engine) == 1 {
            unsafe {
                LLVMDisposeExecutionEngine(*self.execution_engine);
            }
        }
    }
}

/// A wrapper around a function pointer which ensures the symbol being pointed
/// to doesn't accidentally outlive its execution engine.
#[derive(Clone)]
pub struct Symbol<F> {
    pub(crate) execution_engine: Rc<LLVMExecutionEngineRef>,
    inner: F,
}

impl<F: UnsafeFunctionPointer> Deref for Symbol<F> {
    type Target = F;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F> Debug for Symbol<F> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_tuple("Symbol")
            .field(&"<unnamed>")
            .finish()
    }
}

/// Marker trait representing an unsafe function pointer (`unsafe extern "C" fn(A, B, ...) -> Output`).
pub trait UnsafeFunctionPointer: private::Sealed + Copy {}

mod private {
    /// A sealed trait which ensures nobody outside this crate can implement
    /// `UnsafeFunctionPointer`.
    ///
    /// See https://rust-lang-nursery.github.io/api-guidelines/future-proofing.html
    pub trait Sealed {}
}

macro_rules! impl_unsafe_fn {
    ($( $param:ident ),*) => {
        impl<Output, $( $param ),*> private::Sealed for unsafe extern "C" fn($( $param ),*) -> Output {}
        impl<Output, $( $param ),*> UnsafeFunctionPointer for unsafe extern "C" fn($( $param ),*) -> Output {}
    };
}

impl_unsafe_fn!();
impl_unsafe_fn!(A);
impl_unsafe_fn!(A, B);
impl_unsafe_fn!(A, B, C);
impl_unsafe_fn!(A, B, C, D);
impl_unsafe_fn!(A, B, C, D, E);
impl_unsafe_fn!(A, B, C, D, E, F);
impl_unsafe_fn!(A, B, C, D, E, F, G);
impl_unsafe_fn!(A, B, C, D, E, F, G, H);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K, L);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K, L, M);

<commit_msg>Missed a doc test<commit_after>use llvm_sys::core::LLVMDisposeMessage;
use llvm_sys::execution_engine::{LLVMGetExecutionEngineTargetData, LLVMExecutionEngineRef, LLVMRunFunction, LLVMRunFunctionAsMain, LLVMDisposeExecutionEngine, LLVMGetFunctionAddress, LLVMAddModule, LLVMFindFunction, LLVMLinkInMCJIT, LLVMLinkInInterpreter, LLVMRemoveModule, LLVMGenericValueRef, LLVMFreeMachineCodeForFunction, LLVMAddGlobalMapping};

use module::Module;
use targets::TargetData;
use values::{AnyValue, AsValueRef, FunctionValue, GenericValue};

use std::rc::Rc;
use std::ops::Deref;
use std::ffi::{CStr, CString};
use std::mem::{forget, uninitialized, zeroed, transmute_copy, size_of};
use std::fmt::{self, Debug, Formatter};

#[derive(Debug, PartialEq, Eq)]
pub enum FunctionLookupError {
    JITNotEnabled,
    FunctionNotFound, // 404!
}

#[derive(PartialEq, Eq, Debug)]
pub struct ExecutionEngine {
    pub(crate) execution_engine: Rc<LLVMExecutionEngineRef>,
    target_data: Option<TargetData>,
    jit_mode: bool,
}

impl ExecutionEngine {
    pub(crate) fn new(execution_engine: Rc<LLVMExecutionEngineRef>, jit_mode: bool) -> ExecutionEngine {
        assert!(!execution_engine.is_null());

        // REVIEW: Will we have to do this for LLVMGetExecutionEngineTargetMachine too?
        let target_data = unsafe {
            LLVMGetExecutionEngineTargetData(*execution_engine)
        };

        ExecutionEngine {
            execution_engine: execution_engine,
            target_data: Some(TargetData::new(target_data)),
            jit_mode: jit_mode,
        }
    }

    pub fn link_in_mc_jit() {
        unsafe {
            LLVMLinkInMCJIT()
        }
    }

    pub fn link_in_interpreter() {
        unsafe {
            LLVMLinkInInterpreter();
        }
    }

    /// Maps the specified value to an address.
    ///
    /// # Example
    /// ```no_run
    /// use inkwell::targets::{InitializationConfig, Target};
    /// use inkwell::context::Context;
    /// use inkwell::OptimizationLevel;
    ///
    /// Target::initialize_native(&InitializationConfig::default()).unwrap();
    ///
    /// extern fn sumf(a: f64, b: f64) -> f64 {
    ///     a + b
    /// }
    ///
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let builder = context.create_builder();
    ///
    /// let ft = context.f64_type();
    /// let fnt = ft.fn_type(&[], false);
    ///
    /// let f = module.add_function("test_fn", &fnt, None);
    /// let b = context.append_basic_block(&f, "entry");
    /// builder.position_at_end(&b);
    ///
    /// let extf = module.add_function("sumf", &ft.fn_type(&[ &ft, &ft ], false), None);
    ///
    /// let argf = ft.const_float(64.);
    /// let retv = builder.build_call(&extf, &[ &argf, &argf ], "retv", false).left().unwrap().into_float_value();
    ///
    /// builder.build_return(Some(&retv));
    ///
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    /// ee.add_global_mapping(&extf, sumf as usize);
    ///
    /// let result = unsafe { ee.run_function(&f, &[]) }.as_float(&ft);
    ///
    /// assert_eq!(result, 128.);
    /// ```
    pub fn add_global_mapping(&self, value: &AnyValue, addr: usize) {
        unsafe {
            LLVMAddGlobalMapping(*self.execution_engine, value.as_value_ref(), addr as *mut _)
        }
    }

    /// Adds a module to an `ExecutionEngine`.
    ///
    /// The method will be `Ok(())` if the module does not belong to an `ExecutionEngine` already and `Err(())` otherwise.
    ///
    /// ```rust,no_run
    /// use inkwell::targets::{InitializationConfig, Target};
    /// use inkwell::context::Context;
    /// use inkwell::OptimizationLevel;
    ///
    /// Target::initialize_native(&InitializationConfig::default()).unwrap();
    ///
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    ///
    /// assert!(ee.add_module(&module).is_err());
    /// ```
    pub fn add_module(&self, module: &Module) -> Result<(), ()> {
        unsafe {
            LLVMAddModule(*self.execution_engine, module.module.get())
        }

        if module.owned_by_ee.borrow().is_some() {
            return Err(());
        }

        *module.owned_by_ee.borrow_mut() = Some(ExecutionEngine::new(self.execution_engine.clone(), self.jit_mode));

        Ok(())
    }

    pub fn remove_module(&self, module: &Module) -> Result<(), String> {
        match *module.owned_by_ee.borrow() {
            Some(ref ee) if *ee.execution_engine != *self.execution_engine => return Err("Module is not owned by this Execution Engine".into()),
            None => return Err("Module is not owned by an Execution Engine".into()),
            _ => ()
        }

        let mut new_module = unsafe { uninitialized() };
        let mut err_str = unsafe { zeroed() };

        let code = unsafe {
            LLVMRemoveModule(*self.execution_engine, module.module.get(), &mut new_module, &mut err_str)
        };

        if code == 1 {
            let rust_str = unsafe {
                let rust_str = CStr::from_ptr(err_str).to_string_lossy().into_owned();

                LLVMDisposeMessage(err_str);

                rust_str
            };

            return Err(rust_str);
        }

        module.module.set(new_module);
        *module.owned_by_ee.borrow_mut() = None;

        Ok(())
    }

    /// Try to load a function from the execution engine.
    ///
    /// If a target hasn't already been initialized, spurious "function not
    /// found" errors may be encountered.
    ///
    /// The [`UnsafeFunctionPointer`] trait is designed so only `unsafe extern
    /// "C"` functions can be retrieved via the `get_function()` method. If you
    /// get funny type errors then it's probably because you have specified the
    /// wrong calling convention or forgotten to specify the retrieved function
    /// is `unsafe`.
    ///
    /// # Examples
    ///
    ///
    /// ```rust
    /// # use inkwell::targets::{InitializationConfig, Target};
    /// # use inkwell::context::Context;
    /// # use inkwell::OptimizationLevel;
    /// # Target::initialize_native(&InitializationConfig::default()).unwrap();
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let builder = context.create_builder();
    ///
    /// // Set up the function signature
    /// let double = context.f64_type();
    /// let sig = double.fn_type(&[], false);
    ///
    /// // Add the function to our module
    /// let f = module.add_function("test_fn", &sig, None);
    /// let b = context.append_basic_block(&f, "entry");
    /// builder.position_at_end(&b);
    ///
    /// // Insert a return statement
    /// let ret = double.const_float(64.0);
    /// builder.build_return(Some(&ret));
    ///
    /// // create the JIT engine
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    ///
    /// // fetch our JIT'd function and execute it
    /// unsafe {
    ///     let test_fn = ee.get_function::<unsafe extern "C" fn() -> f64>("test_fn").unwrap();
    ///     let return_value = test_fn();
    ///     assert_eq!(return_value, 64.0);
    /// }
    /// ```
    ///
    /// # Safety
    ///
    /// It is the caller's responsibility to ensure they call the function with
    /// the correct signature and calling convention.
    ///
    /// The `Symbol` wrapper ensures a function won't accidentally outlive the
    /// execution engine it came from, but adding functions after calling this
    /// method *may* invalidate the function pointer.
    ///
    /// [`UnsafeFunctionPointer`]: trait.UnsafeFunctionPointer.html
    pub unsafe fn get_function<F>(&self, fn_name: &str) -> Result<Symbol<F>, FunctionLookupError>
    where F: UnsafeFunctionPointer
    {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");

        let address = LLVMGetFunctionAddress(*self.execution_engine, c_string.as_ptr());

        // REVIEW: Can also return 0 if no targets are initialized.
        // One option might be to set a global to true if any at all of the targets have been
        // initialized (maybe we could figure out which config in particular is the trigger)
        // and if not return an "NoTargetsInitialized" error, instead of not found.
        if address == 0 {
            return Err(FunctionLookupError::FunctionNotFound);
        }

        assert_eq!(size_of::<F>(), size_of::<usize>(),
            "The type `F` must have the same size as a function pointer");

        Ok(Symbol {
            execution_engine: self.execution_engine.clone(),
            inner: transmute_copy(&address),
        })
    }

    // REVIEW: Not sure if an EE's target data can change.. if so we might want to update the value
    // when making this call
    pub fn get_target_data(&self) -> &TargetData {
        self.target_data.as_ref().expect("TargetData should always exist until Drop")
    }

    // REVIEW: Can also find nothing if no targeting is initialized. Maybe best to
    // do have a global flag for anything initialized. Catch is that it must be initialized
    // before EE is created
    pub fn get_function_value(&self, fn_name: &str) -> Result<FunctionValue, FunctionLookupError> {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");
        let mut function = unsafe { zeroed() };

        let code = unsafe {
            LLVMFindFunction(*self.execution_engine, c_string.as_ptr(), &mut function)
        };

        if code == 0 {
            return FunctionValue::new(function).ok_or(FunctionLookupError::FunctionNotFound)
        };

        Err(FunctionLookupError::FunctionNotFound)
    }

    // TODOC: Marked as unsafe because input function could very well do something unsafe. It's up to the caller
    // to ensure that doesn't happen by defining their function correctly.
    pub unsafe fn run_function(&self, function: &FunctionValue, args: &[&GenericValue]) -> GenericValue {
        let mut args: Vec<LLVMGenericValueRef> = args.iter()
                                                     .map(|val| val.generic_value)
                                                     .collect();

        let value = LLVMRunFunction(*self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_mut_ptr()); // REVIEW: usize to u32 ok??

        GenericValue::new(value)
    }

    pub fn run_function_as_main(&self, function: &FunctionValue) {
        let args = vec![]; // TODO: Support argc, argv
        let env_p = vec![]; // REVIEW: No clue what this is

        unsafe {
            LLVMRunFunctionAsMain(*self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_ptr(), env_p.as_ptr()); // REVIEW: usize to u32 cast ok??
        }
    }

    pub fn free_fn_machine_code(&self, function: &FunctionValue) {
        unsafe {
            LLVMFreeMachineCodeForFunction(*self.execution_engine, function.as_value_ref())
        }
    }
}

// Modules owned by the EE will be discarded by the EE so we don't
// want owned modules to drop.
impl Drop for ExecutionEngine {
    fn drop(&mut self) {
        forget(self.target_data.take().expect("TargetData should always exist until Drop"));

        if Rc::strong_count(&self.execution_engine) == 1 {
            unsafe {
                LLVMDisposeExecutionEngine(*self.execution_engine);
            }
        }
    }
}

/// A wrapper around a function pointer which ensures the symbol being pointed
/// to doesn't accidentally outlive its execution engine.
#[derive(Clone)]
pub struct Symbol<F> {
    pub(crate) execution_engine: Rc<LLVMExecutionEngineRef>,
    inner: F,
}

impl<F: UnsafeFunctionPointer> Deref for Symbol<F> {
    type Target = F;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F> Debug for Symbol<F> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_tuple("Symbol")
            .field(&"<unnamed>")
            .finish()
    }
}

/// Marker trait representing an unsafe function pointer (`unsafe extern "C" fn(A, B, ...) -> Output`).
pub trait UnsafeFunctionPointer: private::Sealed + Copy {}

mod private {
    /// A sealed trait which ensures nobody outside this crate can implement
    /// `UnsafeFunctionPointer`.
    ///
    /// See https://rust-lang-nursery.github.io/api-guidelines/future-proofing.html
    pub trait Sealed {}
}

macro_rules! impl_unsafe_fn {
    ($( $param:ident ),*) => {
        impl<Output, $( $param ),*> private::Sealed for unsafe extern "C" fn($( $param ),*) -> Output {}
        impl<Output, $( $param ),*> UnsafeFunctionPointer for unsafe extern "C" fn($( $param ),*) -> Output {}
    };
}

impl_unsafe_fn!();
impl_unsafe_fn!(A);
impl_unsafe_fn!(A, B);
impl_unsafe_fn!(A, B, C);
impl_unsafe_fn!(A, B, C, D);
impl_unsafe_fn!(A, B, C, D, E);
impl_unsafe_fn!(A, B, C, D, E, F);
impl_unsafe_fn!(A, B, C, D, E, F, G);
impl_unsafe_fn!(A, B, C, D, E, F, G, H);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K, L);
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K, L, M);

<|endoftext|>


<commit_before>use std::{convert, io};
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection {
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl PeerConnection {
    fn connect(peer: &Peer, metainfo: &Metainfo) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                self.send_interested();
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                self.send_interested();
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    let piece_index = self.next_piece_to_request();
                    self.send_request(piece_index);
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }
    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            let mut message = vec![];
            message.push(0);
            message.push(0);
            message.push(0);
            message.push(1);
            message.push(2);
            try!(self.stream.write_all(&message));
            self.am_i_interested = true;
        }
        Ok(())
    }

    fn next_piece_to_request(&self) -> usize {
        0
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        let mut message = vec![];
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(13);
        message.push(6);
        message.extend(convert_usize_to_bytes(piece).iter().cloned());
        message.extend(convert_usize_to_bytes(0).iter().cloned());
        message.extend(convert_usize_to_bytes(BLOCK_SIZE).iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }
}

#[derive(Debug)]
enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Downloading all the pieces of the file<commit_after>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();

        println!("num_pieces = {}", num_pieces);
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                  try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            let mut message = vec![];
            message.push(0);
            message.push(0);
            message.push(0);
            message.push(1);
            message.push(2);
            try!(self.stream.write_all(&message));
            self.am_i_interested = true;
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        println!("Sending request for {}", piece);
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        println!("Request size: {}", request_size);

        let mut message = vec![];
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(13);
        message.push(6);
        message.extend(convert_usize_to_bytes(piece).iter().cloned());
        message.extend(convert_usize_to_bytes(0).iter().cloned());
        message.extend(convert_usize_to_bytes(request_size).iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request => write!(f, "Request"),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>