<commit_before>use crate::test_utils::TestRandom;
use crate::{BeaconBlockBody, ChainSpec, Eth1Data, Hash256, ProposalSignedData, Slot};
use bls::Signature;
use rand::RngCore;
use serde_derive::Serialize;
use ssz::{hash, Decodable, DecodeError, Encodable, SszStream, TreeHash};

#[derive(Debug, PartialEq, Clone, Serialize)]
pub struct BeaconBlock {
    pub slot: Slot,
    pub parent_root: Hash256,
    pub state_root: Hash256,
    pub randao_reveal: Signature,
    pub eth1_data: Eth1Data,
    pub signature: Signature,
    pub body: BeaconBlockBody,
}

impl BeaconBlock {
    /// Produce the first block of the Beacon Chain.
    pub fn genesis(state_root: Hash256, spec: &ChainSpec) -> BeaconBlock {
        BeaconBlock {
            slot: spec.genesis_slot,
            parent_root: spec.zero_hash,
            state_root,
            randao_reveal: spec.empty_signature.clone(),
            eth1_data: Eth1Data {
                deposit_root: spec.zero_hash,
                block_hash: spec.zero_hash,
            },
            signature: spec.empty_signature.clone(),
            body: BeaconBlockBody {
                proposer_slashings: vec![],
                attester_slashings: vec![],
                attestations: vec![],
                deposits: vec![],
                exits: vec![],
            },
        }
    }

    pub fn canonical_root(&self) -> Hash256 {
        Hash256::from(&self.hash_tree_root()[..])
    }

    pub fn proposal_root(&self, spec: &ChainSpec) -> Hash256 {
        let block_without_signature_root = {
            let mut block_without_signature = self.clone();
            block_without_signature.signature = spec.empty_signature.clone();
            block_without_signature.canonical_root()
        };

        let proposal = ProposalSignedData {
            slot: self.slot,
            shard: spec.beacon_chain_shard_number,
            block_root: block_without_signature_root,
        };
        Hash256::from_slice(&proposal.hash_tree_root()[..])
    }
}

impl Encodable for BeaconBlock {
    fn ssz_append(&self, s: &mut SszStream) {
        s.append(&self.slot);
        s.append(&self.parent_root);
        s.append(&self.state_root);
        s.append(&self.randao_reveal);
        s.append(&self.eth1_data);
        s.append(&self.signature);
        s.append(&self.body);
    }
}

impl Decodable for BeaconBlock {
    fn ssz_decode(bytes: &[u8], i: usize) -> Result<(Self, usize), DecodeError> {
        let (slot, i) = <_>::ssz_decode(bytes, i)?;
        let (parent_root, i) = <_>::ssz_decode(bytes, i)?;
        let (state_root, i) = <_>::ssz_decode(bytes, i)?;
        let (randao_reveal, i) = <_>::ssz_decode(bytes, i)?;
        let (eth1_data, i) = <_>::ssz_decode(bytes, i)?;
        let (signature, i) = <_>::ssz_decode(bytes, i)?;
        let (body, i) = <_>::ssz_decode(bytes, i)?;

        Ok((
            Self {
                slot,
                parent_root,
                state_root,
                randao_reveal,
                eth1_data,
                signature,
                body,
            },
            i,
        ))
    }
}

impl TreeHash for BeaconBlock {
    fn hash_tree_root(&self) -> Vec<u8> {
        let mut result: Vec<u8> = vec![];
        result.append(&mut self.slot.hash_tree_root());
        result.append(&mut self.parent_root.hash_tree_root());
        result.append(&mut self.state_root.hash_tree_root());
        result.append(&mut self.randao_reveal.hash_tree_root());
        result.append(&mut self.eth1_data.hash_tree_root());
        result.append(&mut self.signature.hash_tree_root());
        result.append(&mut self.body.hash_tree_root());
        hash(&result)
    }
}

impl<T: RngCore> TestRandom<T> for BeaconBlock {
    fn random_for_test(rng: &mut T) -> Self {
        Self {
            slot: <_>::random_for_test(rng),
            parent_root: <_>::random_for_test(rng),
            state_root: <_>::random_for_test(rng),
            randao_reveal: <_>::random_for_test(rng),
            eth1_data: <_>::random_for_test(rng),
            signature: <_>::random_for_test(rng),
            body: <_>::random_for_test(rng),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
    use ssz::ssz_encode;

    #[test]
    pub fn test_ssz_round_trip() {
        let mut rng = XorShiftRng::from_seed([42; 16]);
        let original = BeaconBlock::random_for_test(&mut rng);

        let bytes = ssz_encode(&original);
        let (decoded, _) = <_>::ssz_decode(&bytes, 0).unwrap();

        assert_eq!(original, decoded);
    }

    #[test]
    pub fn test_hash_tree_root() {
        let mut rng = XorShiftRng::from_seed([42; 16]);
        let original = BeaconBlock::random_for_test(&mut rng);

        let result = original.hash_tree_root();

        assert_eq!(result.len(), 32);
        // TODO: Add further tests
        // https://github.com/sigp/lighthouse/issues/170
    }
}
<commit_msg>Fix clippy lint<commit_after>use crate::test_utils::TestRandom;
use crate::{BeaconBlockBody, ChainSpec, Eth1Data, Hash256, ProposalSignedData, Slot};
use bls::Signature;
use rand::RngCore;
use serde_derive::Serialize;
use ssz::{hash, Decodable, DecodeError, Encodable, SszStream, TreeHash};

#[derive(Debug, PartialEq, Clone, Serialize)]
pub struct BeaconBlock {
    pub slot: Slot,
    pub parent_root: Hash256,
    pub state_root: Hash256,
    pub randao_reveal: Signature,
    pub eth1_data: Eth1Data,
    pub signature: Signature,
    pub body: BeaconBlockBody,
}

impl BeaconBlock {
    /// Produce the first block of the Beacon Chain.
    pub fn genesis(state_root: Hash256, spec: &ChainSpec) -> BeaconBlock {
        BeaconBlock {
            slot: spec.genesis_slot,
            parent_root: spec.zero_hash,
            state_root,
            randao_reveal: spec.empty_signature.clone(),
            eth1_data: Eth1Data {
                deposit_root: spec.zero_hash,
                block_hash: spec.zero_hash,
            },
            signature: spec.empty_signature.clone(),
            body: BeaconBlockBody {
                proposer_slashings: vec![],
                attester_slashings: vec![],
                attestations: vec![],
                deposits: vec![],
                exits: vec![],
            },
        }
    }

    pub fn canonical_root(&self) -> Hash256 {
        Hash256::from(&self.hash_tree_root()[..])
    }

    pub fn proposal_root(&self, spec: &ChainSpec) -> Hash256 {
        let block_without_signature_root = {
            let mut block_without_signature = self.clone();
            block_without_signature.signature = spec.empty_signature.clone();
            block_without_signature.canonical_root()
        };

        let proposal = ProposalSignedData {
            slot: self.slot,
            shard: spec.beacon_chain_shard_number,
            block_root: block_without_signature_root,
        };
        Hash256::from(&proposal.hash_tree_root()[..])
    }
}

impl Encodable for BeaconBlock {
    fn ssz_append(&self, s: &mut SszStream) {
        s.append(&self.slot);
        s.append(&self.parent_root);
        s.append(&self.state_root);
        s.append(&self.randao_reveal);
        s.append(&self.eth1_data);
        s.append(&self.signature);
        s.append(&self.body);
    }
}

impl Decodable for BeaconBlock {
    fn ssz_decode(bytes: &[u8], i: usize) -> Result<(Self, usize), DecodeError> {
        let (slot, i) = <_>::ssz_decode(bytes, i)?;
        let (parent_root, i) = <_>::ssz_decode(bytes, i)?;
        let (state_root, i) = <_>::ssz_decode(bytes, i)?;
        let (randao_reveal, i) = <_>::ssz_decode(bytes, i)?;
        let (eth1_data, i) = <_>::ssz_decode(bytes, i)?;
        let (signature, i) = <_>::ssz_decode(bytes, i)?;
        let (body, i) = <_>::ssz_decode(bytes, i)?;

        Ok((
            Self {
                slot,
                parent_root,
                state_root,
                randao_reveal,
                eth1_data,
                signature,
                body,
            },
            i,
        ))
    }
}

impl TreeHash for BeaconBlock {
    fn hash_tree_root(&self) -> Vec<u8> {
        let mut result: Vec<u8> = vec![];
        result.append(&mut self.slot.hash_tree_root());
        result.append(&mut self.parent_root.hash_tree_root());
        result.append(&mut self.state_root.hash_tree_root());
        result.append(&mut self.randao_reveal.hash_tree_root());
        result.append(&mut self.eth1_data.hash_tree_root());
        result.append(&mut self.signature.hash_tree_root());
        result.append(&mut self.body.hash_tree_root());
        hash(&result)
    }
}

impl<T: RngCore> TestRandom<T> for BeaconBlock {
    fn random_for_test(rng: &mut T) -> Self {
        Self {
            slot: <_>::random_for_test(rng),
            parent_root: <_>::random_for_test(rng),
            state_root: <_>::random_for_test(rng),
            randao_reveal: <_>::random_for_test(rng),
            eth1_data: <_>::random_for_test(rng),
            signature: <_>::random_for_test(rng),
            body: <_>::random_for_test(rng),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
    use ssz::ssz_encode;

    #[test]
    pub fn test_ssz_round_trip() {
        let mut rng = XorShiftRng::from_seed([42; 16]);
        let original = BeaconBlock::random_for_test(&mut rng);

        let bytes = ssz_encode(&original);
        let (decoded, _) = <_>::ssz_decode(&bytes, 0).unwrap();

        assert_eq!(original, decoded);
    }

    #[test]
    pub fn test_hash_tree_root() {
        let mut rng = XorShiftRng::from_seed([42; 16]);
        let original = BeaconBlock::random_for_test(&mut rng);

        let result = original.hash_tree_root();

        assert_eq!(result.len(), 32);
        // TODO: Add further tests
        // https://github.com/sigp/lighthouse/issues/170
    }
}
<|endoftext|>


<commit_before>use autograd::{Function, FuncIntf, FuncDelegate, FIWrap, OptVarKindList};
use tensor::{TensorKindList, OptTensorKindList};

pub struct Threshold {
    delegate: Function,
    threshold: f64,
    value: f64,
    inplace: bool,
}

impl Threshold {
    pub fn new(threshold: f64, value: f64, inplace: bool) -> FIWrap<Self> {
        if inplace {
            panic!("in-place processing requires value ({}) to not \
                    exceed threshold ({})",
                   value,
                   threshold);
        }

        FIWrap::new(Threshold {
                        delegate: Function::new(),
                        threshold: threshold,
                        value: value,
                        inplace: inplace,
                    })
    }
}

impl_func_delegate!(Threshold);

impl FuncIntf for Threshold {
    fn forward(&mut self, input_: &mut TensorKindList) -> TensorKindList {
        let mut input = input_.remove(0);
        let mut backend = input.backend();
        let mut output = if self.inplace {
            self.mark_dirty(&vec![input.clone()]);
            input.clone()
        } else {
            input.new(()).resize_as_(&input)
        };
        // XXX check if training
        self.save_for_backward(&vec![input.clone()]);
        backend.Threshold_updateOutput(&mut input,
                                       &mut output,
                                       self.threshold,
                                       self.value,
                                       self.inplace);
        vec![output]
    }
    fn backward(&mut self, input: &mut OptTensorKindList) -> OptTensorKindList {
        /* Why are they doing backprop on a volatile variable? */
        unimplemented!()
    }
    fn backward_var(&mut self, input: &mut OptVarKindList) -> OptVarKindList {
        unimplemented!()
    }
}
<commit_msg>implement backward for Threshold<commit_after>use autograd::{Function, FuncIntf, FuncDelegate, FIWrap, OptVarKindList, Variable, VariableArgs, VarKind};
use tensor::{TensorKindList, OptTensorKindList};

pub struct Threshold {
    delegate: Function,
    saved_tensors: Vec<::tensor::TensorKind>,
    threshold: f64,
    value: f64,
    inplace: bool,
}

impl Threshold {
    pub fn new(threshold: f64, value: f64, inplace: bool) -> FIWrap<Self> {
        if inplace {
            panic!("in-place processing requires value ({}) to not \
                    exceed threshold ({})",
                   value,
                   threshold);
        }

        FIWrap::new(Threshold {
                        delegate: Function::new(),
                        saved_tensors: Vec::new(),
                        threshold: threshold,
                        value: value,
                        inplace: inplace,
                    })
    }
}

impl_func_delegate!(Threshold);

impl FuncIntf for Threshold {
    fn forward(&mut self, input_: &mut TensorKindList) -> TensorKindList {
        let mut input = input_.remove(0);
        let mut backend = input.backend();
        let mut output = if self.inplace {
            self.mark_dirty(&vec![input.clone()]);
            input.clone()
        } else {
            input.new(()).resize_as_(&input)
        };
        // XXX check if training
        self.saved_tensors.push(input.clone());
        backend.Threshold_updateOutput(&mut input,
                                       &mut output,
                                       self.threshold,
                                       self.value,
                                       self.inplace);
        vec![output]
    }
    fn backward(&mut self, grad_output_list: &mut OptTensorKindList) -> OptTensorKindList {
        let mut grad_output  = grad_output_list.remove(0).unwrap();
        let mut input = self.saved_tensors.remove(0);
        let needs_input_grad = self.needs_input_grad().clone();

        println!("Threshold backward");
        let grad_input = if needs_input_grad[0] {
            let mut grad_input = input.new(());
            let mut backend = input.backend();
            backend.Threshold_updateGradInput(
                &mut input,
                &mut grad_output,
                &mut grad_input,
                self.threshold,
                self.value,
                false
            );
            Some(grad_input)
        } else {
            None
        };
        vec![grad_input, None, None, None]
    }
}
<|endoftext|>


<commit_before>/// Stub implementation of a systick timer since the NRF doesn't have the
/// Cortex-M0 Systick. This will need to be replaced with one of the other
/// timers on the NRF, or maybe we don't care if only one process will ever run
/// on the NRF51

use common::VolatileCell;
use core::mem;

struct CLOCK {
    pub tasks_hfclkstart    : VolatileCell<u32>,
    pub tasks_hfclkstop     : VolatileCell<u32>,
    pub tasks_lfclkstart    : VolatileCell<u32>,
    pub tasks_lfclkstop     : VolatileCell<u32>,
    pub tasks_cal           : VolatileCell<u32>,
    pub tasks_cstart        : VolatileCell<u32>,
    pub tasks_cstop         : VolatileCell<u32>,
    _reserved1: [VolatileCell<u32>; 57],
    pub events_hfclkstarted : VolatileCell<u32>,
    pub events_lfclkstarted : VolatileCell<u32>,
    pub done                : VolatileCell<u32>,
    pub ctto                : VolatileCell<u32>,
    _reserved2: [VolatileCell<u32>; 125],
    pub intenset            : VolatileCell<u32>,
    pub intenclr            : VolatileCell<u32>,
    _reserved3: [VolatileCell<u32>; 63],
    pub hfclkrun            : VolatileCell<u32>,
    pub hfclkstat           : VolatileCell<u32>,
    _reserved4: [VolatileCell<u32>; 1],
    pub lfclkrun            : VolatileCell<u32>,
    pub lfclkstat           : VolatileCell<u32>,
    pub lfclksrccopy        : VolatileCell<u32>,
    _reserved5: [VolatileCell<u32>; 62],
    pub lfclksrc            : VolatileCell<u32>,
    _reserved6: [VolatileCell<u32>; 7],
    pub ctiv                : VolatileCell<u32>,
    _reserved7: [VolatileCell<u32>; 5],
    pub xtalfreq            : VolatileCell<u32>,
}

struct RTC {
    pub tasks_start: VolatileCell<u32>,
    pub tasks_stop: VolatileCell<u32>,
    pub tasks_clear: VolatileCell<u32>,
    pub tasks_trigovrflw: VolatileCell<u32>,
    _reserved1: [VolatileCell<u32>; 60],
    pub events_tick: VolatileCell<u32>,
    pub events_ovrflw: VolatileCell<u32>,
    _reserved2: [VolatileCell<u32>; 14],
    pub events_compare: [VolatileCell<u32>; 4],
    _reserved3: [VolatileCell<u32>; 108],
    pub inten:    VolatileCell<u32>,
    pub intenset: VolatileCell<u32>,
    pub intenclr: VolatileCell<u32>,
    _reserved4: [VolatileCell<u32>; 13],
    pub evten: VolatileCell<u32>,
    pub evtenset: VolatileCell<u32>,
    pub evtenclr: VolatileCell<u32>,
    _reserved5: [VolatileCell<u32>; 110],
    pub counter: VolatileCell<u32>,
    pub prescaler: VolatileCell<u32>,
    _reserved6: [VolatileCell<u32>; 13],
    pub cc: [VolatileCell<u32>; 4],
    _reserved7: [VolatileCell<u32>; 683],
    pub power: VolatileCell<u32>,
}
const CLOCK_BASE: usize = 0x40000000;
const RTC0_BASE: usize =  0x4000B000;

static mut VAL: usize = 0;

#[allow(non_snake_case)]
fn RTC0() -> &'static RTC {
        unsafe { mem::transmute(RTC0_BASE as usize) }
}

#[allow(non_snake_case)]
fn CLOCK() -> &'static CLOCK {
        unsafe { mem::transmute(CLOCK_BASE as usize) }
}

pub unsafe fn reset() {
    RTC0().tasks_clear.set(1);
}

pub unsafe fn set_timer(val: usize) {
    // Clock is 32768Hz
    // Each tick is 30.51us, approximate as 30.5
    let ticks: u32 = (val as u32 * 10) / 305;
    RTC0().cc[0].set(ticks);
    VAL = val
}

pub unsafe fn enable(_: bool) {
    CLOCK().xtalfreq.set(0);
    CLOCK().tasks_lfclkstop.set(1);
    CLOCK().events_lfclkstarted.set(0);
    CLOCK().lfclksrc.set(0);
    CLOCK().tasks_lfclkstart.set(1);
    while CLOCK().events_lfclkstarted.get() == 0  {}
    RTC0().prescaler.set(0);
    RTC0().tasks_stop.set(1);
    RTC0().tasks_clear.set(1);
    RTC0().evtenclr.set(0xffffffff);
    RTC0().intenclr.set(0xffffffff);
    RTC0().tasks_start.set(1);
}

pub unsafe fn overflowed() -> bool {
//    value() >= (RTC0().cc[0].get() as usize)
    false
}

pub unsafe fn value() -> usize {
    // Clock is 32768Hz
    // Each tick is 30.51us, approximate as 30.5
    let counter = RTC0().counter.get() as usize;
    let val = counter * 305 / 10;
    VAL * 1000
}

<commit_msg>Add some comments<commit_after>/// Stub implementation of a systick timer since the NRF doesn't have the
/// Cortex-M0 Systick. This will need to be replaced with one of the other
/// timers on the NRF, or maybe we don't care if only one process will ever run
/// on the NRF51

use common::VolatileCell;
use core::mem;

struct CLOCK {
    pub tasks_hfclkstart    : VolatileCell<u32>,
    pub tasks_hfclkstop     : VolatileCell<u32>,
    pub tasks_lfclkstart    : VolatileCell<u32>,
    pub tasks_lfclkstop     : VolatileCell<u32>,
    pub tasks_cal           : VolatileCell<u32>,
    pub tasks_cstart        : VolatileCell<u32>,
    pub tasks_cstop         : VolatileCell<u32>,
    _reserved1: [VolatileCell<u32>; 57],
    pub events_hfclkstarted : VolatileCell<u32>,
    pub events_lfclkstarted : VolatileCell<u32>,
    pub done                : VolatileCell<u32>,
    pub ctto                : VolatileCell<u32>,
    _reserved2: [VolatileCell<u32>; 125],
    pub intenset            : VolatileCell<u32>,
    pub intenclr            : VolatileCell<u32>,
    _reserved3: [VolatileCell<u32>; 63],
    pub hfclkrun            : VolatileCell<u32>,
    pub hfclkstat           : VolatileCell<u32>,
    _reserved4: [VolatileCell<u32>; 1],
    pub lfclkrun            : VolatileCell<u32>,
    pub lfclkstat           : VolatileCell<u32>,
    pub lfclksrccopy        : VolatileCell<u32>,
    _reserved5: [VolatileCell<u32>; 62],
    pub lfclksrc            : VolatileCell<u32>,
    _reserved6: [VolatileCell<u32>; 7],
    pub ctiv                : VolatileCell<u32>,
    _reserved7: [VolatileCell<u32>; 5],
    pub xtalfreq            : VolatileCell<u32>,
}

struct RTC {
    pub tasks_start: VolatileCell<u32>,
    pub tasks_stop: VolatileCell<u32>,
    pub tasks_clear: VolatileCell<u32>,
    pub tasks_trigovrflw: VolatileCell<u32>,
    _reserved1: [VolatileCell<u32>; 60],
    pub events_tick: VolatileCell<u32>,
    pub events_ovrflw: VolatileCell<u32>,
    _reserved2: [VolatileCell<u32>; 14],
    pub events_compare: [VolatileCell<u32>; 4],
    _reserved3: [VolatileCell<u32>; 108],
    pub inten:    VolatileCell<u32>,
    pub intenset: VolatileCell<u32>,
    pub intenclr: VolatileCell<u32>,
    _reserved4: [VolatileCell<u32>; 13],
    pub evten: VolatileCell<u32>,
    pub evtenset: VolatileCell<u32>,
    pub evtenclr: VolatileCell<u32>,
    _reserved5: [VolatileCell<u32>; 110],
    pub counter: VolatileCell<u32>,
    pub prescaler: VolatileCell<u32>,
    _reserved6: [VolatileCell<u32>; 13],
    pub cc: [VolatileCell<u32>; 4],
    _reserved7: [VolatileCell<u32>; 683],
    pub power: VolatileCell<u32>,
}
const CLOCK_BASE: usize = 0x40000000;
const RTC0_BASE: usize =  0x4000B000;

static mut VAL: usize = 0;

#[allow(non_snake_case)]
fn RTC0() -> &'static RTC {
        unsafe { mem::transmute(RTC0_BASE as usize) }
}

#[allow(non_snake_case)]
fn CLOCK() -> &'static CLOCK {
        unsafe { mem::transmute(CLOCK_BASE as usize) }
}

pub unsafe fn reset() {
    RTC0().tasks_clear.set(1);
}

pub unsafe fn set_timer(val: usize) {
    // Clock is 32768Hz
    // Each tick is 30.51us, approximate as 30.5
    let ticks: u32 = (val as u32 * 10) / 305;
    RTC0().cc[0].set(ticks); // Compare value
    VAL = val
}

pub unsafe fn enable(_: bool) {
    // Cribbed from mynewt; some of this overkill but be safe
    CLOCK().xtalfreq.set(0);
    CLOCK().tasks_lfclkstop.set(1);
    CLOCK().events_lfclkstarted.set(0);
    CLOCK().lfclksrc.set(0);
    CLOCK().tasks_lfclkstart.set(1);
    while CLOCK().events_lfclkstarted.get() == 0  {}
    RTC0().prescaler.set(0);
    RTC0().tasks_stop.set(1);
    RTC0().tasks_clear.set(1);
    RTC0().evtenclr.set(0xffffffff);
    RTC0().intenclr.set(0xffffffff);
    RTC0().tasks_start.set(1);
}

pub unsafe fn overflowed() -> bool {
//    value() >= (RTC0().cc[0].get() as usize)
    false
}

pub unsafe fn value() -> usize {
    // Clock is 32768Hz
    // Each tick is 30.51us, approximate as 30.5
    let counter = RTC0().counter.get() as usize;
    let val = counter * 305 / 10;
    // Keep this hardcoded value because counter seems to
    // not work correctly
    VAL * 1000
}

<|endoftext|>


<commit_before>#![cfg_attr(not(feature = "clippy"), allow(needless_range_loop))]

use super::{
    LENGTH_BYTES,
};

#[derive(Debug, PartialEq)]
pub enum DecodeError {
    TooShort,
    TooLong,
}

pub trait Decodable: Sized {
    fn ssz_decode(bytes: &[u8], index: usize) -> Result<(Self, usize), DecodeError>;
}

/// Decode the given bytes for the given type
///
/// The single ssz encoded value will be decoded as the given type at the
/// given index.
pub fn decode_ssz<T>(ssz_bytes: &[u8], index: usize)
    -> Result<(T, usize), DecodeError>
    where T: Decodable
{
    if index >= ssz_bytes.len() {
        return Err(DecodeError::TooShort)
    }
    T::ssz_decode(ssz_bytes, index)
}

/// Decode a vector (list) of encoded bytes.
///
/// Each element in the list will be decoded and placed into the vector.
pub fn decode_ssz_list<T>(ssz_bytes: &[u8], index: usize)
    -> Result<(Vec<T>, usize), DecodeError>
    where T: Decodable
{

    if index + LENGTH_BYTES > ssz_bytes.len() {
        return Err(DecodeError::TooShort);
    };

    // get the length
    let serialized_length = match decode_length(ssz_bytes, index, LENGTH_BYTES) {
        Err(v) => return Err(v),
        Ok(v) => v,
    };

    let final_len: usize = index + LENGTH_BYTES + serialized_length;

    if final_len > ssz_bytes.len() {
        return Err(DecodeError::TooShort);
    };

    let mut tmp_index = index + LENGTH_BYTES;
    let mut res_vec: Vec<T> = Vec::new();

    while tmp_index < final_len {
        match T::ssz_decode(ssz_bytes, tmp_index) {
            Err(v) => return Err(v),
            Ok(v) => {
                tmp_index = v.1;
                res_vec.push(v.0);
            },
        };

    };

    Ok((res_vec, final_len))
}

/// Given some number of bytes, interpret the first four
/// bytes as a 32-bit big-endian integer and return the
/// result.
pub fn decode_length(bytes: &[u8], index: usize, length_bytes: usize)
    -> Result<usize, DecodeError>
{
    if bytes.len() < length_bytes {
        return Err(DecodeError::TooShort);
    };
    let mut len: usize = 0;
    for i in index..index+length_bytes {
        let offset = (index+length_bytes - i - 1) * 8;
        len |= (bytes[i] as usize) << offset;
    };
    Ok(len)
}

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

    #[test]
    fn test_ssz_decode_length() {
        let decoded = decode_length(
            &vec![0, 0, 0, 1],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 1);

        let decoded = decode_length(
            &vec![0, 0, 1, 0],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 256);

        let decoded = decode_length(
            &vec![0, 0, 1, 255],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 511);

        let decoded = decode_length(
            &vec![255, 255, 255, 255],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 4294967295);
    }

    #[test]
    fn test_encode_decode_length() {
        let params: Vec<usize> = vec![
            0, 1, 2, 3, 7, 8, 16,
            2^8,  2^8  + 1,
            2^16, 2^16 + 1,
            2^24, 2^24 + 1,
            2^32,
        ];
        for i in params {
            let decoded = decode_length(
                &encode_length(i, LENGTH_BYTES),
                0,
                LENGTH_BYTES).unwrap();
            assert_eq!(i, decoded);
        }
    }

    #[test]
    fn test_decode_ssz_list() {
        // u16
        let v: Vec<u16> = vec![10, 10, 10, 10];
        let decoded: (Vec<u16>, usize) = decode_ssz_list(
            &vec![0, 0, 0, 8, 0, 10, 0, 10, 0, 10, 0, 10],
            0
        ).unwrap();

        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 12);

        // u32
        let v: Vec<u32> = vec![10, 10, 10, 10];
        let decoded: (Vec<u32>, usize) = decode_ssz_list(
            &vec![
                0, 0, 0, 16,
                0, 0, 0, 10, 0, 0, 0, 10, 0, 0, 0, 10, 0, 0, 0, 10
            ],
            0
        ).unwrap();
        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 20);

        // u64
        let v: Vec<u64> = vec![10,10,10,10];
        let decoded: (Vec<u64>, usize) = decode_ssz_list(
            &vec![0, 0, 0, 32,
                0, 0, 0, 0, 0, 0, 0, 10,
                0, 0, 0, 0, 0, 0, 0, 10,
                0, 0, 0, 0, 0, 0, 0, 10,
                0, 0, 0, 0, 0, 0, 0, 10,
            ],
            0
        ).unwrap();
        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 36);

        // Check that it can accept index
        let v: Vec<usize> = vec![15,15,15,15];
        let decoded: (Vec<usize>, usize) = decode_ssz_list(
            &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
                0, 0, 0, 32,
                0, 0, 0, 0, 0, 0, 0, 15,
                0, 0, 0, 0, 0, 0, 0, 15,
                0, 0, 0, 0, 0, 0, 0, 15,
                0, 0, 0, 0, 0, 0, 0, 15,
            ],
            10
        ).unwrap();
        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 46);

        // Check that length > bytes throws error
        let decoded: Result<(Vec<usize>, usize), DecodeError> = decode_ssz_list(
            &vec![0, 0, 0, 32,
                0, 0, 0, 0, 0, 0, 0, 15,
            ],
            0
        );
        assert_eq!(decoded, Err(DecodeError::TooShort));

        // Check that incorrect index throws error
        let decoded: Result<(Vec<usize>, usize), DecodeError> = decode_ssz_list(
            &vec![
                0, 0, 0, 0, 0, 0, 0, 15,
            ],
            16
        );
        assert_eq!(decoded, Err(DecodeError::TooShort));
    }
}
<commit_msg>another loop fixed<commit_after>use super::{
    LENGTH_BYTES,
};

#[derive(Debug, PartialEq)]
pub enum DecodeError {
    TooShort,
    TooLong,
}

pub trait Decodable: Sized {
    fn ssz_decode(bytes: &[u8], index: usize) -> Result<(Self, usize), DecodeError>;
}

/// Decode the given bytes for the given type
///
/// The single ssz encoded value will be decoded as the given type at the
/// given index.
pub fn decode_ssz<T>(ssz_bytes: &[u8], index: usize)
    -> Result<(T, usize), DecodeError>
    where T: Decodable
{
    if index >= ssz_bytes.len() {
        return Err(DecodeError::TooShort)
    }
    T::ssz_decode(ssz_bytes, index)
}

/// Decode a vector (list) of encoded bytes.
///
/// Each element in the list will be decoded and placed into the vector.
pub fn decode_ssz_list<T>(ssz_bytes: &[u8], index: usize)
    -> Result<(Vec<T>, usize), DecodeError>
    where T: Decodable
{

    if index + LENGTH_BYTES > ssz_bytes.len() {
        return Err(DecodeError::TooShort);
    };

    // get the length
    let serialized_length = match decode_length(ssz_bytes, index, LENGTH_BYTES) {
        Err(v) => return Err(v),
        Ok(v) => v,
    };

    let final_len: usize = index + LENGTH_BYTES + serialized_length;

    if final_len > ssz_bytes.len() {
        return Err(DecodeError::TooShort);
    };

    let mut tmp_index = index + LENGTH_BYTES;
    let mut res_vec: Vec<T> = Vec::new();

    while tmp_index < final_len {
        match T::ssz_decode(ssz_bytes, tmp_index) {
            Err(v) => return Err(v),
            Ok(v) => {
                tmp_index = v.1;
                res_vec.push(v.0);
            },
        };

    };

    Ok((res_vec, final_len))
}

/// Given some number of bytes, interpret the first four
/// bytes as a 32-bit big-endian integer and return the
/// result.
pub fn decode_length(bytes: &[u8], index: usize, length_bytes: usize)
    -> Result<usize, DecodeError>
{
    if bytes.len() < length_bytes {
        return Err(DecodeError::TooShort);
    };
    let mut len: usize = 0;
    for (i, byte) in bytes.iter().enumerate().take(index+length_bytes).skip(index) {
        let offset = (index+length_bytes - i - 1) * 8;
        len |= (*byte as usize) << offset;
    }
    Ok(len)
}

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

    #[test]
    fn test_ssz_decode_length() {
        let decoded = decode_length(
            &vec![0, 0, 0, 1],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 1);

        let decoded = decode_length(
            &vec![0, 0, 1, 0],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 256);

        let decoded = decode_length(
            &vec![0, 0, 1, 255],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 511);

        let decoded = decode_length(
            &vec![255, 255, 255, 255],
            0,
            LENGTH_BYTES);
        assert_eq!(decoded.unwrap(), 4294967295);
    }

    #[test]
    fn test_encode_decode_length() {
        let params: Vec<usize> = vec![
            0, 1, 2, 3, 7, 8, 16,
            2^8,  2^8  + 1,
            2^16, 2^16 + 1,
            2^24, 2^24 + 1,
            2^32,
        ];
        for i in params {
            let decoded = decode_length(
                &encode_length(i, LENGTH_BYTES),
                0,
                LENGTH_BYTES).unwrap();
            assert_eq!(i, decoded);
        }
    }

    #[test]
    fn test_decode_ssz_list() {
        // u16
        let v: Vec<u16> = vec![10, 10, 10, 10];
        let decoded: (Vec<u16>, usize) = decode_ssz_list(
            &vec![0, 0, 0, 8, 0, 10, 0, 10, 0, 10, 0, 10],
            0
        ).unwrap();

        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 12);

        // u32
        let v: Vec<u32> = vec![10, 10, 10, 10];
        let decoded: (Vec<u32>, usize) = decode_ssz_list(
            &vec![
                0, 0, 0, 16,
                0, 0, 0, 10, 0, 0, 0, 10, 0, 0, 0, 10, 0, 0, 0, 10
            ],
            0
        ).unwrap();
        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 20);

        // u64
        let v: Vec<u64> = vec![10,10,10,10];
        let decoded: (Vec<u64>, usize) = decode_ssz_list(
            &vec![0, 0, 0, 32,
                0, 0, 0, 0, 0, 0, 0, 10,
                0, 0, 0, 0, 0, 0, 0, 10,
                0, 0, 0, 0, 0, 0, 0, 10,
                0, 0, 0, 0, 0, 0, 0, 10,
            ],
            0
        ).unwrap();
        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 36);

        // Check that it can accept index
        let v: Vec<usize> = vec![15,15,15,15];
        let decoded: (Vec<usize>, usize) = decode_ssz_list(
            &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
                0, 0, 0, 32,
                0, 0, 0, 0, 0, 0, 0, 15,
                0, 0, 0, 0, 0, 0, 0, 15,
                0, 0, 0, 0, 0, 0, 0, 15,
                0, 0, 0, 0, 0, 0, 0, 15,
            ],
            10
        ).unwrap();
        assert_eq!(decoded.0, v);
        assert_eq!(decoded.1, 46);

        // Check that length > bytes throws error
        let decoded: Result<(Vec<usize>, usize), DecodeError> = decode_ssz_list(
            &vec![0, 0, 0, 32,
                0, 0, 0, 0, 0, 0, 0, 15,
            ],
            0
        );
        assert_eq!(decoded, Err(DecodeError::TooShort));

        // Check that incorrect index throws error
        let decoded: Result<(Vec<usize>, usize), DecodeError> = decode_ssz_list(
            &vec![
                0, 0, 0, 0, 0, 0, 0, 15,
            ],
            16
        );
        assert_eq!(decoded, Err(DecodeError::TooShort));
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use store::Store;
use types::{BeaconBlock, BeaconState, BeaconStateError, EthSpec, Hash256, Slot};

/// Extends `BlockRootsIterator`, returning `BeaconBlock` instances, instead of their roots.
pub struct BlockIterator<T: EthSpec, U> {
    roots: BlockRootsIterator<T, U>,
}

impl<T: EthSpec, U: Store> BlockIterator<T, U> {
    /// Create a new iterator over all blocks in the given `beacon_state` and prior states.
    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
        Self {
            roots: BlockRootsIterator::new(store, beacon_state, start_slot),
        }
    }
}

impl<T: EthSpec, U: Store> Iterator for BlockIterator<T, U> {
    type Item = BeaconBlock;

    fn next(&mut self) -> Option<Self::Item> {
        let root = self.roots.next()?;
        self.roots.store.get(&root).ok()?
    }
}

/// Iterates backwards through block roots.
///
/// Uses the `latest_block_roots` field of `BeaconState` to as the source of block roots and will
/// perform a lookup on the `Store` for a prior `BeaconState` if `latest_block_roots` has been
/// exhausted.
///
/// Returns `None` for roots prior to genesis or when there is an error reading from `Store`.
pub struct BlockRootsIterator<T: EthSpec, U> {
    store: Arc<U>,
    beacon_state: BeaconState<T>,
    slot: Slot,
}

impl<T: EthSpec, U: Store> BlockRootsIterator<T, U> {
    /// Create a new iterator over all block roots in the given `beacon_state` and prior states.
    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
        Self {
            slot: start_slot,
            beacon_state,
            store,
        }
    }
}

impl<T: EthSpec, U: Store> Iterator for BlockRootsIterator<T, U> {
    type Item = Hash256;

    fn next(&mut self) -> Option<Self::Item> {
        if self.slot == 0 {
            return None;
        }

        let slot = self.slot - 1;

        match self.beacon_state.get_block_root(slot) {
            Ok(root) => Some(*root),
            Err(BeaconStateError::SlotOutOfBounds) => {
                // Read a `BeaconState` from the store that has access to prior historical root.
                self.beacon_state = {
                    // Read the earliest historic state in the current slot.
                    let earliest_historic_slot =
                        self.beacon_state.slot - Slot::from(T::slots_per_historical_root());

                    // Load the earlier state from disk.
                    let new_state_root = self
                        .beacon_state
                        .get_state_root(earliest_historic_slot)
                        .ok()?;

                    let state_option = self.store.get(&new_state_root).ok()?;
                    state_option?
                };

                self.beacon_state.get_block_root(slot).ok().cloned()
            }
            _ => return None,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use store::MemoryStore;
    use types::{test_utils::TestingBeaconStateBuilder, FoundationEthSpec, Keypair};

    fn get_state<T: EthSpec>() -> BeaconState<T> {
        let builder =
            TestingBeaconStateBuilder::from_single_keypair(0, &Keypair::random(), &T::spec());
        let (state, _keypairs) = builder.build();
        state
    }

    #[test]
    fn root_iter() {
        let store = Arc::new(MemoryStore::open());

        let mut state_a: BeaconState<FoundationEthSpec> = get_state();
        let mut state_b: BeaconState<FoundationEthSpec> = get_state();

        let mut hashes = (0..).into_iter().map(|i| Hash256::from(i));

        for root in &mut state_a.latest_block_roots[..] {
            *root = hashes.next().unwrap()
        }
        for root in &mut state_b.latest_block_roots[..] {
            *root = hashes.next().unwrap()
        }

        let state_a_root = hashes.next().unwrap();
        state_b.latest_state_roots[0] = state_a_root;
        store.put(&state_a_root, &state_a).unwrap();

        let iter = BlockRootsIterator::new(store.clone(), state_b.clone(), state_b.slot);
        let mut collected: Vec<Hash256> = iter.collect();
        collected.reverse();

        for (i, item) in collected.iter().enumerate() {
            assert_eq!(*item, Hash256::from(i as u64));
        }
    }
}
<commit_msg>Fix bug in rev block iter<commit_after>use std::sync::Arc;
use store::Store;
use types::{BeaconBlock, BeaconState, BeaconStateError, EthSpec, Hash256, Slot};

/// Extends `BlockRootsIterator`, returning `BeaconBlock` instances, instead of their roots.
pub struct BlockIterator<T: EthSpec, U> {
    roots: BlockRootsIterator<T, U>,
}

impl<T: EthSpec, U: Store> BlockIterator<T, U> {
    /// Create a new iterator over all blocks in the given `beacon_state` and prior states.
    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
        Self {
            roots: BlockRootsIterator::new(store, beacon_state, start_slot),
        }
    }
}

impl<T: EthSpec, U: Store> Iterator for BlockIterator<T, U> {
    type Item = BeaconBlock;

    fn next(&mut self) -> Option<Self::Item> {
        let root = self.roots.next()?;
        self.roots.store.get(&root).ok()?
    }
}

/// Iterates backwards through block roots.
///
/// Uses the `latest_block_roots` field of `BeaconState` to as the source of block roots and will
/// perform a lookup on the `Store` for a prior `BeaconState` if `latest_block_roots` has been
/// exhausted.
///
/// Returns `None` for roots prior to genesis or when there is an error reading from `Store`.
pub struct BlockRootsIterator<T: EthSpec, U> {
    store: Arc<U>,
    beacon_state: BeaconState<T>,
    slot: Slot,
}

impl<T: EthSpec, U: Store> BlockRootsIterator<T, U> {
    /// Create a new iterator over all block roots in the given `beacon_state` and prior states.
    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
        Self {
            slot: start_slot,
            beacon_state,
            store,
        }
    }
}

impl<T: EthSpec, U: Store> Iterator for BlockRootsIterator<T, U> {
    type Item = Hash256;

    fn next(&mut self) -> Option<Self::Item> {
        if (self.slot == 0) || (self.slot > self.beacon_state.slot) {
            return None;
        }

        let slot = self.slot - 1;

        match self.beacon_state.get_block_root(slot) {
            Ok(root) => Some(*root),
            Err(BeaconStateError::SlotOutOfBounds) => {
                // Read a `BeaconState` from the store that has access to prior historical root.
                self.beacon_state = {
                    // Read the earliest historic state in the current slot.
                    let earliest_historic_slot =
                        self.beacon_state.slot - Slot::from(T::slots_per_historical_root());

                    // Load the earlier state from disk.
                    let new_state_root = self
                        .beacon_state
                        .get_state_root(earliest_historic_slot)
                        .ok()?;

                    let state_option = self.store.get(&new_state_root).ok()?;
                    state_option?
                };

                self.beacon_state.get_block_root(slot).ok().cloned()
            }
            _ => return None,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use store::MemoryStore;
    use types::{test_utils::TestingBeaconStateBuilder, FoundationEthSpec, Keypair};

    fn get_state<T: EthSpec>() -> BeaconState<T> {
        let builder =
            TestingBeaconStateBuilder::from_single_keypair(0, &Keypair::random(), &T::spec());
        let (state, _keypairs) = builder.build();
        state
    }

    #[test]
    fn root_iter() {
        let store = Arc::new(MemoryStore::open());

        let mut state_a: BeaconState<FoundationEthSpec> = get_state();
        let mut state_b: BeaconState<FoundationEthSpec> = get_state();

        let mut hashes = (0..).into_iter().map(|i| Hash256::from(i));

        for root in &mut state_a.latest_block_roots[..] {
            *root = hashes.next().unwrap()
        }
        for root in &mut state_b.latest_block_roots[..] {
            *root = hashes.next().unwrap()
        }

        let state_a_root = hashes.next().unwrap();
        state_b.latest_state_roots[0] = state_a_root;
        store.put(&state_a_root, &state_a).unwrap();

        let iter = BlockRootsIterator::new(store.clone(), state_b.clone(), state_b.slot);
        let mut collected: Vec<Hash256> = iter.collect();
        collected.reverse();

        for (i, item) in collected.iter().enumerate() {
            assert_eq!(*item, Hash256::from(i as u64));
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! An implementation of `MemoryDB` database.
use std::sync::{Arc, RwLock, RwLockReadGuard};
use std::clone::Clone;
use std::collections::btree_map::{BTreeMap, Range};
use std::collections::HashMap;
use std::iter::Peekable;

use super::{Database, Snapshot, Patch, Change, Iterator, Iter, Result};

const DEFAULT_NAME: &'static str = "default";

type DB = HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>;

/// Database implementation that stores all the data in memory.
///
/// It's mainly used for testing and not designed to be efficient.
#[derive(Default, Clone, Debug)]
pub struct MemoryDB {
    map: Arc<RwLock<DB>>,
}

/// An iterator over the entries of a `MemoryDB`.
struct MemoryDBIter<'a> {
    iter: Peekable<Range<'a, Vec<u8>, Vec<u8>>>,
    _guard: RwLockReadGuard<'a, DB>,
}

impl MemoryDB {
    /// Creates a new, empty database.
    pub fn new() -> MemoryDB {
        let mut tables = HashMap::new();
        debug_assert_eq!(
            tables.insert(DEFAULT_NAME.to_string(), BTreeMap::new()),
            None
        );
        MemoryDB { map: Arc::new(RwLock::new(tables)) }
    }
}

impl Database for MemoryDB {
    fn clone(&self) -> Box<Database> {
        Box::new(Clone::clone(self))
    }

    fn snapshot(&self) -> Box<Snapshot> {
        Box::new(MemoryDB {
            map: Arc::new(RwLock::new(self.map.read().unwrap().clone())),
        })
    }

    fn merge(&mut self, patch: Patch) -> Result<()> {
        for (cf_name, changes) in patch {
            let mut guard = self.map.write().unwrap();
            if !guard.contains_key(&cf_name) {
                guard.insert(cf_name.clone(), BTreeMap::new());
            }
            let table = guard.get_mut(&cf_name).unwrap();
            for (key, change) in changes {
                match change {
                    Change::Put(ref value) => {
                        table.insert(key, value.to_vec());
                    }
                    Change::Delete => {
                        table.remove(&key);
                    }
                }
            }
        }
        Ok(())
    }
}

impl Snapshot for MemoryDB {
    fn get(&self, name: &str, key: &[u8]) -> Option<Vec<u8>> {
        self.map.read().unwrap().get(name).and_then(|table| {
            table.get(key).cloned()
        })
    }

    fn contains(&self, name: &str, key: &[u8]) -> bool {
        self.map.read().unwrap().get(name).map_or(false, |table| {
            table.contains_key(key)
        })
    }

    fn iter<'a>(&'a self, name: &str, from: &[u8]) -> Iter<'a> {
        use std::collections::Bound::{Included, Unbounded};
        use std::mem::transmute;
        let map_guard = self.map.read().unwrap();

        Box::new(MemoryDBIter {
            iter: unsafe {
                transmute(match map_guard.get(name) {
                    Some(table) => {
                        table
                            .range::<[u8], _>((Included(from), Unbounded))
                            .peekable()
                    }
                    None => {
                        map_guard
                            .get(DEFAULT_NAME)
                            .unwrap()
                            .range::<[u8], _>((Unbounded, Unbounded))
                            .peekable()
                    }
                })
            },
            _guard: map_guard,
        })
    }
}

impl<'a> Iterator for MemoryDBIter<'a> {
    fn next(&mut self) -> Option<(&[u8], &[u8])> {
        self.iter.next().map(|(k, v)| (k.as_slice(), v.as_slice()))
    }

    fn peek(&mut self) -> Option<(&[u8], &[u8])> {
        self.iter.peek().map(|&(k, v)| (k.as_slice(), v.as_slice()))
    }
}

#[test]
fn test_memorydb_snapshot() {
    let mut db = MemoryDB::new();

    {
        let mut fork = db.fork();
        fork.put(DEFAULT_NAME, vec![1, 2, 3], vec![123]);
        let _ = db.merge(fork.into_patch());
    }

    let snapshot = db.snapshot();
    assert!(snapshot.contains(DEFAULT_NAME, vec![1, 2, 3].as_slice()));

    {
        let mut fork = db.fork();
        fork.put(DEFAULT_NAME, vec![2, 3, 4], vec![234]);
        let _ = db.merge(fork.into_patch());
    }

    assert!(!snapshot.contains(DEFAULT_NAME, vec![2, 3, 4].as_slice()));

    {
        let db_clone = Clone::clone(&db);
        let snap_clone = db_clone.snapshot();
        assert!(snap_clone.contains(DEFAULT_NAME, vec![2, 3, 4].as_slice()));
    }

    let snapshot = db.snapshot();
    assert!(snapshot.contains(DEFAULT_NAME, vec![2, 3, 4].as_slice()));
}
<commit_msg>Fixed memorydb<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! An implementation of `MemoryDB` database.
use std::sync::{Arc, RwLock};
use std::clone::Clone;
use std::collections::btree_map::BTreeMap;
use std::collections::HashMap;

use super::{Database, Snapshot, Patch, Change, Iterator, Iter, Result};

type DB = HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>;

/// Database implementation that stores all the data in memory.
///
/// It's mainly used for testing and not designed to be efficient.
#[derive(Default, Clone, Debug)]
pub struct MemoryDB {
    map: Arc<RwLock<DB>>,
}

/// An iterator over the entries of a `MemoryDB`.
struct MemoryDBIter {
    data: Vec<(Vec<u8>, Vec<u8>)>,
    index: usize,
}

impl MemoryDB {
    /// Creates a new, empty database.
    pub fn new() -> MemoryDB {
        MemoryDB { map: Arc::new(RwLock::new(HashMap::new())) }
    }
}

impl Database for MemoryDB {
    fn clone(&self) -> Box<Database> {
        Box::new(Clone::clone(self))
    }

    fn snapshot(&self) -> Box<Snapshot> {
        Box::new(MemoryDB {
            map: Arc::new(RwLock::new(self.map.read().unwrap().clone())),
        })
    }

    fn merge(&mut self, patch: Patch) -> Result<()> {
        for (cf_name, changes) in patch {
            let mut guard = self.map.write().unwrap();
            if !guard.contains_key(&cf_name) {
                guard.insert(cf_name.clone(), BTreeMap::new());
            }
            let table = guard.get_mut(&cf_name).unwrap();
            for (key, change) in changes {
                match change {
                    Change::Put(ref value) => {
                        table.insert(key, value.to_vec());
                    }
                    Change::Delete => {
                        table.remove(&key);
                    }
                }
            }
        }
        Ok(())
    }
}

impl Snapshot for MemoryDB {
    fn get(&self, name: &str, key: &[u8]) -> Option<Vec<u8>> {
        self.map.read().unwrap().get(name).and_then(|table| {
            table.get(key).cloned()
        })
    }

    fn contains(&self, name: &str, key: &[u8]) -> bool {
        self.map.read().unwrap().get(name).map_or(false, |table| {
            table.contains_key(key)
        })
    }

    fn iter(&self, name: &str, from: &[u8]) -> Iter {
        let map_guard = self.map.read().unwrap();
        let data = match map_guard.get(name) {
            Some(table) => {
                table
                    .iter()
                    .skip_while(|&(k, _)| k.as_slice() < from)
                    .map(|(k, v)| (k.to_vec(), v.to_vec()))
                    .collect()
            }
            None => Vec::new(),
        };

        Box::new(MemoryDBIter { data, index: 0 })
    }
}

impl Iterator for MemoryDBIter {
    fn next(&mut self) -> Option<(&[u8], &[u8])> {
        if self.index < self.data.len() {
            self.index += 1;
            self.data.get(self.index - 1).map(|&(ref k, ref v)| {
                (k.as_slice(), v.as_slice())
            })
        } else {
            None
        }
    }

    fn peek(&mut self) -> Option<(&[u8], &[u8])> {
        if self.index < self.data.len() {
            self.data.get(self.index).map(|&(ref k, ref v)| {
                (k.as_slice(), v.as_slice())
            })
        } else {
            None
        }
    }
}

#[test]
fn test_memorydb_snapshot() {
    let mut db = MemoryDB::new();
    let idx_name = "idx_name";
    {
        let mut fork = db.fork();
        fork.put(idx_name, vec![1, 2, 3], vec![123]);
        let _ = db.merge(fork.into_patch());
    }

    let snapshot = db.snapshot();
    assert!(snapshot.contains(idx_name, vec![1, 2, 3].as_slice()));

    {
        let mut fork = db.fork();
        fork.put(idx_name, vec![2, 3, 4], vec![234]);
        let _ = db.merge(fork.into_patch());
    }

    assert!(!snapshot.contains(idx_name, vec![2, 3, 4].as_slice()));

    {
        let db_clone = Clone::clone(&db);
        let snap_clone = db_clone.snapshot();
        assert!(snap_clone.contains(idx_name, vec![2, 3, 4].as_slice()));
    }

    let snapshot = db.snapshot();
    assert!(snapshot.contains(idx_name, vec![2, 3, 4].as_slice()));
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

extern crate toml;
extern crate rustc_serialize;

use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::io::{self, Read, Write};
use std::path::{PathBuf, Path};
use std::process::{Command, Stdio};

static HOSTS: &'static [&'static str] = &[
    "aarch64-unknown-linux-gnu",
    "arm-unknown-linux-gnueabi",
    "arm-unknown-linux-gnueabihf",
    "armv7-unknown-linux-gnueabihf",
    "i686-apple-darwin",
    "i686-pc-windows-gnu",
    "i686-pc-windows-msvc",
    "i686-unknown-linux-gnu",
    "mips-unknown-linux-gnu",
    "mips64-unknown-linux-gnuabi64",
    "mips64el-unknown-linux-gnuabi64",
    "mipsel-unknown-linux-gnu",
    "powerpc-unknown-linux-gnu",
    "powerpc64-unknown-linux-gnu",
    "powerpc64le-unknown-linux-gnu",
    "s390x-unknown-linux-gnu",
    "x86_64-apple-darwin",
    "x86_64-pc-windows-gnu",
    "x86_64-pc-windows-msvc",
    "x86_64-unknown-freebsd",
    "x86_64-unknown-linux-gnu",
    "x86_64-unknown-netbsd",
];

static TARGETS: &'static [&'static str] = &[
    "aarch64-apple-ios",
    "aarch64-unknown-fuchsia",
    "aarch64-linux-android",
    "aarch64-unknown-linux-gnu",
    "arm-linux-androideabi",
    "arm-unknown-linux-gnueabi",
    "arm-unknown-linux-gnueabihf",
    "arm-unknown-linux-musleabi",
    "arm-unknown-linux-musleabihf",
    "armv7-apple-ios",
    "armv7-linux-androideabi",
    "armv7-unknown-linux-gnueabihf",
    "armv7-unknown-linux-musleabihf",
    "armv7s-apple-ios",
    "asmjs-unknown-emscripten",
    "i386-apple-ios",
    "i586-pc-windows-msvc",
    "i586-unknown-linux-gnu",
    "i686-apple-darwin",
    "i686-linux-android",
    "i686-pc-windows-gnu",
    "i686-pc-windows-msvc",
    "i686-unknown-freebsd",
    "i686-unknown-linux-gnu",
    "i686-unknown-linux-musl",
    "mips-unknown-linux-gnu",
    "mips-unknown-linux-musl",
    "mips64-unknown-linux-gnuabi64",
    "mips64el-unknown-linux-gnuabi64",
    "mipsel-unknown-linux-gnu",
    "mipsel-unknown-linux-musl",
    "powerpc-unknown-linux-gnu",
    "powerpc64-unknown-linux-gnu",
    "powerpc64le-unknown-linux-gnu",
    "s390x-unknown-linux-gnu",
    "sparc64-unknown-linux-gnu",
    "wasm32-unknown-emscripten",
    "x86_64-apple-darwin",
    "x86_64-apple-ios",
    "x86_64-pc-windows-gnu",
    "x86_64-pc-windows-msvc",
    "x86_64-rumprun-netbsd",
    "x86_64-unknown-freebsd",
    "x86_64-unknown-fuchsia",
    "x86_64-unknown-linux-gnu",
    "x86_64-unknown-linux-musl",
    "x86_64-unknown-netbsd",
];

static MINGW: &'static [&'static str] = &[
    "i686-pc-windows-gnu",
    "x86_64-pc-windows-gnu",
];

struct Manifest {
    manifest_version: String,
    date: String,
    pkg: HashMap<String, Package>,
}

#[derive(RustcEncodable)]
struct Package {
    version: String,
    target: HashMap<String, Target>,
}

#[derive(RustcEncodable)]
struct Target {
    available: bool,
    url: Option<String>,
    hash: Option<String>,
    components: Option<Vec<Component>>,
    extensions: Option<Vec<Component>>,
}

#[derive(RustcEncodable)]
struct Component {
    pkg: String,
    target: String,
}

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {}", stringify!($e), e),
    })
}

struct Builder {
    rust_release: String,
    cargo_release: String,
    input: PathBuf,
    output: PathBuf,
    gpg_passphrase: String,
    digests: HashMap<String, String>,
    s3_address: String,
    date: String,
    rust_version: String,
    cargo_version: String,
}

fn main() {
    let mut args = env::args().skip(1);
    let input = PathBuf::from(args.next().unwrap());
    let output = PathBuf::from(args.next().unwrap());
    let date = args.next().unwrap();
    let rust_release = args.next().unwrap();
    let cargo_release = args.next().unwrap();
    let s3_address = args.next().unwrap();
    let mut passphrase = String::new();
    t!(io::stdin().read_to_string(&mut passphrase));

    Builder {
        rust_release: rust_release,
        cargo_release: cargo_release,
        input: input,
        output: output,
        gpg_passphrase: passphrase,
        digests: HashMap::new(),
        s3_address: s3_address,
        date: date,
        rust_version: String::new(),
        cargo_version: String::new(),
    }.build();
}

impl Builder {
    fn build(&mut self) {
        self.rust_version = self.version("rust", "x86_64-unknown-linux-gnu");
        self.cargo_version = self.version("cargo", "x86_64-unknown-linux-gnu");

        self.digest_and_sign();
        let Manifest { manifest_version, date, pkg } = self.build_manifest();

        // Unfortunately we can't use derive(RustcEncodable) here because the
        // version field is called `manifest-version`, not `manifest_version`.
        // In lieu of that just create the table directly here with a `BTreeMap`
        // and wrap it up in a `Value::Table`.
        let mut manifest = BTreeMap::new();
        manifest.insert("manifest-version".to_string(),
                        toml::Value::String(manifest_version));
        manifest.insert("date".to_string(), toml::Value::String(date.clone()));
        manifest.insert("pkg".to_string(), toml::encode(&pkg));
        let manifest = toml::Value::Table(manifest).to_string();

        let filename = format!("channel-rust-{}.toml", self.rust_release);
        self.write_manifest(&manifest, &filename);

        let filename = format!("channel-rust-{}-date.txt", self.rust_release);
        self.write_date_stamp(&date, &filename);

        if self.rust_release != "beta" && self.rust_release != "nightly" {
            self.write_manifest(&manifest, "channel-rust-stable.toml");
            self.write_date_stamp(&date, "channel-rust-stable-date.txt");
        }
    }

    fn digest_and_sign(&mut self) {
        for file in t!(self.input.read_dir()).map(|e| t!(e).path()) {
            let filename = file.file_name().unwrap().to_str().unwrap();
            let digest = self.hash(&file);
            self.sign(&file);
            assert!(self.digests.insert(filename.to_string(), digest).is_none());
        }
    }

    fn build_manifest(&mut self) -> Manifest {
        let mut manifest = Manifest {
            manifest_version: "2".to_string(),
            date: self.date.to_string(),
            pkg: HashMap::new(),
        };

        self.package("rustc", &mut manifest.pkg, HOSTS);
        self.package("cargo", &mut manifest.pkg, HOSTS);
        self.package("rust-mingw", &mut manifest.pkg, MINGW);
        self.package("rust-std", &mut manifest.pkg, TARGETS);
        self.package("rust-docs", &mut manifest.pkg, TARGETS);
        self.package("rust-src", &mut manifest.pkg, &["*"]);

        if self.rust_release == "nightly" {
            self.package("rust-analysis", &mut manifest.pkg, TARGETS);
        }

        let mut pkg = Package {
            version: self.cached_version("rust").to_string(),
            target: HashMap::new(),
        };
        for host in HOSTS {
            let filename = self.filename("rust", host);
            let digest = match self.digests.remove(&filename) {
                Some(digest) => digest,
                None => {
                    pkg.target.insert(host.to_string(), Target {
                        available: false,
                        url: None,
                        hash: None,
                        components: None,
                        extensions: None,
                    });
                    continue
                }
            };
            let mut components = Vec::new();
            let mut extensions = Vec::new();

            // rustc/rust-std/cargo are all required, and so is rust-mingw if it's
            // available for the target.
            components.extend(vec![
                Component { pkg: "rustc".to_string(), target: host.to_string() },
                Component { pkg: "rust-std".to_string(), target: host.to_string() },
                Component { pkg: "cargo".to_string(), target: host.to_string() },
            ]);
            if host.contains("pc-windows-gnu") {
                components.push(Component {
                    pkg: "rust-mingw".to_string(),
                    target: host.to_string(),
                });
            }

            // Docs, other standard libraries, and the source package are all
            // optional.
            extensions.push(Component {
                pkg: "rust-docs".to_string(),
                target: host.to_string(),
            });
            for target in TARGETS {
                if target != host {
                    extensions.push(Component {
                        pkg: "rust-std".to_string(),
                        target: target.to_string(),
                    });
                }
                if self.rust_release == "nightly" {
                    extensions.push(Component {
                        pkg: "rust-analysis".to_string(),
                        target: target.to_string(),
                    });
                }
            }
            extensions.push(Component {
                pkg: "rust-src".to_string(),
                target: "*".to_string(),
            });

            pkg.target.insert(host.to_string(), Target {
                available: true,
                url: Some(self.url("rust", host)),
                hash: Some(digest),
                components: Some(components),
                extensions: Some(extensions),
            });
        }
        manifest.pkg.insert("rust".to_string(), pkg);

        return manifest
    }

    fn package(&mut self,
               pkgname: &str,
               dst: &mut HashMap<String, Package>,
               targets: &[&str]) {
        let targets = targets.iter().map(|name| {
            let filename = self.filename(pkgname, name);
            let digest = match self.digests.remove(&filename) {
                Some(digest) => digest,
                None => {
                    return (name.to_string(), Target {
                        available: false,
                        url: None,
                        hash: None,
                        components: None,
                        extensions: None,
                    })
                }
            };

            (name.to_string(), Target {
                available: true,
                url: Some(self.url(pkgname, name)),
                hash: Some(digest),
                components: None,
                extensions: None,
            })
        }).collect();

        dst.insert(pkgname.to_string(), Package {
            version: self.cached_version(pkgname).to_string(),
            target: targets,
        });
    }

    fn url(&self, component: &str, target: &str) -> String {
        format!("{}/{}/{}",
                self.s3_address,
                self.date,
                self.filename(component, target))
    }

    fn filename(&self, component: &str, target: &str) -> String {
        if component == "rust-src" {
            format!("rust-src-{}.tar.gz", self.rust_release)
        } else if component == "cargo" {
            format!("cargo-{}-{}.tar.gz", self.cargo_release, target)
        } else {
            format!("{}-{}-{}.tar.gz", component, self.rust_release, target)
        }
    }

    fn cached_version(&self, component: &str) -> &str {
        if component == "cargo" {
            &self.cargo_version
        } else {
            &self.rust_version
        }
    }

    fn version(&self, component: &str, target: &str) -> String {
        let mut cmd = Command::new("tar");
        let filename = self.filename(component, target);
        cmd.arg("xf")
           .arg(self.input.join(&filename))
           .arg(format!("{}/version", filename.replace(".tar.gz", "")))
           .arg("-O");
        let version = t!(cmd.output());
        if !version.status.success() {
            panic!("failed to learn version:\n\n{:?}\n\n{}\n\n{}",
                   cmd,
                   String::from_utf8_lossy(&version.stdout),
                   String::from_utf8_lossy(&version.stderr));
        }
        String::from_utf8_lossy(&version.stdout).trim().to_string()
    }

    fn hash(&self, path: &Path) -> String {
        let sha = t!(Command::new("shasum")
                        .arg("-a").arg("256")
                        .arg(path.file_name().unwrap())
                        .current_dir(path.parent().unwrap())
                        .output());
        assert!(sha.status.success());

        let filename = path.file_name().unwrap().to_str().unwrap();
        let sha256 = self.output.join(format!("{}.sha256", filename));
        t!(t!(File::create(&sha256)).write_all(&sha.stdout));

        let stdout = String::from_utf8_lossy(&sha.stdout);
        stdout.split_whitespace().next().unwrap().to_string()
    }

    fn sign(&self, path: &Path) {
        let filename = path.file_name().unwrap().to_str().unwrap();
        let asc = self.output.join(format!("{}.asc", filename));
        println!("signing: {:?}", path);
        let mut cmd = Command::new("gpg");
        cmd.arg("--no-tty")
            .arg("--yes")
            .arg("--passphrase-fd").arg("0")
            .arg("--armor")
            .arg("--output").arg(&asc)
            .arg("--detach-sign").arg(path)
            .stdin(Stdio::piped());
        let mut child = t!(cmd.spawn());
        t!(child.stdin.take().unwrap().write_all(self.gpg_passphrase.as_bytes()));
        assert!(t!(child.wait()).success());
    }

    fn write_manifest(&self, manifest: &str, name: &str) {
        let dst = self.output.join(name);
        t!(t!(File::create(&dst)).write_all(manifest.as_bytes()));
        self.hash(&dst);
        self.sign(&dst);
    }

    fn write_date_stamp(&self, date: &str, name: &str) {
        let dst = self.output.join(name);
        t!(t!(File::create(&dst)).write_all(date.as_bytes()));
        self.hash(&dst);
        self.sign(&dst);
    }
}
<commit_msg>Make docs required again<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

extern crate toml;
extern crate rustc_serialize;

use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::io::{self, Read, Write};
use std::path::{PathBuf, Path};
use std::process::{Command, Stdio};

static HOSTS: &'static [&'static str] = &[
    "aarch64-unknown-linux-gnu",
    "arm-unknown-linux-gnueabi",
    "arm-unknown-linux-gnueabihf",
    "armv7-unknown-linux-gnueabihf",
    "i686-apple-darwin",
    "i686-pc-windows-gnu",
    "i686-pc-windows-msvc",
    "i686-unknown-linux-gnu",
    "mips-unknown-linux-gnu",
    "mips64-unknown-linux-gnuabi64",
    "mips64el-unknown-linux-gnuabi64",
    "mipsel-unknown-linux-gnu",
    "powerpc-unknown-linux-gnu",
    "powerpc64-unknown-linux-gnu",
    "powerpc64le-unknown-linux-gnu",
    "s390x-unknown-linux-gnu",
    "x86_64-apple-darwin",
    "x86_64-pc-windows-gnu",
    "x86_64-pc-windows-msvc",
    "x86_64-unknown-freebsd",
    "x86_64-unknown-linux-gnu",
    "x86_64-unknown-netbsd",
];

static TARGETS: &'static [&'static str] = &[
    "aarch64-apple-ios",
    "aarch64-unknown-fuchsia",
    "aarch64-linux-android",
    "aarch64-unknown-linux-gnu",
    "arm-linux-androideabi",
    "arm-unknown-linux-gnueabi",
    "arm-unknown-linux-gnueabihf",
    "arm-unknown-linux-musleabi",
    "arm-unknown-linux-musleabihf",
    "armv7-apple-ios",
    "armv7-linux-androideabi",
    "armv7-unknown-linux-gnueabihf",
    "armv7-unknown-linux-musleabihf",
    "armv7s-apple-ios",
    "asmjs-unknown-emscripten",
    "i386-apple-ios",
    "i586-pc-windows-msvc",
    "i586-unknown-linux-gnu",
    "i686-apple-darwin",
    "i686-linux-android",
    "i686-pc-windows-gnu",
    "i686-pc-windows-msvc",
    "i686-unknown-freebsd",
    "i686-unknown-linux-gnu",
    "i686-unknown-linux-musl",
    "mips-unknown-linux-gnu",
    "mips-unknown-linux-musl",
    "mips64-unknown-linux-gnuabi64",
    "mips64el-unknown-linux-gnuabi64",
    "mipsel-unknown-linux-gnu",
    "mipsel-unknown-linux-musl",
    "powerpc-unknown-linux-gnu",
    "powerpc64-unknown-linux-gnu",
    "powerpc64le-unknown-linux-gnu",
    "s390x-unknown-linux-gnu",
    "sparc64-unknown-linux-gnu",
    "wasm32-unknown-emscripten",
    "x86_64-apple-darwin",
    "x86_64-apple-ios",
    "x86_64-pc-windows-gnu",
    "x86_64-pc-windows-msvc",
    "x86_64-rumprun-netbsd",
    "x86_64-unknown-freebsd",
    "x86_64-unknown-fuchsia",
    "x86_64-unknown-linux-gnu",
    "x86_64-unknown-linux-musl",
    "x86_64-unknown-netbsd",
];

static MINGW: &'static [&'static str] = &[
    "i686-pc-windows-gnu",
    "x86_64-pc-windows-gnu",
];

struct Manifest {
    manifest_version: String,
    date: String,
    pkg: HashMap<String, Package>,
}

#[derive(RustcEncodable)]
struct Package {
    version: String,
    target: HashMap<String, Target>,
}

#[derive(RustcEncodable)]
struct Target {
    available: bool,
    url: Option<String>,
    hash: Option<String>,
    components: Option<Vec<Component>>,
    extensions: Option<Vec<Component>>,
}

#[derive(RustcEncodable)]
struct Component {
    pkg: String,
    target: String,
}

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {}", stringify!($e), e),
    })
}

struct Builder {
    rust_release: String,
    cargo_release: String,
    input: PathBuf,
    output: PathBuf,
    gpg_passphrase: String,
    digests: HashMap<String, String>,
    s3_address: String,
    date: String,
    rust_version: String,
    cargo_version: String,
}

fn main() {
    let mut args = env::args().skip(1);
    let input = PathBuf::from(args.next().unwrap());
    let output = PathBuf::from(args.next().unwrap());
    let date = args.next().unwrap();
    let rust_release = args.next().unwrap();
    let cargo_release = args.next().unwrap();
    let s3_address = args.next().unwrap();
    let mut passphrase = String::new();
    t!(io::stdin().read_to_string(&mut passphrase));

    Builder {
        rust_release: rust_release,
        cargo_release: cargo_release,
        input: input,
        output: output,
        gpg_passphrase: passphrase,
        digests: HashMap::new(),
        s3_address: s3_address,
        date: date,
        rust_version: String::new(),
        cargo_version: String::new(),
    }.build();
}

impl Builder {
    fn build(&mut self) {
        self.rust_version = self.version("rust", "x86_64-unknown-linux-gnu");
        self.cargo_version = self.version("cargo", "x86_64-unknown-linux-gnu");

        self.digest_and_sign();
        let Manifest { manifest_version, date, pkg } = self.build_manifest();

        // Unfortunately we can't use derive(RustcEncodable) here because the
        // version field is called `manifest-version`, not `manifest_version`.
        // In lieu of that just create the table directly here with a `BTreeMap`
        // and wrap it up in a `Value::Table`.
        let mut manifest = BTreeMap::new();
        manifest.insert("manifest-version".to_string(),
                        toml::Value::String(manifest_version));
        manifest.insert("date".to_string(), toml::Value::String(date.clone()));
        manifest.insert("pkg".to_string(), toml::encode(&pkg));
        let manifest = toml::Value::Table(manifest).to_string();

        let filename = format!("channel-rust-{}.toml", self.rust_release);
        self.write_manifest(&manifest, &filename);

        let filename = format!("channel-rust-{}-date.txt", self.rust_release);
        self.write_date_stamp(&date, &filename);

        if self.rust_release != "beta" && self.rust_release != "nightly" {
            self.write_manifest(&manifest, "channel-rust-stable.toml");
            self.write_date_stamp(&date, "channel-rust-stable-date.txt");
        }
    }

    fn digest_and_sign(&mut self) {
        for file in t!(self.input.read_dir()).map(|e| t!(e).path()) {
            let filename = file.file_name().unwrap().to_str().unwrap();
            let digest = self.hash(&file);
            self.sign(&file);
            assert!(self.digests.insert(filename.to_string(), digest).is_none());
        }
    }

    fn build_manifest(&mut self) -> Manifest {
        let mut manifest = Manifest {
            manifest_version: "2".to_string(),
            date: self.date.to_string(),
            pkg: HashMap::new(),
        };

        self.package("rustc", &mut manifest.pkg, HOSTS);
        self.package("cargo", &mut manifest.pkg, HOSTS);
        self.package("rust-mingw", &mut manifest.pkg, MINGW);
        self.package("rust-std", &mut manifest.pkg, TARGETS);
        self.package("rust-docs", &mut manifest.pkg, TARGETS);
        self.package("rust-src", &mut manifest.pkg, &["*"]);

        if self.rust_release == "nightly" {
            self.package("rust-analysis", &mut manifest.pkg, TARGETS);
        }

        let mut pkg = Package {
            version: self.cached_version("rust").to_string(),
            target: HashMap::new(),
        };
        for host in HOSTS {
            let filename = self.filename("rust", host);
            let digest = match self.digests.remove(&filename) {
                Some(digest) => digest,
                None => {
                    pkg.target.insert(host.to_string(), Target {
                        available: false,
                        url: None,
                        hash: None,
                        components: None,
                        extensions: None,
                    });
                    continue
                }
            };
            let mut components = Vec::new();
            let mut extensions = Vec::new();

            // rustc/rust-std/cargo/docs are all required, and so is rust-mingw
            // if it's available for the target.
            components.extend(vec![
                Component { pkg: "rustc".to_string(), target: host.to_string() },
                Component { pkg: "rust-std".to_string(), target: host.to_string() },
                Component { pkg: "cargo".to_string(), target: host.to_string() },
                Component { pkg: "rust-docs".to_string(), target: host.to_string() },
            ]);
            if host.contains("pc-windows-gnu") {
                components.push(Component {
                    pkg: "rust-mingw".to_string(),
                    target: host.to_string(),
                });
            }

            for target in TARGETS {
                if target != host {
                    extensions.push(Component {
                        pkg: "rust-std".to_string(),
                        target: target.to_string(),
                    });
                }
                if self.rust_release == "nightly" {
                    extensions.push(Component {
                        pkg: "rust-analysis".to_string(),
                        target: target.to_string(),
                    });
                }
            }
            extensions.push(Component {
                pkg: "rust-src".to_string(),
                target: "*".to_string(),
            });

            pkg.target.insert(host.to_string(), Target {
                available: true,
                url: Some(self.url("rust", host)),
                hash: Some(digest),
                components: Some(components),
                extensions: Some(extensions),
            });
        }
        manifest.pkg.insert("rust".to_string(), pkg);

        return manifest
    }

    fn package(&mut self,
               pkgname: &str,
               dst: &mut HashMap<String, Package>,
               targets: &[&str]) {
        let targets = targets.iter().map(|name| {
            let filename = self.filename(pkgname, name);
            let digest = match self.digests.remove(&filename) {
                Some(digest) => digest,
                None => {
                    return (name.to_string(), Target {
                        available: false,
                        url: None,
                        hash: None,
                        components: None,
                        extensions: None,
                    })
                }
            };

            (name.to_string(), Target {
                available: true,
                url: Some(self.url(pkgname, name)),
                hash: Some(digest),
                components: None,
                extensions: None,
            })
        }).collect();

        dst.insert(pkgname.to_string(), Package {
            version: self.cached_version(pkgname).to_string(),
            target: targets,
        });
    }

    fn url(&self, component: &str, target: &str) -> String {
        format!("{}/{}/{}",
                self.s3_address,
                self.date,
                self.filename(component, target))
    }

    fn filename(&self, component: &str, target: &str) -> String {
        if component == "rust-src" {
            format!("rust-src-{}.tar.gz", self.rust_release)
        } else if component == "cargo" {
            format!("cargo-{}-{}.tar.gz", self.cargo_release, target)
        } else {
            format!("{}-{}-{}.tar.gz", component, self.rust_release, target)
        }
    }

    fn cached_version(&self, component: &str) -> &str {
        if component == "cargo" {
            &self.cargo_version
        } else {
            &self.rust_version
        }
    }

    fn version(&self, component: &str, target: &str) -> String {
        let mut cmd = Command::new("tar");
        let filename = self.filename(component, target);
        cmd.arg("xf")
           .arg(self.input.join(&filename))
           .arg(format!("{}/version", filename.replace(".tar.gz", "")))
           .arg("-O");
        let version = t!(cmd.output());
        if !version.status.success() {
            panic!("failed to learn version:\n\n{:?}\n\n{}\n\n{}",
                   cmd,
                   String::from_utf8_lossy(&version.stdout),
                   String::from_utf8_lossy(&version.stderr));
        }
        String::from_utf8_lossy(&version.stdout).trim().to_string()
    }

    fn hash(&self, path: &Path) -> String {
        let sha = t!(Command::new("shasum")
                        .arg("-a").arg("256")
                        .arg(path.file_name().unwrap())
                        .current_dir(path.parent().unwrap())
                        .output());
        assert!(sha.status.success());

        let filename = path.file_name().unwrap().to_str().unwrap();
        let sha256 = self.output.join(format!("{}.sha256", filename));
        t!(t!(File::create(&sha256)).write_all(&sha.stdout));

        let stdout = String::from_utf8_lossy(&sha.stdout);
        stdout.split_whitespace().next().unwrap().to_string()
    }

    fn sign(&self, path: &Path) {
        let filename = path.file_name().unwrap().to_str().unwrap();
        let asc = self.output.join(format!("{}.asc", filename));
        println!("signing: {:?}", path);
        let mut cmd = Command::new("gpg");
        cmd.arg("--no-tty")
            .arg("--yes")
            .arg("--passphrase-fd").arg("0")
            .arg("--armor")
            .arg("--output").arg(&asc)
            .arg("--detach-sign").arg(path)
            .stdin(Stdio::piped());
        let mut child = t!(cmd.spawn());
        t!(child.stdin.take().unwrap().write_all(self.gpg_passphrase.as_bytes()));
        assert!(t!(child.wait()).success());
    }

    fn write_manifest(&self, manifest: &str, name: &str) {
        let dst = self.output.join(name);
        t!(t!(File::create(&dst)).write_all(manifest.as_bytes()));
        self.hash(&dst);
        self.sign(&dst);
    }

    fn write_date_stamp(&self, date: &str, name: &str) {
        let dst = self.output.join(name);
        t!(t!(File::create(&dst)).write_all(date.as_bytes()));
        self.hash(&dst);
        self.sign(&dst);
    }
}
<|endoftext|>