<commit_before>use std::num::Wrapping as W;

use digest;
use utils::buffer::{FixedBuffer64, FixedBuffer, StandardPadding};

use byteorder::{ByteOrder, BigEndian};

// sboxes.c: Tiger S boxeszz
const SBOXES: [[u64; 256]; 4] = include!("tiger.sboxes");
const ROUNDS: usize = 3;

#[derive(Debug, Clone, Copy)]
struct State {
    a: W<u64>,
    b: W<u64>,
    c: W<u64>,
}

macro_rules! round {
    ($a:expr, $b:expr, $c:expr, $x:expr, $mul:expr) => {
        $c = $c ^ $x;
        $a = $a - W(
            SBOXES[0][($c.0 >> (0*8)) as usize & 0xff] ^
            SBOXES[1][($c.0 >> (2*8)) as usize & 0xff] ^
            SBOXES[2][($c.0 >> (4*8)) as usize & 0xff] ^
            SBOXES[3][($c.0 >> (6*8)) as usize & 0xff]);
        $b = $b + W(
            SBOXES[3][($c.0 >> (1*8)) as usize & 0xff] ^
            SBOXES[2][($c.0 >> (3*8)) as usize & 0xff] ^
            SBOXES[1][($c.0 >> (5*8)) as usize & 0xff] ^
            SBOXES[0][($c.0 >> (7*8)) as usize & 0xff]);
        $b = $b * W($mul);
    };
}

impl State {
    fn new() -> Self {
        State {
            a: W(0x0123456789abcdef),
            b: W(0xfedcba9876543210),
            c: W(0xf096a5b4c3b2e187),
        }
    }

    fn pass(&mut self, block: &[W<u64>], mul: u64) {
        round!(self.a, self.b, self.c, block[0], mul);
        round!(self.b, self.c, self.a, block[1], mul);
        round!(self.c, self.a, self.b, block[2], mul);
        round!(self.a, self.b, self.c, block[3], mul);
        round!(self.b, self.c, self.a, block[4], mul);
        round!(self.c, self.a, self.b, block[5], mul);
        round!(self.a, self.b, self.c, block[6], mul);
        round!(self.b, self.c, self.a, block[7], mul);
    }

    fn key_schedule(x: &mut [W<u64>]) {
        x[0] = x[0] - (x[7] ^ W(0xa5a5a5a5a5a5a5a5));
        x[1] = x[1] ^ x[0];
        x[2] = x[2] + x[1];
        x[3] = x[3] - (x[2] ^ (!x[1] << 19));
        x[4] = x[4] ^ x[3];
        x[5] = x[5] + x[4];
        x[6] = x[6] - (x[5] ^ (!x[4] >> 23));
        x[7] = x[7] ^ x[6];

        x[0] = x[0] + x[7];
        x[1] = x[1] - (x[0] ^ (!x[7] << 19));
        x[2] = x[2] ^ x[1];
        x[3] = x[3] + x[2];
        x[4] = x[4] - (x[3] ^ (!x[2] >> 23));
        x[5] = x[5] ^ x[4];
        x[6] = x[6] + x[5];
        x[7] = x[7] - (x[6] ^ W(0x0123456789abcdef));
    }

    fn rotate(&mut self) {
        let tmp = self.a;
        self.a = self.c;
        self.c = self.b;
        self.b = tmp;
    }

    fn compress(&mut self, block: &[u8]) {
        assert_eq!(block.len(), 64);
        let mut wblock = [W(0); 8];

        for (v, c) in wblock.iter_mut().zip(block.chunks(8)) {
            *v = W(BigEndian::read_u64(c));
        }

        let tmp = *self; // save abc
        for i in 0..ROUNDS {
            if i != 0 {
                Self::key_schedule(&mut wblock);
            }
            let mul = match i {
                0 => 5,
                1 => 7,
                _ => 9,
            };
            self.pass(&mut wblock, mul);
            self.rotate();
        }

        // feedforward
        self.a = self.a ^ tmp.a;
        self.b = self.b - tmp.b;
        self.c = self.c + tmp.c;
    }
}

#[derive(Clone)]
pub struct Tiger {
    state: State,
    buffer: FixedBuffer64,
    length: u64,
}

impl Default for Tiger {
    fn default() -> Self {
        Tiger {
            state: State::new(),
            buffer: FixedBuffer64::new(),
            length: 0,
        }
    }
}

impl digest::Digest for Tiger {
    fn update<T>(&mut self, update: T)
        where T: AsRef<[u8]>
    {
        let update = update.as_ref();
        self.length += update.len() as u64;

        let state = &mut self.state;
        self.buffer.input(update, |d| state.compress(d));
    }

    fn output_bits() -> usize {
        192
    }
    fn block_size() -> usize {
        64
    }

    fn result<T>(mut self, mut out: T)
        where T: AsMut<[u8]>
    {
        let state = &mut self.state;

        self.buffer.pad(0x01, 8, |d| state.compress(d));
        BigEndian::write_u64(self.buffer.next(8), self.length << 3);
        state.compress(self.buffer.full_buffer());

        let mut out = out.as_mut();
        assert!(out.len() >= Self::output_bytes());
        BigEndian::write_u64(&mut out[0..8], state.a.0);
        BigEndian::write_u64(&mut out[8..16], state.b.0);
        BigEndian::write_u64(&mut out[16..24], state.c.0);
    }
}

#[cfg(test)]
mod tests {
    use digest::test::Test;
    use super::Tiger;

    const TESTS: &'static [Test<'static>] = &[
        Test { input: b"", output: &[0x60, 0xef, 0x6c, 0x0d, 0xbc, 0x07, 0x7b, 0x9c, 0x17, 0x5f, 0xfb, 0x77, 0x71, 0x00, 0x8c, 0x25, 0x3b, 0xac, 0xea, 0x02, 0x4c, 0x9d, 0x01, 0xab] },
        Test { input: b"abc", output: &[0xc7, 0x9e, 0x79, 0x9e, 0x14, 0xb5, 0x3e, 0x7d, 0xf9, 0x35, 0xd8, 0x34, 0x77, 0xfa, 0x4d, 0xf9, 0x39, 0xd1, 0x8c, 0x44, 0xf7, 0x6b, 0x73, 0xcd] },
        Test { input: b"Tiger", output: &[0xa2, 0x4e, 0xe9, 0x54, 0x0c, 0x41, 0xd7, 0x1b, 0x6a, 0x62, 0x6f, 0x9d, 0xdf, 0x41, 0xd1, 0x2e, 0x30, 0x31, 0x27, 0x2b, 0x6a, 0xab, 0xbd, 0x9a] },
        Test { input: b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-", output: &[0xaf, 0x58, 0xf9, 0xc0, 0x5b, 0x88, 0x60, 0x48, 0xc1, 0x6f, 0x48, 0xbc, 0x90, 0x4b, 0xef, 0xfd, 0xa8, 0x38, 0xf0, 0x05, 0xff, 0x74, 0x03, 0x5e] },
        Test { input: b"ABCDEFGHIJKLMNOPQRSTUVWXYZ=abcdefghijklmnopqrstuvwxyz+0123456789", output: &[0x76, 0xaa, 0x09, 0x2e, 0x58, 0x9b, 0x0f, 0x0b, 0x8b, 0x78, 0x09, 0x44, 0xad, 0x1f, 0x0c, 0x41, 0xc9, 0xa5, 0xce, 0xb3, 0x07, 0xd2, 0xfd, 0xe0] },
        Test { input: b"Tiger - A Fast New Hash Function, by Ross Anderson and Eli Biham", output: &[0x82, 0x81, 0x5d, 0x89, 0x24, 0xe2, 0xf1, 0x1a, 0x88, 0xc0, 0xf2, 0x92, 0x91, 0x4c, 0x6c, 0xfd, 0xbf, 0xd8, 0xe7, 0x8e, 0x2c, 0xf2, 0x9a, 0xd0] },
        Test { input: b"Tiger - A Fast New Hash Function, by Ross Anderson and Eli Biham, proceedings of Fast Software Encryption 3, Cambridge.", output: &[0x48, 0x1f, 0x6d, 0xd0, 0xdf, 0x57, 0xe1, 0x03, 0xd7, 0xde, 0x0b, 0x66, 0x3a, 0xdb, 0x05, 0xf5, 0x03, 0xb8, 0x77, 0xf2, 0x76, 0xad, 0xb2, 0x86] },
        Test { input: b"Tiger - A Fast New Hash Function, by Ross Anderson and Eli Biham, proceedings of Fast Software Encryption 3, Cambridge, 1996.", output: &[0x69, 0x6d, 0x99, 0x60, 0x61, 0x88, 0x47, 0x25, 0x9f, 0x2f, 0xa5, 0x80, 0xdb, 0x2f, 0x95, 0x55, 0x96, 0xd7, 0xbe, 0xa2, 0x04, 0x6f, 0x46, 0xd7] },
        Test { input: b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-", output: &[0x27, 0x20, 0xa3, 0x8d, 0x7f, 0x08, 0x75, 0x54, 0x5f, 0xd9, 0x61, 0x4f, 0x48, 0xf7, 0xe6, 0x7a, 0x92, 0xd3, 0x74, 0x69, 0xda, 0x2e, 0xf3, 0x78] },
    ];

    #[test]
    fn example_implementation_vectors() {
        for test in TESTS {
            test.test(Tiger::default());
        }
    }

    #[test]
    fn hash_of_64k_bytes_string() {
        use digest::Digest;

        let mut hash = Tiger::default();

        for i in 0..65536 {
            hash.update(&[(i & 0xff) as u8]);
        }
        let mut result = [0; 24];
        hash.result(&mut result[..]);

        assert_eq!(&result[..],
                   &[0xcd, 0x7e, 0xb9, 0x64, 0x5f, 0xb4, 0x05, 0xc6, 0x48, 0x5d, 0xd1, 0xaa, 0x14,
                     0x59, 0x6a, 0x63, 0xe5, 0x70, 0x4c, 0xc2, 0xff, 0x28, 0xf2, 0x4a])
    }
}
<commit_msg>Auto merge of #58 - libOctavo:tiger2, r=hauleth<commit_after>use std::num::Wrapping as W;

use digest;
use utils::buffer::{FixedBuffer64, FixedBuffer, StandardPadding};

use byteorder::{ByteOrder, BigEndian};

// sboxes.c: Tiger S boxeszz
const SBOXES: [[u64; 256]; 4] = include!("tiger.sboxes");
const ROUNDS: usize = 3;

#[derive(Debug, Clone, Copy)]
struct State {
    a: W<u64>,
    b: W<u64>,
    c: W<u64>,
}

macro_rules! round {
    ($a:expr, $b:expr, $c:expr, $x:expr, $mul:expr) => {
        $c = $c ^ $x;
        $a = $a - W(
            SBOXES[0][($c.0 >> (0*8)) as usize & 0xff] ^
            SBOXES[1][($c.0 >> (2*8)) as usize & 0xff] ^
            SBOXES[2][($c.0 >> (4*8)) as usize & 0xff] ^
            SBOXES[3][($c.0 >> (6*8)) as usize & 0xff]);
        $b = $b + W(
            SBOXES[3][($c.0 >> (1*8)) as usize & 0xff] ^
            SBOXES[2][($c.0 >> (3*8)) as usize & 0xff] ^
            SBOXES[1][($c.0 >> (5*8)) as usize & 0xff] ^
            SBOXES[0][($c.0 >> (7*8)) as usize & 0xff]);
        $b = $b * W($mul);
    };
}

impl State {
    fn new() -> Self {
        State {
            a: W(0x0123456789abcdef),
            b: W(0xfedcba9876543210),
            c: W(0xf096a5b4c3b2e187),
        }
    }

    fn pass(&mut self, block: &[W<u64>], mul: u64) {
        round!(self.a, self.b, self.c, block[0], mul);
        round!(self.b, self.c, self.a, block[1], mul);
        round!(self.c, self.a, self.b, block[2], mul);
        round!(self.a, self.b, self.c, block[3], mul);
        round!(self.b, self.c, self.a, block[4], mul);
        round!(self.c, self.a, self.b, block[5], mul);
        round!(self.a, self.b, self.c, block[6], mul);
        round!(self.b, self.c, self.a, block[7], mul);
    }

    fn key_schedule(x: &mut [W<u64>]) {
        x[0] = x[0] - (x[7] ^ W(0xa5a5a5a5a5a5a5a5));
        x[1] = x[1] ^ x[0];
        x[2] = x[2] + x[1];
        x[3] = x[3] - (x[2] ^ (!x[1] << 19));
        x[4] = x[4] ^ x[3];
        x[5] = x[5] + x[4];
        x[6] = x[6] - (x[5] ^ (!x[4] >> 23));
        x[7] = x[7] ^ x[6];

        x[0] = x[0] + x[7];
        x[1] = x[1] - (x[0] ^ (!x[7] << 19));
        x[2] = x[2] ^ x[1];
        x[3] = x[3] + x[2];
        x[4] = x[4] - (x[3] ^ (!x[2] >> 23));
        x[5] = x[5] ^ x[4];
        x[6] = x[6] + x[5];
        x[7] = x[7] - (x[6] ^ W(0x0123456789abcdef));
    }

    fn rotate(&mut self) {
        let tmp = self.a;
        self.a = self.c;
        self.c = self.b;
        self.b = tmp;
    }

    fn compress(&mut self, block: &[u8]) {
        assert_eq!(block.len(), 64);
        let mut wblock = [W(0); 8];

        for (v, c) in wblock.iter_mut().zip(block.chunks(8)) {
            *v = W(BigEndian::read_u64(c));
        }

        let tmp = *self; // save abc
        for i in 0..ROUNDS {
            if i != 0 {
                Self::key_schedule(&mut wblock);
            }
            let mul = match i {
                0 => 5,
                1 => 7,
                _ => 9,
            };
            self.pass(&mut wblock, mul);
            self.rotate();
        }

        // feedforward
        self.a = self.a ^ tmp.a;
        self.b = self.b - tmp.b;
        self.c = self.c + tmp.c;
    }
}

macro_rules! tiger_impl {
    ($name:ident, $padding:expr) => {
        #[derive(Clone)]
        pub struct $name {
            state: State,
            buffer: FixedBuffer64,
            length: u64,
        }

        impl Default for $name {
            fn default() -> Self {
                $name {
                    state: State::new(),
                    buffer: FixedBuffer64::new(),
                    length: 0,
                }
            }
        }

        impl digest::Digest for $name {
            fn update<T>(&mut self, update: T)
                where T: AsRef<[u8]>
                {
                    let update = update.as_ref();
                    self.length += update.len() as u64;

                    let state = &mut self.state;
                    self.buffer.input(update, |d| state.compress(d));
                }

            fn output_bits() -> usize {
                192
            }
            fn block_size() -> usize {
                64
            }

            fn result<T>(mut self, mut out: T)
                where T: AsMut<[u8]>
                {
                    let state = &mut self.state;

                    self.buffer.pad($padding, 8, |d| state.compress(d));
                    BigEndian::write_u64(self.buffer.next(8), self.length << 3);
                    state.compress(self.buffer.full_buffer());

                    let mut out = out.as_mut();
                    assert!(out.len() >= Self::output_bytes());
                    BigEndian::write_u64(&mut out[0..8], state.a.0);
                    BigEndian::write_u64(&mut out[8..16], state.b.0);
                    BigEndian::write_u64(&mut out[16..24], state.c.0);
                }
        }
    };
}

tiger_impl!(Tiger,  0x01);
tiger_impl!(Tiger2, 0x80);

#[cfg(test)]
mod tests {
    use digest::test::Test;
    use super::Tiger;

    const TESTS: &'static [Test<'static>] = &[
        Test { input: b"", output: &[0x60, 0xef, 0x6c, 0x0d, 0xbc, 0x07, 0x7b, 0x9c, 0x17, 0x5f, 0xfb, 0x77, 0x71, 0x00, 0x8c, 0x25, 0x3b, 0xac, 0xea, 0x02, 0x4c, 0x9d, 0x01, 0xab] },
        Test { input: b"abc", output: &[0xc7, 0x9e, 0x79, 0x9e, 0x14, 0xb5, 0x3e, 0x7d, 0xf9, 0x35, 0xd8, 0x34, 0x77, 0xfa, 0x4d, 0xf9, 0x39, 0xd1, 0x8c, 0x44, 0xf7, 0x6b, 0x73, 0xcd] },
        Test { input: b"Tiger", output: &[0xa2, 0x4e, 0xe9, 0x54, 0x0c, 0x41, 0xd7, 0x1b, 0x6a, 0x62, 0x6f, 0x9d, 0xdf, 0x41, 0xd1, 0x2e, 0x30, 0x31, 0x27, 0x2b, 0x6a, 0xab, 0xbd, 0x9a] },
        Test { input: b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-", output: &[0xaf, 0x58, 0xf9, 0xc0, 0x5b, 0x88, 0x60, 0x48, 0xc1, 0x6f, 0x48, 0xbc, 0x90, 0x4b, 0xef, 0xfd, 0xa8, 0x38, 0xf0, 0x05, 0xff, 0x74, 0x03, 0x5e] },
        Test { input: b"ABCDEFGHIJKLMNOPQRSTUVWXYZ=abcdefghijklmnopqrstuvwxyz+0123456789", output: &[0x76, 0xaa, 0x09, 0x2e, 0x58, 0x9b, 0x0f, 0x0b, 0x8b, 0x78, 0x09, 0x44, 0xad, 0x1f, 0x0c, 0x41, 0xc9, 0xa5, 0xce, 0xb3, 0x07, 0xd2, 0xfd, 0xe0] },
        Test { input: b"Tiger - A Fast New Hash Function, by Ross Anderson and Eli Biham", output: &[0x82, 0x81, 0x5d, 0x89, 0x24, 0xe2, 0xf1, 0x1a, 0x88, 0xc0, 0xf2, 0x92, 0x91, 0x4c, 0x6c, 0xfd, 0xbf, 0xd8, 0xe7, 0x8e, 0x2c, 0xf2, 0x9a, 0xd0] },
        Test { input: b"Tiger - A Fast New Hash Function, by Ross Anderson and Eli Biham, proceedings of Fast Software Encryption 3, Cambridge.", output: &[0x48, 0x1f, 0x6d, 0xd0, 0xdf, 0x57, 0xe1, 0x03, 0xd7, 0xde, 0x0b, 0x66, 0x3a, 0xdb, 0x05, 0xf5, 0x03, 0xb8, 0x77, 0xf2, 0x76, 0xad, 0xb2, 0x86] },
        Test { input: b"Tiger - A Fast New Hash Function, by Ross Anderson and Eli Biham, proceedings of Fast Software Encryption 3, Cambridge, 1996.", output: &[0x69, 0x6d, 0x99, 0x60, 0x61, 0x88, 0x47, 0x25, 0x9f, 0x2f, 0xa5, 0x80, 0xdb, 0x2f, 0x95, 0x55, 0x96, 0xd7, 0xbe, 0xa2, 0x04, 0x6f, 0x46, 0xd7] },
        Test { input: b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-", output: &[0x27, 0x20, 0xa3, 0x8d, 0x7f, 0x08, 0x75, 0x54, 0x5f, 0xd9, 0x61, 0x4f, 0x48, 0xf7, 0xe6, 0x7a, 0x92, 0xd3, 0x74, 0x69, 0xda, 0x2e, 0xf3, 0x78] },
    ];

    #[test]
    fn example_implementation_vectors() {
        for test in TESTS {
            test.test(Tiger::default());
        }
    }

    #[test]
    fn hash_of_64k_bytes_string() {
        use digest::Digest;

        let mut hash = Tiger::default();

        for i in 0..65536 {
            hash.update(&[(i & 0xff) as u8]);
        }
        let mut result = [0; 24];
        hash.result(&mut result[..]);

        assert_eq!(&result[..],
                   &[0xcd, 0x7e, 0xb9, 0x64, 0x5f, 0xb4, 0x05, 0xc6, 0x48, 0x5d, 0xd1, 0xaa, 0x14,
                   0x59, 0x6a, 0x63, 0xe5, 0x70, 0x4c, 0xc2, 0xff, 0x28, 0xf2, 0x4a])
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
// 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.

//! Common storage-related types
use memmap;

use std::cmp;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufRead, BufReader, ErrorKind, Write};
use std::os::unix::io::AsRawFd;
use std::io::Read;
use std::path::Path;

#[cfg(any(target_os = "linux"))]
use libc::{ftruncate64, off64_t};
#[cfg(not(any(target_os = "linux", target_os = "android")))]
use libc::{ftruncate as ftruncate64, off_t as off64_t};

use core::ser;

/// A no-op function for doing nothing with some pruned data.
pub fn prune_noop(_pruned_data: &[u8]) {}

/// Wrapper for a file that can be read at any position (random read) but for
/// which writes are append only. Reads are backed by a memory map (mmap(2)),
/// relying on the operating system for fast access and caching. The memory
/// map is reallocated to expand it when new writes are flushed.
///
/// Despite being append-only, the file can still be pruned and truncated. The
/// former simply happens by rewriting it, ignoring some of the data. The
/// latter by truncating the underlying file and re-creating the mmap.
pub struct AppendOnlyFile {
	path: String,
	file: File,
	mmap: Option<memmap::Mmap>,
	buffer_start: usize,
	buffer: Vec<u8>,
	buffer_start_bak: usize,
}

impl AppendOnlyFile {
	/// Open a file (existing or not) as append-only, backed by a mmap.
	pub fn open(path: String) -> io::Result<AppendOnlyFile> {
		let file = OpenOptions::new()
			.read(true)
			.append(true)
			.create(true)
			.open(path.clone())?;
		let mut aof = AppendOnlyFile {
			path: path.clone(),
			file: file,
			mmap: None,
			buffer_start: 0,
			buffer: vec![],
			buffer_start_bak: 0,
		};
		// if we have a non-empty file then mmap it.
		if let Ok(sz) = aof.size() {
			if sz > 0 {
				aof.buffer_start = sz as usize;
				aof.mmap = Some(unsafe { memmap::Mmap::map(&aof.file)? });
			}
		}
		Ok(aof)
	}

	/// Append data to the file. Until the append-only file is synced, data is
	/// only written to memory.
	pub fn append(&mut self, buf: &mut Vec<u8>) {
		self.buffer.append(buf);
	}

	/// Rewinds the data file back to a lower position. The new position needs
	/// to be the one of the first byte the next time data is appended.
	pub fn rewind(&mut self, pos: u64) {
		if self.buffer_start_bak > 0 || self.buffer.len() > 0 {
			panic!("Can't rewind on a dirty state.");
		}
		self.buffer_start_bak = self.buffer_start;
		self.buffer_start = pos as usize;
	}

	/// Syncs all writes (fsync), reallocating the memory map to make the newly
	/// written data accessible.
	pub fn flush(&mut self) -> io::Result<()> {
		if self.buffer_start_bak > 0 {
			// flushing a rewound state, we need to truncate before applying
			self.truncate(self.buffer_start)?;
			self.buffer_start_bak = 0;
		}

		self.buffer_start += self.buffer.len();
		self.file.write(&self.buffer[..])?;
		self.file.sync_all()?;

		self.buffer = vec![];

		// Note: file must be non-empty to memory map it
		if self.file.metadata()?.len() == 0 {
			self.mmap = None;
		} else {
			self.mmap = Some(unsafe { memmap::Mmap::map(&self.file)? });
		}

		Ok(())
	}

	/// Returns the last position (in bytes), taking into account whether data
	/// has been rewound
	pub fn last_buffer_pos(&self) -> usize {
		self.buffer_start
	}

	/// Discard the current non-flushed data.
	pub fn discard(&mut self) {
		if self.buffer_start_bak > 0 {
			// discarding a rewound state, restore the buffer start
			self.buffer_start = self.buffer_start_bak;
			self.buffer_start_bak = 0;
		}
		self.buffer = vec![];
	}

	/// Read length bytes of data at offset from the file.
	/// Leverages the memory map.
	pub fn read(&self, offset: usize, length: usize) -> Vec<u8> {
		if offset >= self.buffer_start {
			if self.buffer.is_empty() {
				return vec![];
			}
			let offset = offset - self.buffer_start;
			return self.buffer[offset..(offset + length)].to_vec();
		}
		if let None = self.mmap {
			return vec![];
		}
		let mmap = self.mmap.as_ref().unwrap();

		if mmap.len() < (offset + length) {
			return vec![];
		}

		(&mmap[offset..(offset + length)]).to_vec()
	}

	/// Truncates the underlying file to the provided offset
	pub fn truncate(&self, offs: usize) -> io::Result<()> {
		let fd = self.file.as_raw_fd();
		let res = unsafe { ftruncate64(fd, offs as off64_t) };
		if res == -1 {
			Err(io::Error::last_os_error())
		} else {
			Ok(())
		}
	}

	/// Saves a copy of the current file content, skipping data at the provided
	/// prune indices. The prune Vec must be ordered.
	pub fn save_prune<T>(
		&self,
		target: String,
		prune_offs: Vec<u64>,
		prune_len: u64,
		prune_cb: T,
	) -> io::Result<()>
	where
		T: Fn(&[u8]),
	{
		if prune_offs.is_empty() {
			fs::copy(self.path.clone(), target.clone())?;
			Ok(())
		} else {
			let mut reader = File::open(self.path.clone())?;
			let mut writer = File::create(target.clone())?;

			// align the buffer on prune_len to avoid misalignments
			let mut buf = vec![0; (prune_len * 256) as usize];
			let mut read = 0;
			let mut prune_pos = 0;
			loop {
				// fill our buffer
				let len = match reader.read(&mut buf) {
					Ok(0) => return Ok(()),
					Ok(len) => len,
					Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
					Err(e) => return Err(e),
				} as u64;

				// write the buffer, except if we prune offsets in the current span,
				// in which case we skip
				let mut buf_start = 0;
				while prune_offs[prune_pos] >= read && prune_offs[prune_pos] < read + len {
					let prune_at = (prune_offs[prune_pos] - read) as usize;
					if prune_at != buf_start {
						writer.write_all(&buf[buf_start..prune_at])?;
					} else {
						prune_cb(&buf[buf_start..prune_at]);
					}
					buf_start = prune_at + (prune_len as usize);
					if prune_offs.len() > prune_pos + 1 {
						prune_pos += 1;
					} else {
						break;
					}
				}
				writer.write_all(&mut buf[buf_start..(len as usize)])?;
				read += len;
			}
		}
	}

	/// Current size of the file in bytes.
	pub fn size(&self) -> io::Result<u64> {
		fs::metadata(&self.path).map(|md| md.len())
	}

	/// Path of the underlying file
	pub fn path(&self) -> String {
		self.path.clone()
	}
}

/// Log file fully cached in memory containing all positions that should be
/// eventually removed from the MMR append-only data file. Allows quick
/// checking of whether a piece of data has been marked for deletion. When the
/// log becomes too long, the MMR backend will actually remove chunks from the
/// MMR data file and truncate the remove log.
pub struct RemoveLog {
	path: String,
	/// Ordered vector of MMR positions that should get eventually removed.
	pub removed: Vec<(u64, u32)>,
	// Holds positions temporarily until flush is called.
	removed_tmp: Vec<(u64, u32)>,
	// Holds truncated removed temporarily until discarded or committed
	removed_bak: Vec<(u64, u32)>,
}

impl RemoveLog {
	/// Open the remove log file.
	/// The content of the file will be read in memory for fast checking.
	pub fn open(path: String) -> io::Result<RemoveLog> {
		let removed = read_ordered_vec(path.clone(), 12)?;
		Ok(RemoveLog {
			path: path,
			removed: removed,
			removed_tmp: vec![],
			removed_bak: vec![],
		})
	}

	/// Rewinds the remove log back to the provided index.
	/// We keep everything in the rm_log from that index and earlier.
	/// In practice the index is a block height, so we rewind back to that block
	/// keeping everything in the rm_log up to and including that block.
	pub fn rewind(&mut self, idx: u32) -> io::Result<()> {
		// simplifying assumption: we always remove older than what's in tmp
		self.removed_tmp = vec![];
		// backing it up before truncating
		self.removed_bak = self.removed.clone();

		if idx == 0 {
			self.removed = vec![];
		} else {
			// retain rm_log entries up to and including those at the provided index
			self.removed.retain(|&(_, x)| x <= idx);
		}
		Ok(())
	}

	/// Append a set of new positions to the remove log. Both adds those
	/// positions the ordered in-memory set and to the file.
	pub fn append(&mut self, elmts: Vec<u64>, index: u32) -> io::Result<()> {
		for elmt in elmts {
			match self.removed_tmp.binary_search(&(elmt, index)) {
				Ok(_) => continue,
				Err(idx) => {
					self.removed_tmp.insert(idx, (elmt, index));
				}
			}
		}
		Ok(())
	}

	/// Flush the positions to remove to file.
	pub fn flush(&mut self) -> io::Result<()> {
		let mut file = File::create(self.path.clone())?;
		for elmt in &self.removed_tmp {
			match self.removed.binary_search(&elmt) {
				Ok(_) => continue,
				Err(idx) => {
					self.removed.insert(idx, *elmt);
				}
			}
		}
		for elmt in &self.removed {
			file.write_all(&ser::ser_vec(&elmt).unwrap()[..])?;
		}
		self.removed_tmp = vec![];
		self.removed_bak = vec![];
		file.sync_data()
	}

	/// Discard pending changes
	pub fn discard(&mut self) {
		if self.removed_bak.len() > 0 {
			self.removed = self.removed_bak.clone();
			self.removed_bak = vec![];
		}
		self.removed_tmp = vec![];
	}

	/// Whether the remove log currently includes the provided position.
	pub fn includes(&self, elmt: u64) -> bool {
		include_tuple(&self.removed, elmt) || include_tuple(&self.removed_tmp, elmt)
	}

	/// Number of positions stored in the remove log.
	pub fn len(&self) -> usize {
		self.removed.len()
	}

	/// Return vec of pos for removed elements before the provided cutoff index.
	/// Useful for when we prune and compact an MMR.
	pub fn removed_pre_cutoff(&self, cutoff_idx: u32) -> Vec<u64> {
		self.removed
			.iter()
			.filter_map(
				|&(pos, idx)| {
					if idx < cutoff_idx {
						Some(pos)
					} else {
						None
					}
				},
			)
			.collect()
	}
}

fn include_tuple(v: &Vec<(u64, u32)>, e: u64) -> bool {
	if let Err(pos) = v.binary_search(&(e, 0)) {
		if pos < v.len() && v[pos].0 == e {
			return true;
		}
	}
	false
}

/// Read an ordered vector of scalars from a file.
pub fn read_ordered_vec<T>(path: String, elmt_len: usize) -> io::Result<Vec<T>>
where
	T: ser::Readable + cmp::Ord,
{
	let file_path = Path::new(&path);
	let mut ovec = Vec::with_capacity(1000);
	if file_path.exists() {
		let mut file = BufReader::with_capacity(elmt_len * 1000, File::open(path.clone())?);
		loop {
			// need a block to end mutable borrow before consume
			let buf_len = {
				let buf = file.fill_buf()?;
				if buf.len() == 0 {
					break;
				}
				let elmts_res: Result<Vec<T>, ser::Error> = ser::deserialize(&mut &buf[..]);
				match elmts_res {
					Ok(elmts) => for elmt in elmts {
						if let Err(idx) = ovec.binary_search(&elmt) {
							ovec.insert(idx, elmt);
						}
					},
					Err(_) => {
						return Err(io::Error::new(
							io::ErrorKind::InvalidData,
							format!("Corrupted storage, could not read file at {}", path),
						));
					}
				}
				buf.len()
			};
			file.consume(buf_len);
		}
	}
	Ok(ovec)
}

/// Writes an ordered vector to a file
pub fn write_vec<T>(path: String, v: &Vec<T>) -> io::Result<()>
where
	T: ser::Writeable,
{
	let mut file_path = File::create(&path)?;
	ser::serialize(&mut file_path, v).map_err(|_| {
		io::Error::new(
			io::ErrorKind::InvalidInput,
			format!("Failed to serialize data when writing to {}", path),
		)
	})?;
	Ok(())
}
<commit_msg>Buffers writes in the MMR store types (#1050)<commit_after>// Copyright 2018 The Grin Developers
// 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.

//! Common storage-related types
use memmap;

use std::cmp;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufRead, BufReader, BufWriter, ErrorKind, Write};
use std::os::unix::io::AsRawFd;
use std::io::Read;
use std::path::Path;

#[cfg(any(target_os = "linux"))]
use libc::{ftruncate64, off64_t};
#[cfg(not(any(target_os = "linux", target_os = "android")))]
use libc::{ftruncate as ftruncate64, off_t as off64_t};

use core::ser;

/// A no-op function for doing nothing with some pruned data.
pub fn prune_noop(_pruned_data: &[u8]) {}

/// Wrapper for a file that can be read at any position (random read) but for
/// which writes are append only. Reads are backed by a memory map (mmap(2)),
/// relying on the operating system for fast access and caching. The memory
/// map is reallocated to expand it when new writes are flushed.
///
/// Despite being append-only, the file can still be pruned and truncated. The
/// former simply happens by rewriting it, ignoring some of the data. The
/// latter by truncating the underlying file and re-creating the mmap.
pub struct AppendOnlyFile {
	path: String,
	file: File,
	mmap: Option<memmap::Mmap>,
	buffer_start: usize,
	buffer: Vec<u8>,
	buffer_start_bak: usize,
}

impl AppendOnlyFile {
	/// Open a file (existing or not) as append-only, backed by a mmap.
	pub fn open(path: String) -> io::Result<AppendOnlyFile> {
		let file = OpenOptions::new()
			.read(true)
			.append(true)
			.create(true)
			.open(path.clone())?;
		let mut aof = AppendOnlyFile {
			path: path.clone(),
			file: file,
			mmap: None,
			buffer_start: 0,
			buffer: vec![],
			buffer_start_bak: 0,
		};
		// if we have a non-empty file then mmap it.
		if let Ok(sz) = aof.size() {
			if sz > 0 {
				aof.buffer_start = sz as usize;
				aof.mmap = Some(unsafe { memmap::Mmap::map(&aof.file)? });
			}
		}
		Ok(aof)
	}

	/// Append data to the file. Until the append-only file is synced, data is
	/// only written to memory.
	pub fn append(&mut self, buf: &mut Vec<u8>) {
		self.buffer.append(buf);
	}

	/// Rewinds the data file back to a lower position. The new position needs
	/// to be the one of the first byte the next time data is appended.
	pub fn rewind(&mut self, pos: u64) {
		if self.buffer_start_bak > 0 || self.buffer.len() > 0 {
			panic!("Can't rewind on a dirty state.");
		}
		self.buffer_start_bak = self.buffer_start;
		self.buffer_start = pos as usize;
	}

	/// Syncs all writes (fsync), reallocating the memory map to make the newly
	/// written data accessible.
	pub fn flush(&mut self) -> io::Result<()> {
		if self.buffer_start_bak > 0 {
			// flushing a rewound state, we need to truncate before applying
			self.truncate(self.buffer_start)?;
			self.buffer_start_bak = 0;
		}

		self.buffer_start += self.buffer.len();
		self.file.write(&self.buffer[..])?;
		self.file.sync_all()?;

		self.buffer = vec![];

		// Note: file must be non-empty to memory map it
		if self.file.metadata()?.len() == 0 {
			self.mmap = None;
		} else {
			self.mmap = Some(unsafe { memmap::Mmap::map(&self.file)? });
		}

		Ok(())
	}

	/// Returns the last position (in bytes), taking into account whether data
	/// has been rewound
	pub fn last_buffer_pos(&self) -> usize {
		self.buffer_start
	}

	/// Discard the current non-flushed data.
	pub fn discard(&mut self) {
		if self.buffer_start_bak > 0 {
			// discarding a rewound state, restore the buffer start
			self.buffer_start = self.buffer_start_bak;
			self.buffer_start_bak = 0;
		}
		self.buffer = vec![];
	}

	/// Read length bytes of data at offset from the file.
	/// Leverages the memory map.
	pub fn read(&self, offset: usize, length: usize) -> Vec<u8> {
		if offset >= self.buffer_start {
			if self.buffer.is_empty() {
				return vec![];
			}
			let offset = offset - self.buffer_start;
			return self.buffer[offset..(offset + length)].to_vec();
		}
		if let None = self.mmap {
			return vec![];
		}
		let mmap = self.mmap.as_ref().unwrap();

		if mmap.len() < (offset + length) {
			return vec![];
		}

		(&mmap[offset..(offset + length)]).to_vec()
	}

	/// Truncates the underlying file to the provided offset
	pub fn truncate(&self, offs: usize) -> io::Result<()> {
		let fd = self.file.as_raw_fd();
		let res = unsafe { ftruncate64(fd, offs as off64_t) };
		if res == -1 {
			Err(io::Error::last_os_error())
		} else {
			Ok(())
		}
	}

	/// Saves a copy of the current file content, skipping data at the provided
	/// prune indices. The prune Vec must be ordered.
	pub fn save_prune<T>(
		&self,
		target: String,
		prune_offs: Vec<u64>,
		prune_len: u64,
		prune_cb: T,
	) -> io::Result<()>
	where
		T: Fn(&[u8]),
	{
		if prune_offs.is_empty() {
			fs::copy(self.path.clone(), target.clone())?;
			Ok(())
		} else {
			let mut reader = File::open(self.path.clone())?;
			let mut writer = BufWriter::new(File::create(target.clone())?);

			// align the buffer on prune_len to avoid misalignments
			let mut buf = vec![0; (prune_len * 256) as usize];
			let mut read = 0;
			let mut prune_pos = 0;
			loop {
				// fill our buffer
				let len = match reader.read(&mut buf) {
					Ok(0) => return Ok(()),
					Ok(len) => len,
					Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
					Err(e) => return Err(e),
				} as u64;

				// write the buffer, except if we prune offsets in the current span,
				// in which case we skip
				let mut buf_start = 0;
				while prune_offs[prune_pos] >= read && prune_offs[prune_pos] < read + len {
					let prune_at = (prune_offs[prune_pos] - read) as usize;
					if prune_at != buf_start {
						writer.write_all(&buf[buf_start..prune_at])?;
					} else {
						prune_cb(&buf[buf_start..prune_at]);
					}
					buf_start = prune_at + (prune_len as usize);
					if prune_offs.len() > prune_pos + 1 {
						prune_pos += 1;
					} else {
						break;
					}
				}
				writer.write_all(&mut buf[buf_start..(len as usize)])?;
				read += len;
			}
		}
	}

	/// Current size of the file in bytes.
	pub fn size(&self) -> io::Result<u64> {
		fs::metadata(&self.path).map(|md| md.len())
	}

	/// Path of the underlying file
	pub fn path(&self) -> String {
		self.path.clone()
	}
}

/// Log file fully cached in memory containing all positions that should be
/// eventually removed from the MMR append-only data file. Allows quick
/// checking of whether a piece of data has been marked for deletion. When the
/// log becomes too long, the MMR backend will actually remove chunks from the
/// MMR data file and truncate the remove log.
pub struct RemoveLog {
	path: String,
	/// Ordered vector of MMR positions that should get eventually removed.
	pub removed: Vec<(u64, u32)>,
	// Holds positions temporarily until flush is called.
	removed_tmp: Vec<(u64, u32)>,
	// Holds truncated removed temporarily until discarded or committed
	removed_bak: Vec<(u64, u32)>,
}

impl RemoveLog {
	/// Open the remove log file.
	/// The content of the file will be read in memory for fast checking.
	pub fn open(path: String) -> io::Result<RemoveLog> {
		let removed = read_ordered_vec(path.clone(), 12)?;
		Ok(RemoveLog {
			path: path,
			removed: removed,
			removed_tmp: vec![],
			removed_bak: vec![],
		})
	}

	/// Rewinds the remove log back to the provided index.
	/// We keep everything in the rm_log from that index and earlier.
	/// In practice the index is a block height, so we rewind back to that block
	/// keeping everything in the rm_log up to and including that block.
	pub fn rewind(&mut self, idx: u32) -> io::Result<()> {
		// simplifying assumption: we always remove older than what's in tmp
		self.removed_tmp = vec![];
		// backing it up before truncating
		self.removed_bak = self.removed.clone();

		if idx == 0 {
			self.removed = vec![];
		} else {
			// retain rm_log entries up to and including those at the provided index
			self.removed.retain(|&(_, x)| x <= idx);
		}
		Ok(())
	}

	/// Append a set of new positions to the remove log. Both adds those
	/// positions the ordered in-memory set and to the file.
	pub fn append(&mut self, elmts: Vec<u64>, index: u32) -> io::Result<()> {
		for elmt in elmts {
			match self.removed_tmp.binary_search(&(elmt, index)) {
				Ok(_) => continue,
				Err(idx) => {
					self.removed_tmp.insert(idx, (elmt, index));
				}
			}
		}
		Ok(())
	}

	/// Flush the positions to remove to file.
	pub fn flush(&mut self) -> io::Result<()> {
		for elmt in &self.removed_tmp {
			match self.removed.binary_search(&elmt) {
				Ok(_) => continue,
				Err(idx) => {
					self.removed.insert(idx, *elmt);
				}
			}
		}
		let mut file = BufWriter::new(File::create(self.path.clone())?);
		for elmt in &self.removed {
			file.write_all(&ser::ser_vec(&elmt).unwrap()[..])?;
		}
		self.removed_tmp = vec![];
		self.removed_bak = vec![];
		file.flush()
	}

	/// Discard pending changes
	pub fn discard(&mut self) {
		if self.removed_bak.len() > 0 {
			self.removed = self.removed_bak.clone();
			self.removed_bak = vec![];
		}
		self.removed_tmp = vec![];
	}

	/// Whether the remove log currently includes the provided position.
	pub fn includes(&self, elmt: u64) -> bool {
		include_tuple(&self.removed, elmt) || include_tuple(&self.removed_tmp, elmt)
	}

	/// Number of positions stored in the remove log.
	pub fn len(&self) -> usize {
		self.removed.len()
	}

	/// Return vec of pos for removed elements before the provided cutoff index.
	/// Useful for when we prune and compact an MMR.
	pub fn removed_pre_cutoff(&self, cutoff_idx: u32) -> Vec<u64> {
		self.removed
			.iter()
			.filter_map(
				|&(pos, idx)| {
					if idx < cutoff_idx {
						Some(pos)
					} else {
						None
					}
				},
			)
			.collect()
	}
}

fn include_tuple(v: &Vec<(u64, u32)>, e: u64) -> bool {
	if let Err(pos) = v.binary_search(&(e, 0)) {
		if pos < v.len() && v[pos].0 == e {
			return true;
		}
	}
	false
}

/// Read an ordered vector of scalars from a file.
pub fn read_ordered_vec<T>(path: String, elmt_len: usize) -> io::Result<Vec<T>>
where
	T: ser::Readable + cmp::Ord,
{
	let file_path = Path::new(&path);
	let mut ovec = Vec::with_capacity(1000);
	if file_path.exists() {
		let mut file = BufReader::with_capacity(elmt_len * 1000, File::open(path.clone())?);
		loop {
			// need a block to end mutable borrow before consume
			let buf_len = {
				let buf = file.fill_buf()?;
				if buf.len() == 0 {
					break;
				}
				let elmts_res: Result<Vec<T>, ser::Error> = ser::deserialize(&mut &buf[..]);
				match elmts_res {
					Ok(elmts) => for elmt in elmts {
						if let Err(idx) = ovec.binary_search(&elmt) {
							ovec.insert(idx, elmt);
						}
					},
					Err(_) => {
						return Err(io::Error::new(
							io::ErrorKind::InvalidData,
							format!("Corrupted storage, could not read file at {}", path),
						));
					}
				}
				buf.len()
			};
			file.consume(buf_len);
		}
	}
	Ok(ovec)
}

/// Writes an ordered vector to a file
pub fn write_vec<T>(path: String, v: &Vec<T>) -> io::Result<()>
where
	T: ser::Writeable,
{
	let mut file_path = File::create(&path)?;
	ser::serialize(&mut file_path, v).map_err(|_| {
		io::Error::new(
			io::ErrorKind::InvalidInput,
			format!("Failed to serialize data when writing to {}", path),
		)
	})?;
	Ok(())
}
<|endoftext|>


<commit_before>// run-pass

// edition:2018

#![feature(async_await)]

extern crate arc_wake;

use std::pin::Pin;
use std::future::Future;
use std::sync::{
    Arc,
    atomic::{self, AtomicUsize},
};
use std::task::{Context, Poll};
use arc_wake::ArcWake;

struct Counter {
    wakes: AtomicUsize,
}

impl ArcWake for Counter {
    fn wake(self: Arc<Self>) {
        Self::wake_by_ref(&self)
    }
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
    }
}

struct WakeOnceThenComplete(bool, u8);

impl Future for WakeOnceThenComplete {
    type Output = u8;
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u8> {
        if self.0 {
            Poll::Ready(self.1)
        } else {
            cx.waker().wake_by_ref();
            self.0 = true;
            Poll::Pending
        }
    }
}

fn wait(fut: impl Future<Output = u8>) -> u8 {
    let mut fut = Box::pin(fut);
    let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
    let waker = ArcWake::into_waker(counter.clone());
    let mut cx = Context::from_waker(&waker);
    loop {
        match fut.as_mut().poll(&mut cx) {
            Poll::Ready(out) => return out,
            Poll::Pending => (),
        }
    }
}

fn base() -> WakeOnceThenComplete { WakeOnceThenComplete(false, 1) }

async fn await1_level1() -> u8 {
    base().await
}

async fn await2_level1() -> u8 {
    base().await + base().await
}

async fn await3_level1() -> u8 {
    base().await + base().await + base().await
}

async fn await3_level2() -> u8 {
    await3_level1().await + await3_level1().await + await3_level1().await
}

async fn await3_level3() -> u8 {
    await3_level2().await + await3_level2().await + await3_level2().await
}

async fn await3_level4() -> u8 {
    await3_level3().await + await3_level3().await + await3_level3().await
}

async fn await3_level5() -> u8 {
    await3_level4().await + await3_level4().await + await3_level4().await
}

fn main() {
    assert_eq!(2, std::mem::size_of_val(&base()));
    assert_eq!(8, std::mem::size_of_val(&await1_level1()));
    assert_eq!(12, std::mem::size_of_val(&await2_level1()));
    assert_eq!(12, std::mem::size_of_val(&await3_level1()));
    assert_eq!(20, std::mem::size_of_val(&await3_level2()));
    assert_eq!(28, std::mem::size_of_val(&await3_level3()));
    assert_eq!(36, std::mem::size_of_val(&await3_level4()));
    assert_eq!(44, std::mem::size_of_val(&await3_level5()));

    assert_eq!(1,   wait(base()));
    assert_eq!(1,   wait(await1_level1()));
    assert_eq!(2,   wait(await2_level1()));
    assert_eq!(3,   wait(await3_level1()));
    assert_eq!(9,   wait(await3_level2()));
    assert_eq!(27,  wait(await3_level3()));
    assert_eq!(81,  wait(await3_level4()));
    assert_eq!(243, wait(await3_level5()));
}
<commit_msg>Add missing aux-build directive<commit_after>// run-pass
// aux-build:arc_wake.rs
// edition:2018

#![feature(async_await)]

extern crate arc_wake;

use std::pin::Pin;
use std::future::Future;
use std::sync::{
    Arc,
    atomic::{self, AtomicUsize},
};
use std::task::{Context, Poll};
use arc_wake::ArcWake;

struct Counter {
    wakes: AtomicUsize,
}

impl ArcWake for Counter {
    fn wake(self: Arc<Self>) {
        Self::wake_by_ref(&self)
    }
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
    }
}

struct WakeOnceThenComplete(bool, u8);

impl Future for WakeOnceThenComplete {
    type Output = u8;
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u8> {
        if self.0 {
            Poll::Ready(self.1)
        } else {
            cx.waker().wake_by_ref();
            self.0 = true;
            Poll::Pending
        }
    }
}

fn wait(fut: impl Future<Output = u8>) -> u8 {
    let mut fut = Box::pin(fut);
    let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
    let waker = ArcWake::into_waker(counter.clone());
    let mut cx = Context::from_waker(&waker);
    loop {
        match fut.as_mut().poll(&mut cx) {
            Poll::Ready(out) => return out,
            Poll::Pending => (),
        }
    }
}

fn base() -> WakeOnceThenComplete { WakeOnceThenComplete(false, 1) }

async fn await1_level1() -> u8 {
    base().await
}

async fn await2_level1() -> u8 {
    base().await + base().await
}

async fn await3_level1() -> u8 {
    base().await + base().await + base().await
}

async fn await3_level2() -> u8 {
    await3_level1().await + await3_level1().await + await3_level1().await
}

async fn await3_level3() -> u8 {
    await3_level2().await + await3_level2().await + await3_level2().await
}

async fn await3_level4() -> u8 {
    await3_level3().await + await3_level3().await + await3_level3().await
}

async fn await3_level5() -> u8 {
    await3_level4().await + await3_level4().await + await3_level4().await
}

fn main() {
    assert_eq!(2, std::mem::size_of_val(&base()));
    assert_eq!(8, std::mem::size_of_val(&await1_level1()));
    assert_eq!(12, std::mem::size_of_val(&await2_level1()));
    assert_eq!(12, std::mem::size_of_val(&await3_level1()));
    assert_eq!(20, std::mem::size_of_val(&await3_level2()));
    assert_eq!(28, std::mem::size_of_val(&await3_level3()));
    assert_eq!(36, std::mem::size_of_val(&await3_level4()));
    assert_eq!(44, std::mem::size_of_val(&await3_level5()));

    assert_eq!(1,   wait(base()));
    assert_eq!(1,   wait(await1_level1()));
    assert_eq!(2,   wait(await2_level1()));
    assert_eq!(3,   wait(await3_level1()));
    assert_eq!(9,   wait(await3_level2()));
    assert_eq!(27,  wait(await3_level3()));
    assert_eq!(81,  wait(await3_level4()));
    assert_eq!(243, wait(await3_level5()));
}
<|endoftext|>


<commit_before>use alloc::{Allocator, SliceWrapper};
use super::BrotliAlloc;
use super::backward_references::BrotliEncoderParams;
pub trait Joinable<T:Send+'static>:Sized {
  fn join(self) -> Result<T, ()>;
}

pub trait Spawnable<T:Send+'static> {
  type JoinHandle: Joinable<T>;
  fn spawn<F: FnOnce() -> T + Send + 'static>(f: F) -> Self::JoinHandle;
}

pub struct CompressionThreadResult<Alloc:BrotliAlloc+Send+'static> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  compressed: <Alloc as Allocator<u8>>::AllocatedMemory,
  compressed_size: usize,
  alloc: Alloc,
}
enum InternalSendAlloc<Alloc:BrotliAlloc+Send+'static, Join: Joinable<CompressionThreadResult<Alloc>>>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  A(Alloc),
  Join(Join),
  Spawning,
}
pub struct SendAlloc<Alloc:BrotliAlloc +Send+'static,
                     Join:Joinable<CompressionThreadResult<Alloc>>>(InternalSendAlloc<Alloc, Join>)
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send;

impl<Alloc:BrotliAlloc+Send+'static,Join:Joinable<CompressionThreadResult<Alloc>>> SendAlloc<Alloc, Join>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  pub fn new(alloc: Alloc) -> Self {
    SendAlloc::<Alloc, Join>(InternalSendAlloc::A(alloc))
  }
  pub fn unwrap_or(self, other: Alloc) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::Spawning | InternalSendAlloc::Join(_) => {
        other
      },
    }
  }
  pub fn unwrap(self) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::Join(_) | InternalSendAlloc::Spawning => panic!("Item permanently borrowed/leaked"),
    }
  }
}

enum InternalOwned<T> {
  Item(T),
  Borrowed,
}

pub struct Owned<T>(InternalOwned<T>);
impl<T> Owned<T> {
  pub fn new(data:T) -> Self {
    Owned::<T>(InternalOwned::Item(data))
  }
  pub fn unwrap_or(self, other: T) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      other
    }
  }
  pub fn unwrap(self) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      panic!("Item permanently borrowed")
    }
  }
}

pub fn CompressMulti<Alloc:BrotliAlloc+Send+'static, SliceW: SliceWrapper<u8>, Spawner:Spawnable<CompressionThreadResult<Alloc>>> (
  params:&BrotliEncoderParams,
  input: &mut Owned<SliceW>,
  output: &mut [u8],
  alloc_per_thread:&mut [SendAlloc<Alloc, Spawner::JoinHandle>],
  thread_spawner: Spawner,
) -> Result<usize, ()> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  Err(())
}

<commit_msg>making batch_spawn a little less general--maybe it's too specific now--may generalize it a bit more now<commit_after>use alloc::{Allocator, SliceWrapper};
use super::BrotliAlloc;
use super::backward_references::BrotliEncoderParams;
pub trait Joinable<T:Send+'static>:Sized {
  fn join(self) -> Result<T, ()>;
}

pub struct CompressionThreadResult<Alloc:BrotliAlloc+Send+'static> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  compressed: <Alloc as Allocator<u8>>::AllocatedMemory,
  compressed_size: usize,
  alloc: Alloc,
}
enum InternalSendAlloc<Alloc:BrotliAlloc+Send+'static, Join: Joinable<CompressionThreadResult<Alloc>>>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  A(Alloc),
  Join(Join),
  Spawning,
}
pub struct SendAlloc<Alloc:BrotliAlloc +Send+'static,
                     Join:Joinable<CompressionThreadResult<Alloc>>>(InternalSendAlloc<Alloc, Join>)
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send;

impl<Alloc:BrotliAlloc+Send+'static,Join:Joinable<CompressionThreadResult<Alloc>>> SendAlloc<Alloc, Join>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  pub fn new(alloc: Alloc) -> Self {
    SendAlloc::<Alloc, Join>(InternalSendAlloc::A(alloc))
  }
  pub fn unwrap_or(self, other: Alloc) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::Spawning | InternalSendAlloc::Join(_) => {
        other
      },
    }
  }
  pub fn unwrap(self) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::Join(_) | InternalSendAlloc::Spawning => panic!("Item permanently borrowed/leaked"),
    }
  }
}

enum InternalOwned<T> {
  Item(T),
  Borrowed,
}

pub struct Owned<T>(InternalOwned<T>);
impl<T> Owned<T> {
  pub fn new(data:T) -> Self {
    Owned::<T>(InternalOwned::Item(data))
  }
  pub fn unwrap_or(self, other: T) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      other
    }
  }
  pub fn unwrap(self) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      panic!("Item permanently borrowed")
    }
  }
}

pub trait OwnedRetriever<SliceW: SliceWrapper<u8>+Send+'static> {
  fn unwrap(self) -> SliceW;
}

pub trait BatchSpawnable<Alloc:BrotliAlloc+Send+'static, SliceW:SliceWrapper<u8>+Send+'static>
  where <Alloc as Allocator<u8>>::AllocatedMemory:Send+'static
{
  type JoinHandle: Joinable<CompressionThreadResult<Alloc>>;
  type FinalJoinHandle: OwnedRetriever<SliceW>;
  // this function takes in a SendAlloc per thread and converts them all into JoinHandle
  // the input is borrowed until the joins complete
  // owned is set to borrowed
  // the final join handle is a r/w lock which will return the SliceW to the owner
  // the FinalJoinHandle is only to be called when each individual JoinHandle has been examined
  fn batch_spawn<F: FnOnce() -> CompressionThreadResult<Alloc> + Send + 'static>(
    input: &mut Owned<SliceW>,
    alloc_per_thread:&mut [SendAlloc<Alloc, Self::JoinHandle>],
    f: F,
  ) -> Self::FinalJoinHandle;
}

pub fn CompressMulti<Alloc:BrotliAlloc+Send+'static, SliceW: SliceWrapper<u8>+Send+'static, Spawner:BatchSpawnable<Alloc, SliceW>> (
  params:&BrotliEncoderParams,
  input: &mut Owned<SliceW>,
  output: &mut [u8],
  alloc_per_thread:&mut [SendAlloc<Alloc, Spawner::JoinHandle>],
  thread_spawner: Spawner,
) -> Result<usize, ()> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {

  Err(())
}

<|endoftext|>


<commit_before>use crate::convert::*;
use core::hash::{Hasher};

///This constant come from Kunth's prng (Empirically it works better than those from splitmix32).
const MULTIPLE: u64 = 6364136223846793005;
const INCREMENT: u64 = 1442695040888963407;
const ROT: u32 = 23;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [AHasher] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: u64
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    #[inline]
    pub(crate) fn new_with_key(key: u64, loc: u64) -> AHasher {
        AHasher { buffer: key ^ (loc.rotate_left(ROT)) }
    }

    /// This update function has the goal of updating the buffer with a single multiply
    /// FxHash does this but is venerable to attack. To avoid this input needs to be masked to with an unpredictable value.
    /// However other hashes such as murmurhash have taken that approach but were found venerable to attack.
    /// The attack was based on the idea of reversing the pre-mixing (Which is necessarily reversible otherwise
    /// bits would be lost) then placing a difference in the highest bit before the multiply. Because a multiply
    /// can never affect the bits to the right of it. This version avoids this vulnerability by rotating and
    /// performing a second multiply. This makes it impossible for an attacker to place a single bit
    /// difference between two blocks so as to cancel each other. (While the transform is still reversible if you know the key)
    ///
    /// The update of the buffer to perform the second multiply is moved from the end to the beginning of the method.
    /// This has the effect of causing the next call to update to perform he second multiply. For the final
    /// update this is performed in the finalize method. This might seem wasteful, but its actually an optimization.
    /// If the method get's inlined into the caller where it is being invoked on a single primitive, the first call
    /// to update the buffer will be operating on constants and the compiler will optimize it out, by replacing it with
    /// the result.
    #[inline(always)]
    fn update(&mut self, new_data: u64) {
        let existing = self.buffer.wrapping_mul(MULTIPLE).rotate_left(ROT).wrapping_mul(MULTIPLE);
        self.buffer = existing ^ new_data;
    }

    /// This is similar to the above update function (see it's description). But is designed to run in a loop
    /// that will be unrolled and vectorized. So instead of using the buffer, it uses a 'key' that it updates
    /// and returns. The buffer is only xored at the end. This structure is so that when the method is inlined,
    /// the compiler will unroll any loop this gets placed in and the loop can be automatically vectorized
    /// and the rotates, xors, and multiplies can be paralleled.
    ///
    /// The key needs to be incremented between consecutive calls to prevent (a,b) from hashing the same as (b,a).
    /// The adding of the increment is moved to the bottom rather than the top. This allows one less add to be
    /// performed overall, but more importantly, it follows the multiply, which is expensive. So the CPU can
    /// run another operation afterwords if does not depend on the output of the multiply operation.
    #[inline(always)]
    fn ordered_update(&mut self, new_data: u64, key: u64) -> u64 {
        self.buffer ^= (new_data ^ key).wrapping_mul(MULTIPLE).rotate_left(ROT).wrapping_mul(MULTIPLE);
        key.wrapping_add(INCREMENT)
    }
}

#[inline(never)]
#[no_mangle]
fn hash_test(input: &[u8]) -> u64 {
    let mut a = AHasher::new_with_key(67, 87);
    a.write(input);
    a.finish()
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {

    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        let data: [u64;2] = i.convert();
        self.update(data[0]);
        self.update(data[1]);
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //Needs to be an add rather than an xor because otherwise it could be canceled with carefully formed input.
        self.buffer = self.buffer.wrapping_add(length);
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() > 8 {
            if data.len() > 16 {
                let tail = data.read_last_u64();
                let mut key: u64 = self.buffer;
                while data.len() > 8 {
                    let (val, rest) = data.read_u64();
                    key = self.ordered_update(val, key);
                    data = rest;
                }
                self.update(tail);
            } else {
                self.update(data.read_u64().0);
                self.update(data.read_last_u64());
            }
        } else {
            if data.len() >= 2 {
                if data.len() >= 4 {
                    let block: [u32; 2] = [data.read_u32().0, data.read_last_u32()];
                    self.update(block.convert());
                } else {
                    let block: [u16; 2] = [data.read_u16().0, data.read_last_u16()];
                    let val: u32 = block.convert();
                    self.update(val as u64);
                }
            } else {
                let value;
                if data.len() > 0 {
                    value = data[0]; //len 1
                } else {
                    value = 0;
                }
                self.update(value as u64);
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        self.buffer.wrapping_mul(MULTIPLE).rotate_left(9).wrapping_mul(MULTIPLE)
    }
}

#[cfg(test)]
mod tests {
    use crate::convert::Convert;
    use crate::fallback_hash::*;

    #[test]
    fn test_hash() {
        let mut hasher = AHasher::new_with_key(0,0);
        let value: u64 = 1 << 32;
        hasher.update(value);
        let result = hasher.buffer;
        let mut hasher = AHasher::new_with_key(0,0);
        let value2: u64 = 1;
        hasher.update(value2);
        let result2 = hasher.buffer;
        let result: [u8; 8] = result.convert();
        let result2: [u8; 8] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<commit_msg>Change fallback to eagerly mix in new data. (In prep for moving away for const_random)<commit_after>use crate::convert::*;
use core::hash::{Hasher};

///This constant come from Kunth's prng (Empirically it works better than those from splitmix32).
const MULTIPLE: u64 = 6364136223846793005;
const INCREMENT: u64 = 1442695040888963407;
const ROT: u32 = 23;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [AHasher] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: u64
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    #[inline]
    pub(crate) fn new_with_key(key: u64, loc: u64) -> AHasher {
        AHasher { buffer: key ^ (loc.rotate_left(ROT)) }
    }

    /// This update function has the goal of updating the buffer with a single multiply
    /// FxHash does this but is venerable to attack. To avoid this input needs to be masked to with an unpredictable value.
    /// However other hashes such as murmurhash have taken that approach but were found venerable to attack.
    /// The attack was based on the idea of reversing the pre-mixing (Which is necessarily reversible otherwise
    /// bits would be lost) then placing a difference in the highest bit before the multiply. Because a multiply
    /// can never affect the bits to the right of it. This version avoids this vulnerability by rotating and
    /// performing a second multiply. This makes it impossible for an attacker to place a single bit
    /// difference between two blocks so as to cancel each other. (While the transform is still reversible if you know the key)
    #[inline(always)]
    fn update(&mut self, new_data: u64) {
        self.buffer = (self.buffer ^ new_data).wrapping_mul(MULTIPLE).rotate_left(ROT).wrapping_mul(MULTIPLE);
    }

    /// This is similar to the above update function (see it's description). But is designed to run in a loop
    /// that will be unrolled and vectorized. So instead of using the buffer, it uses a 'key' that it updates
    /// and returns. The buffer is only xored at the end. This structure is so that when the method is inlined,
    /// the compiler will unroll any loop this gets placed in and the loop can be automatically vectorized
    /// and the rotates, xors, and multiplies can be paralleled.
    ///
    /// The key needs to be incremented between consecutive calls to prevent (a,b) from hashing the same as (b,a).
    /// The adding of the increment is moved to the bottom rather than the top. This allows one less add to be
    /// performed overall, but more importantly, it follows the multiply, which is expensive. So the CPU can
    /// run another operation afterwords if does not depend on the output of the multiply operation.
    #[inline(always)]
    fn ordered_update(&mut self, new_data: u64, key: u64) -> u64 {
        self.buffer ^= (new_data ^ key).wrapping_mul(MULTIPLE).rotate_left(ROT).wrapping_mul(MULTIPLE);
        key.wrapping_add(INCREMENT)
    }
}

#[inline(never)]
#[no_mangle]
fn hash_test(input: &[u8]) -> u64 {
    let mut a = AHasher::new_with_key(67, 87);
    a.write(input);
    a.finish()
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {

    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.update(i as u64);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        let data: [u64;2] = i.convert();
        self.update(data[0]);
        self.update(data[1]);
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //Needs to be an add rather than an xor because otherwise it could be canceled with carefully formed input.
        self.buffer = self.buffer.wrapping_add(length);
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() > 8 {
            if data.len() > 16 {
                let tail = data.read_last_u64();
                let mut key: u64 = self.buffer;
                while data.len() > 8 {
                    let (val, rest) = data.read_u64();
                    key = self.ordered_update(val, key);
                    data = rest;
                }
                self.update(tail);
            } else {
                self.update(data.read_u64().0);
                self.update(data.read_last_u64());
            }
        } else {
            if data.len() >= 2 {
                if data.len() >= 4 {
                    let block: [u32; 2] = [data.read_u32().0, data.read_last_u32()];
                    self.update(block.convert());
                } else {
                    let block: [u16; 2] = [data.read_u16().0, data.read_last_u16()];
                    let val: u32 = block.convert();
                    self.update(val as u64);
                }
            } else {
                let value;
                if data.len() > 0 {
                    value = data[0]; //len 1
                } else {
                    value = 0;
                }
                self.update(value as u64);
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        return self.buffer;
    }
}

#[cfg(test)]
mod tests {
    use crate::convert::Convert;
    use crate::fallback_hash::*;

    #[test]
    fn test_hash() {
        let mut hasher = AHasher::new_with_key(0,0);
        let value: u64 = 1 << 32;
        hasher.update(value);
        let result = hasher.buffer;
        let mut hasher = AHasher::new_with_key(0,0);
        let value2: u64 = 1;
        hasher.update(value2);
        let result2 = hasher.buffer;
        let result: [u8; 8] = result.convert();
        let result2: [u8; 8] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 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.

/*!

The iteration traits and common implementation

*/

use cmp::{Eq, Ord};
use kinds::Copy;
use option::{None, Option, Some};
use vec;

/// A function used to initialize the elements of a sequence
pub type InitOp<'self,T> = &'self fn(uint) -> T;

pub trait BaseIter<A> {
    fn each(&self, blk: &fn(v: &A) -> bool);
    fn size_hint(&self) -> Option<uint>;
}

pub trait ReverseIter<A>: BaseIter<A> {
    fn each_reverse(&self, blk: &fn(&A) -> bool);
}

pub trait MutableIter<A>: BaseIter<A> {
    fn each_mut(&mut self, blk: &fn(&mut A) -> bool);
}

pub trait ExtendedIter<A> {
    fn eachi(&self, blk: &fn(uint, v: &A) -> bool);
    fn all(&self, blk: &fn(&A) -> bool) -> bool;
    fn any(&self, blk: &fn(&A) -> bool) -> bool;
    fn foldl<B>(&self, b0: B, blk: &fn(&B, &A) -> B) -> B;
    fn position(&self, f: &fn(&A) -> bool) -> Option<uint>;
    fn map_to_vec<B>(&self, op: &fn(&A) -> B) -> ~[B];
    fn flat_map_to_vec<B,IB: BaseIter<B>>(&self, op: &fn(&A) -> IB) -> ~[B];
}

pub trait EqIter<A:Eq> {
    fn contains(&self, x: &A) -> bool;
    fn count(&self, x: &A) -> uint;
}

pub trait Times {
    fn times(&self, it: &fn() -> bool);
}

pub trait CopyableIter<A:Copy> {
    fn filter_to_vec(&self, pred: &fn(&A) -> bool) -> ~[A];
    fn to_vec(&self) -> ~[A];
    fn find(&self, p: &fn(&A) -> bool) -> Option<A>;
}

pub trait CopyableOrderedIter<A:Copy + Ord> {
    fn min(&self) -> A;
    fn max(&self) -> A;
}

pub trait CopyableNonstrictIter<A:Copy> {
    // Like "each", but copies out the value. If the receiver is mutated while
    // iterating over it, the semantics must not be memory-unsafe but are
    // otherwise undefined.
    fn each_val(&const self, f: &fn(A) -> bool);
}

// A trait for sequences that can be built by imperatively pushing elements
// onto them.
pub trait Buildable<A> {
    /**
     * Builds a buildable sequence by calling a provided function with
     * an argument function that pushes an element onto the back of
     * the sequence.
     * This version takes an initial size for the sequence.
     *
     * # Arguments
     *
     * * size - A hint for an initial size of the sequence
     * * builder - A function that will construct the sequence. It receives
     *             as an argument a function that will push an element
     *             onto the sequence being constructed.
     */
     fn build_sized(size: uint, builder: &fn(push: &fn(A))) -> Self;
}

#[inline(always)]
pub fn eachi<A,IA:BaseIter<A>>(self: &IA, blk: &fn(uint, &A) -> bool) {
    let mut i = 0;
    for self.each |a| {
        if !blk(i, a) { break; }
        i += 1;
    }
}

#[inline(always)]
pub fn all<A,IA:BaseIter<A>>(self: &IA, blk: &fn(&A) -> bool) -> bool {
    for self.each |a| {
        if !blk(a) { return false; }
    }
    return true;
}

#[inline(always)]
pub fn any<A,IA:BaseIter<A>>(self: &IA, blk: &fn(&A) -> bool) -> bool {
    for self.each |a| {
        if blk(a) { return true; }
    }
    return false;
}

#[inline(always)]
pub fn filter_to_vec<A:Copy,IA:BaseIter<A>>(self: &IA,
                                            prd: &fn(&A) -> bool)
                                         -> ~[A] {
    do vec::build_sized_opt(self.size_hint()) |push| {
        for self.each |a| {
            if prd(a) { push(*a); }
        }
    }
}

#[inline(always)]
pub fn map_to_vec<A,B,IA:BaseIter<A>>(self: &IA, op: &fn(&A) -> B) -> ~[B] {
    do vec::build_sized_opt(self.size_hint()) |push| {
        for self.each |a| {
            push(op(a));
        }
    }
}

#[inline(always)]
pub fn flat_map_to_vec<A,B,IA:BaseIter<A>,IB:BaseIter<B>>(self: &IA,
                                                          op: &fn(&A) -> IB)
                                                       -> ~[B] {
    do vec::build |push| {
        for self.each |a| {
            for op(a).each |&b| {
                push(b);
            }
        }
    }
}

#[inline(always)]
pub fn foldl<A,B,IA:BaseIter<A>>(self: &IA, b0: B, blk: &fn(&B, &A) -> B)
                              -> B {
    let mut b = b0;
    for self.each |a| {
        b = blk(&b, a);
    }
    b
}

#[inline(always)]
pub fn to_vec<A:Copy,IA:BaseIter<A>>(self: &IA) -> ~[A] {
    foldl::<A,~[A],IA>(self, ~[], |r, a| vec::append(copy (*r), ~[*a]))
}

#[inline(always)]
pub fn contains<A:Eq,IA:BaseIter<A>>(self: &IA, x: &A) -> bool {
    for self.each |a| {
        if *a == *x { return true; }
    }
    return false;
}

#[inline(always)]
pub fn count<A:Eq,IA:BaseIter<A>>(self: &IA, x: &A) -> uint {
    do foldl(self, 0) |count, value| {
        if *value == *x {
            *count + 1
        } else {
            *count
        }
    }
}

#[inline(always)]
pub fn position<A,IA:BaseIter<A>>(self: &IA, f: &fn(&A) -> bool)
                               -> Option<uint> {
    let mut i = 0;
    for self.each |a| {
        if f(a) { return Some(i); }
        i += 1;
    }
    return None;
}

// note: 'rposition' would only make sense to provide with a bidirectional
// iter interface, such as would provide "reach" in addition to "each". As is,
// it would have to be implemented with foldr, which is too inefficient.

#[inline(always)]
pub fn repeat(times: uint, blk: &fn() -> bool) {
    let mut i = 0;
    while i < times {
        if !blk() { break }
        i += 1;
    }
}

#[inline(always)]
pub fn min<A:Copy + Ord,IA:BaseIter<A>>(self: &IA) -> A {
    match do foldl::<A,Option<A>,IA>(self, None) |a, b| {
        match a {
          &Some(ref a_) if *a_ < *b => {
             *(a)
          }
          _ => Some(*b)
        }
    } {
        Some(val) => val,
        None => fail!(~"min called on empty iterator")
    }
}

#[inline(always)]
pub fn max<A:Copy + Ord,IA:BaseIter<A>>(self: &IA) -> A {
    match do foldl::<A,Option<A>,IA>(self, None) |a, b| {
        match a {
          &Some(ref a_) if *a_ > *b => {
              *(a)
          }
          _ => Some(*b)
        }
    } {
        Some(val) => val,
        None => fail!(~"max called on empty iterator")
    }
}

#[inline(always)]
pub fn find<A:Copy,IA:BaseIter<A>>(self: &IA, f: &fn(&A) -> bool)
                                -> Option<A> {
    for self.each |i| {
        if f(i) { return Some(*i) }
    }
    return None;
}

// Some functions for just building

/**
 * Builds a sequence by calling a provided function with an argument
 * function that pushes an element to the back of a sequence.
 *
 * # Arguments
 *
 * * builder - A function that will construct the sequence. It receives
 *             as an argument a function that will push an element
 *             onto the sequence being constructed.
 */
#[inline(always)]
pub fn build<A,B: Buildable<A>>(builder: &fn(push: &fn(A))) -> B {
    Buildable::build_sized(4, builder)
}

/**
 * Builds a sequence by calling a provided function with an argument
 * function that pushes an element to the back of the sequence.
 * This version takes an initial size for the sequence.
 *
 * # Arguments
 *
 * * size - An option, maybe containing initial size of the sequence
 *          to reserve.
 * * builder - A function that will construct the sequence. It receives
 *             as an argument a function that will push an element
 *             onto the sequence being constructed.
 */
#[inline(always)]
pub fn build_sized_opt<A,B: Buildable<A>>(size: Option<uint>,
                                          builder: &fn(push: &fn(A))) -> B {
    Buildable::build_sized(size.get_or_default(4), builder)
}

// Functions that combine iteration and building

/// Applies a function to each element of an iterable and returns the results
/// in a sequence built via `BU`.  See also `map_to_vec`.
#[inline(always)]
pub fn map<T,IT: BaseIter<T>,U,BU: Buildable<U>>(v: &IT, f: &fn(&T) -> U)
    -> BU {
    do build_sized_opt(v.size_hint()) |push| {
        for v.each() |elem| {
            push(f(elem));
        }
    }
}

/**
 * Creates and initializes a generic sequence from a function.
 *
 * Creates a generic sequence of size `n_elts` and initializes the elements
 * to the value returned by the function `op`.
 */
#[inline(always)]
pub fn from_fn<T,BT: Buildable<T>>(n_elts: uint, op: InitOp<T>) -> BT {
    do Buildable::build_sized(n_elts) |push| {
        let mut i: uint = 0u;
        while i < n_elts { push(op(i)); i += 1u; }
    }
}

/**
 * Creates and initializes a generic sequence with some elements.
 *
 * Creates an immutable vector of size `n_elts` and initializes the elements
 * to the value `t`.
 */
#[inline(always)]
pub fn from_elem<T:Copy,BT:Buildable<T>>(n_elts: uint, t: T) -> BT {
    do Buildable::build_sized(n_elts) |push| {
        let mut i: uint = 0;
        while i < n_elts { push(t); i += 1; }
    }
}

/// Appends two generic sequences.
#[inline(always)]
pub fn append<T:Copy,IT:BaseIter<T>,BT:Buildable<T>>(lhs: &IT, rhs: &IT)
                                                  -> BT {
    let size_opt = lhs.size_hint().chain_ref(
        |sz1| rhs.size_hint().map(|sz2| *sz1+*sz2));
    do build_sized_opt(size_opt) |push| {
        for lhs.each |x| { push(*x); }
        for rhs.each |x| { push(*x); }
    }
}

/// Copies a generic sequence, possibly converting it to a different
/// type of sequence.
#[inline(always)]
pub fn copy_seq<T:Copy,IT:BaseIter<T>,BT:Buildable<T>>(v: &IT) -> BT {
    do build_sized_opt(v.size_hint()) |push| {
        for v.each |x| { push(*x); }
    }
}

/**
 * Helper function to transform an internal iterator into an owned vector.
 *
 * # Example:
 *
 * ~~~
 * let v = ~[1, 2, 3];
 * let v2 = do iter_to_vec |f| { v.each(|e| f(*e)) };
 * if v != v2 { fail!() }
 * ~~~
 */
#[inline(always)]
pub fn iter_to_vec<T>(pusher: &fn(it: &fn(T) -> bool)) -> ~[T] {
    let mut v = ~[];
    let pushf = |e| {v.push(e); true};
    pusher(pushf);
    v
}

#[test]
fn test_iter_to_vec() {
    let v = ~[1, 2, 3];
    let v2 = do iter_to_vec |f| { v.each(|e| f(*e)) };
    if v != v2 { fail!() }
}
<commit_msg>Avoid excessive allocations and copies in iter::to_vec<commit_after>// Copyright 2012-2013 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.

/*!

The iteration traits and common implementation

*/

use cmp::{Eq, Ord};
use kinds::Copy;
use option::{None, Option, Some};
use vec;

/// A function used to initialize the elements of a sequence
pub type InitOp<'self,T> = &'self fn(uint) -> T;

pub trait BaseIter<A> {
    fn each(&self, blk: &fn(v: &A) -> bool);
    fn size_hint(&self) -> Option<uint>;
}

pub trait ReverseIter<A>: BaseIter<A> {
    fn each_reverse(&self, blk: &fn(&A) -> bool);
}

pub trait MutableIter<A>: BaseIter<A> {
    fn each_mut(&mut self, blk: &fn(&mut A) -> bool);
}

pub trait ExtendedIter<A> {
    fn eachi(&self, blk: &fn(uint, v: &A) -> bool);
    fn all(&self, blk: &fn(&A) -> bool) -> bool;
    fn any(&self, blk: &fn(&A) -> bool) -> bool;
    fn foldl<B>(&self, b0: B, blk: &fn(&B, &A) -> B) -> B;
    fn position(&self, f: &fn(&A) -> bool) -> Option<uint>;
    fn map_to_vec<B>(&self, op: &fn(&A) -> B) -> ~[B];
    fn flat_map_to_vec<B,IB: BaseIter<B>>(&self, op: &fn(&A) -> IB) -> ~[B];
}

pub trait EqIter<A:Eq> {
    fn contains(&self, x: &A) -> bool;
    fn count(&self, x: &A) -> uint;
}

pub trait Times {
    fn times(&self, it: &fn() -> bool);
}

pub trait CopyableIter<A:Copy> {
    fn filter_to_vec(&self, pred: &fn(&A) -> bool) -> ~[A];
    fn to_vec(&self) -> ~[A];
    fn find(&self, p: &fn(&A) -> bool) -> Option<A>;
}

pub trait CopyableOrderedIter<A:Copy + Ord> {
    fn min(&self) -> A;
    fn max(&self) -> A;
}

pub trait CopyableNonstrictIter<A:Copy> {
    // Like "each", but copies out the value. If the receiver is mutated while
    // iterating over it, the semantics must not be memory-unsafe but are
    // otherwise undefined.
    fn each_val(&const self, f: &fn(A) -> bool);
}

// A trait for sequences that can be built by imperatively pushing elements
// onto them.
pub trait Buildable<A> {
    /**
     * Builds a buildable sequence by calling a provided function with
     * an argument function that pushes an element onto the back of
     * the sequence.
     * This version takes an initial size for the sequence.
     *
     * # Arguments
     *
     * * size - A hint for an initial size of the sequence
     * * builder - A function that will construct the sequence. It receives
     *             as an argument a function that will push an element
     *             onto the sequence being constructed.
     */
     fn build_sized(size: uint, builder: &fn(push: &fn(A))) -> Self;
}

#[inline(always)]
pub fn eachi<A,IA:BaseIter<A>>(self: &IA, blk: &fn(uint, &A) -> bool) {
    let mut i = 0;
    for self.each |a| {
        if !blk(i, a) { break; }
        i += 1;
    }
}

#[inline(always)]
pub fn all<A,IA:BaseIter<A>>(self: &IA, blk: &fn(&A) -> bool) -> bool {
    for self.each |a| {
        if !blk(a) { return false; }
    }
    return true;
}

#[inline(always)]
pub fn any<A,IA:BaseIter<A>>(self: &IA, blk: &fn(&A) -> bool) -> bool {
    for self.each |a| {
        if blk(a) { return true; }
    }
    return false;
}

#[inline(always)]
pub fn filter_to_vec<A:Copy,IA:BaseIter<A>>(self: &IA,
                                            prd: &fn(&A) -> bool)
                                         -> ~[A] {
    do vec::build_sized_opt(self.size_hint()) |push| {
        for self.each |a| {
            if prd(a) { push(*a); }
        }
    }
}

#[inline(always)]
pub fn map_to_vec<A,B,IA:BaseIter<A>>(self: &IA, op: &fn(&A) -> B) -> ~[B] {
    do vec::build_sized_opt(self.size_hint()) |push| {
        for self.each |a| {
            push(op(a));
        }
    }
}

#[inline(always)]
pub fn flat_map_to_vec<A,B,IA:BaseIter<A>,IB:BaseIter<B>>(self: &IA,
                                                          op: &fn(&A) -> IB)
                                                       -> ~[B] {
    do vec::build |push| {
        for self.each |a| {
            for op(a).each |&b| {
                push(b);
            }
        }
    }
}

#[inline(always)]
pub fn foldl<A,B,IA:BaseIter<A>>(self: &IA, b0: B, blk: &fn(&B, &A) -> B)
                              -> B {
    let mut b = b0;
    for self.each |a| {
        b = blk(&b, a);
    }
    b
}

#[inline(always)]
pub fn to_vec<A:Copy,IA:BaseIter<A>>(self: &IA) -> ~[A] {
    map_to_vec(self, |&x| x)
}

#[inline(always)]
pub fn contains<A:Eq,IA:BaseIter<A>>(self: &IA, x: &A) -> bool {
    for self.each |a| {
        if *a == *x { return true; }
    }
    return false;
}

#[inline(always)]
pub fn count<A:Eq,IA:BaseIter<A>>(self: &IA, x: &A) -> uint {
    do foldl(self, 0) |count, value| {
        if *value == *x {
            *count + 1
        } else {
            *count
        }
    }
}

#[inline(always)]
pub fn position<A,IA:BaseIter<A>>(self: &IA, f: &fn(&A) -> bool)
                               -> Option<uint> {
    let mut i = 0;
    for self.each |a| {
        if f(a) { return Some(i); }
        i += 1;
    }
    return None;
}

// note: 'rposition' would only make sense to provide with a bidirectional
// iter interface, such as would provide "reach" in addition to "each". As is,
// it would have to be implemented with foldr, which is too inefficient.

#[inline(always)]
pub fn repeat(times: uint, blk: &fn() -> bool) {
    let mut i = 0;
    while i < times {
        if !blk() { break }
        i += 1;
    }
}

#[inline(always)]
pub fn min<A:Copy + Ord,IA:BaseIter<A>>(self: &IA) -> A {
    match do foldl::<A,Option<A>,IA>(self, None) |a, b| {
        match a {
          &Some(ref a_) if *a_ < *b => {
             *(a)
          }
          _ => Some(*b)
        }
    } {
        Some(val) => val,
        None => fail!(~"min called on empty iterator")
    }
}

#[inline(always)]
pub fn max<A:Copy + Ord,IA:BaseIter<A>>(self: &IA) -> A {
    match do foldl::<A,Option<A>,IA>(self, None) |a, b| {
        match a {
          &Some(ref a_) if *a_ > *b => {
              *(a)
          }
          _ => Some(*b)
        }
    } {
        Some(val) => val,
        None => fail!(~"max called on empty iterator")
    }
}

#[inline(always)]
pub fn find<A:Copy,IA:BaseIter<A>>(self: &IA, f: &fn(&A) -> bool)
                                -> Option<A> {
    for self.each |i| {
        if f(i) { return Some(*i) }
    }
    return None;
}

// Some functions for just building

/**
 * Builds a sequence by calling a provided function with an argument
 * function that pushes an element to the back of a sequence.
 *
 * # Arguments
 *
 * * builder - A function that will construct the sequence. It receives
 *             as an argument a function that will push an element
 *             onto the sequence being constructed.
 */
#[inline(always)]
pub fn build<A,B: Buildable<A>>(builder: &fn(push: &fn(A))) -> B {
    Buildable::build_sized(4, builder)
}

/**
 * Builds a sequence by calling a provided function with an argument
 * function that pushes an element to the back of the sequence.
 * This version takes an initial size for the sequence.
 *
 * # Arguments
 *
 * * size - An option, maybe containing initial size of the sequence
 *          to reserve.
 * * builder - A function that will construct the sequence. It receives
 *             as an argument a function that will push an element
 *             onto the sequence being constructed.
 */
#[inline(always)]
pub fn build_sized_opt<A,B: Buildable<A>>(size: Option<uint>,
                                          builder: &fn(push: &fn(A))) -> B {
    Buildable::build_sized(size.get_or_default(4), builder)
}

// Functions that combine iteration and building

/// Applies a function to each element of an iterable and returns the results
/// in a sequence built via `BU`.  See also `map_to_vec`.
#[inline(always)]
pub fn map<T,IT: BaseIter<T>,U,BU: Buildable<U>>(v: &IT, f: &fn(&T) -> U)
    -> BU {
    do build_sized_opt(v.size_hint()) |push| {
        for v.each() |elem| {
            push(f(elem));
        }
    }
}

/**
 * Creates and initializes a generic sequence from a function.
 *
 * Creates a generic sequence of size `n_elts` and initializes the elements
 * to the value returned by the function `op`.
 */
#[inline(always)]
pub fn from_fn<T,BT: Buildable<T>>(n_elts: uint, op: InitOp<T>) -> BT {
    do Buildable::build_sized(n_elts) |push| {
        let mut i: uint = 0u;
        while i < n_elts { push(op(i)); i += 1u; }
    }
}

/**
 * Creates and initializes a generic sequence with some elements.
 *
 * Creates an immutable vector of size `n_elts` and initializes the elements
 * to the value `t`.
 */
#[inline(always)]
pub fn from_elem<T:Copy,BT:Buildable<T>>(n_elts: uint, t: T) -> BT {
    do Buildable::build_sized(n_elts) |push| {
        let mut i: uint = 0;
        while i < n_elts { push(t); i += 1; }
    }
}

/// Appends two generic sequences.
#[inline(always)]
pub fn append<T:Copy,IT:BaseIter<T>,BT:Buildable<T>>(lhs: &IT, rhs: &IT)
                                                  -> BT {
    let size_opt = lhs.size_hint().chain_ref(
        |sz1| rhs.size_hint().map(|sz2| *sz1+*sz2));
    do build_sized_opt(size_opt) |push| {
        for lhs.each |x| { push(*x); }
        for rhs.each |x| { push(*x); }
    }
}

/// Copies a generic sequence, possibly converting it to a different
/// type of sequence.
#[inline(always)]
pub fn copy_seq<T:Copy,IT:BaseIter<T>,BT:Buildable<T>>(v: &IT) -> BT {
    do build_sized_opt(v.size_hint()) |push| {
        for v.each |x| { push(*x); }
    }
}

/**
 * Helper function to transform an internal iterator into an owned vector.
 *
 * # Example:
 *
 * ~~~
 * let v = ~[1, 2, 3];
 * let v2 = do iter_to_vec |f| { v.each(|e| f(*e)) };
 * if v != v2 { fail!() }
 * ~~~
 */
#[inline(always)]
pub fn iter_to_vec<T>(pusher: &fn(it: &fn(T) -> bool)) -> ~[T] {
    let mut v = ~[];
    let pushf = |e| {v.push(e); true};
    pusher(pushf);
    v
}

#[test]
fn test_iter_to_vec() {
    let v = ~[1, 2, 3];
    let v2 = do iter_to_vec |f| { v.each(|e| f(*e)) };
    if v != v2 { fail!() }
}
<|endoftext|>